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