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