#134 Update SMT based on pull request comments

Author: martti007

src/smt/LeafOutOfBoundsError.ts

@@ -42,7 +42,6 @@ export function regionFromKey(key: Uint8Array, depth: number): Uint8Array {
  * 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);
   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.`);
   }

src/smt/SparseMerkleTreePathUtils.ts

@@ -4,19 +4,20 @@ 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 } from '../SparseMerkleTreePathUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Finalized leaf in a radix sparse Merkle tree.
  */
 export class FinalizedLeafBranch {
+  public readonly depth: number;
+
   private constructor(
     private readonly _key: Uint8Array,
     private readonly _data: Uint8Array,
     public readonly hash: DataHash,
   ) {
-    this._key = new Uint8Array(_key);
-    this._data = new Uint8Array(_data);
+    this.depth = _key.length * 8;
   }
 
   /**
@@ -66,6 +67,17 @@ export class FinalizedLeafBranch {
     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);
+  }
+
   /**
    * @returns {Promise<FinalizedLeafBranch>} This branch (already finalized).
    */
@@ -78,7 +90,7 @@ export class FinalizedLeafBranch {
    */
   public toString(): string {
     return dedent`
-      FinalizedLeaf[${bitsToString(this._key, this._key.length * 8)}]
+      FinalizedLeaf[${bitsToString(this._key, this.depth)}]
         Key: ${HexConverter.encode(this._key)}
         Data: ${HexConverter.encode(this._data)}
         Hash: ${this.hash.toString()}`;

src/smt/radix/FinalizedLeafBranch.ts

@@ -1,10 +1,11 @@
 import { FinalizedBranch } from './FinalizedBranch.js';
+import { PendingBranch } from './PendingBranch.js';
 import { PendingNodeBranch } from './PendingNodeBranch.js';
 import { DataHash } from '../../crypto/hash/DataHash.js';
 import { IDataHasher } from '../../crypto/hash/IDataHasher.js';
 import { IDataHasherFactory } from '../../crypto/hash/IDataHasherFactory.js';
 import { dedent } from '../../util/StringUtils.js';
-import { bitsToString } from '../SparseMerkleTreePathUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Finalized interior node in a radix sparse Merkle tree.
@@ -48,6 +49,16 @@ export class FinalizedNodeBranch {
     return new FinalizedNodeBranch(node.path, node.depth, left, right, hash);
   }
 
+  /**
+   * Depth at which `key` diverges from this node's committed region, capped at the node's own depth.
+   *
+   * @param {Uint8Array} key Key being inserted.
+   * @returns {number} Common-prefix depth.
+   */
+  public calculateSplitDepth(key: Uint8Array): number {
+    return commonPrefixLength(key, this._path, this.depth);
+  }
+
   /**
    * @returns {Promise<FinalizedNodeBranch>} This branch (already finalized).
    */
@@ -66,4 +77,24 @@ export class FinalizedNodeBranch {
         Left: ${this.left.toString()}
         Right: ${this.right.toString()}`;
   }
+
+  /**
+   * Derive a pending node with `left` as its left child, reusing this node's committed region.
+   *
+   * @param {PendingBranch} left Replacement left child.
+   * @returns {PendingNodeBranch} New pending node.
+   */
+  public withLeftBranch(left: PendingBranch): PendingNodeBranch {
+    return PendingNodeBranch.create(this._path, this.depth, left, this.right);
+  }
+
+  /**
+   * Derive a pending node with `right` as its right child, reusing this node's committed region.
+   *
+   * @param {PendingBranch} right Replacement right child.
+   * @returns {PendingNodeBranch} New pending node.
+   */
+  public withRightBranch(right: PendingBranch): PendingNodeBranch {
+    return PendingNodeBranch.create(this._path, this.depth, this.left, right);
+  }
 }

src/smt/radix/FinalizedNodeBranch.ts

@@ -3,18 +3,19 @@ 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 } from '../SparseMerkleTreePathUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Pending leaf in a radix sparse Merkle tree, awaiting hashing.
  */
 export class PendingLeafBranch {
+  public readonly depth: number;
+
   public constructor(
     private readonly _key: Uint8Array,
     private readonly _data: Uint8Array,
   ) {
-    this._key = new Uint8Array(_key);
-    this._data = new Uint8Array(_data);
+    this.depth = _key.length * 8;
   }
 
   /**
@@ -40,6 +41,17 @@ export class PendingLeafBranch {
     return this._key.slice();
   }
 
+  /**
+   * 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);
+  }
+
   /**
    * Hash this leaf to produce a finalized branch.
    *
@@ -55,7 +67,7 @@ export class PendingLeafBranch {
    */
   public toString(): string {
     return dedent`
-      PendingLeaf[${bitsToString(this._key, this._key.length * 8)}]
+      PendingLeaf[${bitsToString(this._key, this.depth)}]
         Key: ${HexConverter.encode(this._key)}
         Data: ${HexConverter.encode(this._data)}`;
   }

src/smt/radix/PendingLeafBranch.ts

@@ -3,20 +3,18 @@ import { PendingBranch } from './PendingBranch.js';
 import { IDataHasher } from '../../crypto/hash/IDataHasher.js';
 import { IDataHasherFactory } from '../../crypto/hash/IDataHasherFactory.js';
 import { dedent } from '../../util/StringUtils.js';
-import { bitsToString } from '../SparseMerkleTreePathUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Pending interior node in a radix sparse Merkle tree, awaiting hashing.
  */
 export class PendingNodeBranch {
-  public constructor(
+  private constructor(
     private readonly _path: Uint8Array,
     public readonly depth: number,
     public readonly left: PendingBranch,
     public readonly right: PendingBranch,
-  ) {
-    this._path = new Uint8Array(_path);
-  }
+  ) {}
 
   /**
    * @returns {Uint8Array} Copy of the node's committed region (its `depth`-bit prefix, suffix zeroed).
@@ -25,6 +23,29 @@ export class PendingNodeBranch {
     return this._path.slice();
   }
 
+  /**
+   * Create a pending node, defensively copying the committed region.
+   *
+   * @param {Uint8Array} path Committed region (its `depth`-bit prefix, suffix zeroed).
+   * @param {number} depth Node depth.
+   * @param {PendingBranch} left Left child.
+   * @param {PendingBranch} right Right child.
+   * @returns {PendingNodeBranch} New pending node.
+   */
+  public static create(path: Uint8Array, depth: number, left: PendingBranch, right: PendingBranch): PendingNodeBranch {
+    return new PendingNodeBranch(new Uint8Array(path), depth, left, right);
+  }
+
+  /**
+   * Depth at which `key` diverges from this node's committed region, capped at the node's own depth.
+   *
+   * @param {Uint8Array} key Key being inserted.
+   * @returns {number} Common-prefix depth.
+   */
+  public calculateSplitDepth(key: Uint8Array): number {
+    return commonPrefixLength(key, this._path, this.depth);
+  }
+
   /**
    * Hash this node (after finalizing children).
    *
@@ -44,4 +65,24 @@ export class PendingNodeBranch {
         Left: ${this.left.toString()}
         Right: ${this.right.toString()}`;
   }
+
+  /**
+   * Derive a node with `left` as its left child, reusing this node's committed region without copying it.
+   *
+   * @param {PendingBranch} left Replacement left child.
+   * @returns {PendingNodeBranch} New pending node.
+   */
+  public withLeftBranch(left: PendingBranch): PendingNodeBranch {
+    return new PendingNodeBranch(this._path, this.depth, left, this.right);
+  }
+
+  /**
+   * Derive a node with `right` as its right child, reusing this node's committed region without copying it.
+   *
+   * @param {PendingBranch} right Replacement right child.
+   * @returns {PendingNodeBranch} New pending node.
+   */
+  public withRightBranch(right: PendingBranch): PendingNodeBranch {
+    return new PendingNodeBranch(this._path, this.depth, this.left, right);
+  }
 }

src/smt/radix/PendingNodeBranch.ts

@@ -4,7 +4,7 @@ import { LeafExistsError } from '../LeafExistsError.js';
 import { PendingBranch } from './PendingBranch.js';
 import { PendingLeafBranch } from './PendingLeafBranch.js';
 import { PendingNodeBranch } from './PendingNodeBranch.js';
-import { commonPrefixLength, getBitAtDepth, regionFromKey } from '../SparseMerkleTreePathUtils.js';
+import { getBitAtDepth, regionFromKey } from '../SparseMerkleTreePathUtils.js';
 import { SparseMerkleTreeRootNode } from './SparseMerkleTreeRootNode.js';
 import { IDataHasher } from '../../crypto/hash/IDataHasher.js';
 import { IDataHasherFactory } from '../../crypto/hash/IDataHasherFactory.js';
@@ -73,29 +73,29 @@ export class SparseMerkleTree {
 
   private buildTree(branch: PendingBranch, key: Uint8Array, data: Uint8Array): PendingBranch {
     if (branch instanceof PendingLeafBranch || branch instanceof FinalizedLeafBranch) {
-      const depth = commonPrefixLength(key, branch.path, 256);
-      if (depth === 256) {
+      const depth = branch.calculateSplitDepth(key);
+      if (depth === branch.depth) {
         throw new LeafExistsError();
       }
 
       // If a leaf must be split from the middle
       const isRight = getBitAtDepth(key, depth);
       const newBranch = new PendingLeafBranch(key, data);
-      return new PendingNodeBranch(
+      return PendingNodeBranch.create(
         regionFromKey(key, depth),
         depth,
         isRight ? branch : newBranch,
         isRight ? newBranch : branch,
       );
     }
 
-    const depth = commonPrefixLength(key, branch.path, branch.depth);
+    const depth = branch.calculateSplitDepth(key);
     const isRight = getBitAtDepth(key, depth);
 
     // If node branch is split in the middle
     if (depth < branch.depth) {
       const newBranch = new PendingLeafBranch(key, data);
-      return new PendingNodeBranch(
+      return PendingNodeBranch.create(
         regionFromKey(key, depth),
         depth,
         isRight ? branch : newBranch,
@@ -105,9 +105,9 @@ export class SparseMerkleTree {
 
     // Key belongs inside current node
     if (isRight) {
-      return new PendingNodeBranch(branch.path, branch.depth, branch.left, this.buildTree(branch.right, key, data));
+      return branch.withRightBranch(this.buildTree(branch.right, key, data));
     }
 
-    return new PendingNodeBranch(branch.path, branch.depth, this.buildTree(branch.left, key, data), branch.right);
+    return branch.withLeftBranch(this.buildTree(branch.left, key, data));
   }
 }

src/smt/radix/SparseMerkleTree.ts

@@ -41,7 +41,7 @@ export class SparseMerkleTreeRootNode {
     }
 
     if (left != null && right != null) {
-      const node = await new PendingNodeBranch(new Uint8Array(32), 0, left, right).finalize(factory);
+      const node = await PendingNodeBranch.create(new Uint8Array(32), 0, left, right).finalize(factory);
       return new SparseMerkleTreeRootNode(node.left, node.right, node.hash);
     }
 

src/smt/radix/SparseMerkleTreeRootNode.ts

@@ -5,22 +5,23 @@ import { IDataHasherFactory } from '../../crypto/hash/IDataHasherFactory.js';
 import { BigintConverter } from '../../util/BigintConverter.js';
 import { HexConverter } from '../../util/HexConverter.js';
 import { dedent } from '../../util/StringUtils.js';
-import { bitsToString } from '../SparseMerkleTreePathUtils.js';
+import { bitsToString, commonPrefixLength } from '../SparseMerkleTreePathUtils.js';
 
 /**
  * Finalized leaf in a radix sparse Merkle sum tree. The leaf hash is
  * `SHA-256(0x10 || key || data || u256(value))`, where `u256` is the 32-byte
  * big-endian encoding of the leaf amount.
  */
 export class FinalizedLeafBranch {
+  public readonly depth: number;
+
   private constructor(
     private readonly _key: Uint8Array,
     private readonly _data: Uint8Array,
     public readonly value: bigint,
     public readonly hash: DataHash,
   ) {
-    this._key = new Uint8Array(_key);
-    this._data = new Uint8Array(_data);
+    this.depth = _key.length * 8;
   }
 
   /**
@@ -73,6 +74,17 @@ export class FinalizedLeafBranch {
     return new FinalizedLeafBranch(key, data, leaf.value, 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);
+  }
+
   /**
    * @returns {Promise<FinalizedLeafBranch>} This branch (already finalized).
    */
@@ -85,7 +97,7 @@ export class FinalizedLeafBranch {
    */
   public toString(): string {
     return dedent`
-      FinalizedLeaf[${bitsToString(this._key, this._key.length * 8)}]
+      FinalizedLeaf[${bitsToString(this._key, this.depth)}]
         Key: ${HexConverter.encode(this._key)}
         Data: ${HexConverter.encode(this._data)}
         Value: ${this.value}