Improve Word jet CMR computation

Author: roconnor-blockstream

bitstring.h

@@ -32,4 +32,21 @@ static inline bool getBit(const bitstring *s, size_t n) {
   return 1 & (s->arr[total_offset / CHAR_BIT] >> (CHAR_BIT - 1 - (total_offset % CHAR_BIT)));
 }
 
+/* Return 8 bits from a bitstring staring from the nth bit.
+ *
+ * Precondition: NULL != s
+ *               n + 8 <= s->len;
+ */
+static inline uint_fast8_t getByte(const bitstring *s, size_t n) {
+  simplicity_assert(8 <= s->len);
+  simplicity_assert(n <= s->len - 8);
+  size_t total_offset = s->offset + n;
+  if (total_offset % CHAR_BIT <= CHAR_BIT - 8) {
+    return (uint_fast8_t)(0xff & (s->arr[total_offset / CHAR_BIT] >> (CHAR_BIT - 8 - (total_offset % CHAR_BIT))));
+  } else {
+    /* CHAR_BIT < total_offset % CHAR_BIT + 8 */
+    return (uint_fast8_t)(0xff & (((1U * s->arr[total_offset / CHAR_BIT]) << (total_offset % CHAR_BIT - (CHAR_BIT - 8)))
+                                  | (s->arr[total_offset / CHAR_BIT + 1] >> (CHAR_BIT - (total_offset % CHAR_BIT - (CHAR_BIT - 8))))));
+  }
+}
 #endif

dag.c

@@ -97,31 +97,39 @@ static sha256_midstate mkJetCMR(uint32_t *imr, uint_fast64_t weight) {
  * Precondition: 2^n == value->len
  */
 sha256_midstate computeWordCMR(const bitstring* value, size_t n) {
-  /* 'stack' is an array of 33 hashes consisting of 8 'uint32_t's each. */
-  uint32_t stack[8*33] = {0};
+  /* 'stack' is an array of 30 hashes consisting of 8 'uint32_t's each. */
+  uint32_t stack[8*30] = {0};
   uint32_t *stack_ptr = stack;
   sha256_midstate imr = identityIV;
   simplicity_assert(n < 32);
   simplicity_assert((size_t)1 << n == value->len);
   /* Pass 1: Compute the CMR for the expression that writes 'value'.
    * This expression consists of deeply nested PAIRs of expressions that write one bit each.
-   *
-   * :TODO: This can be optimized by a constant factor by precomputing a table of CMRs of expressions
-   * that, for example, write out every possible byte sequence.
    */
   /* stack[0..7] (8 bytes) is kept as all zeros for later.
    * We start the stack_ptr at the second item.
    */
-  for(size_t i = 0; i < value->len; ++i) {
-    /* stack_ptr == stack + 8*<count of the number of set bits in the value i> */
+  if (n < 3) {
     stack_ptr += 8;
-    memcpy(stack_ptr, &bit_cmr[getBit(value, i)], sizeof(uint32_t[8]));
-    /* This inner for loop runs in ammortized constant time. */
-    for (size_t j = i; j & 1; j = j >> 1) {
-      sha256_midstate pair = cmrIV(PAIR);
-      stack_ptr -= 8;
-      sha256_compression(pair.s, stack_ptr);
-      memcpy(stack_ptr, pair.s, sizeof(uint32_t[8]));
+    size_t i;
+    switch(n) {
+     case 0: i = getBit(value, 0); break;
+     case 1: i = 2 + ((1U * getBit(value, 0) << 1) | getBit(value, 1)); break;
+     case 2: i = 6 + ((1U * getBit(value, 0) << 3) | (1U * getBit(value, 1) << 2) | (1U * getBit(value, 2) << 1) | getBit(value, 3)); break;
+    }
+    memcpy(stack_ptr, &word_cmr[i], sizeof(uint32_t[8]));
+  } else {
+    for (size_t i = 0; i < value->len >> 3; ++i) {
+      /* stack_ptr == stack + 8*<count of the number of set bits in the value i> */
+      stack_ptr += 8;
+      memcpy(stack_ptr, &word_cmr[22 + getByte(value, 8*i)], sizeof(uint32_t[8]));
+      /* This inner for loop runs in amortized constant time. */
+      for (size_t j = i; j & 1; j = j >> 1) {
+        sha256_midstate pair = cmrIV(PAIR);
+        stack_ptr -= 8;
+        sha256_compression(pair.s, stack_ptr);
+        memcpy(stack_ptr, pair.s, sizeof(uint32_t[8]));
+      }
     }
   }
   /* value->len is a power of 2.*/