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