Extract ImmutableSharedKeysMap class from TransitiveInfoProviderMap.

PiperOrigin-RevId: 159498323

6 files changed, 320 insertions, 129 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/BUILD b/src/main/java/com/google/devtools/build/lib/BUILD
index a9fcb0c043..ca9f117d3e 100644
--- a/src/main/java/com/google/devtools/build/lib/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/BUILD
@@ -127,6 +127,7 @@ java_library(
         "collect/nestedset/*.java",
     ]),
     deps = [
+        ":concurrent",
         ":preconditions",
         "//third_party:guava",
         "//third_party:jsr305",
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 06a7726114..89ca1f13bd 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
@@ -14,7 +14,6 @@
 
 package com.google.devtools.build.lib.analysis;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.util.Preconditions;
 import java.util.Arrays;
 import java.util.LinkedHashMap;
@@ -74,6 +73,6 @@ public class TransitiveInfoProviderMapBuilder {
   }
 
   public TransitiveInfoProviderMap build() {
-    return new TransitiveInfoProviderMapOffsetBased(ImmutableMap.copyOf(providers));
+    return new TransitiveInfoProviderMapImpl(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
new file mode 100644
index 0000000000..ca1b9c854a
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapImpl.java
@@ -0,0 +1,57 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis;
+
+import com.google.devtools.build.lib.collect.ImmutableSharedKeyMap;
+import java.util.Map;
+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<Class<? extends TransitiveInfoProvider>, TransitiveInfoProvider>
+    implements TransitiveInfoProviderMap {
+
+  TransitiveInfoProviderMapImpl(
+      Map<Class<? extends TransitiveInfoProvider>, TransitiveInfoProvider> map) {
+    super(map);
+  }
+
+  @SuppressWarnings("unchecked")
+  @Nullable
+  @Override
+  public <P extends TransitiveInfoProvider> P getProvider(Class<P> providerClass) {
+    Class<? extends TransitiveInfoProvider> effectiveClass =
+        TransitiveInfoProviderEffectiveClassHelper.get(providerClass);
+    return (P) get(effectiveClass);
+  }
+
+  @Override
+  public int getProviderCount() {
+    return size();
+  }
+
+  @Override
+  public Class<? extends TransitiveInfoProvider> getProviderClassAt(int i) {
+    return keyAt(i);
+  }
+
+  @Override
+  public TransitiveInfoProvider getProviderAt(int i) {
+    return valueAt(i);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapOffsetBased.java b/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapOffsetBased.java
deleted file mode 100644
index a36a3e14be..0000000000
--- a/src/main/java/com/google/devtools/build/lib/analysis/TransitiveInfoProviderMapOffsetBased.java
+++ /dev/null
@@ -1,127 +0,0 @@
-// Copyright 2017 The Bazel Authors. All rights reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//    http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.google.devtools.build.lib.analysis;
-
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Interner;
-import com.google.devtools.build.lib.concurrent.BlazeInterners;
-import java.util.Arrays;
-import java.util.Map.Entry;
-import javax.annotation.Nullable;
-import javax.annotation.concurrent.Immutable;
-
-/**
- * Provides a mapping between a TransitiveInfoProvider class and an instance.
- *
- * <p>This class implements a map where it is expected that a lot of the key sets will be the same.
- * These key sets are shared and an offset table of indices is computed. Each provider map instance
- * thus contains only a reference to the shared offset table, and a plain array of providers.
- */
-@Immutable
-final class TransitiveInfoProviderMapOffsetBased implements TransitiveInfoProviderMap {
-  private static final Interner<OffsetTable> offsetTables = BlazeInterners.newWeakInterner();
-
-  private final OffsetTable offsetTable;
-  private final TransitiveInfoProvider[] providers;
-
-  private static final class OffsetTable {
-    private final Class<? extends TransitiveInfoProvider>[] providerClasses;
-    // 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
-    // during interning anyway (the majority).
-    private volatile ImmutableMap<Class<? extends TransitiveInfoProvider>, Integer> indexMap;
-
-    OffsetTable(Class<? extends TransitiveInfoProvider>[] providerClasses) {
-      this.providerClasses = providerClasses;
-    }
-
-    private ImmutableMap<Class<? extends TransitiveInfoProvider>, Integer> getIndexMap() {
-      if (indexMap == null) {
-        synchronized (this) {
-          if (indexMap == null) {
-            ImmutableMap.Builder<Class<? extends TransitiveInfoProvider>, Integer> builder =
-                ImmutableMap.builder();
-            for (int i = 0; i < providerClasses.length; ++i) {
-              builder.put(providerClasses[i], i);
-            }
-            this.indexMap = builder.build();
-          }
-        }
-      }
-      return indexMap;
-    }
-
-    int getOffsetForClass(Class effectiveClass) {
-      return getIndexMap().getOrDefault(effectiveClass, -1);
-    }
-
-    @Override
-    public boolean equals(Object o) {
-      if (this == o) {
-        return true;
-      }
-      if (!(o instanceof OffsetTable)) {
-        return false;
-      }
-      OffsetTable that = (OffsetTable) o;
-      return Arrays.equals(this.providerClasses, that.providerClasses);
-    }
-
-    @Override
-    public int hashCode() {
-      return Arrays.hashCode(providerClasses);
-    }
-  }
-
-  TransitiveInfoProviderMapOffsetBased(
-      ImmutableMap<Class<? extends TransitiveInfoProvider>, TransitiveInfoProvider> map) {
-    int count = map.size();
-    Class<? extends TransitiveInfoProvider>[] providerClasses = new Class[count];
-    this.providers = new TransitiveInfoProvider[count];
-    int i = 0;
-    for (Entry<Class<? extends TransitiveInfoProvider>, TransitiveInfoProvider> entry :
-        map.entrySet()) {
-      providerClasses[i] = entry.getKey();
-      providers[i] = entry.getValue();
-      ++i;
-    }
-    OffsetTable offsetTable = new OffsetTable(providerClasses);
-    this.offsetTable = offsetTables.intern(offsetTable);
-  }
-
-  /** Returns the instance for the provided providerClass, or <tt>null</tt> if not present. */
-  @Override
-  @Nullable
-  public <P extends TransitiveInfoProvider> P getProvider(Class<P> providerClass) {
-    Class effectiveClass = TransitiveInfoProviderEffectiveClassHelper.get(providerClass);
-    int offset = offsetTable.getOffsetForClass(effectiveClass);
-    return offset >= 0 ? (P) providers[offset] : null;
-  }
-
-  @Override
-  public int getProviderCount() {
-    return providers.length;
-  }
-
-  @Override
-  public Class<? extends TransitiveInfoProvider> getProviderClassAt(int i) {
-    return offsetTable.providerClasses[i];
-  }
-
-  @Override
-  public TransitiveInfoProvider getProviderAt(int i) {
-    return providers[i];
-  }
-}
diff --git a/src/main/java/com/google/devtools/build/lib/collect/CompactImmutableMap.java b/src/main/java/com/google/devtools/build/lib/collect/CompactImmutableMap.java
new file mode 100644
index 0000000000..d410d0cdba
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/collect/CompactImmutableMap.java
@@ -0,0 +1,59 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.collect;
+
+import java.util.Iterator;
+
+/**
+ * A minimal map interface that avoids methods whose implementation tends to force GC churn, or
+ * otherwise overly constrain implementation freedom.
+ *
+ * <p>TODO: Convert to interface once we move to Java 8.
+ */
+public abstract class CompactImmutableMap<K, V> implements Iterable<K> {
+
+  public boolean containsKey(K key) {
+    return get(key) != null;
+  }
+
+  public abstract V get(K key);
+
+  public abstract int size();
+
+  public abstract K keyAt(int index);
+
+  public abstract V valueAt(int index);
+
+  @Override
+  public Iterator<K> iterator() {
+    return new ImmutableMapSharedKeysIterator();
+  }
+
+  private class ImmutableMapSharedKeysIterator implements Iterator<K> {
+    int index = 0;
+
+    @Override
+    public boolean hasNext() {
+      return index < size();
+    }
+
+    @Override
+    public K next() {
+      K key = keyAt(index);
+      ++index;
+      return key;
+    }
+  }
+}
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
new file mode 100644
index 0000000000..792c671f7e
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/collect/ImmutableSharedKeyMap.java
@@ -0,0 +1,202 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.collect;
+
+import com.google.common.base.Objects;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Interner;
+import com.google.devtools.build.lib.concurrent.BlazeInterners;
+import com.google.devtools.build.lib.util.Preconditions;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.concurrent.Immutable;
+
+/**
+ * Provides a memory-efficient map when the key sets are likely to be shared between multiple
+ * instances of this class.
+ *
+ * <p>This class is appropriate where it is expected that a lot of the key sets will be the same.
+ * These key sets are shared and an offset table of indices is computed. Each map instance thus
+ * contains only a reference to the shared offset table, and a plain array of instances.
+ *
+ * <p>The map is sensitive to insertion order. Two maps with different insertion orders are *not*
+ * considered equal, and will not share keys.
+ *
+ * <p>This class explicitly does *not* implement the Map interface, as use of that would lead to a
+ * lot of GC churn.
+ */
+@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;
+
+  private 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
+    // during interning anyway (the majority).
+    private volatile ImmutableMap<K, Integer> indexMap;
+
+    private OffsetTable(Object[] keys) {
+      this.keys = keys;
+    }
+
+    void initIndexMap() {
+      if (indexMap == null) {
+        synchronized (this) {
+          if (indexMap == null) {
+            ImmutableMap.Builder<K, Integer> builder = ImmutableMap.builder();
+            for (int i = 0; i < keys.length; ++i) {
+              @SuppressWarnings("unchecked")
+              K key = (K) keys[i];
+              builder.put(key, i);
+            }
+            this.indexMap = builder.build();
+          }
+        }
+      }
+    }
+
+    private ImmutableMap<K, Integer> getIndexMap() {
+      return indexMap;
+    }
+
+    int offsetForKey(K key) {
+      return getIndexMap().getOrDefault(key, -1);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof OffsetTable)) {
+        return false;
+      }
+      OffsetTable that = (OffsetTable) o;
+      return Arrays.equals(this.keys, that.keys);
+    }
+
+    @Override
+    public int hashCode() {
+      return Arrays.hashCode(keys);
+    }
+  }
+
+  protected ImmutableSharedKeyMap(Object[] keys, Object[] values) {
+    Preconditions.checkArgument(keys.length == values.length);
+    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) {
+    OffsetTable<K> offsetTable = new OffsetTable<>(keys);
+    OffsetTable<K> internedTable = (OffsetTable<K>) offsetTables.intern(offsetTable);
+    internedTable.initIndexMap();
+    return internedTable;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public V get(K key) {
+    int offset = offsetTable.offsetForKey(key);
+    return offset != -1 ? (V) values[offset] : null;
+  }
+
+  @Override
+  public int size() {
+    return values.length;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public K keyAt(int index) {
+    return (K) offsetTable.keys[index];
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public V valueAt(int index) {
+    return (V) values[index];
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    ImmutableSharedKeyMap<?, ?> that = (ImmutableSharedKeyMap<?, ?>) o;
+    // We can use object identity for the offset table due to
+    // it being interned
+    return offsetTable == that.offsetTable && Arrays.equals(values, that.values);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(offsetTable, Arrays.hashCode(values));
+  }
+
+  public static <K, V> Builder<K, V> builder() {
+    return new Builder<>();
+  }
+
+  /** Builder for {@link ImmutableSharedKeyMap}. */
+  public static class Builder<K, V> {
+    private final List<Object> entries = new ArrayList<>();
+
+    private Builder() {}
+
+    public Builder<K, V> put(K key, V value) {
+      entries.add(key);
+      entries.add(value);
+      return this;
+    }
+
+    public ImmutableSharedKeyMap<K, V> build() {
+      int count = entries.size() / 2;
+      Object[] keys = new Object[count];
+      Object[] values = new Object[count];
+      int entryIndex = 0;
+      for (int i = 0; i < count; ++i) {
+        keys[i] = entries.get(entryIndex++);
+        values[i] = entries.get(entryIndex++);
+      }
+      return new ImmutableSharedKeyMap<>(keys, values);
+    }
+  }
+}