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