com.facebook.presto.jdbc.internal.guava.collect.TreeRangeMap Maven / Gradle / Ivy
/*
* Copyright (C) 2012 The Guava Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.facebook.presto.jdbc.internal.guava.collect;
import static com.facebook.presto.jdbc.internal.guava.base.Preconditions.checkArgument;
import static com.facebook.presto.jdbc.internal.guava.base.Preconditions.checkNotNull;
import static com.facebook.presto.jdbc.internal.guava.base.Predicates.compose;
import static com.facebook.presto.jdbc.internal.guava.base.Predicates.in;
import static com.facebook.presto.jdbc.internal.guava.base.Predicates.not;
import com.facebook.presto.jdbc.internal.guava.annotations.Beta;
import com.facebook.presto.jdbc.internal.guava.annotations.GwtIncompatible;
import com.facebook.presto.jdbc.internal.guava.base.MoreObjects;
import com.facebook.presto.jdbc.internal.guava.base.Predicate;
import com.facebook.presto.jdbc.internal.guava.collect.Maps.IteratorBasedAbstractMap;
import java.util.AbstractMap;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.NavigableMap;
import java.util.NoSuchElementException;
import java.util.Set;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
* An implementation of {@code RangeMap} based on a {@code TreeMap}, supporting all optional
* operations.
*
* Like all {@code RangeMap} implementations, this supports neither null keys nor null values.
*
* @author Louis Wasserman
* @since 14.0
*/
@Beta
@GwtIncompatible // NavigableMap
public final class TreeRangeMap implements RangeMap {
private final NavigableMap, RangeMapEntry> entriesByLowerBound;
public static TreeRangeMap create() {
return new TreeRangeMap<>();
}
private TreeRangeMap() {
this.entriesByLowerBound = Maps.newTreeMap();
}
private static final class RangeMapEntry
extends AbstractMapEntry, V> {
private final Range range;
private final V value;
RangeMapEntry(Cut lowerBound, Cut upperBound, V value) {
this(Range.create(lowerBound, upperBound), value);
}
RangeMapEntry(Range range, V value) {
this.range = range;
this.value = value;
}
@Override
public Range getKey() {
return range;
}
@Override
public V getValue() {
return value;
}
public boolean contains(K value) {
return range.contains(value);
}
Cut getLowerBound() {
return range.lowerBound;
}
Cut getUpperBound() {
return range.upperBound;
}
}
@Override
public @Nullable V get(K key) {
Entry, V> entry = getEntry(key);
return (entry == null) ? null : entry.getValue();
}
@Override
public @Nullable Entry, V> getEntry(K key) {
Entry, RangeMapEntry> mapEntry =
entriesByLowerBound.floorEntry(Cut.belowValue(key));
if (mapEntry != null && mapEntry.getValue().contains(key)) {
return mapEntry.getValue();
} else {
return null;
}
}
@Override
public void put(Range range, V value) {
// don't short-circuit if the range is empty - it may be between two ranges we can coalesce.
if (!range.isEmpty()) {
checkNotNull(value);
remove(range);
entriesByLowerBound.put(range.lowerBound, new RangeMapEntry(range, value));
}
}
@Override
public void putCoalescing(Range range, V value) {
if (entriesByLowerBound.isEmpty()) {
put(range, value);
return;
}
Range coalescedRange = coalescedRange(range, checkNotNull(value));
put(coalescedRange, value);
}
/** Computes the coalesced range for the given range+value - does not mutate the map. */
private Range coalescedRange(Range range, V value) {
Range coalescedRange = range;
Entry, RangeMapEntry> lowerEntry =
entriesByLowerBound.lowerEntry(range.lowerBound);
coalescedRange = coalesce(coalescedRange, value, lowerEntry);
Entry, RangeMapEntry> higherEntry =
entriesByLowerBound.floorEntry(range.upperBound);
coalescedRange = coalesce(coalescedRange, value, higherEntry);
return coalescedRange;
}
/** Returns the range that spans the given range and entry, if the entry can be coalesced. */
private static Range coalesce(
Range range, V value, @Nullable Entry, RangeMapEntry> entry) {
if (entry != null
&& entry.getValue().getKey().isConnected(range)
&& entry.getValue().getValue().equals(value)) {
return range.span(entry.getValue().getKey());
}
return range;
}
@Override
public void putAll(RangeMap rangeMap) {
for (Entry, V> entry : rangeMap.asMapOfRanges().entrySet()) {
put(entry.getKey(), entry.getValue());
}
}
@Override
public void clear() {
entriesByLowerBound.clear();
}
@Override
public Range span() {
Entry, RangeMapEntry> firstEntry = entriesByLowerBound.firstEntry();
Entry, RangeMapEntry> lastEntry = entriesByLowerBound.lastEntry();
if (firstEntry == null) {
throw new NoSuchElementException();
}
return Range.create(
firstEntry.getValue().getKey().lowerBound, lastEntry.getValue().getKey().upperBound);
}
private void putRangeMapEntry(Cut lowerBound, Cut upperBound, V value) {
entriesByLowerBound.put(lowerBound, new RangeMapEntry(lowerBound, upperBound, value));
}
@Override
public void remove(Range rangeToRemove) {
if (rangeToRemove.isEmpty()) {
return;
}
/*
* The comments for this method will use [ ] to indicate the bounds of rangeToRemove and ( ) to
* indicate the bounds of ranges in the range map.
*/
Entry, RangeMapEntry> mapEntryBelowToTruncate =
entriesByLowerBound.lowerEntry(rangeToRemove.lowerBound);
if (mapEntryBelowToTruncate != null) {
// we know ( [
RangeMapEntry rangeMapEntry = mapEntryBelowToTruncate.getValue();
if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.lowerBound) > 0) {
// we know ( [ )
if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.upperBound) > 0) {
// we know ( [ ] ), so insert the range ] ) back into the map --
// it's being split apart
putRangeMapEntry(
rangeToRemove.upperBound,
rangeMapEntry.getUpperBound(),
mapEntryBelowToTruncate.getValue().getValue());
}
// overwrite mapEntryToTruncateBelow with a truncated range
putRangeMapEntry(
rangeMapEntry.getLowerBound(),
rangeToRemove.lowerBound,
mapEntryBelowToTruncate.getValue().getValue());
}
}
Entry, RangeMapEntry> mapEntryAboveToTruncate =
entriesByLowerBound.lowerEntry(rangeToRemove.upperBound);
if (mapEntryAboveToTruncate != null) {
// we know ( ]
RangeMapEntry rangeMapEntry = mapEntryAboveToTruncate.getValue();
if (rangeMapEntry.getUpperBound().compareTo(rangeToRemove.upperBound) > 0) {
// we know ( ] ), and since we dealt with truncating below already,
// we know [ ( ] )
putRangeMapEntry(
rangeToRemove.upperBound,
rangeMapEntry.getUpperBound(),
mapEntryAboveToTruncate.getValue().getValue());
}
}
entriesByLowerBound.subMap(rangeToRemove.lowerBound, rangeToRemove.upperBound).clear();
}
@Override
public Map, V> asMapOfRanges() {
return new AsMapOfRanges(entriesByLowerBound.values());
}
@Override
public Map, V> asDescendingMapOfRanges() {
return new AsMapOfRanges(entriesByLowerBound.descendingMap().values());
}
private final class AsMapOfRanges extends IteratorBasedAbstractMap, V> {
final Iterable, V>> entryIterable;
@SuppressWarnings("unchecked") // it's safe to upcast iterables
AsMapOfRanges(Iterable> entryIterable) {
this.entryIterable = (Iterable) entryIterable;
}
@Override
public boolean containsKey(@Nullable Object key) {
return get(key) != null;
}
@Override
public V get(@Nullable Object key) {
if (key instanceof Range) {
Range range = (Range) key;
RangeMapEntry rangeMapEntry = entriesByLowerBound.get(range.lowerBound);
if (rangeMapEntry != null && rangeMapEntry.getKey().equals(range)) {
return rangeMapEntry.getValue();
}
}
return null;
}
@Override
public int size() {
return entriesByLowerBound.size();
}
@Override
Iterator, V>> entryIterator() {
return entryIterable.iterator();
}
}
@Override
public RangeMap subRangeMap(Range subRange) {
if (subRange.equals(Range.all())) {
return this;
} else {
return new SubRangeMap(subRange);
}
}
@SuppressWarnings("unchecked")
private RangeMap emptySubRangeMap() {
return EMPTY_SUB_RANGE_MAP;
}
private static final RangeMap EMPTY_SUB_RANGE_MAP =
new RangeMap() {
@Override
public @Nullable Object get(Comparable key) {
return null;
}
@Override
public @Nullable Entry getEntry(Comparable key) {
return null;
}
@Override
public Range span() {
throw new NoSuchElementException();
}
@Override
public void put(Range range, Object value) {
checkNotNull(range);
throw new IllegalArgumentException(
"Cannot insert range " + range + " into an empty subRangeMap");
}
@Override
public void putCoalescing(Range range, Object value) {
checkNotNull(range);
throw new IllegalArgumentException(
"Cannot insert range " + range + " into an empty subRangeMap");
}
@Override
public void putAll(RangeMap rangeMap) {
if (!rangeMap.asMapOfRanges().isEmpty()) {
throw new IllegalArgumentException(
"Cannot putAll(nonEmptyRangeMap) into an empty subRangeMap");
}
}
@Override
public void clear() {}
@Override
public void remove(Range range) {
checkNotNull(range);
}
@Override
public Map asMapOfRanges() {
return Collections.emptyMap();
}
@Override
public Map asDescendingMapOfRanges() {
return Collections.emptyMap();
}
@Override
public RangeMap subRangeMap(Range range) {
checkNotNull(range);
return this;
}
};
private class SubRangeMap implements RangeMap {
private final Range subRange;
SubRangeMap(Range subRange) {
this.subRange = subRange;
}
@Override
public @Nullable V get(K key) {
return subRange.contains(key) ? TreeRangeMap.this.get(key) : null;
}
@Override
public @Nullable Entry, V> getEntry(K key) {
if (subRange.contains(key)) {
Entry, V> entry = TreeRangeMap.this.getEntry(key);
if (entry != null) {
return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
}
}
return null;
}
@Override
public Range span() {
Cut lowerBound;
Entry, RangeMapEntry> lowerEntry =
entriesByLowerBound.floorEntry(subRange.lowerBound);
if (lowerEntry != null
&& lowerEntry.getValue().getUpperBound().compareTo(subRange.lowerBound) > 0) {
lowerBound = subRange.lowerBound;
} else {
lowerBound = entriesByLowerBound.ceilingKey(subRange.lowerBound);
if (lowerBound == null || lowerBound.compareTo(subRange.upperBound) >= 0) {
throw new NoSuchElementException();
}
}
Cut upperBound;
Entry, RangeMapEntry> upperEntry =
entriesByLowerBound.lowerEntry(subRange.upperBound);
if (upperEntry == null) {
throw new NoSuchElementException();
} else if (upperEntry.getValue().getUpperBound().compareTo(subRange.upperBound) >= 0) {
upperBound = subRange.upperBound;
} else {
upperBound = upperEntry.getValue().getUpperBound();
}
return Range.create(lowerBound, upperBound);
}
@Override
public void put(Range range, V value) {
checkArgument(
subRange.encloses(range), "Cannot put range %s into a subRangeMap(%s)", range, subRange);
TreeRangeMap.this.put(range, value);
}
@Override
public void putCoalescing(Range range, V value) {
if (entriesByLowerBound.isEmpty() || range.isEmpty() || !subRange.encloses(range)) {
put(range, value);
return;
}
Range coalescedRange = coalescedRange(range, checkNotNull(value));
// only coalesce ranges within the subRange
put(coalescedRange.intersection(subRange), value);
}
@Override
public void putAll(RangeMap rangeMap) {
if (rangeMap.asMapOfRanges().isEmpty()) {
return;
}
Range span = rangeMap.span();
checkArgument(
subRange.encloses(span),
"Cannot putAll rangeMap with span %s into a subRangeMap(%s)",
span,
subRange);
TreeRangeMap.this.putAll(rangeMap);
}
@Override
public void clear() {
TreeRangeMap.this.remove(subRange);
}
@Override
public void remove(Range range) {
if (range.isConnected(subRange)) {
TreeRangeMap.this.remove(range.intersection(subRange));
}
}
@Override
public RangeMap subRangeMap(Range range) {
if (!range.isConnected(subRange)) {
return emptySubRangeMap();
} else {
return TreeRangeMap.this.subRangeMap(range.intersection(subRange));
}
}
@Override
public Map, V> asMapOfRanges() {
return new SubRangeMapAsMap();
}
@Override
public Map, V> asDescendingMapOfRanges() {
return new SubRangeMapAsMap() {
@Override
Iterator, V>> entryIterator() {
if (subRange.isEmpty()) {
return Iterators.emptyIterator();
}
final Iterator> backingItr =
entriesByLowerBound
.headMap(subRange.upperBound, false)
.descendingMap()
.values()
.iterator();
return new AbstractIterator, V>>() {
@Override
protected Entry, V> computeNext() {
if (backingItr.hasNext()) {
RangeMapEntry entry = backingItr.next();
if (entry.getUpperBound().compareTo(subRange.lowerBound) <= 0) {
return endOfData();
}
return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
}
return endOfData();
}
};
}
};
}
@Override
public boolean equals(@Nullable Object o) {
if (o instanceof RangeMap) {
RangeMap rangeMap = (RangeMap) o;
return asMapOfRanges().equals(rangeMap.asMapOfRanges());
}
return false;
}
@Override
public int hashCode() {
return asMapOfRanges().hashCode();
}
@Override
public String toString() {
return asMapOfRanges().toString();
}
class SubRangeMapAsMap extends AbstractMap, V> {
@Override
public boolean containsKey(Object key) {
return get(key) != null;
}
@Override
public V get(Object key) {
try {
if (key instanceof Range) {
@SuppressWarnings("unchecked") // we catch ClassCastExceptions
Range r = (Range) key;
if (!subRange.encloses(r) || r.isEmpty()) {
return null;
}
RangeMapEntry candidate = null;
if (r.lowerBound.compareTo(subRange.lowerBound) == 0) {
// r could be truncated on the left
Entry, RangeMapEntry> entry =
entriesByLowerBound.floorEntry(r.lowerBound);
if (entry != null) {
candidate = entry.getValue();
}
} else {
candidate = entriesByLowerBound.get(r.lowerBound);
}
if (candidate != null
&& candidate.getKey().isConnected(subRange)
&& candidate.getKey().intersection(subRange).equals(r)) {
return candidate.getValue();
}
}
} catch (ClassCastException e) {
return null;
}
return null;
}
@Override
public V remove(Object key) {
V value = get(key);
if (value != null) {
@SuppressWarnings("unchecked") // it's definitely in the map, so safe
Range range = (Range) key;
TreeRangeMap.this.remove(range);
return value;
}
return null;
}
@Override
public void clear() {
SubRangeMap.this.clear();
}
private boolean removeEntryIf(Predicate, V>> predicate) {
List> toRemove = Lists.newArrayList();
for (Entry, V> entry : entrySet()) {
if (predicate.apply(entry)) {
toRemove.add(entry.getKey());
}
}
for (Range range : toRemove) {
TreeRangeMap.this.remove(range);
}
return !toRemove.isEmpty();
}
@Override
public Set> keySet() {
return new Maps.KeySet, V>(SubRangeMapAsMap.this) {
@Override
public boolean remove(@Nullable Object o) {
return SubRangeMapAsMap.this.remove(o) != null;
}
@Override
public boolean retainAll(Collection c) {
return removeEntryIf(compose(not(in(c)), Maps.>keyFunction()));
}
};
}
@Override
public Set, V>> entrySet() {
return new Maps.EntrySet, V>() {
@Override
Map, V> map() {
return SubRangeMapAsMap.this;
}
@Override
public Iterator, V>> iterator() {
return entryIterator();
}
@Override
public boolean retainAll(Collection c) {
return removeEntryIf(not(in(c)));
}
@Override
public int size() {
return Iterators.size(iterator());
}
@Override
public boolean isEmpty() {
return !iterator().hasNext();
}
};
}
Iterator, V>> entryIterator() {
if (subRange.isEmpty()) {
return Iterators.emptyIterator();
}
Cut cutToStart =
MoreObjects.firstNonNull(
entriesByLowerBound.floorKey(subRange.lowerBound), subRange.lowerBound);
final Iterator> backingItr =
entriesByLowerBound.tailMap(cutToStart, true).values().iterator();
return new AbstractIterator, V>>() {
@Override
protected Entry, V> computeNext() {
while (backingItr.hasNext()) {
RangeMapEntry entry = backingItr.next();
if (entry.getLowerBound().compareTo(subRange.upperBound) >= 0) {
return endOfData();
} else if (entry.getUpperBound().compareTo(subRange.lowerBound) > 0) {
// this might not be true e.g. at the start of the iteration
return Maps.immutableEntry(entry.getKey().intersection(subRange), entry.getValue());
}
}
return endOfData();
}
};
}
@Override
public Collection values() {
return new Maps.Values, V>(this) {
@Override
public boolean removeAll(Collection c) {
return removeEntryIf(compose(in(c), Maps.valueFunction()));
}
@Override
public boolean retainAll(Collection c) {
return removeEntryIf(compose(not(in(c)), Maps.valueFunction()));
}
};
}
}
}
@Override
public boolean equals(@Nullable Object o) {
if (o instanceof RangeMap) {
RangeMap rangeMap = (RangeMap) o;
return asMapOfRanges().equals(rangeMap.asMapOfRanges());
}
return false;
}
@Override
public int hashCode() {
return asMapOfRanges().hashCode();
}
@Override
public String toString() {
return entriesByLowerBound.values().toString();
}
}