Have Dict#contents switch to an ImmutableMap when the Dict is frozen.

This implements the idea I documented in a TODO in 24ec910. The approach is to introduce an optional mechanism for Starlark values to get notified when their `Mutability` gets frozen, and to use that mechanism in `Dict`. In implementing this idea, I noticed yet another idea for saving memory. I added a TODO for that, and I also added a TODO for a code unification noticed by arostovtsev@ during code review.

For a `bazel build --nobuild //very/large:target` invocation with many retained `Dict` instances, this approach reduces retained heap by ~1%. There is a ~0.65% increase in CPU usage but no measurable increase in wall time, so I think this tradeoff is very much so worth it.

PiperOrigin-RevId: 433563032

Author: haxorz

src/main/java/net/starlark/java/eval/Dict.java

@@ -14,6 +14,7 @@
 
 package net.starlark.java.eval;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -111,8 +112,21 @@ public class Dict<K, V>
         StarlarkIndexable,
         StarlarkIterable<K> {
 
-  private final Map<K, V> contents;
-  private int iteratorCount; // number of active iterators (unused once frozen)
+  /**
+   * The contents of the Dict.
+   *
+   * <p>This will either be a {@link LinkedHashMap} (for a mutable Dict that hasn't yet been frozen)
+   * or a {@link ImmutableMap} (for a Dict that started frozen or has been frozen; see {@link
+   * #Dict(ImmutableMap)} and {@link #onFreeze}, respectively).
+   */
+  private Map<K, V> contents;
+
+  /**
+   * The number of active iterators.
+   *
+   * <p>This is unused after the Dict is frozen.
+   */
+  private int iteratorCount;
 
   /** Final except for {@link #unsafeShallowFreeze}; must not be modified any other way. */
   private Mutability mutability;
@@ -121,8 +135,7 @@ private Dict(Mutability mutability, LinkedHashMap<K, V> contents) {
     Preconditions.checkNotNull(mutability);
     Preconditions.checkState(mutability != Mutability.IMMUTABLE);
     this.mutability = mutability;
-    // TODO(bazel-team): Memory optimization opportunity: Make it so that a call to
-    //  `mutability.freeze()` causes `contents` here to become an ImmutableMap.
+    mutability.notifyOnFreeze(this);
     this.contents = contents;
   }
 
@@ -560,6 +573,35 @@ public Mutability mutability() {
   public void unsafeShallowFreeze() {
     Mutability.Freezable.checkUnsafeShallowFreezePrecondition(this);
     this.mutability = Mutability.IMMUTABLE;
+    this.contents = ImmutableMap.copyOf(contents);
+  }
+
+  /**
+   * Switches {@link #contents} to be a {@link ImmutableMap} in order to save memory.
+   *
+   * <p>See the comment in {@link #Dict(ImmutableMap)} for details.
+   */
+  @Override
+  public void onFreeze() {
+    if (contents instanceof LinkedHashMap) {
+      this.contents = ImmutableMap.copyOf(contents);
+      // TODO(bazel-team): Make #mutability IMMUTABLE, causing the old Mutability instance to be
+      // eligible for GC. Do this in StarlarkList too.
+    } else {
+      // Assert #onFreeze hasn't been called before. But also hackily support the usage in
+      // ScriptTest.
+      //
+      // ScriptTest extends the language with a `freeze` feature, allowing #unsafeShallowFreeze to
+      // be called. When the bzl code has already caused that method to be called, #contents will
+      // already be an ImmutableMap, but #mutability will still not be.
+      //
+      // And we definitely want #unsafeShallowFreeze to cause #contents to become an ImmutableMap.
+      // That method is used at the end of deserialization, which uses a mutable #mutability that
+      // became IMMUTABLE in that method.
+      //
+      // TODO(bazel-team): Combine the two methods.
+      Preconditions.checkState(mutability == Mutability.IMMUTABLE);
+    }
   }
 
   /**
@@ -623,6 +665,11 @@ public String toString() {
     return Starlark.repr(this);
   }
 
+  @VisibleForTesting
+  Map<K, V> getContentsForTesting() {
+    return contents;
+  }
+
   /**
    * Casts a non-null Starlark value {@code x} to a {@code Dict<K, V>} after checking that all keys
    * and values are instances of {@code keyType} and {@code valueType}, respectively. On error, it

src/main/java/net/starlark/java/eval/Mutability.java

@@ -14,6 +14,7 @@
 package net.starlark.java.eval;
 
 import com.google.common.base.Joiner;
+import java.util.ArrayList;
 import java.util.IdentityHashMap;
 
 /**
@@ -105,6 +106,8 @@ public final class Mutability implements AutoCloseable {
   /** Controls access to {@link Freezable#unsafeShallowFreeze}. */
   private final boolean allowsUnsafeShallowFreeze;
 
+  private ArrayList<Freezable> freezablesToNotifyOnFreeze;
+
   private Mutability(Object[] annotation, boolean allowsUnsafeShallowFreeze) {
     this.annotation = annotation;
     this.allowsUnsafeShallowFreeze = allowsUnsafeShallowFreeze;
@@ -171,13 +174,22 @@ private boolean updateIteratorCount(Freezable x, int delta) {
    * Freezes this {@code Mutability}, rendering all {@link Freezable} objects that refer to it
    * immutable.
    *
-   * Note that freezing does not directly touch all the {@code Freezables}, so this operation is
-   * constant-time.
+   * <p>Note that freezing directly touches only the {@code Freezables} for which there was a call
+   * to {@code thisMutability.notifyOnFreeze(thatFreezable)}, so this is linear-time in the number
+   * of such {@code Freezables}s.
    *
    * @return this object, in the fluent style
    */
   public Mutability freeze() {
     this.iteratorCount = null;
+
+    if (freezablesToNotifyOnFreeze != null) {
+      for (Freezable freezable : freezablesToNotifyOnFreeze) {
+        freezable.onFreeze();
+      }
+      freezablesToNotifyOnFreeze = null;
+    }
+
     return this;
   }
 
@@ -188,12 +200,22 @@ public void close() {
 
   /**
    * Returns whether {@link Freezable}s having this {@code Mutability} allow the {@link
-   * #unsafeShallowFreeze} operation.
+   * Freezable#unsafeShallowFreeze} operation.
    */
   public boolean allowsUnsafeShallowFreeze() {
     return allowsUnsafeShallowFreeze;
   }
 
+  /**
+   * Causes {@code freezable.onFreeze()} to be called in the future when {@link #freeze} is called.
+   */
+  public void notifyOnFreeze(Freezable freezable) {
+    if (freezablesToNotifyOnFreeze == null) {
+      freezablesToNotifyOnFreeze = new ArrayList<>();
+    }
+    freezablesToNotifyOnFreeze.add(freezable);
+  }
+
   /**
    * An object that refers to a {@link Mutability} to decide whether to allow mutation. All {@link
    * Freezable} Starlark objects created within a given {@link StarlarkThread} will share the same
@@ -207,6 +229,14 @@ public interface Freezable {
      */
     Mutability mutability();
 
+    /**
+     * If {@code mutability().notifyOnFreeze(this)} has been called, this method gets called when
+     * {@code mutability().freeze()} gets called.
+     *
+     * <p>Do not call this method from outside {@link Mutability#freeze}.
+     */
+    default void onFreeze() {}
+
     /**
      * Registers a change to this Freezable's iterator count and reports whether it is temporarily
      * immutable.

src/test/java/net/starlark/java/eval/MutabilityTest.java

@@ -127,4 +127,44 @@ public void checkUnsafeShallowFreezePrecondition_SucceedsWhenAllowed() throws Ex
     Mutability mutability = Mutability.createAllowingShallowFreeze("test");
     Freezable.checkUnsafeShallowFreezePrecondition(new DummyFreezable(mutability));
   }
+
+  @Test
+  public void notifyOnFreeze() {
+    class FreezableWithNotify implements Freezable {
+      private final Mutability mutability;
+      private boolean onFreezeCalled = false;
+
+      FreezableWithNotify(Mutability mutability) {
+        this.mutability = mutability;
+      }
+
+      @Override
+      public Mutability mutability() {
+        return mutability;
+      }
+
+      @Override
+      public void onFreeze() {
+        onFreezeCalled = true;
+      }
+    }
+
+    // When we have an unfrozen Mutability,
+    Mutability mutability = Mutability.create("test");
+
+    // And we create two Freezables using the Mutability,
+    FreezableWithNotify freezable1 = new FreezableWithNotify(mutability);
+    FreezableWithNotify freezable2 = new FreezableWithNotify(mutability);
+
+    // And we tell the Mutability to notify one of the two Freezables when it has been frozen,
+    mutability.notifyOnFreeze(freezable1);
+
+    // And we freeze the Mutability,
+    mutability.freeze();
+
+    // Then the Freezable that was supposed to be notified was notified,
+    assertThat(freezable1.onFreezeCalled).isTrue();
+    // And the other one wasn't.
+    assertThat(freezable2.onFreezeCalled).isFalse();
+  }
 }

src/test/java/net/starlark/java/eval/StarlarkMutableTest.java

@@ -20,13 +20,14 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Tests for {@link StarlarkMutable}. */
+/** Tests for how Starlark value implementations interact with {@link Mutability#freeze}. */
 @RunWith(JUnit4.class)
 public final class StarlarkMutableTest {
 
@@ -142,4 +143,33 @@ public void testDictBuilder() throws Exception {
         .isEqualTo(
             "{\"one\": \"1\", \"two\": \"22\", \"three\": \"33\", \"four\": \"4\"}"); // unchanged
   }
+
+  @Test
+  public void testImmutableDictUsesImmutableMap() {
+    Mutability mu = Mutability.IMMUTABLE;
+    Dict<String, String> dict = Dict.<String, String>builder().put("cat", "dog").build(mu);
+    assertThat(dict.getContentsForTesting()).isInstanceOf(ImmutableMap.class);
+    assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog");
+  }
+
+  @Test
+  public void testMutableDictSwitchesToImmutableMapWhenFrozen() {
+    Mutability mu = Mutability.create("test");
+    Dict<String, String> dict = Dict.<String, String>builder().put("cat", "dog").build(mu);
+    assertThat(dict.getContentsForTesting()).isInstanceOf(LinkedHashMap.class);
+    assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog");
+    mu.freeze();
+    assertThat(dict.getContentsForTesting()).isInstanceOf(ImmutableMap.class);
+    assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog");
+  }
+
+  @Test
+  public void testMutableDictSwitchesToImmutableMapWhenFrozen_usesCanonicalEmptyImmutableMap() {
+    Mutability mu = Mutability.create("test");
+    Dict<String, String> dict = Dict.<String, String>builder().build(mu);
+    assertThat(dict.getContentsForTesting()).isInstanceOf(LinkedHashMap.class);
+    assertThat(dict.getContentsForTesting()).isEmpty();
+    mu.freeze();
+    assertThat(dict.getContentsForTesting()).isSameInstanceAs(ImmutableMap.of());
+  }
 }