Optimize memory use of AttributeContainer:

- Relatively few locations are set relative to the number of attributes.
  Replace the sparse array with a dense one.
- BitSet requires a minimum of two objects (48 bytes in our current JVM),
  but we set an average of just six bits.  Replace it with a list of
  indices packed into a byte[], shared with the locations array.

Also add some assertions to help the next reader.

--
MOS_MIGRATED_REVID=124830576

2 files changed, 180 insertions, 18 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/AttributeContainer.java b/src/main/java/com/google/devtools/build/lib/packages/AttributeContainer.java
index 8c5bfffaa6..8da08274f5 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/AttributeContainer.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/AttributeContainer.java
@@ -13,10 +13,11 @@
 // limitations under the License.
 package com.google.devtools.build.lib.packages;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.devtools.build.lib.events.Location;
 
-import java.util.BitSet;
+import java.util.Arrays;
 
 /**
  * Provides attribute setting and retrieval for a Rule. Encapsulating attribute access
@@ -38,26 +39,42 @@ public class AttributeContainer {
   // Attribute values, keyed by attribute index:
   private final Object[] attributeValues;
 
-  // Whether an attribute value has been set explicitly in the BUILD file, keyed by attribute index.
-  private final BitSet attributeValueExplicitlySpecified;
+  // Holds two lists of attribute indices.
+  // The first byte gives the length of the first list.
+  // The first list records which attributes were set explicitly in the BUILD file.
+  // The second list ends at the end of the array.
+  // The second list records which attributes have Locations, in reverse order
+  // from the attributeLocations array.
+  // Between the lists there may be unused zero bytes (zeros are forbidden within each list).
+  private byte[] state;
+
+  // Attribute locations, packed:
+  private Location[] attributeLocations;
 
-  // Attribute locations, keyed by attribute index:
-  private final Location[] attributeLocations;
 
   /**
    * Create a container for a rule of the given rule class.
    */
   public AttributeContainer(RuleClass ruleClass) {
-   this(ruleClass, new Location[ruleClass.getAttributeCount()]);
+   this(ruleClass, EMPTY_LOCATIONS);
   }
-  
+
   AttributeContainer(RuleClass ruleClass, Location[] locations) {
+    int n = ruleClass.getAttributeCount();
+    if (n > 254) {
+      // We reserve the zero byte as a hole/sentinel inside state[].
+      // If you hit this limit, replace byte with char and remove the masking with 0xff.
+      throw new AssertionError("can't pack " + n + " rule indices into bytes");
+    }
     this.ruleClass = ruleClass;
-    this.attributeValues = new Object[ruleClass.getAttributeCount()];
-    this.attributeValueExplicitlySpecified = new BitSet(ruleClass.getAttributeCount());
+    this.attributeValues = new Object[n];
+    this.state = EMPTY_STATE;
     this.attributeLocations = locations;
   }
 
+  private static final byte[] EMPTY_STATE = {0};
+  private static final Location[] EMPTY_LOCATIONS = {};
+
   /**
    * Returns an attribute value by instance, or null on no match.
    */
@@ -83,7 +100,78 @@ public class AttributeContainer {
 
   public boolean isAttributeValueExplicitlySpecified(String attributeName) {
     Integer idx = ruleClass.getAttributeIndex(attributeName);
-    return idx != null && attributeValueExplicitlySpecified.get(idx);
+    return idx != null && getExplicit(idx);
+  }
+
+ /**
+  * Returns the number of elements of state[] currently used to store
+  * indices of "explicitly set" attributes.
+  */
+  private int explicitCount() {
+    return 0xff & state[0];
+  }
+
+  private boolean getExplicit(int index) {
+    int n = explicitCount();
+    for (int i = 1; i <= n; ++i) {
+      if ((0xff & state[i]) == index + 1) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private int getLocationIndex(int index) {
+    int n = explicitCount();
+    for (int i = state.length - 1; i > n; --i) {
+      if ((0xff & state[i]) == index + 1) {
+        return state.length - 1 - i;
+      }
+    }
+    return -1;
+  }
+
+  private void setExplicit(int index) {
+    if (getExplicit(index)) {
+      return;
+    }
+    ensureSpace();
+    int n = explicitCount() + 1;
+    state[0] = (byte) n;
+    state[n] = (byte) (index + 1);
+  }
+
+  private int addLocationIndex(int index) {
+    ensureSpace();
+    for (int i = state.length - 1; ; --i) {
+      if (i <= explicitCount()) {
+        throw new AssertionError("ensureSpace() did not insert a zero");
+      }
+      if (state[i] == 0) {
+        state[i] = (byte) (index + 1);
+        return state.length - 1 - i;
+      }
+    }
+  }
+
+  /**
+   * Ensures that the state[n] byte is equal to the sentinel value 0, so there is room for another
+   * attribute's explicit bit to be set, or for an attribute's location to be set.
+   */
+  private void ensureSpace() {
+    int n = explicitCount() + 1;
+    if (n < state.length && state[n] == 0) {
+      return;
+    }
+    // Grow up to the next multiple of eight bytes, as the object will be
+    // aligned to eight bytes anyway.  Insert zeros between the two lists.
+    byte[] newState = new byte[(state.length | 7) + 1];
+    // Copy stored explicit attributes to the beginning of the array.
+    System.arraycopy(state, 0, newState, 0, n);
+    // Copy stored attribute locations to the *end* of the array.
+    int oldLocations = state.length - n;
+    System.arraycopy(state, n, newState, newState.length - oldLocations, oldLocations);
+    state = newState;
   }
 
   /**
@@ -91,7 +179,8 @@ public class AttributeContainer {
    */
   public Location getAttributeLocation(String attrName) {
     Integer idx = ruleClass.getAttributeIndex(attrName);
-    return idx != null ? attributeLocations[idx] : null;
+    int locationIndex = idx != null ? getLocationIndex(idx) : -1;
+    return locationIndex >= 0 ? attributeLocations[locationIndex] : null;
   }
 
   Object getAttributeValue(int index) {
@@ -99,24 +188,43 @@ public class AttributeContainer {
   }
 
   void setAttributeValue(Attribute attribute, Object value, boolean explicit) {
-    Integer index = ruleClass.getAttributeIndex(attribute.getName());
+    String name = attribute.getName();
+    Integer index = ruleClass.getAttributeIndex(name);
+    if (!explicit && getExplicit(index)) {
+      throw new IllegalArgumentException("attribute " + name + " already explicitly set");
+    }
     attributeValues[index] = value;
-    attributeValueExplicitlySpecified.set(index, explicit);
+    if (explicit) {
+      setExplicit(index);
+    }
   }
 
+  // This sets the attribute "explicitly" as if it came from the BUILD file.
+  // At present, the sole use of this is for the test_suite.$implicit_tests
+  // attribute, which is synthesized during package loading.  We do want to
+  // consider that "explicitly set" so that it appears in query output.
   void setAttributeValueByName(String attrName, Object value) {
-    Integer index = ruleClass.getAttributeIndex(attrName);
-    attributeValues[index] = value;
-    attributeValueExplicitlySpecified.set(index);
+    setAttributeValue(ruleClass.getAttributeByName(attrName), value, true);
   }
 
   void setAttributeLocation(int attrIndex, Location location) {
-    attributeLocations[attrIndex] = location;
+    int locationIndex = getLocationIndex(attrIndex);
+    if (locationIndex >= 0) {
+      throw new IllegalArgumentException("already have a location for attribute "
+          + ruleClass.getAttribute(attrIndex).getName() + ": " + attributeLocations[locationIndex]);
+    }
+    locationIndex = addLocationIndex(attrIndex);
+    if (locationIndex >= attributeLocations.length) {
+      // Grow by two references, as the object will be aligned to eight bytes anyway.
+      attributeLocations = Arrays.copyOf(attributeLocations, attributeLocations.length + 2);
+    }
+    attributeLocations[locationIndex] = location;
   }
 
+  @VisibleForTesting
   void setAttributeLocation(Attribute attribute, Location location) {
     Integer index = ruleClass.getAttributeIndex(attribute.getName());
-    attributeLocations[index] = location;
+    setAttributeLocation(index, location);
   }
 
   public static final Function<RuleClass, AttributeContainer> ATTRIBUTE_CONTAINER_FACTORY =
diff --git a/src/test/java/com/google/devtools/build/lib/packages/AttributeContainerTest.java b/src/test/java/com/google/devtools/build/lib/packages/AttributeContainerTest.java
index 1aa0790cd9..8bf923b589 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/AttributeContainerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/AttributeContainerTest.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.packages;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 import com.google.devtools.build.lib.events.Location;
@@ -27,6 +28,10 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Random;
+
 /**
  * Unit tests for {@link AttributeContainer}.
  */
@@ -99,4 +104,53 @@ public class AttributeContainerTest {
     assertEquals(location2, container.getAttributeLocation(attribute2.getName()));
     assertNull(container.getAttributeLocation("nomatch"));
   }
+
+  @Test
+  public void testPackedState() throws Exception {
+    Random rng = new Random();
+    // The state packing machinery has special behavior at multiples of 8,
+    // so set enough explicit values and locations to exercise that.
+    final int N = 17;
+    Attribute[] attributes = new Attribute[N];
+    for (int i = 0; i < N; ++i) {
+      attributes[i] = ruleClass.getAttribute(i);
+    }
+    Object someValue = new Object();
+    Location[] locations = new Location[N];
+    for (int i = 0; i < N; ++i) {
+      locations[i] = newLocation();
+    }
+    assertTrue(locations[0] != locations[1]);  // test relies on checking reference inequality
+    for (int explicitCount = 0; explicitCount <= N; ++explicitCount) {
+      for (int locationCount = 0; locationCount <= N; ++locationCount) {
+        AttributeContainer container = new AttributeContainer(ruleClass);
+        // Shuffle the attributes each time through, to exercise
+        // different stored indices and orderings.
+        Collections.shuffle(Arrays.asList(attributes));
+        // Also randomly interleave calls to the two setters.
+        int valuePassKey = rng.nextInt(1 << N);
+        int locationPassKey = rng.nextInt(1 << N);
+        for (int pass = 0; pass <= 1; ++pass) {
+          for (int i = 0; i < explicitCount; ++i) {
+            if (pass == ((valuePassKey >> i) & 1)) {
+              container.setAttributeValue(attributes[i], someValue, true);
+            }
+          }
+          for (int i = 0; i < locationCount; ++i) {
+            if (pass == ((locationPassKey >> i) & 1)) {
+              container.setAttributeLocation(attributes[i], locations[i]);
+            }
+          }
+        }
+        for (int i = 0; i < N; ++i) {
+          boolean expected = i < explicitCount;
+          assertEquals(expected, container.isAttributeValueExplicitlySpecified(attributes[i]));
+        }
+        for (int i = 0; i < N; ++i) {
+          Location expected = i < locationCount ? locations[i] : null;
+          assertSame(expected, container.getAttributeLocation(attributes[i].getName()));
+        }
+      }
+    }
+  }
 }