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 8a0fc00585f5f6..8c9f5003bf594c 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
@@ -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,6 +153,25 @@ 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.
@@ -161,7 +179,7 @@ private static int nonNullCount(Object[] attrValues) {
   @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,7 +307,7 @@ 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);
     }
@@ -316,8 +315,8 @@ private static void setBit(byte[] bits, int bitIndex) {
     /** 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
index de8be7432a238a..9e1295d6cfb25d 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
@@ -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);
+  }
 }