001/*
002 * Copyright (C) 2012 The Guava Authors
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package com.google.common.collect;
018
019import static com.google.common.base.Preconditions.checkArgument;
020import static com.google.common.base.Preconditions.checkNotNull;
021import static com.google.common.base.Predicates.compose;
022import static com.google.common.base.Predicates.in;
023import static com.google.common.base.Predicates.not;
024
025import com.google.common.annotations.Beta;
026import com.google.common.annotations.GwtIncompatible;
027import com.google.common.base.MoreObjects;
028import com.google.common.base.Predicate;
029import com.google.common.collect.Maps.IteratorBasedAbstractMap;
030import java.util.AbstractMap;
031import java.util.Collection;
032import java.util.Collections;
033import java.util.Iterator;
034import java.util.List;
035import java.util.Map;
036import java.util.Map.Entry;
037import java.util.NavigableMap;
038import java.util.NoSuchElementException;
039import java.util.Set;
040import javax.annotation.Nullable;
041
042/**
043 * An implementation of {@code RangeMap} based on a {@code TreeMap}, supporting
044 * all optional operations.
045 *
046 * <p>Like all {@code RangeMap} implementations, this supports neither null
047 * keys nor null values.
048 *
049 * @author Louis Wasserman
050 * @since 14.0
051 */
052@Beta
053@GwtIncompatible // NavigableMap
054public final class TreeRangeMap<K extends Comparable, V> implements RangeMap<K, V> {
055
056  private final NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound;
057
058  public static <K extends Comparable, V> TreeRangeMap<K, V> create() {
059    return new TreeRangeMap<K, V>();
060  }
061
062  private TreeRangeMap() {
063    this.entriesByLowerBound = Maps.newTreeMap();
064  }
065
066  private static final class RangeMapEntry<K extends Comparable, V>
067      extends AbstractMapEntry<Range<K>, V> {
068    private final Range<K> range;
069    private final V value;
070
071    RangeMapEntry(Cut<K> lowerBound, Cut<K> upperBound, V value) {
072      this(Range.create(lowerBound, upperBound), value);
073    }
074
075    RangeMapEntry(Range<K> range, V value) {
076      this.range = range;
077      this.value = value;
078    }
079
080    @Override
081    public Range<K> getKey() {
082      return range;
083    }
084
085    @Override
086    public V getValue() {
087      return value;
088    }
089
090    public boolean contains(K value) {
091      return range.contains(value);
092    }
093
094    Cut<K> getLowerBound() {
095      return range.lowerBound;
096    }
097
098    Cut<K> getUpperBound() {
099      return range.upperBound;
100    }
101  }
102
103  @Override
104  @Nullable
105  public V get(K key) {
106    Entry<Range<K>, V> entry = getEntry(key);
107    return (entry == null) ? null : entry.getValue();
108  }
109
110  @Override
111  @Nullable
112  public Entry<Range<K>, V> getEntry(K key) {
113    Map.Entry<Cut<K>, RangeMapEntry<K, V>> mapEntry =
114        entriesByLowerBound.floorEntry(Cut.belowValue(key));
115    if (mapEntry != null && mapEntry.getValue().contains(key)) {
116      return mapEntry.getValue();
117    } else {
118      return null;
119    }
120  }
121
122  @Override
123  public void put(Range<K> range, V value) {
124    // don't short-circuit if the range is empty - it may be between two ranges we can coalesce.
125    if (!range.isEmpty()) {
126      checkNotNull(value);
127      remove(range);
128      entriesByLowerBound.put(range.lowerBound, new RangeMapEntry<K, V>(range, value));
129    }
130  }
131
132  @Override
133  public void putCoalescing(Range<K> range, V value) {
134    if (entriesByLowerBound.isEmpty()) {
135      put(range, value);
136      return;
137    }
138
139    Range<K> coalescedRange = coalescedRange(range, checkNotNull(value));
140    put(coalescedRange, value);
141  }
142
143  /** Computes the coalesced range for the given range+value - does not mutate the map. */
144  private Range<K> coalescedRange(Range<K> range, V value) {
145    Range<K> coalescedRange = range;
146    Map.Entry<Cut<K>, RangeMapEntry<K, V>> lowerEntry =
147        entriesByLowerBound.lowerEntry(range.lowerBound);
148    coalescedRange = coalesce(coalescedRange, value, lowerEntry);
149
150    Map.Entry<Cut<K>, RangeMapEntry<K, V>> higherEntry =
151        entriesByLowerBound.floorEntry(range.upperBound);
152    coalescedRange = coalesce(coalescedRange, value, higherEntry);
153
154    return coalescedRange;
155  }
156
157  /** Returns the range that spans the given range and entry, if the entry can be coalesced. */
158  private static <K extends Comparable, V> Range<K> coalesce(
159      Range<K> range, V value, @Nullable Map.Entry<Cut<K>, RangeMapEntry<K, V>> entry) {
160    if (entry != null
161        && entry.getValue().getKey().isConnected(range)
162        && entry.getValue().getValue().equals(value)) {
163      return range.span(entry.getValue().getKey());
164    }
165    return range;
166  }
167
168  @Override
169  public void putAll(RangeMap<K, V> rangeMap) {
170    for (Map.Entry<Range<K>, V> entry : rangeMap.asMapOfRanges().entrySet()) {
171      put(entry.getKey(), entry.getValue());
172    }
173  }
174
175  @Override
176  public void clear() {
177    entriesByLowerBound.clear();
178  }
179
180  @Override
181  public Range<K> span() {
182    Entry<Cut<K>, RangeMapEntry<K, V>> firstEntry = entriesByLowerBound.firstEntry();
183    Entry<Cut<K>, RangeMapEntry<K, V>> lastEntry = entriesByLowerBound.lastEntry();
184    if (firstEntry == null) {
185      throw new NoSuchElementException();
186    }
187    return Range.create(
188        firstEntry.getValue().getKey().lowerBound, lastEntry.getValue().getKey().upperBound);
189  }
190
191  private void putRangeMapEntry(Cut<K> lowerBound, Cut<K> upperBound, V value) {
192    entriesByLowerBound.put(lowerBound, new RangeMapEntry<K, V>(lowerBound, upperBound, value));
193  }
194
195  @Override
196  public void remove(Range<K> rangeToRemove) {
197    if (rangeToRemove.isEmpty()) {
198      return;
199    }
200
201    /*
202     * The comments for this method will use [ ] to indicate the bounds of rangeToRemove and ( ) to
203     * indicate the bounds of ranges in the range map.
204     */
205    Map.Entry<Cut<K>, RangeMapEntry<K, V>> mapEntryBelowToTruncate =
206        entriesByLowerBound.lowerEntry(rangeToRemove.lowerBound);
207    if (mapEntryBelowToTruncate != null) {
208      // we know ( [
209      RangeMapEntry<K, V> rangeMapEntry = mapEntryBelowToTruncate.getValue();
210      if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.lowerBound) > 0) {
211        // we know ( [ )
212        if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.upperBound) > 0) {
213          // we know ( [ ] ), so insert the range ] ) back into the map --
214          // it's being split apart
215          putRangeMapEntry(
216              rangeToRemove.upperBound,
217              rangeMapEntry.getUpperBound(),
218              mapEntryBelowToTruncate.getValue().getValue());
219        }
220        // overwrite mapEntryToTruncateBelow with a truncated range
221        putRangeMapEntry(
222            rangeMapEntry.getLowerBound(),
223            rangeToRemove.lowerBound,
224            mapEntryBelowToTruncate.getValue().getValue());
225      }
226    }
227
228    Map.Entry<Cut<K>, RangeMapEntry<K, V>> mapEntryAboveToTruncate =
229        entriesByLowerBound.lowerEntry(rangeToRemove.upperBound);
230    if (mapEntryAboveToTruncate != null) {
231      // we know ( ]
232      RangeMapEntry<K, V> rangeMapEntry = mapEntryAboveToTruncate.getValue();
233      if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.upperBound) > 0) {
234        // we know ( ] ), and since we dealt with truncating below already,
235        // we know [ ( ] )
236        putRangeMapEntry(
237            rangeToRemove.upperBound,
238            rangeMapEntry.getUpperBound(),
239            mapEntryAboveToTruncate.getValue().getValue());
240      }
241    }
242    entriesByLowerBound.subMap(rangeToRemove.lowerBound, rangeToRemove.upperBound).clear();
243  }
244
245  @Override
246  public Map<Range<K>, V> asMapOfRanges() {
247    return new AsMapOfRanges(entriesByLowerBound.values());
248  }
249
250  @Override
251  public Map<Range<K>, V> asDescendingMapOfRanges() {
252    return new AsMapOfRanges(entriesByLowerBound.descendingMap().values());
253  }
254
255  private final class AsMapOfRanges extends IteratorBasedAbstractMap<Range<K>, V> {
256
257    final Iterable<Entry<Range<K>, V>> entryIterable;
258
259    @SuppressWarnings("unchecked") // it's safe to upcast iterables
260    AsMapOfRanges(Iterable<RangeMapEntry<K, V>> entryIterable) {
261      this.entryIterable = (Iterable) entryIterable;
262    }
263
264    @Override
265    public boolean containsKey(@Nullable Object key) {
266      return get(key) != null;
267    }
268
269    @Override
270    public V get(@Nullable Object key) {
271      if (key instanceof Range) {
272        Range<?> range = (Range<?>) key;
273        RangeMapEntry<K, V> rangeMapEntry = entriesByLowerBound.get(range.lowerBound);
274        if (rangeMapEntry != null && rangeMapEntry.getKey().equals(range)) {
275          return rangeMapEntry.getValue();
276        }
277      }
278      return null;
279    }
280
281    @Override
282    public int size() {
283      return entriesByLowerBound.size();
284    }
285
286    @Override
287    Iterator<Entry<Range<K>, V>> entryIterator() {
288      return entryIterable.iterator();
289    }
290  }
291
292  @Override
293  public RangeMap<K, V> subRangeMap(Range<K> subRange) {
294    if (subRange.equals(Range.all())) {
295      return this;
296    } else {
297      return new SubRangeMap(subRange);
298    }
299  }
300
301  @SuppressWarnings("unchecked")
302  private RangeMap<K, V> emptySubRangeMap() {
303    return EMPTY_SUB_RANGE_MAP;
304  }
305
306  private static final RangeMap EMPTY_SUB_RANGE_MAP =
307      new RangeMap() {
308        @Override
309        @Nullable
310        public Object get(Comparable key) {
311          return null;
312        }
313
314        @Override
315        @Nullable
316        public Entry<Range, Object> getEntry(Comparable key) {
317          return null;
318        }
319
320        @Override
321        public Range span() {
322          throw new NoSuchElementException();
323        }
324
325        @Override
326        public void put(Range range, Object value) {
327          checkNotNull(range);
328          throw new IllegalArgumentException(
329              "Cannot insert range " + range + " into an empty subRangeMap");
330        }
331
332        @Override
333        public void putCoalescing(Range range, Object value) {
334          checkNotNull(range);
335          throw new IllegalArgumentException(
336              "Cannot insert range " + range + " into an empty subRangeMap");
337        }
338
339        @Override
340        public void putAll(RangeMap rangeMap) {
341          if (!rangeMap.asMapOfRanges().isEmpty()) {
342            throw new IllegalArgumentException(
343                "Cannot putAll(nonEmptyRangeMap) into an empty subRangeMap");
344          }
345        }
346
347        @Override
348        public void clear() {}
349
350        @Override
351        public void remove(Range range) {
352          checkNotNull(range);
353        }
354
355        @Override
356        public Map<Range, Object> asMapOfRanges() {
357          return Collections.emptyMap();
358        }
359
360        @Override
361        public Map<Range, Object> asDescendingMapOfRanges() {
362          return Collections.emptyMap();
363        }
364
365        @Override
366        public RangeMap subRangeMap(Range range) {
367          checkNotNull(range);
368          return this;
369        }
370      };
371
372  private class SubRangeMap implements RangeMap<K, V> {
373
374    private final Range<K> subRange;
375
376    SubRangeMap(Range<K> subRange) {
377      this.subRange = subRange;
378    }
379
380    @Override
381    @Nullable
382    public V get(K key) {
383      return subRange.contains(key) ? TreeRangeMap.this.get(key) : null;
384    }
385
386    @Override
387    @Nullable
388    public Entry<Range<K>, V> getEntry(K key) {
389      if (subRange.contains(key)) {
390        Entry<Range<K>, V> entry = TreeRangeMap.this.getEntry(key);
391        if (entry != null) {
392          return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
393        }
394      }
395      return null;
396    }
397
398    @Override
399    public Range<K> span() {
400      Cut<K> lowerBound;
401      Entry<Cut<K>, RangeMapEntry<K, V>> lowerEntry =
402          entriesByLowerBound.floorEntry(subRange.lowerBound);
403      if (lowerEntry != null
404          && lowerEntry.getValue().getUpperBound().compareTo(subRange.lowerBound) > 0) {
405        lowerBound = subRange.lowerBound;
406      } else {
407        lowerBound = entriesByLowerBound.ceilingKey(subRange.lowerBound);
408        if (lowerBound == null || lowerBound.compareTo(subRange.upperBound) >= 0) {
409          throw new NoSuchElementException();
410        }
411      }
412
413      Cut<K> upperBound;
414      Entry<Cut<K>, RangeMapEntry<K, V>> upperEntry =
415          entriesByLowerBound.lowerEntry(subRange.upperBound);
416      if (upperEntry == null) {
417        throw new NoSuchElementException();
418      } else if (upperEntry.getValue().getUpperBound().compareTo(subRange.upperBound) >= 0) {
419        upperBound = subRange.upperBound;
420      } else {
421        upperBound = upperEntry.getValue().getUpperBound();
422      }
423      return Range.create(lowerBound, upperBound);
424    }
425
426    @Override
427    public void put(Range<K> range, V value) {
428      checkArgument(
429          subRange.encloses(range), "Cannot put range %s into a subRangeMap(%s)", range, subRange);
430      TreeRangeMap.this.put(range, value);
431    }
432
433    @Override
434    public void putCoalescing(Range<K> range, V value) {
435      if (entriesByLowerBound.isEmpty() || range.isEmpty() || !subRange.encloses(range)) {
436        put(range, value);
437        return;
438      }
439
440      Range<K> coalescedRange = coalescedRange(range, checkNotNull(value));
441      // only coalesce ranges within the subRange
442      put(coalescedRange.intersection(subRange), value);
443    }
444
445    @Override
446    public void putAll(RangeMap<K, V> rangeMap) {
447      if (rangeMap.asMapOfRanges().isEmpty()) {
448        return;
449      }
450      Range<K> span = rangeMap.span();
451      checkArgument(
452          subRange.encloses(span),
453          "Cannot putAll rangeMap with span %s into a subRangeMap(%s)",
454          span,
455          subRange);
456      TreeRangeMap.this.putAll(rangeMap);
457    }
458
459    @Override
460    public void clear() {
461      TreeRangeMap.this.remove(subRange);
462    }
463
464    @Override
465    public void remove(Range<K> range) {
466      if (range.isConnected(subRange)) {
467        TreeRangeMap.this.remove(range.intersection(subRange));
468      }
469    }
470
471    @Override
472    public RangeMap<K, V> subRangeMap(Range<K> range) {
473      if (!range.isConnected(subRange)) {
474        return emptySubRangeMap();
475      } else {
476        return TreeRangeMap.this.subRangeMap(range.intersection(subRange));
477      }
478    }
479
480    @Override
481    public Map<Range<K>, V> asMapOfRanges() {
482      return new SubRangeMapAsMap();
483    }
484
485    @Override
486    public Map<Range<K>, V> asDescendingMapOfRanges() {
487      return new SubRangeMapAsMap() {
488
489        @Override
490        Iterator<Entry<Range<K>, V>> entryIterator() {
491          if (subRange.isEmpty()) {
492            return Iterators.emptyIterator();
493          }
494          final Iterator<RangeMapEntry<K, V>> backingItr =
495              entriesByLowerBound
496                  .headMap(subRange.upperBound, false)
497                  .descendingMap()
498                  .values()
499                  .iterator();
500          return new AbstractIterator<Entry<Range<K>, V>>() {
501
502            @Override
503            protected Entry<Range<K>, V> computeNext() {
504              if (backingItr.hasNext()) {
505                RangeMapEntry<K, V> entry = backingItr.next();
506                if (entry.getUpperBound().compareTo(subRange.lowerBound) <= 0) {
507                  return endOfData();
508                }
509                return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
510              }
511              return endOfData();
512            }
513          };
514        }
515      };
516    }
517
518    @Override
519    public boolean equals(@Nullable Object o) {
520      if (o instanceof RangeMap) {
521        RangeMap<?, ?> rangeMap = (RangeMap<?, ?>) o;
522        return asMapOfRanges().equals(rangeMap.asMapOfRanges());
523      }
524      return false;
525    }
526
527    @Override
528    public int hashCode() {
529      return asMapOfRanges().hashCode();
530    }
531
532    @Override
533    public String toString() {
534      return asMapOfRanges().toString();
535    }
536
537    class SubRangeMapAsMap extends AbstractMap<Range<K>, V> {
538
539      @Override
540      public boolean containsKey(Object key) {
541        return get(key) != null;
542      }
543
544      @Override
545      public V get(Object key) {
546        try {
547          if (key instanceof Range) {
548            @SuppressWarnings("unchecked") // we catch ClassCastExceptions
549            Range<K> r = (Range<K>) key;
550            if (!subRange.encloses(r) || r.isEmpty()) {
551              return null;
552            }
553            RangeMapEntry<K, V> candidate = null;
554            if (r.lowerBound.compareTo(subRange.lowerBound) == 0) {
555              // r could be truncated on the left
556              Entry<Cut<K>, RangeMapEntry<K, V>> entry =
557                  entriesByLowerBound.floorEntry(r.lowerBound);
558              if (entry != null) {
559                candidate = entry.getValue();
560              }
561            } else {
562              candidate = entriesByLowerBound.get(r.lowerBound);
563            }
564
565            if (candidate != null
566                && candidate.getKey().isConnected(subRange)
567                && candidate.getKey().intersection(subRange).equals(r)) {
568              return candidate.getValue();
569            }
570          }
571        } catch (ClassCastException e) {
572          return null;
573        }
574        return null;
575      }
576
577      @Override
578      public V remove(Object key) {
579        V value = get(key);
580        if (value != null) {
581          @SuppressWarnings("unchecked") // it's definitely in the map, so safe
582          Range<K> range = (Range<K>) key;
583          TreeRangeMap.this.remove(range);
584          return value;
585        }
586        return null;
587      }
588
589      @Override
590      public void clear() {
591        SubRangeMap.this.clear();
592      }
593
594      private boolean removeEntryIf(Predicate<? super Entry<Range<K>, V>> predicate) {
595        List<Range<K>> toRemove = Lists.newArrayList();
596        for (Entry<Range<K>, V> entry : entrySet()) {
597          if (predicate.apply(entry)) {
598            toRemove.add(entry.getKey());
599          }
600        }
601        for (Range<K> range : toRemove) {
602          TreeRangeMap.this.remove(range);
603        }
604        return !toRemove.isEmpty();
605      }
606
607      @Override
608      public Set<Range<K>> keySet() {
609        return new Maps.KeySet<Range<K>, V>(SubRangeMapAsMap.this) {
610          @Override
611          public boolean remove(@Nullable Object o) {
612            return SubRangeMapAsMap.this.remove(o) != null;
613          }
614
615          @Override
616          public boolean retainAll(Collection<?> c) {
617            return removeEntryIf(compose(not(in(c)), Maps.<Range<K>>keyFunction()));
618          }
619        };
620      }
621
622      @Override
623      public Set<Entry<Range<K>, V>> entrySet() {
624        return new Maps.EntrySet<Range<K>, V>() {
625          @Override
626          Map<Range<K>, V> map() {
627            return SubRangeMapAsMap.this;
628          }
629
630          @Override
631          public Iterator<Entry<Range<K>, V>> iterator() {
632            return entryIterator();
633          }
634
635          @Override
636          public boolean retainAll(Collection<?> c) {
637            return removeEntryIf(not(in(c)));
638          }
639
640          @Override
641          public int size() {
642            return Iterators.size(iterator());
643          }
644
645          @Override
646          public boolean isEmpty() {
647            return !iterator().hasNext();
648          }
649        };
650      }
651
652      Iterator<Entry<Range<K>, V>> entryIterator() {
653        if (subRange.isEmpty()) {
654          return Iterators.emptyIterator();
655        }
656        Cut<K> cutToStart =
657            MoreObjects.firstNonNull(
658                entriesByLowerBound.floorKey(subRange.lowerBound), subRange.lowerBound);
659        final Iterator<RangeMapEntry<K, V>> backingItr =
660            entriesByLowerBound.tailMap(cutToStart, true).values().iterator();
661        return new AbstractIterator<Entry<Range<K>, V>>() {
662
663          @Override
664          protected Entry<Range<K>, V> computeNext() {
665            while (backingItr.hasNext()) {
666              RangeMapEntry<K, V> entry = backingItr.next();
667              if (entry.getLowerBound().compareTo(subRange.upperBound) >= 0) {
668                return endOfData();
669              } else if (entry.getUpperBound().compareTo(subRange.lowerBound) > 0) {
670                // this might not be true e.g. at the start of the iteration
671                return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
672              }
673            }
674            return endOfData();
675          }
676        };
677      }
678
679      @Override
680      public Collection<V> values() {
681        return new Maps.Values<Range<K>, V>(this) {
682          @Override
683          public boolean removeAll(Collection<?> c) {
684            return removeEntryIf(compose(in(c), Maps.<V>valueFunction()));
685          }
686
687          @Override
688          public boolean retainAll(Collection<?> c) {
689            return removeEntryIf(compose(not(in(c)), Maps.<V>valueFunction()));
690          }
691        };
692      }
693    }
694  }
695
696  @Override
697  public boolean equals(@Nullable Object o) {
698    if (o instanceof RangeMap) {
699      RangeMap<?, ?> rangeMap = (RangeMap<?, ?>) o;
700      return asMapOfRanges().equals(rangeMap.asMapOfRanges());
701    }
702    return false;
703  }
704
705  @Override
706  public int hashCode() {
707    return asMapOfRanges().hashCode();
708  }
709
710  @Override
711  public String toString() {
712    return entriesByLowerBound.values().toString();
713  }
714}