Source code

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