Merge pull request #458 from adriancole/variable-bitset

Reduce allocation for HPACK decompression

okhttp-protocols/src/main/java/com/squareup/okhttp/internal/BitArray.java

@@ -19,101 +19,159 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
+import static java.lang.String.format;
+
 /** A simple bitset which supports left shifting. */
-public final class BitArray {
+public interface BitArray {
 
-  long[] data;
+  void clear();
 
-  // Start offset which allows for cheap shifting. Data is always kept on 64-bit bounds but we
-  // offset the outward facing index to support shifts without having to move the underlying bits.
-  private int start; // Valid values are [0..63]
+  void set(int index);
 
-  public BitArray() {
-    data = new long[1];
-  }
+  void toggle(int index);
 
-  private void growToSize(int size) {
-    long[] newData = new long[size];
-    if (data != null) {
-      System.arraycopy(data, 0, newData, 0, data.length);
+  boolean get(int index);
+
+  void shiftLeft(int count);
+
+  /** Bit set that only supports settings bits 0 - 63. */
+  public final class FixedCapacity implements BitArray {
+    long data = 0x0000000000000000L;
+
+    @Override public void clear() {
+      data = 0x0000000000000000L;
     }
-    data = newData;
-  }
 
-  private int offsetOf(int index) {
-    index += start;
-    int offset = index / 64;
-    if (offset > data.length - 1) {
-      growToSize(offset + 1);
+    @Override public void set(int index) {
+      data |= (1L << checkInput(index));
     }
-    return offset;
-  }
 
-  private int shiftOf(int index) {
-    return (index + start) % 64;
-  }
-
-  public void clear() {
-    Arrays.fill(data, 0);
-  }
-
-  public void set(int index) {
-    if (index < 0) {
-      throw new IllegalArgumentException("index < 0: " + index);
+    @Override public void toggle(int index) {
+      data ^= (1L << checkInput(index));
+    }
+
+    @Override public boolean get(int index) {
+      return ((data >> checkInput(index)) & 1L) == 1;
+    }
+
+    @Override public void shiftLeft(int count) {
+      data = data << checkInput(count);
+    }
+
+    @Override public String toString() {
+      return Long.toBinaryString(data);
+    }
+
+    public BitArray toVariableCapacity() {
+      return new VariableCapacity(this);
+    }
+
+    private static int checkInput(int index) {
+      if (index < 0 || index > 63) {
+        throw new IllegalArgumentException(format("input must be between 0 and 63: %s", index));
+      }
+      return index;
     }
-    int offset = offsetOf(index);
-    data[offset] |= 1L << shiftOf(index);
   }
 
-  public void toggle(int index) {
-    if (index < 0) {
-      throw new IllegalArgumentException("index < 0: " + index);
-    }
-    int offset = offsetOf(index);
-    data[offset] ^= 1L << shiftOf(index);
-  }
+  /** Bit set that grows as needed. */
+  public final class VariableCapacity implements BitArray {
 
-  public boolean get(int index) {
-    if (index < 0) {
-      throw new IllegalArgumentException("index < 0: " + index);
-    }
-    int offset = offsetOf(index);
-    return (data[offset] & (1L << shiftOf(index))) != 0;
-  }
+    long[] data;
 
-  public void shiftLeft(int count) {
-    if (count < 0) {
-      throw new IllegalArgumentException("count < 0: " + count);
+    // Start offset which allows for cheap shifting. Data is always kept on 64-bit bounds but we
+    // offset the outward facing index to support shifts without having to move the underlying bits.
+    private int start; // Valid values are [0..63]
+
+    public VariableCapacity() {
+      data = new long[1];
     }
-    start -= count;
-    if (start < 0) {
-      int arrayShift = (start / -64) + 1;
-      long[] newData = new long[data.length + arrayShift];
-      System.arraycopy(data, 0, newData, arrayShift, data.length);
+
+    private VariableCapacity(FixedCapacity small) {
+      data = new long[] {small.data, 0};
+    }
+
+    private void growToSize(int size) {
+      long[] newData = new long[size];
+      if (data != null) {
+        System.arraycopy(data, 0, newData, 0, data.length);
+      }
       data = newData;
-      start = 64 + (start % 64);
     }
-  }
 
-  @Override public String toString() {
-    StringBuilder builder = new StringBuilder("{");
-    List<Integer> ints = toIntegerList();
-    for (int i = 0, count = ints.size(); i < count; i++) {
-      if (i > 0) {
-        builder.append(',');
+    private int offsetOf(int index) {
+      index += start;
+      int offset = index / 64;
+      if (offset > data.length - 1) {
+        growToSize(offset + 1);
       }
-      builder.append(ints.get(i));
+      return offset;
     }
-    return builder.append('}').toString();
-  }
 
-  List<Integer> toIntegerList() {
-    List<Integer> ints = new ArrayList<Integer>();
-    for (int i = 0, count = data.length * 64 - start; i < count; i++) {
-      if (get(i)) {
-        ints.add(i);
+    private int shiftOf(int index) {
+      return (index + start) % 64;
+    }
+
+    @Override public void clear() {
+      Arrays.fill(data, 0);
+    }
+
+    @Override public void set(int index) {
+      checkInput(index);
+      int offset = offsetOf(index);
+      data[offset] |= 1L << shiftOf(index);
+    }
+
+    @Override public void toggle(int index) {
+      checkInput(index);
+      int offset = offsetOf(index);
+      data[offset] ^= 1L << shiftOf(index);
+    }
+
+    @Override public boolean get(int index) {
+      checkInput(index);
+      int offset = offsetOf(index);
+      return (data[offset] & (1L << shiftOf(index))) != 0;
+    }
+
+    @Override public void shiftLeft(int count) {
+      start -= checkInput(count);
+      if (start < 0) {
+        int arrayShift = (start / -64) + 1;
+        long[] newData = new long[data.length + arrayShift];
+        System.arraycopy(data, 0, newData, arrayShift, data.length);
+        data = newData;
+        start = 64 + (start % 64);
       }
     }
-    return ints;
+
+    @Override public String toString() {
+      StringBuilder builder = new StringBuilder("{");
+      List<Integer> ints = toIntegerList();
+      for (int i = 0, count = ints.size(); i < count; i++) {
+        if (i > 0) {
+          builder.append(',');
+        }
+        builder.append(ints.get(i));
+      }
+      return builder.append('}').toString();
+    }
+
+    List<Integer> toIntegerList() {
+      List<Integer> ints = new ArrayList<Integer>();
+      for (int i = 0, count = data.length * 64 - start; i < count; i++) {
+        if (get(i)) {
+          ints.add(i);
+        }
+      }
+      return ints;
+    }
+
+    private static int checkInput(int index) {
+      if (index < 0) {
+        throw new IllegalArgumentException(format("input must be a positive number: %s", index));
+      }
+      return index;
+    }
   }
 }

okhttp-protocols/src/main/java/com/squareup/okhttp/internal/spdy/HpackDraft05.java

@@ -132,13 +132,14 @@ final class HpackDraft05 {
      * Set bit positions indicate {@code headerTable[pos]} should be emitted.
      */
     // Using a BitArray as it has left-shift operator.
-    BitArray referencedHeaders = new BitArray();
+    BitArray referencedHeaders = new BitArray.FixedCapacity();
 
     /**
      * Set bit positions indicate {@code STATIC_HEADER_TABLE[pos]} should be
      * emitted.
      */
-    BitArray referencedStaticHeaders = new BitArray();
+    // Using a long since the static table < 64 entries.
+    long referencedStaticHeaders = 0L;;
     int headerTableByteCount = 0;
 
     Reader(boolean client, int maxHeaderTableByteCount, DataInputStream in) {
@@ -214,13 +215,13 @@ final class HpackDraft05 {
     }
 
     private void clearReferenceSet() {
-      referencedStaticHeaders.clear();
+      referencedStaticHeaders = 0L;
       referencedHeaders.clear();
     }
 
     public void emitReferenceSet() {
       for (int i = 0; i < STATIC_HEADER_TABLE.length; ++i) {
-        if (referencedStaticHeaders.get(i)) {
+        if (((referencedStaticHeaders >> i) & 1L) == 1) {
           emittedHeaders.add(STATIC_HEADER_TABLE[i].name);
           emittedHeaders.add(STATIC_HEADER_TABLE[i].value);
         }
@@ -246,7 +247,7 @@ final class HpackDraft05 {
     private void readIndexedHeader(int index) {
       if (isStaticHeader(index)) {
         if (maxHeaderTableByteCount == 0) {
-          referencedStaticHeaders.set(index - headerCount);
+          referencedStaticHeaders |= (1L << (index - headerCount));
         } else {
           HeaderEntry staticEntry = STATIC_HEADER_TABLE[index - headerCount];
           insertIntoHeaderTable(-1, staticEntry);
@@ -328,6 +329,9 @@ final class HpackDraft05 {
         if (headerCount + 1 > headerTable.length) {
           HeaderEntry[] doubled = new HeaderEntry[headerTable.length * 2];
           System.arraycopy(headerTable, 0, doubled, headerTable.length, headerTable.length);
+          if (doubled.length == 64) {
+            referencedHeaders = ((BitArray.FixedCapacity) referencedHeaders).toVariableCapacity();
+          }
           referencedHeaders.shiftLeft(headerTable.length);
           nextHeaderIndex = headerTable.length - 1;
           headerTable = doubled;

okhttp-protocols/src/test/java/com/squareup/okhttp/internal/BitArrayTest.java

@@ -22,17 +22,80 @@ import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class BitArrayTest {
 
-  @Test public void setExpandsData() {
-    BitArray b = new BitArray();
+  /** Lazy grow into a variable capacity bit set. */
+  @Test public void hpackUseCase() {
+    BitArray b = new BitArray.FixedCapacity();
+    for (int i = 0; i < 64; i++) {
+      b.set(i);
+    }
+    assertTrue(b.get(0));
+    assertTrue(b.get(1));
+    assertTrue(b.get(63));
+    try {
+      b.get(64);
+      fail();
+    } catch (IllegalArgumentException expected) {
+    }
+    b = ((BitArray.FixedCapacity) b).toVariableCapacity();
+    assertTrue(b.get(0));
+    assertTrue(b.get(1));
+    assertTrue(b.get(63));
+    assertFalse(b.get(64));
+    b.set(64);
+    assertTrue(b.get(64));
+  }
+
+  @Test public void setExpandsData_FixedCapacity() {
+    BitArray.FixedCapacity b = new BitArray.FixedCapacity();
+    b.set(63);
+    assertEquals(b.data, BigInteger.ZERO.setBit(63).longValue());
+  }
+
+  @Test public void toggleBit_FixedCapacity() {
+    BitArray.FixedCapacity b = new BitArray.FixedCapacity();
+    b.set(63);
+    b.toggle(63);
+    assertEquals(b.data, 0l);
+    b.toggle(1);
+    assertEquals(b.data, 2l);
+  }
+
+  @Test public void shiftLeft_FixedCapacity() {
+    BitArray.FixedCapacity b = new BitArray.FixedCapacity();
+    b.set(0);
+    b.shiftLeft(1);
+    assertEquals(b.data, 2l);
+  }
+
+  @Test public void multipleShifts_FixedCapacity() {
+    BitArray.FixedCapacity b = new BitArray.FixedCapacity();
+    b.set(10);
+    b.shiftLeft(2);
+    b.shiftLeft(2);
+    assertEquals(b.data, BigInteger.ZERO.setBit(10).shiftLeft(2).shiftLeft(2).longValue());
+  }
+
+  @Test public void clearBits_FixedCapacity() {
+    BitArray.FixedCapacity b = new BitArray.FixedCapacity();
+    b.set(1);
+    b.set(3);
+    b.set(5);
+    b.clear();
+    assertEquals(b.data, 0l);
+  }
+
+  @Test public void setExpandsData_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(64);
     assertEquals(asList(64), b.toIntegerList());
   }
 
-  @Test public void toggleBit() {
-    BitArray b = new BitArray();
+  @Test public void toggleBit_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(100);
     b.toggle(100);
     assertTrue(b.toIntegerList().isEmpty());
@@ -40,22 +103,22 @@ public class BitArrayTest {
     assertEquals(asList(1), b.toIntegerList());
   }
 
-  @Test public void shiftLeftExpandsData() {
-    BitArray b = new BitArray();
+  @Test public void shiftLeftExpandsData_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(0);
     b.shiftLeft(64);
     assertEquals(asList(64), b.toIntegerList());
   }
 
-  @Test public void shiftLeftFromZero() {
-    BitArray b = new BitArray();
+  @Test public void shiftLeftFromZero_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(0);
     b.shiftLeft(1);
     assertEquals(asList(1), b.toIntegerList());
   }
 
-  @Test public void shiftLeftAcrossOffset() {
-    BitArray b = new BitArray();
+  @Test public void shiftLeftAcrossOffset_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(63);
     assertEquals(1, b.data.length);
     b.shiftLeft(1);
@@ -63,8 +126,8 @@ public class BitArrayTest {
     assertEquals(2, b.data.length);
   }
 
-  @Test public void multipleShiftsLeftAcrossOffset() {
-    BitArray b = new BitArray();
+  @Test public void multipleShiftsLeftAcrossOffset_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(1000);
     b.shiftLeft(67);
     assertEquals(asList(1067), b.toIntegerList());
@@ -72,8 +135,8 @@ public class BitArrayTest {
     assertEquals(asList(1136), b.toIntegerList());
   }
 
-  @Test public void clearBits() {
-    BitArray b = new BitArray();
+  @Test public void clearBits_VariableCapacity() {
+    BitArray.VariableCapacity b = new BitArray.VariableCapacity();
     b.set(10);
     b.set(100);
     b.set(1000);
@@ -81,8 +144,8 @@ public class BitArrayTest {
     assertTrue(b.toIntegerList().isEmpty());
   }
 
-  @Test public void bigIntegerSanityCheck() {
-    BitArray a = new BitArray();
+  @Test public void bigIntegerSanityCheck_VariableCapacity() {
+    BitArray a = new BitArray.VariableCapacity();
     BigInteger b = BigInteger.ZERO;
 
     a.set(64);