more robust bit-trickery functions

4 files changed, 81 insertions, 47 deletions

diff --git a/c/blake3.c b/c/blake3.c
index 17fbc3b..b4a5992 100644
--- a/c/blake3.c
+++ b/c/blake3.c
@@ -484,7 +484,7 @@ void blake3_hasher_update(blake3_hasher *self, const void *input,
   // Because we might need to break up the input to form powers of 2, or to
   // evenly divide what we already have, this part runs in a loop.
   while (input_len > BLAKE3_CHUNK_LEN) {
-    size_t subtree_len = (size_t)round_down_to_power_of_2((uint64_t)input_len);
+    size_t subtree_len = round_down_to_power_of_2(input_len);
     uint64_t count_so_far = self->chunk.chunk_counter * BLAKE3_CHUNK_LEN;
     // Shrink the subtree_len until *half of it* it evenly divides the count so
     // far. Why half? Because compress_subtree_to_parent_node will return a
@@ -522,7 +522,7 @@ void blake3_hasher_update(blake3_hasher *self, const void *input,
                      self->chunk.chunk_counter + (subtree_chunks / 2));
     }
     self->chunk.chunk_counter += subtree_chunks;
-    input_bytes = input_bytes + subtree_len;
+    input_bytes += subtree_len;
     input_len -= subtree_len;
   }
 
diff --git a/c/blake3.h b/c/blake3.h
index 43e6522..a7e369b 100644
--- a/c/blake3.h
+++ b/c/blake3.h
@@ -1,4 +1,5 @@
-#pragma once
+#ifndef BLAKE3_H
+#define BLAKE3_H
 
 #include <stddef.h>
 #include <stdint.h>
@@ -39,3 +40,5 @@ void blake3_hasher_update(blake3_hasher *self, const void *input,
                           size_t input_len);
 void blake3_hasher_finalize(const blake3_hasher *self, uint8_t *out,
                             size_t out_len);
+
+#endif /* BLAKE3_H */
diff --git a/c/blake3_impl.h b/c/blake3_impl.h
index 20ef83d..84caae9 100644
--- a/c/blake3_impl.h
+++ b/c/blake3_impl.h
@@ -1,4 +1,5 @@
-#pragma once
+#ifndef BLAKE3_IMPL_H
+#define BLAKE3_IMPL_H
 
 #include <assert.h>
 #include <stdbool.h>
@@ -6,32 +7,35 @@
 #include <stdint.h>
 #include <string.h>
 
-#if __POPCNT__
-#include <nmmintrin.h>
-#endif
-
 #include "blake3.h"
 
 // internal flags
-#define CHUNK_START 1
-#define CHUNK_END 2
-#define PARENT 4
-#define ROOT 8
-#define KEYED_HASH 16
-#define DERIVE_KEY_CONTEXT 32
-#define DERIVE_KEY_MATERIAL 64
+enum blake3_flags {
+  CHUNK_START         = 1 << 0,
+  CHUNK_END           = 1 << 1,
+  PARENT              = 1 << 2,
+  ROOT                = 1 << 3,
+  KEYED_HASH          = 1 << 4,
+  DERIVE_KEY_CONTEXT  = 1 << 5,
+  DERIVE_KEY_MATERIAL = 1 << 6,
+};
 
 // This C implementation tries to support recent versions of GCC, Clang, and
 // MSVC.
 #if defined(_MSC_VER)
-#define INLINE __forceinline static
+#define INLINE static __forceinline
 #else
-#define INLINE __attribute__((always_inline)) static inline
+#define INLINE static inline __attribute__((always_inline))
+#endif
+
+#if defined(__x86_64__) || defined(_M_X64) 
+#define IS_X86
+#define IS_X86_64
 #endif
 
-#if defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) ||            \
-    defined(_M_X64)
+#if defined(__i386__) || defined(_M_IX86)
 #define IS_X86
+#define IS_X86_32
 #endif
 
 #if defined(__arm__)
@@ -39,6 +43,13 @@
 #endif
 
 #if defined(IS_X86)
+#if defined(_MSC_VER)
+#include <intrin.h>
+#endif
+#include <immintrin.h>
+#endif
+
+#if defined(IS_X86)
 #define MAX_SIMD_DEGREE 16
 #elif defined(BLAKE3_USE_NEON)
 #define MAX_SIMD_DEGREE 4
@@ -48,13 +59,7 @@
 
 // There are some places where we want a static size that's equal to the
 // MAX_SIMD_DEGREE, but also at least 2.
-#if defined(IS_X86)
-#define MAX_SIMD_DEGREE_OR_2 16
-#elif defined(BLAKE3_USE_NEON)
-#define MAX_SIMD_DEGREE_OR_2 4
-#else
-#define MAX_SIMD_DEGREE_OR_2 2
-#endif
+#define MAX_SIMD_DEGREE_OR_2 (MAX_SIMD_DEGREE > 2 ? MAX_SIMD_DEGREE : 2)
 
 static const uint32_t IV[8] = {0x6A09E667UL, 0xBB67AE85UL, 0x3C6EF372UL,
                                0xA54FF53AUL, 0x510E527FUL, 0x9B05688CUL,
@@ -70,13 +75,44 @@ static const uint8_t MSG_SCHEDULE[7][16] = {
     {11, 15, 5, 0, 1, 9, 8, 6, 14, 10, 2, 12, 3, 4, 7, 13},
 };
 
+/* Find index of the highest set bit */
+/* x is assumed to be nonzero.       */
+static unsigned int fls(uint64_t x) {
+#if defined(__GNUC__) || defined(__clang__)
+  return 63 ^ __builtin_clzll(x);
+#elif defined(_MSC_VER) && defined(IS_X86_64)
+  unsigned long index;
+  _BitScanReverse64(&index, x);
+  return index;
+#elif defined(_MSC_VER) && defined(IS_X86_32)
+  if(x >> 32) {
+    unsigned long index;
+    _BitScanReverse(&index, x >> 32);
+    return 32 + index;
+  } else {
+    unsigned long index;
+    _BitScanReverse(&index, x);
+    return index;
+  }
+#else
+  unsigned int c = 0;
+  if(x & 0xffffffff00000000ULL) { x >>= 32; c += 32; }
+  if(x & 0x00000000ffff0000ULL) { x >>= 16; c += 16; }
+  if(x & 0x000000000000ff00ULL) { x >>=  8; c +=  8; }
+  if(x & 0x00000000000000f0ULL) { x >>=  4; c +=  4; }
+  if(x & 0x000000000000000cULL) { x >>=  2; c +=  2; }
+  if(x & 0x0000000000000002ULL) {           c +=  1; }
+  return c;
+#endif
+}
+
 // Count the number of 1 bits.
-INLINE uint8_t popcnt(uint64_t x) {
-#if __POPCNT__
-  return (uint8_t)_mm_popcnt_u64(x);
+INLINE unsigned int popcnt(uint64_t x) {
+#if defined(__GNUC__) || defined(__clang__)
+  return __builtin_popcountll(x);
 #else
-  uint8_t count = 0;
-  while (x > 0) {
+  unsigned int count = 0;
+  while (x != 0) {
     count += 1;
     x &= x - 1;
   }
@@ -85,16 +121,9 @@ INLINE uint8_t popcnt(uint64_t x) {
 }
 
 // Largest power of two less than or equal to x. As a special case, returns 1
-// when x is 0. Based on
-// https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2.
+// when x is 0. 
 INLINE uint64_t round_down_to_power_of_2(uint64_t x) {
-  x |= x >> 1;
-  x |= x >> 2;
-  x |= x >> 4;
-  x |= x >> 8;
-  x |= x >> 16;
-  x |= x >> 32;
-  return (x >> 1) + 1;
+  return 1ULL << fls(x | 1);
 }
 
 INLINE uint32_t counter_low(uint64_t counter) { return (uint32_t)counter; }
@@ -137,3 +166,6 @@ void blake3_hash_many(const uint8_t *const *inputs, size_t num_inputs,
                       uint8_t flags_start, uint8_t flags_end, uint8_t *out);
 
 size_t blake3_simd_degree();
+
+
+#endif /* BLAKE3_IMPL_H */
diff --git a/c/main.c b/c/main.c
index 9ef7bce..9433937 100644
--- a/c/main.c
+++ b/c/main.c
@@ -78,14 +78,11 @@ int main(int argc, char **argv) {
     if (strcmp("--length", argv[1]) == 0) {
       char *endptr = NULL;
       unsigned long long out_len_ll = strtoull(argv[2], &endptr, 10);
-      // TODO: There are so many possible error conditions for parsing a
-      //       non-negative size_t...I probably missed something.
       if (errno != 0 || out_len > SIZE_MAX || endptr == argv[2] ||
           *endptr != 0) {
         fprintf(stderr, "Bad length argument.\n");
         return 1;
       }
-      // TODO: A more sanitary cast?
       out_len = (size_t)out_len_ll;
     } else if (strcmp("--keyed", argv[1]) == 0) {
       mode = KEYED_HASH_MODE;
@@ -104,9 +101,11 @@ int main(int argc, char **argv) {
     argv += 2;
   }
 
-  // We're going to hash the input multiple times, so we need to buffer it all.
-  // This is just for test cases, so go ahead and assume that the input is less
-  // than 1 MiB.
+  /* 
+   * We're going to hash the input multiple times, so we need to buffer it all.
+   * This is just for test cases, so go ahead and assume that the input is less
+   * than 1 MiB. 
+   */
   size_t buf_capacity = 1 << 20;
   uint8_t *buf = malloc(buf_capacity);
   assert(buf != NULL);
@@ -142,7 +141,7 @@ int main(int argc, char **argv) {
 
     blake3_hasher_update(&hasher, buf, buf_len);
 
-    // TODO: An incremental output reader API to avoid this allocation.
+    /* TODO: An incremental output reader API to avoid this allocation. */
     uint8_t *out = malloc(out_len);
     memset(out, 0, out_len);
     if (out_len > 0 && out == NULL) {