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