AttributeContainer.Large now handles more than 127 attributes.

katre · copybara-github · commit e8835c1c221d · 2020-12-02T13:39:23.000-08:00
The previous implementation failed with more than 127 attributes because the call to Arrays.binarySearch treats the byte values as signed, so values of 128 and above wrapped around and became negative, breaking binarySearch's precondition that the array be sorted.

Also adds a small test.

PiperOrigin-RevId: 345302645
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
@@ -16,7 +16,6 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.errorprone.annotations.CheckReturnValue;
-import java.util.Arrays;
 import java.util.BitSet;
 import javax.annotation.Nullable;
 
@@ -154,14 +153,33 @@ private static int nonNullCount(Object[] attrValues) {
     return numSet;
   }
 
+  /** Returns index into state array for attrIndex, or -1 if not found */
+  private static int getStateIndex(byte[] state, int start, int attrIndex, int mask) {
+    // Binary search, treating values as unsigned bytes.
+    int lo = start;
+    int hi = state.length - 1;
+    while (hi >= lo) {
+      int mid = (lo + hi) / 2;
+      int midAttrIndex = state[mid] & mask;
+      if (midAttrIndex == attrIndex) {
+        return mid;
+      } else if (midAttrIndex < attrIndex) {
+        lo = mid + 1;
+      } else {
+        hi = mid - 1;
+      }
+    }
+    return -1;
+  }
+
   /**
    * A frozen implementation of AttributeContainer which supports RuleClasses with up to 126
    * attributes.
    */
   @VisibleForTesting
   static final class Small extends Frozen {
 
-    private int maxAttrCount;
+    private final int maxAttrCount;
 
     // Conceptually an AttributeContainer is an unordered set of triples
     // (attribute, value, explicit).
@@ -219,25 +237,6 @@ private Small(Object[] attrValues, BitSet explicitAttrs) {
       }
     }
 
-    /** Returns index into state array for attrIndex, or -1 if not found */
-    private int getStateIndex(int attrIndex) {
-      // Binary search on the bottom 7 bits.
-      int lo = 0;
-      int hi = state.length - 1;
-      while (hi >= lo) {
-        int mid = (lo + hi) / 2;
-        int midAttrIndex = state[mid] & 0x7f;
-        if (midAttrIndex == attrIndex) {
-          return mid;
-        } else if (midAttrIndex < attrIndex) {
-          lo = mid + 1;
-        } else {
-          hi = mid - 1;
-        }
-      }
-      return -1;
-    }
-
     /**
      * Returns true iff the value of the specified attribute is explicitly set in the BUILD file. In
      * addition, this method return false if the rule has no attribute with the given name.
@@ -247,7 +246,7 @@ boolean isAttributeValueExplicitlySpecified(int attrIndex) {
       if (attrIndex < 0) {
         return false;
       }
-      int stateIndex = getStateIndex(attrIndex);
+      int stateIndex = getStateIndex(state, 0, attrIndex, 0x7f);
       return stateIndex >= 0 && (state[stateIndex] & 0x80) != 0;
     }
 
@@ -258,7 +257,7 @@ Object getAttributeValue(int attrIndex) {
         throw new IndexOutOfBoundsException(
             "Maximum valid attrIndex is " + (maxAttrCount - 1) + ". Given " + attrIndex);
       }
-      int stateIndex = getStateIndex(attrIndex);
+      int stateIndex = getStateIndex(state, 0, attrIndex, 0x7f);
       return stateIndex < 0 ? null : values[stateIndex];
     }
   }
@@ -308,16 +307,16 @@ private static int prefixSize(int attrCount) {
     /** Set the specified bit in the byte array. Assumes bitIndex is a valid index. */
     private static void setBit(byte[] bits, int bitIndex) {
       int idx = (bitIndex + 1);
-      byte explicitByte = bits[idx >> 3];
+      int explicitByte = bits[idx >> 3];
       byte mask = (byte) (1 << (idx & 0x07));
       bits[idx >> 3] = (byte) (explicitByte | mask);
     }
 
     /** Get the specified bit in the byte array. Assumes bitIndex is a valid index. */
     private static boolean getBit(byte[] bits, int bitIndex) {
       int idx = (bitIndex + 1);
-      byte explicitByte = bits[idx >> 3];
-      byte mask = (byte) (1 << (idx & 0x07));
+      int explicitByte = bits[idx >> 3];
+      int mask = (byte) (1 << (idx & 0x07));
       return (explicitByte & mask) != 0;
     }
 
@@ -374,7 +373,7 @@ Object getAttributeValue(int attrIndex) {
             "Maximum valid attrIndex is " + (maxAttrCount - 1) + ". Given " + attrIndex);
       }
       int p = prefixSize(maxAttrCount);
-      int stateIndex = Arrays.binarySearch(state, p, state.length, (byte) attrIndex);
+      int stateIndex = getStateIndex(state, p, attrIndex, 0xff);
       return stateIndex < 0 ? null : values[stateIndex - p];
     }
   }
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
@@ -158,4 +158,29 @@ public void testFreezeWorks_smallImplementation() {
   public void testFreezeWorks_largeImplementation() {
     checkFreezeWorks((short) 150, AttributeContainer.Large.class);
   }
+
+  private void testContainerSize(int size) {
+    AttributeContainer container = new Mutable(size);
+    for (int i = 0; i < size; i++) {
+      container.setAttributeValue(i, "value " + i, i % 2 == 0);
+    }
+    AttributeContainer frozen = container.freeze();
+    // Check values.
+    for (int i = 0; i < size; i++) {
+      assertThat(frozen.getAttributeValue(i)).isEqualTo("value " + i);
+      assertThat(frozen.isAttributeValueExplicitlySpecified(i)).isEqualTo(i % 2 == 0);
+    }
+  }
+
+  @Test
+  public void testSmallContainer() {
+    // At 127 attributes, we shift from AttributeContainer.Small to AttributeContainer.Large.
+    testContainerSize(126);
+  }
+
+  @Test
+  public void testLargeContainer() {
+    // AttributeContainer.Large can handle at max 254 attributes.
+    testContainerSize(254);
+  }
 }
