Shielded pools

/// The shielded pool. Holds token balances and tracks an incremental
/// Merkle tree of note commitments and spent nullifiers.
///
/// The Merkle tree uses a "filled subtrees" approach: we store one hash
/// per level representing the most recently completed subtree at that
/// depth. This allows O(depth) insertion without storing every node.
public struct ShieldedPool<phantom T> has key {
    id: UID,
    /// Pool balance
    balance: Balance<T>,
    /// Current Merkle root of the note commitment tree (BN254 field element)
    merkle_root: u256,
    /// Spent nullifiers (prevents double-spend)
    nullifiers: Table<u256, bool>,
    /// Next leaf index in the Merkle tree
    next_leaf_index: u64,
    /// Filled subtrees: one Poseidon hash per level (index 0 = leaf level).
    /// When a subtree at level i is completed, its root is stored here
    /// and used as the left child when computing the next level up.
    filled_subtrees: vector<u256>,
    /// Merkle tree depth (determines max capacity: 2^depth notes)
    tree_depth: u8,
}

/// Hash two BN254 field elements together using Poseidon.
/// This matches the circomlib Poseidon hash for 2 inputs.
fun hash_pair(left: u256, right: u256): u256 {
    let data = vector[left, right];
    poseidon::poseidon_bn254(&data)
}

/// Insert a leaf into the incremental Merkle tree and return the new root.
///
/// Algorithm: Starting from the leaf level, walk up the tree. At each level,
/// if the current index is even (left child), the sibling is the zero hash
/// at that level; if odd (right child), the sibling is the stored filled
/// subtree at that level. When the current index is even, we update the
/// filled_subtrees entry for that level because we just completed a left subtree.
fun insert_leaf<T>(pool: &mut ShieldedPool<T>, leaf: u256): u256 {
    let mut current_index = pool.next_leaf_index;
    let mut current_hash = leaf;
    let depth = pool.tree_depth as u64;

    // Precompute zero hashes for each level
    let mut zero_hashes = vector[];
    let mut zh = zero_leaf();
    let mut i = 0;
    while (i < depth) {
        zero_hashes.push_back(zh);
        zh = hash_pair(zh, zh);
        i = i + 1;
    };

    i = 0;
    while (i < depth) {
        if (current_index % 2 == 0) {
            // Left child: sibling is the zero hash at this level.
            *pool.filled_subtrees.borrow_mut(i) = current_hash;
            current_hash = hash_pair(current_hash, *zero_hashes.borrow(i));
        } else {
            // Right child: sibling is the stored filled subtree (left neighbor).
            let left = *pool.filled_subtrees.borrow(i);
            current_hash = hash_pair(left, current_hash);
        };
        current_index = current_index / 2;
        i = i + 1;
    };

    current_hash
}

const DEFAULT_TREE_DEPTH: u8 = 20; // 2^20 = ~1M notes