Reuse objects when inserting identical keys/values in PathCopyingPersistentTreeMap

src/org/sosy_lab/common/collect/PathCopyingPersistentTreeMap.java

@@ -54,7 +54,7 @@
  * always compares according to the natural ordering. All methods may throw {@link
  * ClassCastException} is key objects are passed that do not implement {@link Comparable}.
  *
- * <p>The natural ordering of the keys needs to be consistent with equals.
+ * <p>The natural ordering of the keys needs to be consistent with equals and object identity.
  *
  * <p>As for all {@link PersistentMap}s, all collection views and all iterators are immutable. They
  * do not reflect changes made to the map and all their modifying operations throw {@link
@@ -137,6 +137,14 @@ Node<K, V> withColor(boolean color) {
       }
     }
 
+    Node<K, V> withValue(V newValue) {
+      if (newValue == getValue()) {
+        return this;
+      } else {
+        return new Node<>(getKey(), newValue, left, right, isRed);
+      }
+    }
+
     @SuppressWarnings("ReferenceEquality") // cannot use equals() for check whether tree is the same
     Node<K, V> withLeftChild(Node<K, V> newLeft) {
       if (newLeft == left) {
@@ -555,7 +563,14 @@ private static <K extends Comparable<? super K>, V> Node<K, V> putAndCopy0(
       current = current.withRightChild(newRight);
 
     } else {
-      current = new Node<>(key, value, current.left, current.right, current.getColor());
+      // This always keeps the old (equal) key object. This has useful implications:
+      // Because we reuse the old key object, the key instance does not change and potential `==`
+      // comparisons on the key at other locations still work successfully.
+      // We do always use the new value object; but in case that the new value object is identical
+      // to the old value object, we can reuse the existing Node object and the whole map. This also
+      // enables `==` comparisons to succeed and saves some memory.
+      // This behavior also matches what JDK maps do.
+      current = current.withValue(value);
     }
 
     // restore invariants

src/org/sosy_lab/common/collect/PathCopyingPersistentTreeMapTest.java

@@ -12,6 +12,7 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.Assert.assertThrows;
 
+import com.google.common.collect.FluentIterable;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.testing.NavigableMapTestSuiteBuilder;
 import com.google.common.collect.testing.TestStringSortedMapGenerator;
@@ -21,6 +22,7 @@
 import com.google.common.testing.EqualsTester;
 import com.google.errorprone.annotations.Var;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Random;
@@ -434,4 +436,96 @@ public void testEntrySetContains() {
     assertThat(second.entrySet().containsAll(first.entrySet())).isFalse();
     assertThat(first.entrySet().containsAll(second.entrySet())).isFalse();
   }
+
+  @Test
+  public void testRemovingMissingKey() {
+    map = map.putAndCopy("a", "").putAndCopy("b", "").putAndCopy("y", "").putAndCopy("z", "");
+
+    assertWithMessage("Removing missing key should produce same map")
+        .that(map.removeAndCopy("key"))
+        .isSameInstanceAs(map);
+  }
+
+  @Test
+  @SuppressWarnings({"checkstyle:IllegalInstantiation", "StringUselessMethods"})
+  public void testSettingIdenticalObjects() {
+    String k1 = new String("key");
+    String k2 = new String("key");
+    String v1 = new String("value");
+    String v2 = new String("value");
+    map =
+        map.putAndCopy("a", "")
+            .putAndCopy("b", "")
+            .putAndCopy(k1, v1)
+            .putAndCopy("y", "")
+            .putAndCopy("z", "");
+
+    assertWithMessage("Reinserting same k/v pair should produce same map")
+        .that(map.putAndCopy(k1, v1))
+        .isSameInstanceAs(map);
+
+    assertWithMessage("Reinserting same value should produce same map")
+        .that(map.putAndCopy(k2, v1))
+        .isSameInstanceAs(map);
+
+    assertWithMessage("Inserting new value should produce map with new value")
+        .that(map.putAndCopy(k1, v2).get(k1))
+        .isSameInstanceAs(v2);
+
+    assertWithMessage("Inserting new k/v pair should keep old key")
+        .that(
+            FluentIterable.from(map.putAndCopy(k2, v2).keySet())
+                .filter(s -> s.length() > 1)
+                .first()
+                .get())
+        .isSameInstanceAs(k1);
+    assertWithMessage("Inserting new k/v pair should produce map with new value")
+        .that(map.putAndCopy(k2, v2).get(k2))
+        .isSameInstanceAs(v2);
+  }
+
+  @Test
+  public void testSettingIdenticalObjectsInHashMap() {
+    testSettingIdenticalObjectsInStandardMap(new HashMap<>());
+  }
+
+  @Test
+  public void testSettingIdenticalObjectsInTreeMap() {
+    testSettingIdenticalObjectsInStandardMap(new TreeMap<>());
+  }
+
+  @SuppressWarnings({"checkstyle:IllegalInstantiation", "StringUselessMethods"})
+  private static void testSettingIdenticalObjectsInStandardMap(Map<String, String> map) {
+    // not testing own code, but checking expectations of other map implementations
+
+    String k1 = new String("key");
+    String k2 = new String("key");
+    String k3 = new String("key");
+    String v1 = new String("value");
+    String v2 = new String("value");
+    String v3 = new String("value");
+    map.put("a", "");
+    map.put("b", "");
+    map.put(k1, v1);
+    map.put("y", "");
+    map.put("z", "");
+
+    map.put(k2, v1);
+    assertWithMessage("Inserting same value should keep old key")
+        .that(FluentIterable.from(map.keySet()).filter(s -> s.length() > 1).first().get())
+        .isSameInstanceAs(k1);
+
+    map.put(k1, v2);
+    assertWithMessage("Inserting new value should produce map with new value")
+        .that(map.get(k1))
+        .isSameInstanceAs(v2);
+
+    map.put(k3, v3);
+    assertWithMessage("Inserting new k/v pair should keep old key")
+        .that(FluentIterable.from(map.keySet()).filter(s -> s.length() > 1).first().get())
+        .isSameInstanceAs(k1);
+    assertWithMessage("Inserting new k/v pair should produce map with new value")
+        .that(map.get(k1))
+        .isSameInstanceAs(v3);
+  }
 }

src/org/sosy_lab/common/collect/PersistentMap.java

@@ -20,8 +20,8 @@
 /**
  * Interface for persistent map. A persistent data structure is immutable, but provides cheap
  * copy-and-write operations. Thus all write operations ({{@link #putAndCopy(Object, Object)},
- * {{@link #removeAndCopy(Object)}}) will not modify the current instance, but return a new instance
- * instead.
+ * {{@link #removeAndCopy(Object)}}) will not modify the current instance, but return an updated
+ * instance instead.
  *
  * <p>All modifying operations inherited from {@link Map} are not supported and will always throw
  * {@link UnsupportedOperationException}. All collections returned by methods of this interface are
@@ -35,11 +35,11 @@
 @Immutable(containerOf = {"K", "V"})
 public interface PersistentMap<K, V extends @Nullable Object> extends Map<K, V> {
 
-  /** Replacement for {{@link #put(Object, Object)} that returns a fresh instance. */
+  /** Replacement for {{@link #put(Object, Object)} that returns an updated map. */
   @CheckReturnValue
   PersistentMap<K, V> putAndCopy(@CompatibleWith("K") K key, @CompatibleWith("V") V value);
 
-  /** Replacement for {{@link #remove(Object)} that returns a fresh instance. */
+  /** Replacement for {{@link #remove(Object)} that returns an updated map. */
   @CheckReturnValue
   PersistentMap<K, V> removeAndCopy(@CompatibleWith("K") Object key);