From a46565dd1c6ce45f444d707013afb5fda9aa6b06 Mon Sep 17 00:00:00 2001
From: Google Java Core Libraries <java-team-github-bot@google.com>
Date: Sun, 6 Oct 2024 06:00:57 -0700
Subject: [PATCH] Provide an optimized copyOf method for TreeRangeMap

RELNOTES=Provide an optimized copyOf method for TreeRangeMap
PiperOrigin-RevId: 682878547
---
 .../common/collect/TreeRangeMapTest.java      | 48 +++++++++++++++++++
 .../google/common/collect/TreeRangeMap.java   | 21 ++++++++
 .../common/collect/TreeRangeMapTest.java      | 48 +++++++++++++++++++
 .../google/common/collect/TreeRangeMap.java   | 21 ++++++++
 4 files changed, 138 insertions(+)

diff --git a/android/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java b/android/guava-tests/test/com/google/common/collect/TreeRangeMapTest.java
index 75ed02b9cec3..1921cbff35ab 100644
--- 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
index f831f752f618..75bee5e702c2 100644
--- 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
index 382afc042fbb..0e939fd6c802 100644
--- 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
index dbcb82ea2e55..80915a375875 100644
--- 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;