@typeberry/lib 0.5.1-a4d7c93 → 0.5.1-eb84786

package/packages/core/trie/trie.js

@@ -0,0 +1,356 @@
+import { Bytes } from "#@typeberry/bytes";
+import { HASH_SIZE, TRUNCATED_HASH_SIZE } from "#@typeberry/hash";
+import { check } from "#@typeberry/utils";
+import { BranchNode, LeafNode, NodeType, TRUNCATED_KEY_BITS, } from "./nodes.js";
+import { WriteableNodesDb } from "./nodesDb.js";
+/** Compare two trie `LeafNode`s only by their key. */
+export const leafComparator = (x, y) => x.getKey().compare(y.getKey());
+const zero = Bytes.zero(HASH_SIZE).asOpaque();
+export class InMemoryTrie {
+    nodes;
+    root;
+    /** Create an empty in-memory trie. */
+    static empty(hasher) {
+        return new InMemoryTrie(new WriteableNodesDb(hasher));
+    }
+    /** Given a collection of leaves, compute the state root. */
+    static computeStateRoot(hasher, leaves) {
+        const sorted = leaves.slice();
+        const firstSorted = sorted.shift();
+        if (firstSorted === undefined) {
+            return zero;
+        }
+        const nodes = [
+            {
+                leaf: firstSorted,
+                sharedBitsWithPrev: 0,
+            },
+        ];
+        let last = nodes[0];
+        // first we go through all of the sorted leaves and figure out how much in common
+        // they have with the previous node.
+        // If the shared-prefix drops, it means we are going up in depth (i.e. we are in different branch).
+        for (const leaf of sorted) {
+            const sharedBitsCount = findSharedPrefix(leaf.getKey(), last.leaf.getKey());
+            last = {
+                leaf,
+                sharedBitsWithPrev: sharedBitsCount,
+            };
+            nodes.push(last);
+        }
+        // Now we will go backwards and hash them together (or create branch nodes).
+        nodes.reverse();
+        const stack = [];
+        let currentDepth = 0;
+        const lastNode = nodes.length === 1 ? undefined : nodes[nodes.length - 1];
+        for (const node of nodes) {
+            const isLastNode = node === lastNode;
+            const key = node.leaf.getKey();
+            const prevDepth = currentDepth;
+            currentDepth = node.sharedBitsWithPrev;
+            // first push all missing right-hand zero nodes.
+            // Handle the case if all nodes are on the left side and we need one more top-level
+            // extra.
+            const startDepth = isLastNode ? prevDepth : prevDepth + 1;
+            for (let i = startDepth; i <= currentDepth; i++) {
+                if (getBit(key, i) === false) {
+                    stack.push(zero);
+                }
+            }
+            // now let's push the hash of the current leaf
+            const hash = hasher.hashConcat(node.leaf.node.raw);
+            stack.push(hash);
+            // we are going further down, so no need to merge anything
+            if (prevDepth < currentDepth) {
+                continue;
+            }
+            // jumping back to some lower depth, we need to merge what we have on the stack.
+            // we need to handle a case where we have no nodes on the top-most left side.
+            // in such case we just add extra zero on the left.
+            const endDepth = isLastNode ? currentDepth - 1 : currentDepth;
+            for (let i = prevDepth; i > endDepth; i--) {
+                if (getBit(key, i) === true) {
+                    stack.push(zero);
+                }
+                const current = stack.pop() ?? zero;
+                const next = stack.pop() ?? zero;
+                const branchNode = BranchNode.fromSubNodes(current, next);
+                const hash = hasher.hashConcat(branchNode.node.raw);
+                stack.push(hash);
+            }
+        }
+        return stack.pop() ?? zero;
+    }
+    /**
+     * Construct a `LeafNode` from given `key` and `value`.
+     *
+     * NOTE: for large value it WILL NOT be embedded in the leaf node,
+     * and should rather be stored separately.
+     */
+    static constructLeaf(hasher, key, value, maybeValueHash) {
+        const valueHash = () => maybeValueHash ?? hasher.hashConcat(value.raw).asOpaque();
+        return LeafNode.fromValue(key, value, valueHash);
+    }
+    /**
+     * Reconstruct the entire trie from it's leaves.
+     *
+     * Note that if only the state root is needed, this is rather inefficient.
+     */
+    static fromLeaves(hasher, leaves) {
+        // TODO [ToDr] [opti] Pair up the leaves and build upper levels.
+        let root = null;
+        const nodes = new WriteableNodesDb(hasher);
+        for (const leaf of leaves) {
+            root = trieInsert(root, nodes, leaf);
+        }
+        return new InMemoryTrie(nodes, root);
+    }
+    constructor(
+    // Exposed for trie-visualiser
+    nodes, root = null) {
+        this.nodes = nodes;
+        this.root = root;
+    }
+    set(key, value, maybeValueHash) {
+        const leafNode = InMemoryTrie.constructLeaf(this.nodes.hasher, key, value, maybeValueHash);
+        this.root = trieInsert(this.root, this.nodes, leafNode);
+        return leafNode;
+    }
+    remove(_) {
+        // TODO [ToDr] Trie removal is most likely not needed. If we run into this issue,
+        // we should most likely decide to optimize the code that requires removal.
+        // NOTE: we pretty much NEVER need the full trie at all. We only need to compute the
+        // state root (from time to time), but we can easily just operate on the leafs.
+        // Removal should happen on the leaf level as well. Most likey the whole
+        // `InMemoryTrie` stuff should be eradicated in favor of leaves-operating methods.
+        throw new Error("Removing from the trie not implemented yet.");
+    }
+    getRootNode() {
+        return this.root;
+    }
+    getRootHash() {
+        if (this.root === null) {
+            return Bytes.zero(HASH_SIZE).asOpaque();
+        }
+        return this.nodes.hashNode(this.root);
+    }
+    toString() {
+        return trieStringify(this.root, this.nodes);
+    }
+}
+/**
+ * Insert a new leaf node into a trie starting at the given `root` node.
+ *
+ * The function will find a place where the leaf node should be present and update
+ * the entire branch up to the trie root.
+ *
+ * New root node is returned.
+ */
+function trieInsert(root, nodes, leaf) {
+    if (root === null) {
+        nodes.insert(leaf.node);
+        return leaf.node;
+    }
+    // first we look up a good place to insert the node to the tree, based on it's key.
+    const traversedPath = findNodeToReplace(root, nodes, leaf.getKey());
+    // now we analyze two possible situations:
+    // 1. We found a leaf node - that means we need to create a branch node (and possible
+    //    extra branch nodes for a common prefix) with these two leaves. Finally we update the
+    //    traversed path from root.
+    // 2. We found an empty spot (i.e. branch node with zero hash) - we can just update already
+    //    traversed path from root.
+    const nodeToInsert = traversedPath.leafToReplace !== undefined
+        ? createSubtreeForBothLeaves(traversedPath, nodes, traversedPath.leafToReplace, leaf)
+        : [leaf.node, nodes.insert(leaf.node)];
+    // finally update the traversed path from `root` to the insertion location.
+    let historicalBranch = traversedPath.branchingHistory.pop();
+    let [lastNode, lastHash] = nodeToInsert;
+    while (historicalBranch !== undefined) {
+        const [branchNode, branchHash, bit] = historicalBranch;
+        nodes.remove(branchHash);
+        // TODO [ToDr] [opti] Avoid allocation here by re-using the old branch node?
+        const newBranchNode = bit
+            ? BranchNode.fromSubNodes(branchNode.getLeft(), lastHash)
+            : BranchNode.fromSubNodes(lastHash, branchNode.getRight());
+        lastHash = nodes.insert(newBranchNode.node);
+        lastNode = newBranchNode.node;
+        historicalBranch = traversedPath.branchingHistory.pop();
+    }
+    return lastNode;
+}
+/**
+ * Path of branch nodes traversed while looking for the best place to put a new leaf.
+ */
+class TraversedPath {
+    /** history of branch nodes (with their hashes) and the branching bit. */
+    branchingHistory = [];
+    /** last bitIndex */
+    bitIndex = 0;
+    /** in case of a leaf node at destination, details of that leaf node */
+    leafToReplace;
+}
+/**
+ * Traverse the trie starting from root and return the path leading to the destination
+ * where leaf with `key` should be placed.
+ */
+function findNodeToReplace(root, nodes, key) {
+    const traversedPath = new TraversedPath();
+    let currentNode = root;
+    let currentNodeHash = nodes.hashNode(root);
+    while (true) {
+        const kind = currentNode.getNodeType();
+        if (kind !== NodeType.Branch) {
+            // we found a leaf that needs to be merged with the one being inserted.
+            const leaf = currentNode.asLeafNode();
+            traversedPath.leafToReplace = [leaf, currentNodeHash];
+            return traversedPath;
+        }
+        // going down the trie
+        const branch = currentNode.asBranchNode();
+        const currBit = getBit(key, traversedPath.bitIndex);
+        const nextHash = currBit ? branch.getRight() : branch.getLeft();
+        traversedPath.branchingHistory.push([branch, currentNodeHash, currBit]);
+        const nextNode = nodes.get(nextHash);
+        if (nextNode === null) {
+            if (nextHash.isEqualTo(zero)) {
+                return traversedPath;
+            }
+            throw new Error(`Missing trie node '${nextHash}' with key prefix: ${key}[0..${traversedPath.bitIndex}]`);
+        }
+        currentNode = nextNode;
+        currentNodeHash = nextHash;
+        traversedPath.bitIndex += 1;
+    }
+}
+/**
+ * Handle a situation where we replace an existing leaf node at destination.
+ *
+ * In such case we need to create a subtree that will hold both of the leaves.
+ *
+ * The function returns a root of the subtree.
+ */
+function createSubtreeForBothLeaves(traversedPath, nodes, leafToReplace, leaf) {
+    const key = leaf.getKey();
+    let [existingLeaf, existingLeafHash] = leafToReplace;
+    const existingLeafKey = existingLeaf.getKey();
+    // TODO [ToDr] [opti] instead of inserting/removing a bunch of nodes, it might be
+    // better to return a changeset that can be batch-applied to the DB.
+    const leafNodeHash = nodes.insert(leaf.node);
+    if (existingLeafKey.isEqualTo(key)) {
+        // remove only if we are not inserting the same value twice
+        // we compare values only, since the hashes might have a difference at first
+        // bit.
+        if (!existingLeaf.getValueHash().isEqualTo(leaf.getValueHash())) {
+            nodes.remove(existingLeafHash);
+        }
+        // just replacing an existing value
+        return [leaf.node, leafNodeHash];
+    }
+    // In case both keys share a prefix we need to add a bunch of branch
+    // nodes up until the keys start diverging.
+    // Here we identify the common bit prefix that will later be used
+    // in reverse to construct required branch nodes.
+    const commonBits = [];
+    let divergingBit = getBit(key, traversedPath.bitIndex);
+    while (traversedPath.bitIndex < TRUNCATED_KEY_BITS) {
+        divergingBit = getBit(key, traversedPath.bitIndex);
+        const bit2 = getBit(existingLeafKey, traversedPath.bitIndex);
+        if (divergingBit === bit2) {
+            commonBits.push(bit2);
+            traversedPath.bitIndex += 1;
+        }
+        else {
+            break;
+        }
+    }
+    // Now construct the common branches, and insert zero hash in place of other sub-trees.
+    const zero = Bytes.zero(HASH_SIZE).asOpaque();
+    // In case we move the leaf from left to right it's hash needs to be re-calculated (missing bit).
+    // TODO [ToDr] [opti] might be better to store the original bit value instead of recalculating.
+    const leafWasInLeftBranch = (() => {
+        const l = traversedPath.branchingHistory.length;
+        if (l > 0) {
+            return traversedPath.branchingHistory[l - 1][2] === false;
+        }
+        return false;
+    })();
+    if (leafWasInLeftBranch && !divergingBit) {
+        existingLeafHash = nodes.hashNode(existingLeaf.node);
+    }
+    let lastBranch = divergingBit
+        ? BranchNode.fromSubNodes(existingLeafHash, leafNodeHash)
+        : BranchNode.fromSubNodes(leafNodeHash, existingLeafHash);
+    let lastHash = nodes.insert(lastBranch.node);
+    let bit = commonBits.pop();
+    // go up and create branch nodes for the common prefix
+    while (bit !== undefined) {
+        lastBranch = bit ? BranchNode.fromSubNodes(zero, lastHash) : BranchNode.fromSubNodes(lastHash, zero);
+        lastHash = nodes.insert(lastBranch.node);
+        bit = commonBits.pop();
+    }
+    // let's return the top branch to join with the history
+    return [lastBranch.node, lastHash];
+}
+/**
+ * Return a single bit from `key` located at `bitIndex`.
+ */
+function getBit(key, bitIndex) {
+    check `${bitIndex < TRUNCATED_KEY_BITS} invalid bit index passed ${bitIndex}`;
+    const byte = bitIndex >>> 3;
+    const bit = bitIndex - (byte << 3);
+    const mask = 0b10_00_00_00 >>> bit;
+    const val = key.raw[byte] & mask;
+    return val !== 0;
+}
+function trieStringify(root, nodes) {
+    if (root === null) {
+        return "<empty tree>";
+    }
+    const kind = root.getNodeType();
+    if (kind === NodeType.Branch) {
+        const branch = root.asBranchNode();
+        const leftHash = branch.getLeft();
+        const rightHash = branch.getRight();
+        const indent = (v) => v
+            .split("\n")
+            .map((v) => `\t\t${v}`)
+            .join("\n");
+        const left = trieStringify(nodes.get(leftHash), nodes);
+        const right = trieStringify(nodes.get(rightHash), nodes);
+        return `<branch>
+	-- ${leftHash}: ${indent(left)}
+	-- ${rightHash}: ${indent(right)}
+`;
+    }
+    const leaf = root.asLeafNode();
+    const valueLength = leaf.getValueLength();
+    const value = valueLength > 0 ? `'${leaf.getValue()}'(len:${valueLength})` : `'<hash>${leaf.getValueHash()}'`;
+    return `\nLeaf('${leaf.getKey().toString()}',${value})`;
+}
+export function findSharedPrefix(a, b) {
+    for (let i = 0; i < TRUNCATED_HASH_SIZE; i++) {
+        const diff = a.raw[i] ^ b.raw[i];
+        if (diff === 0) {
+            continue;
+        }
+        // check how many bits match
+        for (const [mask, matchingBits] of bitLookup) {
+            if ((mask & diff) !== 0) {
+                return i * 8 + matchingBits;
+            }
+        }
+        return i;
+    }
+    return TRUNCATED_HASH_SIZE * 8;
+}
+const bitLookup = [
+    [0b10000000, 0],
+    [0b01000000, 1],
+    [0b00100000, 2],
+    [0b00010000, 3],
+    [0b00001000, 4],
+    [0b00000100, 5],
+    [0b00000010, 6],
+    [0b00000001, 7],
+    [0b00000000, 8],
+];

package/packages/core/trie/trie.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=trie.test.d.ts.map

package/packages/core/trie/trie.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"trie.test.d.ts","sourceRoot":"","sources":["../../../../../packages/core/trie/trie.test.ts"],"names":[],"mappings":""}

package/packages/core/trie/trie.test.js

@@ -0,0 +1,287 @@
+import assert from "node:assert";
+import { before, describe, it } from "node:test";
+import { Bytes, BytesBlob } from "#@typeberry/bytes";
+import { SortedSet } from "#@typeberry/collections";
+import { Blake2b, HASH_SIZE, TRUNCATED_HASH_SIZE } from "#@typeberry/hash";
+import { deepEqual } from "#@typeberry/utils";
+import { getBlake2bTrieHasher } from "./hasher.js";
+import { LeafNode, parseInputKey } from "./nodes.js";
+import { findSharedPrefix, InMemoryTrie, leafComparator } from "./trie.js";
+let blake2bTrieHasher;
+before(async () => {
+    const blake2b = await Blake2b.createHasher();
+    blake2bTrieHasher = getBlake2bTrieHasher(blake2b);
+});
+describe("Root hash", () => {
+    it("should compute state root hash that is equal to the trie one (complex)", () => {
+        const data = Object.entries(testVector10.input).map(([key, val]) => {
+            return [parseInputKey(key), BytesBlob.parseBlobNoPrefix(val)];
+        });
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of data) {
+            trie.set(key.asOpaque(), val);
+        }
+        const expected = trie.getRootHash();
+        const leaves = SortedSet.fromArray(leafComparator, data.map(([key, value]) => {
+            return InMemoryTrie.constructLeaf(blake2bTrieHasher, key.asOpaque(), value);
+        }));
+        deepEqual(InMemoryTrie.computeStateRoot(blake2bTrieHasher, leaves), expected);
+        deepEqual(`0x${testVector10.output}`, expected.toString());
+    });
+    it("should compute state root hash that is equal to the trie one (simple)", () => {
+        const data = [
+            [Bytes.fill(HASH_SIZE, 5), BytesBlob.blobFromString("five")],
+            [Bytes.fill(HASH_SIZE, 1), BytesBlob.blobFromString("one")],
+            [Bytes.fill(HASH_SIZE, 3), BytesBlob.blobFromString("three")],
+        ];
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of data) {
+            trie.set(key.asOpaque(), val);
+        }
+        const expected = trie.getRootHash();
+        const leaves = SortedSet.fromArray(leafComparator, data.map(([key, value]) => {
+            return InMemoryTrie.constructLeaf(blake2bTrieHasher, key.asOpaque(), value);
+        }));
+        deepEqual(InMemoryTrie.computeStateRoot(blake2bTrieHasher, leaves), expected);
+    });
+    it("should find shared prefix bits number", () => {
+        const examples = [
+            [
+                "0x03345958f90731bce89d07c2722dc693425a541b5230f99a6867882993576a",
+                "0x3e7d409b9037b1fd870120de92ebb7285219ce4526c54701b888c5a13995f7",
+                2,
+            ],
+            [
+                "0x3e7d409b9037b1fd870120de92ebb7285219ce4526c54701b888c5a13995f7",
+                "0x5dffe0e2c9f089d30e50b04ee562445cf2c0e7e7d677580ef0ccf2c6fa3522",
+                1,
+            ],
+            [
+                "0x5dffe0e2c9f089d30e50b04ee562445cf2c0e7e7d677580ef0ccf2c6fa3522",
+                "0x6bf8460545baf5b0af874ebbbd56ae09ee73cd24926b4549238b797b447e05",
+                2,
+            ],
+            [
+                "0x6bf8460545baf5b0af874ebbbd56ae09ee73cd24926b4549238b797b447e05",
+                "0x7723a8383e43a1713eb920bae44880b2ae9225ea2d38c031cf3b22434b4507",
+                3,
+            ],
+            [
+                "0x7723a8383e43a1713eb920bae44880b2ae9225ea2d38c031cf3b22434b4507",
+                "0xc2d3bda8f77cc483d2f4368cf998203097230fd353d2223e5a333eb58f76a4",
+                0,
+            ],
+            [
+                "0xc2d3bda8f77cc483d2f4368cf998203097230fd353d2223e5a333eb58f76a4",
+                "0xc7a04effd2c0cede0279747f58bd210d0cc9d65c2eba265c6b4dfbc058a704",
+                5,
+            ],
+        ];
+        for (const [left, right, len] of examples) {
+            const a = Bytes.parseBytes(left, TRUNCATED_HASH_SIZE).asOpaque();
+            const b = Bytes.parseBytes(right, TRUNCATED_HASH_SIZE).asOpaque();
+            const actual = findSharedPrefix(a, b);
+            assert.strictEqual(actual, len, `Expected shared prefix of ${len}, got: ${actual}. Values:\n${a}\n${b}.`);
+        }
+    });
+});
+describe("Root hash", () => {
+    it("should compute state root hash that is equal to the trie one", () => {
+        const data = [
+            [Bytes.fill(HASH_SIZE, 5), BytesBlob.blobFromString("five")],
+            [Bytes.fill(HASH_SIZE, 1), BytesBlob.blobFromString("one")],
+            [Bytes.fill(HASH_SIZE, 3), BytesBlob.blobFromString("three")],
+            [Bytes.fill(HASH_SIZE, 4), BytesBlob.blobFromString("four")],
+            [Bytes.fill(HASH_SIZE, 2), BytesBlob.blobFromString("two")],
+        ];
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of data) {
+            trie.set(key.asOpaque(), val);
+        }
+        const expected = trie.getRootHash();
+        const leaves = SortedSet.fromArray(leafComparator, data.map(([key, value]) => {
+            return InMemoryTrie.constructLeaf(blake2bTrieHasher, key.asOpaque(), value);
+        }));
+        deepEqual(InMemoryTrie.computeStateRoot(blake2bTrieHasher, leaves), expected);
+    });
+});
+describe("Trie", () => {
+    it("Empty trie", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        assert.deepStrictEqual(trie.getRootHash(), Bytes.parseBytesNoPrefix("0000000000000000000000000000000000000000000000000000000000000000", 32));
+    });
+    it("Leaf Node", () => {
+        const key = parseInputKey("16c72e0c2e0b78157e3a116d86d90461a199e439325317aea160b30347adb8ec");
+        const value = BytesBlob.parseBlob("0x4227b4a465084852cd87d8f23bec0db6fa7766b9685ab5e095ef9cda9e15e49dff");
+        const valueHash = () => blake2bTrieHasher.hashConcat(value.raw).asOpaque();
+        const node = LeafNode.fromValue(key, value, valueHash);
+        assert.deepStrictEqual(node.getKey(), Bytes.parseBytes("0x16c72e0c2e0b78157e3a116d86d90461a199e439325317aea160b30347adb8", 31));
+        assert.deepStrictEqual(node.getValueLength(), 0);
+        assert.deepStrictEqual(node.getValue().raw, Bytes.zero(0).raw);
+        assert.deepStrictEqual(node.getValueHash(), valueHash());
+    });
+    it("Empty value", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        trie.set(parseInputKey("16c72e0c2e0b78157e3a116d86d90461a199e439325317aea160b30347adb8ec"), BytesBlob.blobFromNumbers([]));
+        assert.deepStrictEqual(trie.getRootHash().toString(), Bytes.parseBytesNoPrefix("99ecb1509d2cbc16bab389714e5933932977e742472fcd9277d67f45699e076a", 32).toString());
+    });
+    it("Should import some keys", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        trie.set(parseInputKey("645eece27fdce6fd3852790131a50dc5b2dd655a855421b88700e6eb43279ad9"), BytesBlob.blobFromNumbers([0x72]));
+        assert.deepStrictEqual(trie.getRootHash().toString(), Bytes.parseBytesNoPrefix("e9a89ab2f10d45a46d47127110e8353d6443b635b08e989a743c27bb82740d7d", 32).toString());
+    });
+    it("Non embedded leaf", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        trie.set(parseInputKey("3dbc5f775f6156957139100c343bb5ae6589af7398db694ab6c60630a9ed0fcd"), BytesBlob.parseBlob("0x4227b4a465084852cd87d8f23bec0db6fa7766b9685ab5e095ef9cda9e15e49d"));
+        assert.deepStrictEqual(trie.getRootHash().toString(), Bytes.parseBytesNoPrefix("5fd68f074c914741601931d64c6c772c18ab8a4cd0cd3a4fff0611a5d97ecc94", 32).toString());
+    });
+    it("More complicated trie", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        trie.set(parseInputKey("f2a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x22c62f84ee5775d1e75ba6519f6dfae571eb1888768f2a203281579656b6a29097f7c7e2cf44e38da9a541d9b4c773db8b71e1d3"));
+        trie.set(parseInputKey("f3a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x44d0b26211d9d4a44e375207"));
+        assert.deepStrictEqual(trie.getRootHash().toString(), Bytes.parseBytesNoPrefix("fb4bc560e0c314b09a29fc3f83a7f063ec118ff3fc1fba4430fcc0fbea09a207", 32).toString());
+    });
+    it("Move leaf from left to right branch", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        // left value
+        trie.set(parseInputKey("f2a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x23"));
+        // right value
+        trie.set(parseInputKey("f1a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x1234"));
+        // now insert another leaf, which causes `0xf2..` to move to the right.
+        trie.set(parseInputKey("f0a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x1234"));
+        assert.deepStrictEqual(trie.getRootHash().toString(), "0xf07003b9d508b5b017960d069ada7893146c3425be6293b7ac3d4a9709d33b47");
+    });
+    it("Replace leaf value", () => {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        const insert = {
+            f2a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3: "0x23",
+            f1a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3: "0x1234",
+            f0a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3: "0x1234",
+        };
+        for (const [k, v] of Object.entries(insert)) {
+            trie.set(parseInputKey(k), BytesBlob.parseBlob(v));
+        }
+        assert.deepStrictEqual(trie.getRootHash().toString(), "0xf07003b9d508b5b017960d069ada7893146c3425be6293b7ac3d4a9709d33b47");
+        // now set the same key again
+        trie.set(parseInputKey("f2a9fcaf8ae0ff770b0908ebdee1daf8457c0ef5e1106c89ad364236333c5fb3"), BytesBlob.parseBlob("0x1234"));
+        assert.deepStrictEqual(trie.getRootHash().toString(), "0xd6349d4e62a4288ba3d80b94421f2933b13f6f882b2c69a90508aa92ff88343f");
+    });
+    const testVector9 = {
+        d7f99b746f23411983df92806725af8e5cb66eba9f200737accae4a1ab7f47b9: "24232437f5b3f2380ba9089bdbc45efaffbe386602cb1ecc2c17f1d0",
+        "59ee947b94bcc05634d95efb474742f6cd6531766e44670ec987270a6b5a4211": "72fdb0c99cf47feb85b2dad01ee163139ee6d34a8d893029a200aff76f4be5930b9000a1bbb2dc2b6c79f8f3c19906c94a3472349817af21181c3eef6b",
+        a3dc3bed1b0727caf428961bed11c9998ae2476d8a97fad203171b628363d9a2: "8a0dafa9d6ae6177",
+        "15207c233b055f921701fc62b41a440d01dfa488016a97cc653a84afb5f94fd5": "157b6c821169dacabcf26690df",
+        b05ff8a05bb23c0d7b177d47ce466ee58fd55c6a0351a3040cf3cbf5225aab19: "6a208734106f38b73880684b",
+    };
+    it("should return all leaf nodes", () => {
+        const data = { ...testVector9 };
+        // construct the trie
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of Object.entries(data)) {
+            const stateKey = parseInputKey(key);
+            const value = BytesBlob.parseBlobNoPrefix(val);
+            trie.set(stateKey, value);
+        }
+        // when
+        const leaves = Array.from(trie.nodes.leaves());
+        assert.deepStrictEqual(leaves.map((val) => `${val.getKey()}: ${val.node}`), [
+            "0xd7f99b746f23411983df92806725af8e5cb66eba9f200737accae4a1ab7f47: 0x9cd7f99b746f23411983df92806725af8e5cb66eba9f200737accae4a1ab7f4724232437f5b3f2380ba9089bdbc45efaffbe386602cb1ecc2c17f1d000000000",
+            "0x59ee947b94bcc05634d95efb474742f6cd6531766e44670ec987270a6b5a42: 0xc059ee947b94bcc05634d95efb474742f6cd6531766e44670ec987270a6b5a42a7f2482020023b85b4009884a31aea08f03b4bbdb5efd5e6ff1d63f1a86aaa53",
+            "0xa3dc3bed1b0727caf428961bed11c9998ae2476d8a97fad203171b628363d9: 0x88a3dc3bed1b0727caf428961bed11c9998ae2476d8a97fad203171b628363d98a0dafa9d6ae6177000000000000000000000000000000000000000000000000",
+            "0x15207c233b055f921701fc62b41a440d01dfa488016a97cc653a84afb5f94f: 0x8d15207c233b055f921701fc62b41a440d01dfa488016a97cc653a84afb5f94f157b6c821169dacabcf26690df00000000000000000000000000000000000000",
+            "0xb05ff8a05bb23c0d7b177d47ce466ee58fd55c6a0351a3040cf3cbf5225aab: 0x8cb05ff8a05bb23c0d7b177d47ce466ee58fd55c6a0351a3040cf3cbf5225aab6a208734106f38b73880684b0000000000000000000000000000000000000000",
+        ]);
+    });
+    it("should create trie from leaf nodes", () => {
+        const data = { ...testVector9 };
+        // construct the trie manually
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of Object.entries(data)) {
+            const stateKey = parseInputKey(key);
+            const value = BytesBlob.parseBlobNoPrefix(val);
+            trie.set(stateKey, value);
+        }
+        // when
+        const leaves = Array.from(trie.nodes.leaves());
+        const actual = InMemoryTrie.fromLeaves(blake2bTrieHasher, leaves);
+        assert.deepStrictEqual(`${actual.getRootHash()}`, `${trie.getRootHash()}`);
+        assert.deepStrictEqual(actual.nodes, trie.nodes);
+    });
+    it("should return correct leafs after updates", () => {
+        const data = { ...testVector9 };
+        // construct the trie manually
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of Object.entries(data)) {
+            const stateKey = parseInputKey(key);
+            const value = BytesBlob.parseBlobNoPrefix(val);
+            trie.set(stateKey, value);
+        }
+        const initialLeaves = Array.from(trie.nodes.leaves());
+        // insert again
+        for (const [key, val] of Object.entries(data)) {
+            const stateKey = parseInputKey(key);
+            const value = BytesBlob.parseBlobNoPrefix(val);
+            trie.set(stateKey, value);
+        }
+        // when
+        const leaves = Array.from(trie.nodes.leaves());
+        const actual = InMemoryTrie.fromLeaves(blake2bTrieHasher, leaves);
+        deepEqual(leaves.map((x) => x.getKey().toString()).sort(), initialLeaves.map((x) => x.getKey().toString()).sort());
+        assert.deepStrictEqual(`${actual.getRootHash()}`, `${trie.getRootHash()}`);
+        assert.deepStrictEqual(Array.from(actual.nodes.leaves()), Array.from(trie.nodes.leaves()));
+    });
+    it("Test vector 9", () => {
+        const vector = {
+            input: { ...testVector9 },
+            output: "f6ac87ea6258a68bd3288cf73ac9d03419b548858c5466b1927b796f29db13fc",
+        };
+        runTestVector(vector);
+    });
+    it("Test vector 10", () => {
+        runTestVector(testVector10);
+    });
+    it("should work with shorter keys as well", () => {
+        const vector = {
+            input: {
+                "5dffe0e2c9f089d30e50b04ee562445cf2c0e7e7d677580ef0ccf2c6fa3522": "bb11c256876fe10442213dd78714793394d2016134c28a64eb27376ddc147fc6044df72bdea44d9ec66a3ea1e6d523f7de71db1d05a980e001e9fa",
+                df08871e8a54fde4834d83851469e635713615ab1037128df138a6cd223f12: "b8bded4e1c",
+                "7723a8383e43a1713eb920bae44880b2ae9225ea2d38c031cf3b22434b4507": "e46ddd41a5960807d528f5d9282568e622a023b94b72cb63f0353baff189257d",
+                "3e7d409b9037b1fd870120de92ebb7285219ce4526c54701b888c5a13995f73c": "9bc5d0",
+                c2d3bda8f77cc483d2f4368cf998203097230fd353d2223e5a333eb58f76a4: "9ae1dc59670bd3ef6fb51cbbbc05f1d2635fd548cb31f72500000a",
+                "6bf8460545baf5b0af874ebbbd56ae09ee73cd24926b4549238b797b447e05": "0964801caa928bc8c1869d60dbf1d8233233e0261baf725f2631d2b27574efc0316ce3067b4fccfa607274",
+                "832c15668a451578b4c69974085280b4bac5b01e220398f06e06a1d0aff285": "4881dd3238fd6c8af1090d455e7b449a",
+                c7a04effd2c0cede0279747f58bd210d0cc9d65c2eba265c6b4dfbc058a704: "d1fddfd63fd00cd6749a441b6ceaea1f250982a3a6b6d38f1b40cae00972cce3f9f4eaf7f9d7bc3070bd1e8d088500b10ca72e5ed5956f62",
+                "9e78a15cc0b45c83c83218efadd234cbac22dbffb24a76e2eb5f6a81d32df6": "e8256c6b5a9623cf2b293090f78f8fbceea6fc3991ac5f872400608f14d2a8b3d494fcda1c51d93b9904e3242cdeaa4b227c68cea89cca05ab6b5296edf105",
+                // this one is intentionally left longer to show that we support both mixed.
+                "03345958f90731bce89d07c2722dc693425a541b5230f99a6867882993576a23": "cd759a8d88edb46dda489a45ba6e48a42ce7efd36f1ca31d3bdfa40d2091f27740c5ec5de746d90d9841b986f575d545d0fb642398914eaab5",
+            },
+            output: "7d874bf70ea045e278f5cd2eafb28b74cdaee9c225ca884dee82532caa7bad0f",
+        };
+        runTestVector(vector);
+    });
+    function runTestVector(vector) {
+        const trie = InMemoryTrie.empty(blake2bTrieHasher);
+        for (const [key, val] of Object.entries(vector.input)) {
+            const stateKey = parseInputKey(key);
+            const value = BytesBlob.parseBlobNoPrefix(val);
+            trie.set(stateKey, value);
+        }
+        const expected = Bytes.parseBytesNoPrefix(vector.output, 32);
+        assert.deepStrictEqual(trie.getRootHash().toString(), expected.toString());
+    }
+});
+const testVector10 = {
+    input: {
+        "5dffe0e2c9f089d30e50b04ee562445cf2c0e7e7d677580ef0ccf2c6fa3522dd": "bb11c256876fe10442213dd78714793394d2016134c28a64eb27376ddc147fc6044df72bdea44d9ec66a3ea1e6d523f7de71db1d05a980e001e9fa",
+        df08871e8a54fde4834d83851469e635713615ab1037128df138a6cd223f1242: "b8bded4e1c",
+        "7723a8383e43a1713eb920bae44880b2ae9225ea2d38c031cf3b22434b4507e7": "e46ddd41a5960807d528f5d9282568e622a023b94b72cb63f0353baff189257d",
+        "3e7d409b9037b1fd870120de92ebb7285219ce4526c54701b888c5a13995f73c": "9bc5d0",
+        c2d3bda8f77cc483d2f4368cf998203097230fd353d2223e5a333eb58f76a429: "9ae1dc59670bd3ef6fb51cbbbc05f1d2635fd548cb31f72500000a",
+        "6bf8460545baf5b0af874ebbbd56ae09ee73cd24926b4549238b797b447e050a": "0964801caa928bc8c1869d60dbf1d8233233e0261baf725f2631d2b27574efc0316ce3067b4fccfa607274",
+        "832c15668a451578b4c69974085280b4bac5b01e220398f06e06a1d0aff2859a": "4881dd3238fd6c8af1090d455e7b449a",
+        c7a04effd2c0cede0279747f58bd210d0cc9d65c2eba265c6b4dfbc058a7047b: "d1fddfd63fd00cd6749a441b6ceaea1f250982a3a6b6d38f1b40cae00972cce3f9f4eaf7f9d7bc3070bd1e8d088500b10ca72e5ed5956f62",
+        "9e78a15cc0b45c83c83218efadd234cbac22dbffb24a76e2eb5f6a81d32df616": "e8256c6b5a9623cf2b293090f78f8fbceea6fc3991ac5f872400608f14d2a8b3d494fcda1c51d93b9904e3242cdeaa4b227c68cea89cca05ab6b5296edf105",
+        "03345958f90731bce89d07c2722dc693425a541b5230f99a6867882993576a23": "cd759a8d88edb46dda489a45ba6e48a42ce7efd36f1ca31d3bdfa40d2091f27740c5ec5de746d90d9841b986f575d545d0fb642398914eaab5",
+    },
+    output: "7d874bf70ea045e278f5cd2eafb28b74cdaee9c225ca884dee82532caa7bad0f",
+};