trie: improve node typings and class architecture (#3708)

* trie: add RawNode types and replace EmbeddedNode type

* trie: implement new types in BranchNode class

* trie: improve extension and leaf node class architecture

* trie: small improvements

* trie: improve raw nodes comments

* trie: improve comments for raw nodes

Author: gabrocheleau

packages/trie/src/node/branch.ts

@@ -1,9 +1,9 @@
 import { RLP } from '@ethereumjs/rlp'
 
-import type { EmbeddedNode } from '../types.js'
+import type { BranchNodeBranchValue, NodeReferenceOrRawNode } from '../types.js'
 
 export class BranchNode {
-  _branches: (EmbeddedNode | null)[]
+  _branches: BranchNodeBranchValue[]
   _value: Uint8Array | null
 
   constructor() {
@@ -26,19 +26,19 @@ export class BranchNode {
     return this._value && this._value.length > 0 ? this._value : null
   }
 
-  setBranch(i: number, v: EmbeddedNode | null) {
+  setBranch(i: number, v: BranchNodeBranchValue) {
     this._branches[i] = v
   }
 
-  raw(): (EmbeddedNode | null)[] {
+  raw(): BranchNodeBranchValue[] {
     return [...this._branches, this._value]
   }
 
   serialize(): Uint8Array {
-    return RLP.encode(this.raw() as Uint8Array[])
+    return RLP.encode(this.raw())
   }
 
-  getBranch(i: number) {
+  getBranch(i: number): BranchNodeBranchValue {
     const b = this._branches[i]
     if (b !== null && b.length > 0) {
       return b
@@ -47,8 +47,8 @@ export class BranchNode {
     }
   }
 
-  getChildren(): [number, EmbeddedNode][] {
-    const children: [number, EmbeddedNode][] = []
+  getChildren(): [number, NodeReferenceOrRawNode][] {
+    const children: [number, NodeReferenceOrRawNode][] = []
     for (let i = 0; i < 16; i++) {
       const b = this._branches[i]
       if (b !== null && b.length > 0) {

packages/trie/src/node/extension.ts

@@ -1,15 +1,13 @@
-import { addHexPrefix } from '../util/hex.js'
+import { ExtensionOrLeafNodeBase } from './extensionOrLeafNodeBase.js'
 
-import { Node } from './node.js'
+import type { Nibbles, RawExtensionNode } from '../types.js'
 
-import type { Nibbles } from '../types.js'
-
-export class ExtensionNode extends Node {
+export class ExtensionNode extends ExtensionOrLeafNodeBase {
   constructor(nibbles: Nibbles, value: Uint8Array) {
     super(nibbles, value, false)
   }
 
-  static encodeKey(key: Nibbles): Nibbles {
-    return addHexPrefix(key, false)
+  raw(): RawExtensionNode {
+    return super.raw()
   }
 }

packages/trie/src/node/node.ts renamed to packages/trie/src/node/extensionOrLeafNodeBase.ts

@@ -3,23 +3,27 @@ import { RLP } from '@ethereumjs/rlp'
 import { addHexPrefix, removeHexPrefix } from '../util/hex.js'
 import { nibblesTypeToPackedBytes } from '../util/nibbles.js'
 
-import type { Nibbles } from '../types.js'
+import type { Nibbles, RawExtensionNode, RawLeafNode } from '../types.js'
 
-export class Node {
+export abstract class ExtensionOrLeafNodeBase {
   _nibbles: Nibbles
   _value: Uint8Array
-  _terminator: boolean
+  _isLeaf: boolean
 
-  constructor(nibbles: Nibbles, value: Uint8Array, terminator: boolean) {
+  constructor(nibbles: Nibbles, value: Uint8Array, isLeaf: boolean) {
     this._nibbles = nibbles
     this._value = value
-    this._terminator = terminator
+    this._isLeaf = isLeaf
   }
 
   static decodeKey(key: Nibbles): Nibbles {
     return removeHexPrefix(key)
   }
 
+  encodedKey(): Nibbles {
+    return addHexPrefix(this._nibbles.slice(0), this._isLeaf)
+  }
+
   key(k?: Nibbles): Nibbles {
     if (k !== undefined) {
       this._nibbles = k
@@ -40,11 +44,7 @@ export class Node {
     return this._value
   }
 
-  encodedKey(): Nibbles {
-    return addHexPrefix(this._nibbles.slice(0), this._terminator)
-  }
-
-  raw(): [Uint8Array, Uint8Array] {
+  raw(): RawExtensionNode | RawLeafNode {
     return [nibblesTypeToPackedBytes(this.encodedKey()), this._value]
   }
 

packages/trie/src/node/leaf.ts

@@ -1,15 +1,13 @@
-import { addHexPrefix } from '../util/hex.js'
+import { ExtensionOrLeafNodeBase } from './extensionOrLeafNodeBase.js'
 
-import { Node } from './node.js'
+import type { Nibbles, RawLeafNode } from '../types.js'
 
-import type { Nibbles } from '../types.js'
-
-export class LeafNode extends Node {
+export class LeafNode extends ExtensionOrLeafNodeBase {
   constructor(nibbles: Nibbles, value: Uint8Array) {
     super(nibbles, value, true)
   }
 
-  static encodeKey(key: Nibbles): Nibbles {
-    return addHexPrefix(key, true)
+  raw(): RawLeafNode {
+    return super.raw()
   }
 }

packages/trie/src/trie.ts

@@ -34,9 +34,10 @@ import { bytesToNibbles, matchingNibbleLength, nibblesTypeToPackedBytes } from '
 import { WalkController } from './util/walkController.js'
 
 import type {
-  EmbeddedNode,
+  BranchNodeBranchValue,
   FoundNodeFunction,
   Nibbles,
+  NodeReferenceOrRawNode,
   Path,
   TrieNode,
   TrieOpts,
@@ -357,7 +358,7 @@ export class Trie {
               debugString +=
                 branchNode instanceof Uint8Array
                   ? `NodeHash: ${bytesToHex(branchNode)}`
-                  : `Raw_Node: ${branchNode!.toString()}`
+                  : `Raw_Node: ${branchNode.toString()}`
             }
 
             this.debug(debugString, ['find_path', 'branch_node'])
@@ -564,8 +565,8 @@ export class Trie {
       }
 
       if (
-        matchingNibbleLength(lastNode.key(), key.slice(l)) === lastNode.key().length &&
-        keyRemainder.length === 0
+        keyRemainder.length === 0 &&
+        matchingNibbleLength(lastNode.key(), key.slice(l)) === lastNode.key().length
       ) {
         matchLeaf = true
       }
@@ -574,7 +575,7 @@ export class Trie {
     if (matchLeaf) {
       // just updating a found value
       lastNode.value(value)
-      stack.push(lastNode as TrieNode)
+      stack.push(lastNode)
     } else if (lastNode instanceof BranchNode) {
       stack.push(lastNode)
       if (keyRemainder.length !== 0) {
@@ -609,8 +610,8 @@ export class Trie {
         if (lastKey.length !== 0 || lastNode instanceof LeafNode) {
           // shrinking extension or leaf
           lastNode.key(lastKey)
-          const formattedNode = this._formatNode(lastNode, false, toSave)
-          newBranchNode.setBranch(branchKey, formattedNode as EmbeddedNode)
+          const formattedNode = this._formatNode(lastNode, false, toSave) as NodeReferenceOrRawNode
+          newBranchNode.setBranch(branchKey, formattedNode)
         } else {
           // remove extension or attaching
           this._formatNode(lastNode, false, toSave, true)
@@ -703,7 +704,7 @@ export class Trie {
 
     let key = bytesToNibbles(k)
 
-    if (!parentNode) {
+    if (parentNode === undefined) {
       // the root here has to be a leaf.
       this.root(this.EMPTY_TRIE_ROOT)
       return
@@ -728,7 +729,7 @@ export class Trie {
 
     // nodes on the branch
     // count the number of nodes on the branch
-    const branchNodes: [number, EmbeddedNode][] = lastNode.getChildren()
+    const branchNodes: [number, NodeReferenceOrRawNode][] = lastNode.getChildren()
 
     // if there is only one branch node left, collapse the branch node
     if (branchNodes.length === 1) {
@@ -753,10 +754,8 @@ export class Trie {
 
       // look up node
       const foundNode = await this.lookupNode(branchNode)
-      // if (foundNode) {
       key = processBranchNode(key, branchNodeKey, foundNode, parentNode as TrieNode, stack)
       await this.saveStack(key, stack, opStack)
-      // }
     } else {
       // simple removing a leaf and recalculation the stack
       if (parentNode) {
@@ -819,7 +818,7 @@ export class Trie {
     topLevel: boolean,
     opStack: BatchDBOp[],
     remove: boolean = false,
-  ): Uint8Array | (EmbeddedNode | null)[] {
+  ): Uint8Array | NodeReferenceOrRawNode | BranchNodeBranchValue[] {
     const encoded = node.serialize()
 
     if (encoded.length >= 32 || topLevel) {

packages/trie/src/types.ts

@@ -8,9 +8,18 @@ export type TrieNode = BranchNode | ExtensionNode | LeafNode
 
 export type Nibbles = number[]
 
+// A raw node refers to the non-serialized, array form of the node
+// A raw extension node is a 2-item node, where the first item is the encoded path to the next node, and the second item is the reference to the next node
+// A raw leaf node is a 2-item node, where the first item is the remaining path to the leaf node, and the second item is the value
+// To learn more: https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie/#optimization
+export type RawExtensionNode = [Uint8Array, Uint8Array]
+export type RawLeafNode = [Uint8Array, Uint8Array]
+
 // Branch and extension nodes might store
-// hash to next node, or embed it if its len < 32
-export type EmbeddedNode = Uint8Array | Uint8Array[]
+// hash to next node, or a raw node if its length < 32
+export type NodeReferenceOrRawNode = Uint8Array | RawExtensionNode | RawLeafNode
+
+export type BranchNodeBranchValue = NodeReferenceOrRawNode | null
 
 export type Proof = Uint8Array[]
 

packages/trie/src/util/nibbles.ts

@@ -81,13 +81,3 @@ export function matchingNibbleLength(nib1: Nibbles, nib2: Nibbles): number {
   }
   return i
 }
-
-/**
- * Compare two nibble array keys.
- * @param keyA
- * @param keyB
- */
-export function doKeysMatch(keyA: Nibbles, keyB: Nibbles): boolean {
-  const length = matchingNibbleLength(keyA, keyB)
-  return length === keyA.length && length === keyB.length
-}