Provide an optimized copyOf method for TreeRangeMap

RELNOTES=Provide an optimized copyOf method for TreeRangeMap PiperOrigin-RevId: 682878547
google · Oct 6, 2024 · a46565d · a46565d
1 parent 81be061
commit a46565d
Show file tree

Hide file tree

Showing 4 changed files with 138 additions and 0 deletions.
diff --git a/android/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java b/android/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java
@@ -25,6 +25,7 @@
 import com.google.common.collect.testing.features.CollectionFeature;
 import com.google.common.collect.testing.features.CollectionSize;
 import com.google.common.collect.testing.features.MapFeature;
+import com.google.common.testing.EqualsTester;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -702,6 +703,53 @@ public void testSubRangeMapClear() {
         ImmutableMap.of(Range.open(3, 5), 1, Range.closed(12, 16), 3), rangeMap.asMapOfRanges());
   }
 
+  public void testCopyOfTreeRangeMap() {
+    RangeMap<Integer, Integer> rangeMap = TreeRangeMap.create();
+    rangeMap.put(Range.open(3, 7), 1);
+    rangeMap.put(Range.closed(9, 10), 2);
+    rangeMap.put(Range.closed(12, 16), 3);
+
+    RangeMap<Integer, Integer> copy = TreeRangeMap.copyOf(rangeMap);
+
+    assertEquals(rangeMap.asMapOfRanges(), copy.asMapOfRanges());
+  }
+
+  public void testCopyOfImmutableRangeMap() {
+    ImmutableRangeMap<Integer, Integer> rangeMap =
+        ImmutableRangeMap.<Integer, Integer>builder()
+            .put(Range.open(3, 7), 1)
+            .put(Range.closed(9, 10), 2)
+            .put(Range.closed(12, 16), 3)
+            .build();
+
+    RangeMap<Integer, Integer> copy = TreeRangeMap.copyOf(rangeMap);
+
+    assertEquals(rangeMap.asMapOfRanges(), copy.asMapOfRanges());
+  }
+
+  // Overriding testEquals because it seems that we get spurious failures when it things empty
+  // should be unequal to empty.
+  public void testEquals() {
+    TreeRangeMap<Integer, Integer> empty = TreeRangeMap.create();
+    TreeRangeMap<Integer, Integer> nonEmpty = TreeRangeMap.create();
+    nonEmpty.put(Range.all(), 1);
+    TreeRangeMap<Integer, Integer> coalesced = TreeRangeMap.create();
+    coalesced.put(Range.atLeast(1), 1);
+    coalesced.putCoalescing(Range.atMost(1), 1);
+    TreeRangeMap<Integer, Integer> differentValues = TreeRangeMap.create();
+    differentValues.put(Range.closedOpen(1, 2), 2);
+    differentValues.put(Range.closedOpen(3, 4), 2);
+    TreeRangeMap<Double, Integer> differentTypes = TreeRangeMap.create();
+    differentTypes.put(Range.closedOpen(1.0, 2.0), 2);
+    differentTypes.put(Range.closedOpen(3.0, 4.0), 2);
+    new EqualsTester()
+        .addEqualityGroup(empty, TreeRangeMap.<Integer, Integer>create())
+        .addEqualityGroup(nonEmpty, coalesced)
+        .addEqualityGroup(differentValues)
+        .addEqualityGroup(differentTypes)
+        .testEquals();
+  }
+
   private void verify(Map<Integer, Integer> model, RangeMap<Integer, Integer> test) {
     for (int i = MIN_BOUND - 1; i <= MAX_BOUND + 1; i++) {
       assertEquals(model.get(i), test.get(i));

diff --git a/android/guava/src/com/google/common/collect/TreeRangeMap.java b/android/guava/src/com/google/common/collect/TreeRangeMap.java
@@ -59,10 +59,31 @@ public static <K extends Comparable, V> TreeRangeMap<K, V> create() {
     return new TreeRangeMap<>();
   }
 
+  @SuppressWarnings("unchecked")
+  public static <K extends Comparable<?>, V> TreeRangeMap<K, V> copyOf(
+      RangeMap<K, ? extends V> rangeMap) {
+    if (rangeMap instanceof TreeRangeMap) {
+      NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound = Maps.newTreeMap();
+      entriesByLowerBound.putAll(((TreeRangeMap<K, V>) rangeMap).entriesByLowerBound);
+      return new TreeRangeMap<>(entriesByLowerBound);
+    } else {
+      NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound = Maps.newTreeMap();
+      for (Entry<Range<K>, ? extends V> entry : rangeMap.asMapOfRanges().entrySet()) {
+        entriesByLowerBound.put(
+            entry.getKey().lowerBound(), new RangeMapEntry<K, V>(entry.getKey(), entry.getValue()));
+      }
+      return new TreeRangeMap<>(entriesByLowerBound);
+    }
+  }
+
   private TreeRangeMap() {
     this.entriesByLowerBound = Maps.newTreeMap();
   }
 
+  private TreeRangeMap(NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound) {
+    this.entriesByLowerBound = entriesByLowerBound;
+  }
+
   private static final class RangeMapEntry<K extends Comparable, V>
       extends AbstractMapEntry<Range<K>, V> {
     private final Range<K> range;

diff --git a/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java b/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java
@@ -25,6 +25,7 @@
 import com.google.common.collect.testing.features.CollectionFeature;
 import com.google.common.collect.testing.features.CollectionSize;
 import com.google.common.collect.testing.features.MapFeature;
+import com.google.common.testing.EqualsTester;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -853,6 +854,53 @@ public void testSubRangeMapClear() {
         ImmutableMap.of(Range.open(3, 5), 1, Range.closed(12, 16), 3), rangeMap.asMapOfRanges());
   }
 
+  public void testCopyOfTreeRangeMap() {
+    RangeMap<Integer, Integer> rangeMap = TreeRangeMap.create();
+    rangeMap.put(Range.open(3, 7), 1);
+    rangeMap.put(Range.closed(9, 10), 2);
+    rangeMap.put(Range.closed(12, 16), 3);
+
+    RangeMap<Integer, Integer> copy = TreeRangeMap.copyOf(rangeMap);
+
+    assertEquals(rangeMap.asMapOfRanges(), copy.asMapOfRanges());
+  }
+
+  public void testCopyOfImmutableRangeMap() {
+    ImmutableRangeMap<Integer, Integer> rangeMap =
+        ImmutableRangeMap.<Integer, Integer>builder()
+            .put(Range.open(3, 7), 1)
+            .put(Range.closed(9, 10), 2)
+            .put(Range.closed(12, 16), 3)
+            .build();
+
+    RangeMap<Integer, Integer> copy = TreeRangeMap.copyOf(rangeMap);
+
+    assertEquals(rangeMap.asMapOfRanges(), copy.asMapOfRanges());
+  }
+
+  // Overriding testEquals because it seems that we get spurious failures when it things empty
+  // should be unequal to empty.
+  public void testEquals() {
+    TreeRangeMap<Integer, Integer> empty = TreeRangeMap.create();
+    TreeRangeMap<Integer, Integer> nonEmpty = TreeRangeMap.create();
+    nonEmpty.put(Range.all(), 1);
+    TreeRangeMap<Integer, Integer> coalesced = TreeRangeMap.create();
+    coalesced.put(Range.atLeast(1), 1);
+    coalesced.putCoalescing(Range.atMost(1), 1);
+    TreeRangeMap<Integer, Integer> differentValues = TreeRangeMap.create();
+    differentValues.put(Range.closedOpen(1, 2), 2);
+    differentValues.put(Range.closedOpen(3, 4), 2);
+    TreeRangeMap<Double, Integer> differentTypes = TreeRangeMap.create();
+    differentTypes.put(Range.closedOpen(1.0, 2.0), 2);
+    differentTypes.put(Range.closedOpen(3.0, 4.0), 2);
+    new EqualsTester()
+        .addEqualityGroup(empty, TreeRangeMap.<Integer, Integer>create())
+        .addEqualityGroup(nonEmpty, coalesced)
+        .addEqualityGroup(differentValues)
+        .addEqualityGroup(differentTypes)
+        .testEquals();
+  }
+
   private void verify(Map<Integer, Integer> model, RangeMap<Integer, Integer> test) {
     for (int i = MIN_BOUND - 1; i <= MAX_BOUND + 1; i++) {
       assertEquals(model.get(i), test.get(i));

diff --git a/guava/src/com/google/common/collect/TreeRangeMap.java b/guava/src/com/google/common/collect/TreeRangeMap.java
@@ -61,10 +61,31 @@ public static <K extends Comparable, V> TreeRangeMap<K, V> create() {
     return new TreeRangeMap<>();
   }
 
+  @SuppressWarnings("unchecked")
+  public static <K extends Comparable<?>, V> TreeRangeMap<K, V> copyOf(
+      RangeMap<K, ? extends V> rangeMap) {
+    if (rangeMap instanceof TreeRangeMap) {
+      NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound = Maps.newTreeMap();
+      entriesByLowerBound.putAll(((TreeRangeMap<K, V>) rangeMap).entriesByLowerBound);
+      return new TreeRangeMap<>(entriesByLowerBound);
+    } else {
+      NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound = Maps.newTreeMap();
+      for (Entry<Range<K>, ? extends V> entry : rangeMap.asMapOfRanges().entrySet()) {
+        entriesByLowerBound.put(
+            entry.getKey().lowerBound(), new RangeMapEntry<K, V>(entry.getKey(), entry.getValue()));
+      }
+      return new TreeRangeMap<>(entriesByLowerBound);
+    }
+  }
+
   private TreeRangeMap() {
     this.entriesByLowerBound = Maps.newTreeMap();
   }
 
+  private TreeRangeMap(NavigableMap<Cut<K>, RangeMapEntry<K, V>> entriesByLowerBound) {
+    this.entriesByLowerBound = entriesByLowerBound;
+  }
+
   private static final class RangeMapEntry<K extends Comparable, V>
       extends AbstractMapEntry<Range<K>, V> {
     private final Range<K> range;