Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

trie: improve node typings and class architecture #3708

Merged
merged 7 commits into from
Sep 29, 2024
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions packages/trie/src/node/branch.ts
Original file line number Diff line number Diff line change
@@ -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() {
Expand All @@ -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
Expand All @@ -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) {
Expand Down
12 changes: 5 additions & 7 deletions packages/trie/src/node/extension.ts
Original file line number Diff line number Diff line change
@@ -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()
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
jochem-brouwer marked this conversation as resolved.
Show resolved Hide resolved

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
Expand All @@ -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]
}

Expand Down
12 changes: 5 additions & 7 deletions packages/trie/src/node/leaf.ts
Original file line number Diff line number Diff line change
@@ -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()
}
}
23 changes: 11 additions & 12 deletions packages/trie/src/trie.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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'])
Expand Down Expand Up @@ -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
}
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down
12 changes: 10 additions & 2 deletions packages/trie/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,17 @@ export type TrieNode = BranchNode | ExtensionNode | LeafNode

export type Nibbles = number[]

// [ encodedPath, key ]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, I find this documentation somewhat misleading. The key in this context would be the pointer to whatever hash this ExtensionNode points to. I think this is implied to by "key" but this might be confused with the MPT-keys (so the ones you use as get argument). Not sure what a better term here would be though. referencedHash? Or maybe pointerHash?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've copied what is in the ethereum.org documentation. I agree that it is not immediately obvious what this represents. I will remove these two comments and instead add a link to the docs.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I dunno if I like the updated comment any better. The docs do explain what the extension node is, in that the key is the DB key for the next node.

Can we get rid of the "raw node" terminology and just have extension nodes, which are simply the encoded path to the node and then the DB key for the next lookup, and then lead nodes, which are the encoded path and then the leaf value (which is either an rlp encoded account or null, right)?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah that's fair.

I'm not sure I get the idea behind getting rid of the "raw node" terminology. We could do a renaming, but then we would have to rework the classes since these are already called "ExensionNode" and "LeafNode". Additionally, in comments and various places (e.g. methods like decodeRawNode), we have been referring to "raw nodes" quite often, and I would tend not to make changes that are too intrusive.

I've updated the comment, let me know if that seems clear enough

export type RawExtensionNode = [Uint8Array, Uint8Array]

// [ encodedPath, value ]
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[]

Expand Down
10 changes: 0 additions & 10 deletions packages/trie/src/util/nibbles.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Loading