@AutoCodec TransitiveInfoProviderMapImpl and JavaInfo.

PiperOrigin-RevId: 186789569

5 files changed, 48 insertions, 33 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/BUILD b/src/main/java/com/google/devtools/build/lib/BUILD
index 77ffe93545..6e221a203c 100644
--- a/src/main/java/com/google/devtools/build/lib/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/BUILD
@@ -841,6 +841,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/rules/cpp",
         "//src/main/java/com/google/devtools/build/lib/shell",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:extra_actions_base_java_proto",
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapBuilder.java b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapBuilder.java
index 0c3bd38b48..7a5ed434d9 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapBuilder.java
@@ -109,6 +109,6 @@ public class TransitiveInfoProviderMapBuilder {
   }
 
   public TransitiveInfoProviderMap build() {
-    return new TransitiveInfoProviderMapImpl(providers);
+    return TransitiveInfoProviderMapImpl.create(providers);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapImpl.java b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapImpl.java
index b43129d902..1e13d4286f 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapImpl.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapImpl.java
@@ -14,9 +14,12 @@
 
 package com.google.devtools.build.lib.analysis;
 
+import com.google.common.base.Preconditions;
 import com.google.devtools.build.lib.collect.ImmutableSharedKeyMap;
 import com.google.devtools.build.lib.packages.Info;
 import com.google.devtools.build.lib.packages.Provider;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -24,12 +27,28 @@ import javax.annotation.Nullable;
  * Implementation of {@link TransitiveInfoProvider} that uses {@link ImmutableSharedKeyMap}. For
  * memory efficiency, inheritance is used instead of aggregation as an implementation detail.
  */
-class TransitiveInfoProviderMapImpl
-    extends ImmutableSharedKeyMap<Object, Object>
+@AutoCodec
+class TransitiveInfoProviderMapImpl extends ImmutableSharedKeyMap<Object, Object>
     implements TransitiveInfoProviderMap {
 
-  TransitiveInfoProviderMapImpl(Map<Object, Object> map) {
-    super(map);
+  @VisibleForSerialization
+  @AutoCodec.Instantiator
+  TransitiveInfoProviderMapImpl(OffsetTable<Object> offsetTable, Object[] values) {
+    super(offsetTable, values);
+  }
+
+  static TransitiveInfoProviderMapImpl create(Map<Object, Object> map) {
+    int count = map.size();
+    Object[] keys = new Object[count];
+    Object[] values = new Object[count];
+    int i = 0;
+    for (Map.Entry<Object, Object> entry : map.entrySet()) {
+      keys[i] = entry.getKey();
+      values[i] = entry.getValue();
+      ++i;
+    }
+    Preconditions.checkArgument(keys.length == values.length);
+    return new TransitiveInfoProviderMapImpl(ImmutableSharedKeyMap.createOffsetTable(keys), values);
   }
 
   @SuppressWarnings("unchecked")
diff --git a/src/main/java/com/google/devtools/build/lib/collect/ImmutableSharedKeyMap.java b/src/main/java/com/google/devtools/build/lib/collect/ImmutableSharedKeyMap.java
index 3fdde7fbfe..42a7af8723 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/ImmutableSharedKeyMap.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/ImmutableSharedKeyMap.java
@@ -22,7 +22,6 @@ import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import javax.annotation.concurrent.Immutable;
 
 /**
@@ -43,10 +42,16 @@ import javax.annotation.concurrent.Immutable;
 public class ImmutableSharedKeyMap<K, V> extends CompactImmutableMap<K, V> {
   private static final Interner<OffsetTable> offsetTables = BlazeInterners.newWeakInterner();
 
-  private final OffsetTable<K> offsetTable;
-  private final Object[] values;
+  // Visible only for serialization.
+  protected final OffsetTable<K> offsetTable;
+  protected final Object[] values;
 
-  private static final class OffsetTable<K> {
+  /**
+   * Table storing {@code keys} compactly as an array, with an index lookup map for speed.
+   *
+   * <p>Only externally visible for serialization purposes.
+   */
+  public static final class OffsetTable<K> {
     private final Object[] keys;
     // Keep a map around to speed up get lookups for larger maps.
     // We make this value lazy to avoid computing for values that end up being thrown away
@@ -99,29 +104,14 @@ public class ImmutableSharedKeyMap<K, V> extends CompactImmutableMap<K, V> {
     }
   }
 
-  protected ImmutableSharedKeyMap(Object[] keys, Object[] values) {
-    Preconditions.checkArgument(keys.length == values.length);
+  protected ImmutableSharedKeyMap(OffsetTable<K> offsetTable, Object[] values) {
+    Preconditions.checkArgument(offsetTable.keys.length == values.length);
+    this.offsetTable = offsetTable;
     this.values = values;
-    this.offsetTable = createOffsetTable(keys);
-  }
-
-  protected ImmutableSharedKeyMap(Map<K, V> map) {
-    int count = map.size();
-    Object[] keys = new Object[count];
-    Object[] values = new Object[count];
-    int i = 0;
-    for (Map.Entry<K, V> entry : map.entrySet()) {
-      keys[i] = entry.getKey();
-      values[i] = entry.getValue();
-      ++i;
-    }
-    Preconditions.checkArgument(keys.length == values.length);
-    this.values = values;
-    this.offsetTable = createOffsetTable(keys);
   }
 
   @SuppressWarnings("unchecked")
-  private static <K> OffsetTable<K> createOffsetTable(Object[] keys) {
+  protected static <K> OffsetTable<K> createOffsetTable(Object[] keys) {
     OffsetTable<K> offsetTable = new OffsetTable<>(keys);
     OffsetTable<K> internedTable = (OffsetTable<K>) offsetTables.intern(offsetTable);
     internedTable.initIndexMap();
@@ -197,7 +187,7 @@ public class ImmutableSharedKeyMap<K, V> extends CompactImmutableMap<K, V> {
         keys[i] = entries.get(entryIndex++);
         values[i] = entries.get(entryIndex++);
       }
-      return new ImmutableSharedKeyMap<>(keys, values);
+      return new ImmutableSharedKeyMap<>(createOffsetTable(keys), values);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaInfo.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaInfo.java
index e766a21168..5323dd3add 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaInfo.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaInfo.java
@@ -34,6 +34,8 @@ import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.packages.NativeInfo;
 import com.google.devtools.build.lib.packages.NativeProvider;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkCallable;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkModule;
 import com.google.devtools.build.lib.skylarkinterface.SkylarkModuleCategory;
@@ -53,11 +55,12 @@ import javax.annotation.Nullable;
 
 /** A Skylark declared provider that encapsulates all providers that are needed by Java rules. */
 @SkylarkModule(
-    name = "JavaInfo",
-    doc = "Encapsulates all information provided by Java rules",
-    category = SkylarkModuleCategory.PROVIDER
+  name = "JavaInfo",
+  doc = "Encapsulates all information provided by Java rules",
+  category = SkylarkModuleCategory.PROVIDER
 )
 @Immutable
+@AutoCodec
 public final class JavaInfo extends NativeInfo {
 
   public static final String SKYLARK_NAME = "JavaInfo";
@@ -297,7 +300,9 @@ public final class JavaInfo extends NativeInfo {
     return providersList.build();
   }
 
-  private JavaInfo(TransitiveInfoProviderMap providers, boolean neverlink, Location location) {
+  @VisibleForSerialization
+  @AutoCodec.Instantiator
+  JavaInfo(TransitiveInfoProviderMap providers, boolean neverlink, Location location) {
     super(PROVIDER, ImmutableMap.of(), location);
     this.providers = providers;
     this.neverlink = neverlink;