Implement trie.put (#3473)

* first pass at put implementation [no ci]

* Fix logging [no ci]

* Revert optimization

* Comment and test cleanup [no ci]

* fix child index selection [no ci]

* Finish tests

* Remove .only

* Add helper for `put`

* Fix node copy step

* lint

* verkle: add and implement typeguards

* verkle: add typeguard to tests

* Address feedback

---------

Co-authored-by: Gabriel Rocheleau <[email protected]>

Author: acolytec3

packages/verkle/src/node/util.ts

@@ -57,3 +57,10 @@ export const createCValues = (values: Uint8Array[], deletedValues = new Array(12
   }
   return expandedValues
 }
+export function isLeafNode(node: VerkleNode): node is LeafNode {
+  return node.type === VerkleNodeType.Leaf
+}
+
+export function isInternalNode(node: VerkleNode): node is InternalNode {
+  return node.type === VerkleNodeType.Internal
+}

packages/verkle/src/util/bytes.ts

@@ -10,30 +10,13 @@ export function matchingBytesLength(bytes1: Uint8Array, bytes2: Uint8Array): num
   let count = 0
   const minLength = Math.min(bytes1.length, bytes2.length)
 
-  // Unroll the loop for better performance
-  for (let i = 0; i < minLength - 3; i += 4) {
-    // Compare 4 bytes at a time
-    if (
-      bytes1[i] === bytes2[i] &&
-      bytes1[i + 1] === bytes2[i + 1] &&
-      bytes1[i + 2] === bytes2[i + 2] &&
-      bytes1[i + 3] === bytes2[i + 3]
-    ) {
-      count += 4
-    } else {
-      // Break early if a mismatch is found
-      break
-    }
-  }
-
-  // Handle any remaining elements
-  for (let i = minLength - (minLength % 4); i < minLength; i++) {
+  for (let i = 0; i < minLength; i++) {
     if (bytes1[i] === bytes2[i]) {
       count++
     } else {
+      // Break early if a mismatch is found
       break
     }
   }
-
   return count
 }

packages/verkle/src/verkleTree.ts

@@ -5,6 +5,7 @@ import {
   ValueEncoding,
   bytesToHex,
   equalsBytes,
+  intToHex,
   zeros,
 } from '@ethereumjs/util'
 import debug from 'debug'
@@ -14,7 +15,7 @@ import { CheckpointDB } from './db/checkpoint.js'
 import { InternalNode } from './node/internalNode.js'
 import { LeafNode } from './node/leafNode.js'
 import { type VerkleNode } from './node/types.js'
-import { decodeNode } from './node/util.js'
+import { decodeNode, isLeafNode } from './node/util.js'
 import {
   type Proof,
   ROOT_DB_KEY,
@@ -133,7 +134,9 @@ export class VerkleTree {
       }
     }
 
-    return new VerkleTree(opts)
+    const trie = new VerkleTree(opts)
+    await trie._createRootNode()
+    return trie
   }
 
   database(db?: DB<Uint8Array, Uint8Array>) {
@@ -194,7 +197,6 @@ export class VerkleTree {
     const suffix = key[key.length - 1]
     this.DEBUG && this.debug(`Stem: ${bytesToHex(stem)}; Suffix: ${suffix}`, ['GET'])
     const res = await this.findPath(stem)
-
     if (res.node instanceof LeafNode) {
       // The retrieved leaf node contains an array of 256 possible values.
       // The index of the value we want is at the key's last byte
@@ -213,11 +215,169 @@ export class VerkleTree {
    * @param value - the value to store
    * @returns A Promise that resolves once value is stored.
    */
-  // TODO: Rewrite following logic in verkle.spec.ts "findPath validation" test
-  async put(_key: Uint8Array, _value: Uint8Array): Promise<void> {
-    throw new Error('not implemented')
+  async put(key: Uint8Array, value: Uint8Array): Promise<void> {
+    if (key.length !== 32) throw new Error(`expected key with length 32; got ${key.length}`)
+    const stem = key.slice(0, 31)
+    const suffix = key[key.length - 1]
+    this.DEBUG && this.debug(`Stem: ${bytesToHex(stem)}; Suffix: ${suffix}`, ['PUT'])
+
+    const putStack: [Uint8Array, VerkleNode][] = []
+    // Find path to nearest node
+    const foundPath = await this.findPath(stem)
+
+    // Sanity check - we should at least get the root node back
+    if (foundPath.stack.length === 0) {
+      throw new Error(`Root node not found in trie`)
+    }
+
+    // Step 1) Create or update the leaf node
+    let leafNode: LeafNode
+    // First see if leaf node already exists
+    if (foundPath.node !== null) {
+      // Sanity check to verify we have the right node type
+      if (!isLeafNode(foundPath.node)) {
+        throw new Error(
+          `expected leaf node found at ${bytesToHex(stem)}. Got internal node instead`
+        )
+      }
+      leafNode = foundPath.node
+      // Sanity check to verify we have the right leaf node
+      if (!equalsBytes(leafNode.stem, stem)) {
+        throw new Error(
+          `invalid leaf node found. Expected stem: ${bytesToHex(stem)}; got ${bytesToHex(
+            foundPath.node.stem
+          )}`
+        )
+      }
+    } else {
+      // Leaf node doesn't exist, create a new one
+      leafNode = await LeafNode.create(
+        stem,
+        new Array(256).fill(new Uint8Array(32)),
+        this.verkleCrypto
+      )
+      this.DEBUG && this.debug(`Creating new leaf node at stem: ${bytesToHex(stem)}`, ['PUT'])
+    }
+    // Update value in leaf node and push to putStack
+    leafNode.setValue(suffix, value)
+    this.DEBUG &&
+      this.debug(
+        `Updating value for suffix: ${suffix} at leaf node with stem: ${bytesToHex(stem)}`,
+        ['PUT']
+      )
+    putStack.push([leafNode.hash(), leafNode])
+
+    // `path` is the path to the last node pushed to the `putStack`
+    let lastPath = leafNode.stem
+
+    // Step 2) Determine if a new internal node is needed
+    if (foundPath.stack.length > 1) {
+      // Only insert new internal node if we have more than 1 node in the path
+      // since a single node indicates only the root node is in the path
+      const nearestNodeTuple = foundPath.stack.pop()!
+      const nearestNode = nearestNodeTuple[0]
+      lastPath = nearestNodeTuple[1]
+      const updatedParentTuple = this.updateParent(leafNode, nearestNode, lastPath)
+      putStack.push([updatedParentTuple.node.hash(), updatedParentTuple.node])
+      lastPath = updatedParentTuple.lastPath
+
+      // Step 3) Walk up trie and update child references in parent internal nodes
+      while (foundPath.stack.length > 1) {
+        const [nextNode, nextPath] = foundPath.stack.pop()! as [InternalNode, Uint8Array]
+        // Compute the child index to be updated on `nextNode`
+        const childIndex = lastPath[matchingBytesLength(lastPath, nextPath)]
+        // Update child reference
+        nextNode.setChild(childIndex, {
+          commitment: putStack[putStack.length - 1][1].commitment,
+          path: lastPath,
+        })
+        this.DEBUG &&
+          this.debug(
+            `Updating child reference for node with path: ${bytesToHex(
+              lastPath
+            )} at index ${childIndex} in internal node at path ${bytesToHex(nextPath)}`,
+            ['PUT']
+          )
+        // Hold onto `path` to current node for updating next parent node child index
+        lastPath = nextPath
+        putStack.push([nextNode.hash(), nextNode])
+      }
+    }
+
+    // Step 4) Update root node
+    const rootNode = foundPath.stack.pop()![0] as InternalNode
+    rootNode.setChild(stem[0], {
+      commitment: putStack[putStack.length - 1][1].commitment,
+      path: lastPath,
+    })
+    this.root(this.verkleCrypto.serializeCommitment(rootNode.commitment))
+    this.DEBUG &&
+      this.debug(
+        `Updating child reference for node with path: ${bytesToHex(lastPath)} at index ${
+          lastPath[0]
+        } in root node`,
+        ['PUT']
+      )
+    this.DEBUG && this.debug(`Updating root node hash to ${bytesToHex(this._root)}`, ['PUT'])
+    putStack.push([this._root, rootNode])
+    await this.saveStack(putStack)
   }
 
+  /**
+   * Helper method for updating or creating the parent internal node for a given leaf node
+   * @param leafNode the child leaf node that will be referenced by the new/updated internal node
+   * returned by this method
+   * @param nearestNode the nearest node to the new leaf node
+   * @param pathToNode the path to `nearestNode`
+   * @returns a tuple of the updated parent node and the path to that parent (i.e. the partial stem of the leaf node that leads to the parent)
+   */
+  updateParent(
+    leafNode: LeafNode,
+    nearestNode: VerkleNode,
+    pathToNode: Uint8Array
+  ): { node: InternalNode; lastPath: Uint8Array } {
+    // Compute the portion of leafNode.stem and nearestNode.path that match (i.e. the partial path closest to leafNode.stem)
+    const partialMatchingStemIndex = matchingBytesLength(leafNode.stem, pathToNode)
+    let internalNode: InternalNode
+    if (isLeafNode(nearestNode)) {
+      // We need to create a new internal node and set nearestNode and leafNode as child nodes of it
+      // Create new internal node
+      internalNode = InternalNode.create(this.verkleCrypto)
+      // Set leafNode and nextNode as children of the new internal node
+      internalNode.setChild(leafNode.stem[partialMatchingStemIndex], {
+        commitment: leafNode.commitment,
+        path: leafNode.stem,
+      })
+      internalNode.setChild(nearestNode.stem[partialMatchingStemIndex], {
+        commitment: nearestNode.commitment,
+        path: nearestNode.stem,
+      })
+      // Find the path to the new internal node (the matching portion of the leaf node and next node's stems)
+      pathToNode = leafNode.stem.slice(0, partialMatchingStemIndex)
+      this.DEBUG &&
+        this.debug(`Creating new internal node at path ${bytesToHex(pathToNode)}`, ['PUT'])
+    } else {
+      // Nearest node is an internal node.  We need to update the appropriate child reference
+      // to the new leaf node
+      internalNode = nearestNode
+      internalNode.setChild(leafNode.stem[partialMatchingStemIndex], {
+        commitment: leafNode.commitment,
+        path: leafNode.stem,
+      })
+      this.DEBUG &&
+        this.debug(
+          `Updating child reference for leaf node with stem: ${bytesToHex(
+            leafNode.stem
+          )} at index ${
+            leafNode.stem[partialMatchingStemIndex]
+          } in internal node at path ${bytesToHex(
+            leafNode.stem.slice(0, partialMatchingStemIndex)
+          )}`,
+          ['PUT']
+        )
+    }
+    return { node: internalNode, lastPath: pathToNode }
+  }
   /**
    * Tries to find a path to the node for the given key.
    * It returns a `stack` of nodes to the closest node.
@@ -231,12 +391,13 @@ export class VerkleTree {
       stack: [],
       remaining: key,
     }
-    if (equalsBytes(this.root(), this.EMPTY_TREE_ROOT)) return result
+
+    // TODO: Decide if findPath should return an empty stack if we have an empty trie or a path with just the empty root node
+    // if (equalsBytes(this.root(), this.EMPTY_TREE_ROOT)) return result
 
     // Get root node
     let rawNode = await this._db.get(this.root())
-    if (rawNode === undefined)
-      throw new Error('root node should exist when root not empty tree root')
+    if (rawNode === undefined) throw new Error('root node should exist')
 
     const rootNode = decodeNode(rawNode, this.verkleCrypto) as InternalNode
 
@@ -246,7 +407,7 @@ export class VerkleTree {
 
     // Root node doesn't contain a child node's commitment on the first byte of the path so we're done
     if (equalsBytes(child.commitment, this.verkleCrypto.zeroCommitment)) {
-      this.DEBUG && this.debug(`Partial Path ${key[0]} - found no child.`, ['FIND_PATH'])
+      this.DEBUG && this.debug(`Partial Path ${intToHex(key[0])} - found no child.`, ['FIND_PATH'])
       return result
     }
     let finished = false
@@ -260,7 +421,7 @@ export class VerkleTree {
       // Calculate the index of the last matching byte in the key
       const matchingKeyLength = matchingBytesLength(key, child.path)
       const foundNode = equalsBytes(key, child.path)
-      if (foundNode || child.path.length >= key.length || decodedNode instanceof LeafNode) {
+      if (foundNode || child.path.length >= key.length || isLeafNode(decodedNode)) {
         // If the key and child.path are equal, then we found the node
         // If the child.path is the same length or longer than the key but doesn't match it
         // or the found node is a leaf node, we've found another node where this node should
@@ -282,16 +443,15 @@ export class VerkleTree {
         // We found a different node than the one specified by `key`
         // so the sought node doesn't exist
         result.remaining = key.slice(matchingKeyLength)
+        const pathToNearestNode = isLeafNode(decodedNode) ? decodedNode.stem : child.path
         this.DEBUG &&
           this.debug(
-            `Path ${bytesToHex(
-              key.slice(0, matchingKeyLength)
-            )} - found path to nearest node ${bytesToHex(
+            `Path ${bytesToHex(pathToNearestNode)} - found path to nearest node ${bytesToHex(
               decodedNode.hash()
             )} but target node not found.`,
             ['FIND_PATH']
           )
-        result.stack.push([decodedNode, key.slice(0, matchingKeyLength)])
+        result.stack.push([decodedNode, pathToNearestNode])
         return result
       }
       // Push internal node to path stack

packages/verkle/test/internalNode.spec.ts

@@ -2,7 +2,7 @@ import { type VerkleCrypto, equalsBytes, randomBytes } from '@ethereumjs/util'
 import { loadVerkleCrypto } from 'verkle-cryptography-wasm'
 import { assert, beforeAll, describe, it } from 'vitest'
 
-import { NODE_WIDTH, VerkleNodeType, decodeNode } from '../src/node/index.js'
+import { NODE_WIDTH, VerkleNodeType, decodeNode, isInternalNode } from '../src/node/index.js'
 import { InternalNode } from '../src/node/internalNode.js'
 
 describe('verkle node - internal', () => {
@@ -14,6 +14,7 @@ describe('verkle node - internal', () => {
     const commitment = randomBytes(32)
     const node = new InternalNode({ commitment, verkleCrypto })
 
+    assert.ok(isInternalNode(node), 'typeguard should return true')
     assert.equal(node.type, VerkleNodeType.Internal, 'type should be set')
     assert.ok(equalsBytes(node.commitment, commitment), 'commitment should be set')
 

packages/verkle/test/leafNode.spec.ts

@@ -2,7 +2,7 @@ import { type VerkleCrypto, equalsBytes, randomBytes } from '@ethereumjs/util'
 import { loadVerkleCrypto } from 'verkle-cryptography-wasm'
 import { assert, beforeAll, describe, it } from 'vitest'
 
-import { VerkleNodeType } from '../src/node/index.js'
+import { VerkleNodeType, isLeafNode } from '../src/node/index.js'
 import { LeafNode } from '../src/node/leafNode.js'
 
 describe('verkle node - leaf', () => {
@@ -25,6 +25,7 @@ describe('verkle node - leaf', () => {
       verkleCrypto,
     })
 
+    assert.ok(isLeafNode(node), 'typeguard should return true')
     assert.equal(node.type, VerkleNodeType.Leaf, 'type should be set')
     assert.ok(
       equalsBytes(node.commitment as unknown as Uint8Array, commitment),