Optimise Starlark providers' memory layout

Only the memory layout of Starlark providers returned from the rules is optimised.
It's optimised by converting mutable lists to immutable ones and direct references to other Starlark providers.

Elements of the lists are not optimised.

PiperOrigin-RevId: 510996467
Change-Id: I95a26399e1aa024ead8dabb74a4d57dd04e66979

Author: comius

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleConfiguredTargetUtil.java

@@ -333,6 +333,9 @@ private static void addProviders(
           }
         }
       } else {
+        if (info instanceof StarlarkInfo) {
+          info = ((StarlarkInfo) info).unsafeOptimizeMemoryLayout();
+        }
         Provider.Key providerKey = getProviderKey(info);
         // Single declared provider
         declaredProviders.put(providerKey, info);
@@ -341,6 +344,12 @@ private static void addProviders(
       // Sequence of declared providers
       for (Info provider :
           Sequence.cast(rawProviders, Info.class, "result of rule implementation function")) {
+        if (provider instanceof StarlarkInfo) {
+          // Provider instances are optimised recursively, without optimising elements of the list.
+          // Tradeoff is that some object may be duplicated if they are reachable by more than one
+          // path, but we don't expect that much in practice.
+          provider = ((StarlarkInfo) provider).unsafeOptimizeMemoryLayout();
+        }
         Provider.Key providerKey = getProviderKey(provider);
         if (declaredProviders.put(providerKey, provider) != null) {
           context.ruleError("Multiple conflicting returned providers with key " + providerKey);

src/main/java/com/google/devtools/build/lib/packages/StarlarkInfo.java

@@ -46,4 +46,15 @@ public static StarlarkInfo create(
   @Nullable
   @Override
   public abstract Object getValue(String name);
+
+  /**
+   * Returns an equivalent memory-optimized version of this provider instance (which might be the
+   * same existing instance modified in-place).
+   *
+   * <p>The existing instance must not be accessed concurrently while calling this method.
+   *
+   * <p>The mutability of any values contained in this instance must already be frozen prior to
+   * calling this method.
+   */
+  public abstract StarlarkInfo unsafeOptimizeMemoryLayout();
 }

src/main/java/com/google/devtools/build/lib/packages/StarlarkInfoNoSchema.java

@@ -23,6 +23,7 @@
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkList;
 import net.starlark.java.syntax.Location;
 import net.starlark.java.syntax.TokenKind;
 
@@ -293,4 +294,18 @@ private static StarlarkInfo plus(StarlarkInfoNoSchema x, StarlarkInfoNoSchema y)
 
     return new StarlarkInfoNoSchema(x.provider, ztable, Location.BUILTIN);
   }
+
+  @Override
+  public StarlarkInfoNoSchema unsafeOptimizeMemoryLayout() {
+    for (int i = table.length / 2; i < table.length; i++) {
+      if (table[i] instanceof StarlarkList<?>) {
+        // On duplicated lists, ImmutableStarlarkLists objects are duplicated, but not underlying
+        // Object arrays
+        table[i] = ((StarlarkList<?>) table[i]).unsafeOptimizeMemoryLayout();
+      } else if (table[i] instanceof StarlarkInfo) {
+        table[i] = ((StarlarkInfo) table[i]).unsafeOptimizeMemoryLayout();
+      }
+    }
+    return this;
+  }
 }

src/main/java/com/google/devtools/build/lib/packages/StarlarkInfoWithSchema.java

@@ -24,6 +24,7 @@
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
 import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkList;
 import net.starlark.java.syntax.Location;
 import net.starlark.java.syntax.TokenKind;
 
@@ -153,4 +154,19 @@ private static StarlarkInfoWithSchema plus(StarlarkInfoWithSchema x, StarlarkInf
     }
     return new StarlarkInfoWithSchema(x.provider, ztable, Location.BUILTIN);
   }
+
+  @Override
+  public StarlarkInfoWithSchema unsafeOptimizeMemoryLayout() {
+    int n = table.length;
+    for (int i = 0; i < n; i++) {
+      if (table[i] instanceof StarlarkList<?>) {
+        // On duplicated lists, ImmutableStarlarkLists are duplicated, but not underlying Object
+        // arrays
+        table[i] = ((StarlarkList<?>) table[i]).unsafeOptimizeMemoryLayout();
+      } else if (table[i] instanceof StarlarkInfo) {
+        table[i] = ((StarlarkInfo) table[i]).unsafeOptimizeMemoryLayout();
+      }
+    }
+    return this;
+  }
 }

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

@@ -219,4 +219,17 @@ public <T> T[] toArray(T[] a) {
   Object[] elems() {
     return elems;
   }
+
+  @Override
+  public StarlarkList<E> unsafeOptimizeMemoryLayout() {
+    Preconditions.checkState(mutability.isFrozen());
+    // Canonicalize our Mutability, on the off-chance the old one may be freed.
+    mutability = Mutability.IMMUTABLE;
+    if (elems.length > size) {
+      // shrink the Object array
+      elems = Arrays.copyOf(elems, size);
+    }
+    // Give the caller an immutable specialization of StarlarkList.
+    return wrap(mutability, elems);
+  }
 }

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

@@ -418,4 +418,19 @@ public Object pop(Object i) throws EvalException {
     removeElementAt(index);
     return result;
   }
+
+  /**
+   * Mutates this list in-place to reduce memory usage, and returns an optimized list (which might
+   * be the same as this instance).
+   *
+   * <p>This operation is not protected by the mutability mechanism. It is the caller's
+   * responsibility to ensure this list is not concurrently accessed during this method's execution.
+   *
+   * <p>The mutated list and the returned list are both equivalent to the original list.
+   *
+   * <p>The mutability must be frozen prior to calling this method.
+   */
+  public StarlarkList<E> unsafeOptimizeMemoryLayout() {
+    return this;
+  }
 }