Merge pull request #135 from unicitynetwork/issue-134

#134 Use big endian bit ordering and byte array paths in smt

Author: b3y0urs3lf

src/api/InclusionCertificate.ts

@@ -5,8 +5,7 @@ import { HashAlgorithm } from '../crypto/hash/HashAlgorithm.js';
 import { FinalizedBranch } from '../smt/radix/FinalizedBranch.js';
 import { FinalizedLeafBranch } from '../smt/radix/FinalizedLeafBranch.js';
 import { SparseMerkleTreeRootNode } from '../smt/radix/SparseMerkleTreeRootNode.js';
-import { pathToRegion } from '../smt/SparseMerkleTreePathUtils.js';
-import { BitString } from '../util/BitString.js';
+import { getBitAtDepth, regionFromKey } from '../smt/SparseMerkleTreePathUtils.js';
 import { HexConverter } from '../util/HexConverter.js';
 import { dedent } from '../util/StringUtils.js';
 import { areUint8ArraysEqual } from '../util/TypedArrayUtils.js';
@@ -36,7 +35,6 @@ export class InclusionCertificate {
 
     const siblings: DataHash[] = [];
     const bitmap = new Uint8Array(InclusionCertificate.BITMAP_SIZE);
-    const keyPath = BitString.fromBytesReversedLSB(key).toBigInt();
 
     while (node != null) {
       if (node instanceof FinalizedLeafBranch) {
@@ -47,12 +45,12 @@ export class InclusionCertificate {
         return new InclusionCertificate(bitmap, siblings);
       }
 
-      const isRight: bigint = (keyPath >> BigInt(node.depth)) & 1n;
+      const isRight = getBitAtDepth(key, node.depth);
 
       const sibling = isRight ? node.left : node.right;
 
       if (sibling != null) {
-        bitmap[Math.floor(node.depth / 8)] |= 1 << (node.depth % 8);
+        bitmap[node.depth >> 3] |= 0x80 >> (node.depth & 7);
         siblings.push(sibling.hash);
       }
 
@@ -121,12 +119,9 @@ export class InclusionCertificate {
       .update(value)
       .digest();
 
-    const keyPath = BitString.fromBytesReversedLSB(key).toBigInt();
-    const bitmapPath = BitString.fromBytesReversedLSB(this.bitmap).toBigInt();
-
     let position = this.siblings.length;
     for (let depth = InclusionCertificate.MAX_DEPTH; depth >= 0; depth--) {
-      if (!((bitmapPath >> BigInt(depth)) & 1n)) continue;
+      if (!getBitAtDepth(this.bitmap, depth)) continue;
 
       position -= 1;
       if (position < 0) {
@@ -136,7 +131,7 @@ export class InclusionCertificate {
       const sibling = this.siblings[position];
 
       let left: Uint8Array, right: Uint8Array;
-      if ((keyPath >> BigInt(depth)) & 1n) {
+      if (getBitAtDepth(key, depth)) {
         left = sibling.data;
         right = hash.data;
       } else {
@@ -146,7 +141,7 @@ export class InclusionCertificate {
 
       hash = await new DataHasher(HashAlgorithm.SHA256)
         .update(new Uint8Array([0x01, depth]))
-        .update(pathToRegion(keyPath, depth))
+        .update(regionFromKey(key, depth))
         .update(left)
         .update(right)
         .digest();

src/payment/SplitAllocationProof.ts

@@ -7,8 +7,8 @@ import { CborSerializer } from '../serialization/cbor/CborSerializer.js';
 import { FinalizedBranch } from '../smt/radixsum/FinalizedBranch.js';
 import { FinalizedLeafBranch } from '../smt/radixsum/FinalizedLeafBranch.js';
 import { SparseMerkleSumTreeRootNode } from '../smt/radixsum/SparseMerkleSumTreeRootNode.js';
+import { getBitAtDepth } from '../smt/SparseMerkleTreePathUtils.js';
 import { BigintConverter } from '../util/BigintConverter.js';
-import { BitString } from '../util/BitString.js';
 import { HexConverter } from '../util/HexConverter.js';
 import { dedent } from '../util/StringUtils.js';
 import { areUint8ArraysEqual } from '../util/TypedArrayUtils.js';
@@ -54,7 +54,6 @@ export class SplitAllocationProof {
    * @throws {Error} If the key is not present in the tree.
    */
   public static create(root: SparseMerkleSumTreeRootNode, key: Uint8Array): SplitAllocationProof {
-    const keyPath = BitString.fromBytesReversedLSB(key).toBigInt();
     const siblings: ISibling[] = [];
 
     let node: FinalizedBranch | SparseMerkleSumTreeRootNode | null = root;
@@ -68,7 +67,7 @@ export class SplitAllocationProof {
         return new SplitAllocationProof(siblings);
       }
 
-      const isRight: number = Number((keyPath >> BigInt(node.depth)) & 1n);
+      const isRight = getBitAtDepth(key, node.depth);
       const sibling = isRight ? node.left : node.right;
       if (sibling != null) {
         siblings.push({ depth: node.depth, hash: sibling.hash, sum: sibling.value });
@@ -148,7 +147,8 @@ export class SplitAllocationProof {
       throw new Error('Data must be 32 bytes long.');
     }
 
-    const keyPath = BitString.fromBytesReversedLSB(key).toBigInt();
+    key = new Uint8Array(key);
+    data = new Uint8Array(data);
 
     let hash = await new DataHasher(HashAlgorithm.SHA256)
       .update(new Uint8Array([0x10]))
@@ -164,7 +164,7 @@ export class SplitAllocationProof {
         throw new Error('Reconstructed sum overflows 256 bits.');
       }
 
-      const isRight = Number((keyPath >> BigInt(sibling.depth)) & 1n);
+      const isRight = getBitAtDepth(key, sibling.depth);
       const left = isRight ? { hash: sibling.hash, value: sibling.sum } : { hash, value: sum };
       const right = isRight ? { hash, value: sum } : { hash: sibling.hash, value: sibling.sum };
 

src/smt/LeafExistsError.ts

@@ -0,0 +1,10 @@
+/**
+ * Thrown when a sparse Merkle tree insertion targets a key that is already
+ * present in the tree.
+ */
+export class LeafExistsError extends Error {
+  public constructor() {
+    super('Leaf already exists.');
+    this.name = 'LeafExistsError';
+  }
+}

src/smt/LeafInBranchError.ts

@@ -1,43 +1,71 @@
-type CommonPath = { length: number; path: bigint };
-
-export function calculateCommonPath(path1: bigint, path2: bigint): CommonPath {
-  let path = 1n;
-  let mask = 1n;
-  let length = 0;
+/**
+ * Length of the common big-endian bit prefix shared by keys `a` and `b`, capped at `maxDepth`
+ * (depth 0 is the most significant bit of byte 0). Used to find where a new key bifurcates from an
+ * existing branch: pass `256` for a leaf (compare the whole key) or the node's depth for an interior
+ * branch (its stored region is only meaningful up to that depth).
+ */
+export function commonPrefixLength(a: Uint8Array, b: Uint8Array, maxDepth: number): number {
+  const fullBytes = maxDepth >> 3;
+  for (let i = 0; i < fullBytes; i++) {
+    if (a[i] !== b[i]) {
+      return (i << 3) + Math.clz32(a[i] ^ b[i]) - 24;
+    }
+  }
 
-  while ((path1 & mask) === (path2 & mask) && path < path1 && path < path2) {
-    mask <<= 1n;
-    length += 1;
-    path = mask | ((mask - 1n) & path1);
+  const remainderBits = maxDepth & 7;
+  if (remainderBits > 0) {
+    const diff = (a[fullBytes] ^ b[fullBytes]) & (0xff << (8 - remainderBits));
+    if (diff !== 0) {
+      return (fullBytes << 3) + Math.clz32(diff) - 24;
+    }
   }
 
-  return { length, path };
+  return maxDepth;
 }
 
 /**
- * Region committed by an interior node: the `depth`-bit common prefix of all leaves in the node's
- * sub-tree. The `i`th lowest bit of path will be the `i mod 8`th lowest bit in the `i div 8`th
- * byte of the returned 32-byte array (so the packing is little-endian); the remaining bits of the array
- * are set to zero.
+ * The key's first `depth` bits, with the remaining bits of the 32-byte array zeroed.
  */
-export function pathToRegion(path: bigint, depth: number): Uint8Array {
+export function regionFromKey(key: Uint8Array, depth: number): Uint8Array {
   const region = new Uint8Array(32);
-  const fullBytes = Math.floor(depth / 8);
-  const remainderBits = depth % 8;
-
-  let bits = path;
-  for (let j = 0; j < fullBytes; j++, bits >>= 8n) {
-    region[j] = Number(bits & 0xffn);
-  }
+  const fullBytes = depth >> 3;
+  const remainderBits = depth & 7;
+  region.set(key.subarray(0, fullBytes));
   if (remainderBits > 0) {
-    region[fullBytes] = Number(bits & 0xffn) & ((1 << remainderBits) - 1);
+    region[fullBytes] = key[fullBytes] & ((0xff << (8 - remainderBits)) & 0xff);
   }
   return region;
 }
 
+/**
+ * Big-endian bit of `data` at the given depth per the Yellowpaper: depth 0 is the most significant
+ * bit of `data[0]` (`data[0] & 0x80`) and depth 255 is the least significant bit of `data[31]`.
+ */
 export function getBitAtDepth(data: Uint8Array, depth: number): number {
-  depth = Number(depth);
-  const byteIndex = Math.floor(depth / 8);
-  const bitInByte = depth % 8;
-  return (data[byteIndex] >> bitInByte) & 1;
+  if (!Number.isInteger(depth) || depth < 0 || depth >= data.length * 8) {
+    throw new Error(`Depth ${depth} is out of bounds for a ${data.length}-byte value.`);
+  }
+  const byteIndex = depth >> 3;
+  const bitInByte = depth & 7;
+  return (data[byteIndex] >> (7 - bitInByte)) & 1;
+}
+
+/**
+ * Render the first `length` big-endian bits of `data` as a `'0'`/`'1'` string: whole bytes first,
+ * then the high bits of the partial byte.
+ */
+export function bitsToString(data: Uint8Array, length: number): string {
+  if (!Number.isInteger(length) || length < 0 || length > data.length * 8) {
+    throw new Error(`Length ${length} is out of bounds for a ${data.length}-byte value.`);
+  }
+  const fullBytes = length >> 3;
+  const remainderBits = length & 7;
+  let bits = '';
+  for (let i = 0; i < fullBytes; i++) {
+    bits += data[i].toString(2).padStart(8, '0');
+  }
+  if (remainderBits > 0) {
+    bits += data[fullBytes].toString(2).padStart(8, '0').slice(0, remainderBits);
+  }
+  return bits;
 }

src/smt/LeafOutOfBoundsError.ts

@@ -4,17 +4,21 @@ import { IDataHasher } from '../../crypto/hash/IDataHasher.js';
 import { IDataHasherFactory } from '../../crypto/hash/IDataHasherFactory.js';
 import { HexConverter } from '../../util/HexConverter.js';
 import { dedent } from '../../util/StringUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Finalized leaf in a radix sparse Merkle tree.
  */
 export class FinalizedLeafBranch {
-  public constructor(
-    public readonly path: bigint,
+  public readonly depth: number;
+
+  private constructor(
     private readonly _key: Uint8Array,
     private readonly _data: Uint8Array,
     public readonly hash: DataHash,
-  ) {}
+  ) {
+    this.depth = _key.length * 8;
+  }
 
   /**
    * @returns {Uint8Array} Copy of the leaf data bytes.
@@ -30,6 +34,15 @@ export class FinalizedLeafBranch {
     return this._key.slice();
   }
 
+  /**
+   * Routing key: the leaf's own key, read bit-by-bit during tree construction.
+   *
+   * @returns {Uint8Array} Copy of the routing key.
+   */
+  public get path(): Uint8Array {
+    return this._key.slice();
+  }
+
   /**
    * Hash a {@link PendingLeafBranch} into a finalized leaf.
    *
@@ -51,7 +64,18 @@ export class FinalizedLeafBranch {
       .update(data)
       .digest();
 
-    return new FinalizedLeafBranch(leaf.path, key, data, hash);
+    return new FinalizedLeafBranch(key, data, hash);
+  }
+
+  /**
+   * Depth at which `key` diverges from this leaf's key, capped at the leaf's own depth (a full match
+   * returns `depth`, signalling a duplicate key).
+   *
+   * @param {Uint8Array} key Key being inserted.
+   * @returns {number} Common-prefix depth.
+   */
+  public calculateSplitDepth(key: Uint8Array): number {
+    return commonPrefixLength(key, this._key, this.depth);
   }
 
   /**
@@ -66,7 +90,7 @@ export class FinalizedLeafBranch {
    */
   public toString(): string {
     return dedent`
-      FinalizedLeaf[${this.path.toString(2)}]
+      FinalizedLeaf[${bitsToString(this._key, this.depth)}]
         Key: ${HexConverter.encode(this._key)}
         Data: ${HexConverter.encode(this._data)}
         Hash: ${this.hash.toString()}`;