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