packages/contracts-core/contracts/libs/merkle/MerkleTree.sol
// SPDX-License-Identifier: MIT
pragma solidity 0.8.17;
import {MerkleMath} from "./MerkleMath.sol";
import {AGENT_TREE_HEIGHT, ORIGIN_TREE_HEIGHT} from "../Constants.sol";
import {NotEnoughLeafs, MerkleTreeFull, LeafNotProven} from "../Errors.sol";
/// `BaseTree` is a struct representing incremental merkle tree.
/// - Contains only the current branch.
/// - The number of inserted leaves is stored externally, and is later supplied for tree operations.
/// - Has a fixed height of `ORIGIN_TREE_HEIGHT`.
/// - Can store up to `2 ** ORIGIN_TREE_HEIGHT - 1` leaves.
/// > Note: the hash function for the tree `H(x, y)` is defined as:
/// > - `H(0,0) = 0`
/// > - `H(x,y) = keccak256(x, y), when (x != 0 || y != 0)`
/// ## Invariant for leafs
/// - The leftmost empty leaf has `index == count`, where `count` is the amount of the inserted leafs so far.
/// - Value for any empty leaf or node is bytes32(0).
/// ## Invariant for current branch
/// `branch[i]` is always the value of a node on the i-th level.
/// Levels are numbered from leafs to root: `0 .. ORIGIN_TREE_HEIGHT`.
/// `branch[i]` stores the value for the node, such that:
// - The node is a "left child" (e.g. has an even index).
/// - The node must have two non-empty children.
/// - Out of all level's "left child" nodes with "non-empty children",
/// the one with the biggest index (the rightmost one) is stored as `branch[i]`.
/// > __`branch` could be used to form a proof of inclusion for the first empty leaf (`index == count`).__
/// _Here is how:_
/// - Let's walk along the path from the "first empty leaf" to the root.
/// - i-th bit of the "first empty leaf" index (which is equal to `count`) determines if the path's node
/// for this i-th level is a "left child" or a "right child".
/// - i-th bit in `count` is 0 → we are the left child on this level → sibling is the right child
/// that does not exist yet → `proof[i] = bytes32(0)`.
/// - i-th bit in `count` is 1 → we are the right child on this level → sibling is the left child
/// sibling is the rightmost "left child" node on the level → `proof[i] = branch[i]`.
/// > Therefore `proof[i] = (count & (1 << i)) == 0 ? bytes32(0) : branch[i])`
struct BaseTree {
bytes32[ORIGIN_TREE_HEIGHT] branch;
}
using MerkleTree for BaseTree global;
/// `HistoricalTree` is an incremental merkle tree keeping track of its historical merkle roots.
/// > - `roots[N]` is the root of the tree after `N` leafs were inserted
/// > - `roots[0] == bytes32(0)`
/// @param tree Incremental merkle tree
/// @param roots Historical merkle roots of the tree
struct HistoricalTree {
BaseTree tree;
bytes32[] roots;
}
using MerkleTree for HistoricalTree global;
/// `DynamicTree` is a struct representing a Merkle Tree with `2**AGENT_TREE_HEIGHT` leaves.
/// - A single operation is available: update value for leaf with an arbitrary index (which might be a non-empty leaf).
/// - This is done by requesting the proof of inclusion for the old value, which is used to both
/// verify the old value, and calculate the new root.
/// > Based on Original idea from [ER forum post](https://ethresear.ch/t/efficient-on-chain-dynamic-merkle-tree/11054).
struct DynamicTree {
bytes32 root;
}
using MerkleTree for DynamicTree global;
/// MerkleTree is work based on Nomad's Merkle.sol, which is used under MIT OR Apache-2.0
/// [link](https://github.com/nomad-xyz/monorepo/blob/main/packages/contracts-core/contracts/libs/Merkle.sol).
/// With the following changes:
/// - Adapted for Solidity 0.8.x.
/// - Amount of tree leaves stored externally.
/// - Added thorough documentation.
/// - `H(0,0) = 0` optimization from [ER forum post](https://ethresear.ch/t/optimizing-sparse-merkle-trees/3751/6).
/// > Nomad's Merkle.sol is work based on eth2 deposit contract, which is used under CC0-1.0
///[link](https://github.com/ethereum/deposit_contract/blob/dev/deposit_contract/contracts/validator_registration.v.py).
/// With the following changes:
/// > - Implemented in Solidity 0.7.6 (eth2 deposit contract implemented in Vyper).
/// > - `H() = keccak256()` is used as the hashing function instead of `sha256()`.
library MerkleTree {
/// @dev For root calculation we need at least one empty leaf, thus the minus one in the formula.
uint256 internal constant MAX_LEAVES = 2 ** ORIGIN_TREE_HEIGHT - 1;
// ═════════════════════════════════════════════════ BASE TREE ═════════════════════════════════════════════════════
/**
* @notice Inserts `node` into merkle tree
* @dev Reverts if tree is full
* @param newCount Amount of inserted leaves in the tree after the insertion (i.e. current + 1)
* @param node Element to insert into tree
*/
function insertBase(BaseTree storage tree, uint256 newCount, bytes32 node) internal {
if (newCount > MAX_LEAVES) revert MerkleTreeFull();
// We go up the tree following the branch from the empty leaf AFTER the just inserted one.
// We stop when we find the first "right child" node.
// Its sibling is now the rightmost "left child" node that has two non-empty children.
// Therefore we need to update `tree.branch` value on this level.
// One could see that `tree.branch` value on lower and higher levels remain unchanged.
// Loop invariant: `node` is the current level's value for the branch from JUST INSERTED leaf
for (uint256 i = 0; i < ORIGIN_TREE_HEIGHT;) {
if ((newCount & 1) == 1) {
// Found the first "right child" node on the branch from EMPTY leaf
// `node` is the value for node on branch from JUST INSERTED leaf
// Which in this case is the "left child".
// We update tree.branch and exit
tree.branch[i] = node;
return;
}
// On the branch from EMPTY leaf this is still a "left child".
// Meaning on branch from JUST INSERTED leaf, `node` is a right child.
// We compute value for `node` parent using `tree.branch` invariant:
// This is the rightmost "left child" node, which would be sibling of `node`.
node = MerkleMath.getParent(tree.branch[i], node);
// Get the parent index, and go to the next tree level
newCount >>= 1;
unchecked {
++i;
}
}
// As the loop should always end prematurely with the `return` statement,
// this code should be unreachable. We assert `false` just to be safe.
assert(false);
}
/**
* @notice Calculates and returns current root of the merkle tree.
* @param count Current amount of inserted leaves in the tree
* @return current Calculated root of `tree`
*/
function rootBase(BaseTree storage tree, uint256 count) internal view returns (bytes32 current) {
// To calculate the root we follow the branch of first EMPTY leaf (index == count)
for (uint256 i = 0; i < ORIGIN_TREE_HEIGHT;) {
// Check if we are the left or the right child on the current level
if ((count & 1) == 1) {
// We are the right child. Our sibling is the "rightmost" "left-child" node
// that has two non-empty children → sibling is `tree.branch[i]`
current = MerkleMath.getParent(tree.branch[i], current);
} else {
// We are the left child. Our sibling does not exist yet → sibling is EMPTY
current = MerkleMath.getParent(current, bytes32(0));
}
// Get the parent index, and go to the next tree level
count >>= 1;
unchecked {
++i;
}
}
}
// ══════════════════════════════════════════════ HISTORICAL TREE ══════════════════════════════════════════════════
/// @notice Initializes the historical roots for the tree by inserting an empty root.
// solhint-disable-next-line ordering
function initializeRoots(HistoricalTree storage tree) internal returns (bytes32 savedRoot) {
// This should only be called once, when the contract is initialized
assert(tree.roots.length == 0);
// Save root for empty merkle tree: bytes32(0)
tree.roots.push(savedRoot);
}
/// @notice Inserts a new leaf into the merkle tree.
/// @dev Reverts if tree is full.
/// @param node Element to insert into tree
/// @return newRoot Merkle root after the leaf was inserted
function insert(HistoricalTree storage tree, bytes32 node) internal returns (bytes32 newRoot) {
// Tree count after the new leaf will be inserted (we store roots[0] as root of empty tree)
uint256 newCount = tree.roots.length;
tree.tree.insertBase(newCount, node);
// Save the new root
newRoot = tree.tree.rootBase(newCount);
tree.roots.push(newRoot);
}
/// @notice Returns the historical root of the merkle tree.
/// @dev Reverts if not enough leafs have been inserted.
/// @param count Amount of leafs in the tree at some point of time
/// @return historicalRoot Merkle root after `count` leafs were inserted
function root(HistoricalTree storage tree, uint256 count) internal view returns (bytes32 historicalRoot) {
if (count >= tree.roots.length) revert NotEnoughLeafs();
return tree.roots[count];
}
// ═══════════════════════════════════════════════ DYNAMIC TREE ════════════════════════════════════════════════════
/**
* @notice Updates the value for the leaf with the given index in the Dynamic Merkle Tree.
* @dev Will revert if incorrect proof of inclusion for old value is supplied.
* @param tree Dynamic merkle tree
* @param index Index of the leaf to update
* @param oldValue Previous value of the leaf
* @param branch Proof of inclusion of previous value into the tree
* @param newValue New leaf value to assign
* @return newRoot New value for the Merkle Root after the leaf is updated
*/
function update(
DynamicTree storage tree,
uint256 index,
bytes32 oldValue,
bytes32[] memory branch,
bytes32 newValue
) internal returns (bytes32 newRoot) {
// Check that the old value + proof result in a correct root
if (MerkleMath.proofRoot(index, oldValue, branch, AGENT_TREE_HEIGHT) != tree.root) {
revert LeafNotProven();
}
// New root is new value + the same proof (values for sibling nodes are not updated)
newRoot = MerkleMath.proofRoot(index, newValue, branch, AGENT_TREE_HEIGHT);
// Write the new root
tree.root = newRoot;
}
}