arcium-client 0.13.2

//! Client-side reconstruction of the reward Merkle tree maintained by
//! `arcium_staking`'s `PrimaryStakingAccount::rewards_mt`.
//!
//! `claim_primary_stake_rewards` (and its delegated sibling) require an [`MTProof`] — a Merkle
//! opening of the `RewardClaim` leaf being claimed against the account's compressed tree. The
//! on-chain account only stores the tree *frontier* (`nodes`), its `root`, and `leaf_count`; it
//! does NOT store the historical leaves, and no event surfaces them. So to build a proof a client
//! has to reconstruct the leaves itself.
//!
//! This module mirrors the on-chain hashing exactly (see
//! `programs/arcium_staking/src/state/compressed_mt.rs`) and rebuilds the opening for a contiguous
//! *batch* of leaves appended on top of a known frontier — the situation right after a
//! `finalize_epoch_rewards` call, where the caller already knows the frontier from the
//! pre-finalize account and can compute the freshly-inserted leaf values from the pre-finalize
//! cluster state.
//!
//! Crucially the opening for an appended leaf needs only:
//!   * its **left** siblings — completed subtrees of earlier leaves, which are exactly the frontier
//!     `nodes` the on-chain `insert_node` carries forward (no historical leaf *values* required),
//!     and
//!   * its **right** siblings — subtrees of *later* leaves in the same batch (whose values we have)
//!     or the empty-subtree default hash.
//!
//! Callers must still treat the result as a hypothesis: feed each `(claim, opening)` through
//! [`verify_proof`] against the freshly-fetched on-chain `root` before submitting, so a stale read
//! or a since-changed formula can never produce an invalid on-chain transaction.

/// Height of `PrimaryStakingAccount::rewards_mt` (`REWARDS_MT_HEIGHT` on-chain).
pub const REWARDS_MT_HEIGHT: usize = 15;

/// `compressed_mt::LEAF_PREFIX` — domain separator prepended before hashing a leaf.
const LEAF_PREFIX: u8 = 0;
/// `compressed_mt::INTERMEDIATE_PREFIX` — domain separator for hashing two child nodes.
const INTERMEDIATE_PREFIX: u8 = 1;

/// A single reward-tree leaf, mirroring `arcium_staking::state::RewardClaim`. The fields are
/// borsh-serialized in declaration order (all little-endian `u64`s) to form the leaf preimage,
/// matching `hash_leaf_static`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct LeafClaim {
    pub total_rewards: u64,
    pub primary_stake: u64,
    pub delegated_stake: u64,
    pub epoch: u64,
}

/// A reconstructed claim ready to be turned into a `claim_primary_stake_rewards` instruction.
#[derive(Clone, Debug)]
pub struct ReconstructedProof {
    pub claim: LeafClaim,
    pub leaf_index: u16,
    pub opening: [[u8; 32]; REWARDS_MT_HEIGHT],
}

fn sha256(parts: &[&[u8]]) -> [u8; 32] {
    solana_program::hash::hashv(parts).to_bytes()
}

/// `sha256(LEAF_PREFIX || borsh(claim))`, identical to on-chain `hash_leaf_static` for a
/// `REWARD_BATCHING_RATE == 1` tree (one `RewardClaim` per leaf).
pub fn hash_leaf(claim: &LeafClaim) -> [u8; 32] {
    let mut buf = [0u8; 1 + 32];
    buf[0] = LEAF_PREFIX;
    buf[1..9].copy_from_slice(&claim.total_rewards.to_le_bytes());
    buf[9..17].copy_from_slice(&claim.primary_stake.to_le_bytes());
    buf[17..25].copy_from_slice(&claim.delegated_stake.to_le_bytes());
    buf[25..33].copy_from_slice(&claim.epoch.to_le_bytes());
    sha256(&[&buf])
}

/// `sha256(INTERMEDIATE_PREFIX || left || right)`, identical to on-chain
/// `hash_intermediate_static`.
fn hash_intermediate(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
    sha256(&[&[INTERMEDIATE_PREFIX], left, right])
}

/// Default (all-zero-leaf) subtree hash for each level, computed via the same recurrence as the
/// on-chain `DEFAULT_HASHES` constant: index 0 is the hash of an empty leaf, each higher index is
/// the hash of two copies of the level below. Computed at runtime (15 hashes) instead of hard-coded
/// so it can never drift from the hashing functions above.
fn default_hashes() -> [[u8; 32]; REWARDS_MT_HEIGHT] {
    let mut out = [[0u8; 32]; REWARDS_MT_HEIGHT];
    out[0] = sha256(&[&[LEAF_PREFIX], &[0u8; 32]]);
    for i in 1..REWARDS_MT_HEIGHT {
        out[i] = hash_intermediate(&out[i - 1], &out[i - 1]);
    }
    out
}

/// Whether the node on `leaf_index`'s path at `level` is a left child (mirror of
/// `CompressedMT::node_is_left`).
fn node_is_left(level: usize, leaf_index: u64) -> bool {
    (leaf_index >> level) % 2 == 0
}

/// Whether the path node at `level` will still change after inserting `leaf_index` (mirror of
/// `CompressedMT::inner_node_will_still_change`).
fn inner_node_will_still_change(level: usize, leaf_index: u64) -> bool {
    if leaf_index == 0 {
        return true;
    }
    let chunk_size: u64 = 1 << level;
    (leaf_index + 1) % chunk_size != 0
}

/// Advance the frontier (`CompressedMT::nodes`) by inserting `leaf` at position `leaf_count`,
/// returning the new frontier. A faithful port of the `new_nodes` bookkeeping in
/// `CompressedMT::insert_node` (we don't track the root here — callers verify against the on-chain
/// root instead).
fn advance_frontier(
    nodes: &[[u8; 32]; REWARDS_MT_HEIGHT],
    leaf_count: u64,
    leaf: [u8; 32],
    defaults: &[[u8; 32]; REWARDS_MT_HEIGHT],
) -> [[u8; 32]; REWARDS_MT_HEIGHT] {
    let mut inter = *nodes;
    let mut new_nodes = *nodes;
    for level in 0..REWARDS_MT_HEIGHT {
        if level == 0 {
            inter[0] = leaf;
            new_nodes[0] = leaf;
        } else {
            let (left, right) = if node_is_left(level - 1, leaf_count) {
                // On the left edge the right sibling isn't initialised yet → default.
                (inter[level - 1], defaults[level - 1])
            } else {
                (nodes[level - 1], inter[level - 1])
            };
            inter[level] = hash_intermediate(&left, &right);
            if !inner_node_will_still_change(level, leaf_count) {
                new_nodes[level] = inter[level];
            }
        }
    }
    new_nodes
}

/// Hash of the subtree rooted at (`level`, `index`), where leaf hashes are drawn from `batch`
/// (covering global leaf indices `[leaf_count_pre, leaf_count_pre + batch.len())`) and any index
/// outside that range is treated as an empty leaf. Only ever queried for *right* siblings of batch
/// leaves, whose covered indices are strictly greater than the leaf in question (hence ≥
/// `leaf_count_pre`), so we never need a pre-existing leaf's value.
fn subtree_hash(
    level: usize,
    index: u64,
    batch: &[[u8; 32]],
    leaf_count_pre: u64,
    defaults: &[[u8; 32]; REWARDS_MT_HEIGHT],
) -> [u8; 32] {
    if level == 0 {
        if index >= leaf_count_pre && index < leaf_count_pre + batch.len() as u64 {
            batch[(index - leaf_count_pre) as usize]
        } else {
            defaults[0]
        }
    } else {
        let left = subtree_hash(level - 1, 2 * index, batch, leaf_count_pre, defaults);
        let right = subtree_hash(level - 1, 2 * index + 1, batch, leaf_count_pre, defaults);
        hash_intermediate(&left, &right)
    }
}

/// Reconstruct the leaf index and Merkle opening for each leaf in `batch`, assuming `batch` was
/// appended (in order) on top of a tree whose frontier is `frontier` and which already held
/// `leaf_count_pre` leaves.
///
/// `frontier` is the on-chain `CompressedMT::nodes` array read *before* the inserts. For a tree
/// reconstructed from scratch (`leaf_count_pre == 0`) pass [`empty_frontier`].
pub fn reconstruct_batch(
    frontier: &[[u8; 32]; REWARDS_MT_HEIGHT],
    leaf_count_pre: u16,
    batch: &[LeafClaim],
) -> Vec<ReconstructedProof> {
    let defaults = default_hashes();
    let leaf_count_pre = leaf_count_pre as u64;
    let leaf_hashes: Vec<[u8; 32]> = batch.iter().map(hash_leaf).collect();

    let mut frontier = *frontier;
    let mut out = Vec::with_capacity(batch.len());
    for (k, claim) in batch.iter().enumerate() {
        let j = leaf_count_pre + k as u64;
        let mut opening = [[0u8; 32]; REWARDS_MT_HEIGHT];
        for (level, slot) in opening.iter_mut().enumerate() {
            *slot = if node_is_left(level, j) {
                // Right sibling: a subtree of later batch leaves (or empty).
                subtree_hash(
                    level,
                    (j >> level) + 1,
                    &leaf_hashes,
                    leaf_count_pre,
                    &defaults,
                )
            } else {
                // Left sibling: a completed subtree carried in the frontier.
                frontier[level]
            };
        }
        out.push(ReconstructedProof {
            claim: *claim,
            leaf_index: j as u16,
            opening,
        });
        frontier = advance_frontier(&frontier, j, leaf_hashes[k], &defaults);
    }
    out
}

/// The frontier of an empty `CompressedMT` (`CompressedMT::new`): the per-level default hashes.
pub fn empty_frontier() -> [[u8; 32]; REWARDS_MT_HEIGHT] {
    default_hashes()
}

/// Recompute the root implied by a leaf + opening and compare it to `root`. A faithful port of
/// on-chain `CompressedMT::verify_proof` — the same computation the program runs — so a `true`
/// result means the corresponding `claim_*_rewards` instruction will pass its proof check.
pub fn verify_proof(
    claim: &LeafClaim,
    leaf_index: u16,
    opening: &[[u8; 32]; REWARDS_MT_HEIGHT],
    root: &[u8; 32],
) -> bool {
    let mut acc = hash_leaf(claim);
    let li = leaf_index as u64;
    for (i, sibling) in opening.iter().enumerate() {
        acc = if node_is_left(i, li) {
            hash_intermediate(&acc, sibling)
        } else {
            hash_intermediate(sibling, &acc)
        };
    }
    acc == *root
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The first on-chain `DEFAULT_HASHES` entry (`compressed_mt::DEFAULT_HASHES[0]`), used to
    /// pin our runtime-computed defaults to the program's hard-coded constant.
    const ON_CHAIN_DEFAULT_0: [u8; 32] = [
        127, 156, 158, 49, 172, 130, 86, 202, 47, 37, 133, 131, 223, 38, 45, 188, 125, 111, 104,
        242, 160, 48, 67, 213, 201, 154, 74, 229, 167, 57, 108, 233,
    ];
    /// `compressed_mt::DEFAULT_HASHES[14]` (the top level), pinning the full recurrence.
    const ON_CHAIN_DEFAULT_14: [u8; 32] = [
        70, 11, 7, 114, 152, 108, 34, 153, 83, 53, 232, 228, 247, 56, 169, 118, 164, 180, 150, 145,
        64, 225, 2, 196, 185, 234, 128, 131, 99, 111, 99, 177,
    ];

    fn claim(seed: u64) -> LeafClaim {
        LeafClaim {
            total_rewards: seed,
            primary_stake: seed.wrapping_mul(3),
            delegated_stake: seed.wrapping_add(7),
            epoch: seed,
        }
    }

    #[test]
    fn default_hashes_match_on_chain() {
        let d = default_hashes();
        assert_eq!(d[0], ON_CHAIN_DEFAULT_0);
        assert_eq!(d[REWARDS_MT_HEIGHT - 1], ON_CHAIN_DEFAULT_14);
    }

    /// Independent, allocation-heavy reference tree (mirror of the on-chain test `TestMerkleTree`)
    /// that stores every node so we can derive the root and any opening by brute force, to check
    /// the frontier-only reconstruction against.
    struct RefTree {
        levels: Vec<Vec<[u8; 32]>>,
        defaults: [[u8; 32]; REWARDS_MT_HEIGHT],
    }

    impl RefTree {
        fn new() -> Self {
            RefTree {
                levels: vec![Vec::new(); REWARDS_MT_HEIGHT],
                defaults: default_hashes(),
            }
        }

        fn node(&self, level: usize, index: usize) -> [u8; 32] {
            self.levels[level]
                .get(index)
                .copied()
                .unwrap_or(self.defaults[level])
        }

        fn insert(&mut self, claim: &LeafClaim) {
            self.levels[0].push(hash_leaf(claim));
            let mut idx = self.levels[0].len() - 1;
            for level in 0..REWARDS_MT_HEIGHT - 1 {
                let parent = idx / 2;
                let left = self.node(level, parent * 2);
                let right = self.node(level, parent * 2 + 1);
                let h = hash_intermediate(&left, &right);
                if parent < self.levels[level + 1].len() {
                    self.levels[level + 1][parent] = h;
                } else {
                    self.levels[level + 1].push(h);
                }
                idx = parent;
            }
        }

        fn root(&self) -> [u8; 32] {
            let top = self.node(REWARDS_MT_HEIGHT - 1, 0);
            let top_sib = self.node(REWARDS_MT_HEIGHT - 1, 1);
            hash_intermediate(&top, &top_sib)
        }

        fn opening(&self, leaf_index: usize) -> [[u8; 32]; REWARDS_MT_HEIGHT] {
            let mut opening = [[0u8; 32]; REWARDS_MT_HEIGHT];
            let mut idx = leaf_index;
            for (level, slot) in opening.iter_mut().enumerate() {
                let sib = if idx % 2 == 0 { idx + 1 } else { idx - 1 };
                *slot = self.node(level, sib);
                idx /= 2;
            }
            opening
        }
    }

    #[test]
    fn reconstruct_from_empty_matches_reference_and_verifies() {
        for n in 1..=40usize {
            let batch: Vec<LeafClaim> = (0..n as u64).map(claim).collect();
            let mut reference = RefTree::new();
            for c in &batch {
                reference.insert(c);
            }
            let root = reference.root();

            let proofs = reconstruct_batch(&empty_frontier(), 0, &batch);
            assert_eq!(proofs.len(), n);
            for (i, p) in proofs.iter().enumerate() {
                assert_eq!(p.leaf_index as usize, i);
                assert_eq!(
                    p.opening,
                    reference.opening(i),
                    "opening mismatch n={n} i={i}"
                );
                assert!(
                    verify_proof(&p.claim, p.leaf_index, &p.opening, &root),
                    "verify failed n={n} i={i}"
                );
            }
        }
    }

    #[test]
    fn reconstruct_appended_batch_matches_full_tree() {
        // Build a full reference tree, then reconstruct only the *trailing* batch using the
        // frontier captured after the first `pre` leaves — the exact post-finalize scenario.
        for pre in 0..20usize {
            for add in 1..=12usize {
                let all: Vec<LeafClaim> = (0..(pre + add) as u64).map(claim).collect();

                let mut reference = RefTree::new();
                for c in &all {
                    reference.insert(c);
                }
                let root = reference.root();

                // Frontier after the first `pre` inserts, derived via advance_frontier.
                let defaults = default_hashes();
                let mut frontier = empty_frontier();
                for (i, c) in all.iter().take(pre).enumerate() {
                    frontier = advance_frontier(&frontier, i as u64, hash_leaf(c), &defaults);
                }

                let batch = &all[pre..];
                let proofs = reconstruct_batch(&frontier, pre as u16, batch);
                assert_eq!(proofs.len(), add);
                for (k, p) in proofs.iter().enumerate() {
                    let global = pre + k;
                    assert_eq!(p.leaf_index as usize, global);
                    assert!(
                        verify_proof(&p.claim, p.leaf_index, &p.opening, &root),
                        "verify failed pre={pre} add={add} k={k}"
                    );
                }
            }
        }
    }
}