malt 0.1.1

use crate::tree::largest_pow2_lt;
use crate::TreeHasher;

// ---------------------------------------------------------------------------
// Proof generation algorithms (free functions).
//
// These depend only on a hasher and a leaf slice — they are decoupled
// from Log state.  Moved here from tree.rs to unify the proof concern
// in a single module (generation + verification).
// ---------------------------------------------------------------------------

/// Batch Merkle Tree Hash per formal model §2.1.
///
/// Computes the root hash of an ordered list of leaf hashes using the
/// recursive definition. Used internally for proof generation.
pub(crate) fn mth<H: TreeHasher>(hasher: &H, leaves: &[H::Digest]) -> H::Digest {
    match leaves.len() {
        0 => hasher.empty(),
        1 => leaves[0].clone(),
        n => {
            let k = largest_pow2_lt(n);
            let left = mth(hasher, &leaves[..k]);
            let right = mth(hasher, &leaves[k..]);
            hasher.node(&left, &right)
        }
    }
}

/// PATH algorithm for inclusion proofs (formal model §4.2).
///
/// Recursively computes the sibling hashes from leaf `m` to the root.
/// This is a free function — it depends only on the hasher and the leaf
/// slice, not on `Log` state.
pub(crate) fn gen_path<H: TreeHasher>(
    hasher: &H,
    m: usize,
    leaves: &[H::Digest],
) -> Vec<H::Digest> {
    let n = leaves.len();
    if n == 1 {
        // P-BASE: single leaf, no siblings needed.
        return Vec::new();
    }
    let k = largest_pow2_lt(n);
    if m < k {
        // P-LEFT: leaf is in the left (complete) subtree.
        let mut result = gen_path(hasher, m, &leaves[..k]);
        result.push(mth(hasher, &leaves[k..]));
        result
    } else {
        // P-RIGHT: leaf is in the right subtree.
        let mut result = gen_path(hasher, m - k, &leaves[k..]);
        result.push(mth(hasher, &leaves[..k]));
        result
    }
}

/// SUBPROOF algorithm for consistency proofs (formal model §5.2).
///
/// Recursively computes the intermediate hashes proving that the first
/// `m` leaves are a prefix of the `leaves` slice. This is a free
/// function — it depends only on the hasher and the leaf slice, not on
/// `Log` state.
pub(crate) fn gen_subproof<H: TreeHasher>(
    hasher: &H,
    m: usize,
    leaves: &[H::Digest],
    b: bool,
) -> Vec<H::Digest> {
    let n = leaves.len();
    if m == n {
        if b {
            // C-SAME: old tree equals current subtree, flag is true.
            return Vec::new();
        } else {
            // C-HASH: old tree equals current subtree, flag is false.
            return vec![mth(hasher, leaves)];
        }
    }
    let k = largest_pow2_lt(n);
    if m <= k {
        // C-LEFT: old size fits within left subtree.
        let mut result = gen_subproof(hasher, m, &leaves[..k], b);
        result.push(mth(hasher, &leaves[k..]));
        result
    } else {
        // C-RIGHT: old size exceeds left subtree.
        let mut result = gen_subproof(hasher, m - k, &leaves[k..], false);
        result.push(mth(hasher, &leaves[..k]));
        result
    }
}

// ---------------------------------------------------------------------------
// Proof types and verification (standalone — no Log needed).
// ---------------------------------------------------------------------------

/// Inclusion proof — proves a leaf at `index` exists in a tree of `tree_size`.
///
/// Generated by [`Log::inclusion_proof`](crate::Log::inclusion_proof) and
/// verified by [`verify_inclusion`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InclusionProof<D> {
    /// 0-based leaf index.
    pub index: u64,
    /// Size of the tree for which this proof is valid.
    pub tree_size: u64,
    /// Sibling hashes from leaf to root.
    pub path: Vec<D>,
}

/// Consistency proof — proves a tree of `old_size` is a prefix of `new_size`.
///
/// Generated by [`Log::consistency_proof`](crate::Log::consistency_proof) and
/// verified by [`verify_consistency`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConsistencyProof<D> {
    /// Size of the older tree.
    pub old_size: u64,
    /// Size of the newer tree.
    pub new_size: u64,
    /// Intermediate hashes.
    pub path: Vec<D>,
}

/// Verify an inclusion proof (formal model §4.3).
///
/// Returns `true` if the proof demonstrates that `leaf_hash` is the leaf at
/// `proof.index` in a tree of `proof.tree_size` with the given `root`.
///
/// This is a standalone verifier — no [`Log`](crate::Log) is needed.
pub fn verify_inclusion<H: TreeHasher>(
    hasher: &H,
    leaf_hash: &H::Digest,
    proof: &InclusionProof<H::Digest>,
    root: &H::Digest,
) -> bool {
    if proof.index >= proof.tree_size {
        return false;
    }

    let mut fn_ = proof.index;
    let mut sn = proof.tree_size - 1;
    let mut r = leaf_hash.clone();

    for p in &proof.path {
        if sn == 0 {
            return false;
        }
        if (fn_ & 1 == 1) || fn_ == sn {
            r = hasher.node(p, &r);
            while fn_ & 1 == 0 && fn_ != 0 {
                fn_ >>= 1;
                sn >>= 1;
            }
        } else {
            r = hasher.node(&r, p);
        }
        fn_ >>= 1;
        sn >>= 1;
    }

    sn == 0 && r == *root
}

/// Verify a consistency proof (formal model §5.3).
///
/// Returns `true` if the proof demonstrates that the tree of `proof.old_size`
/// with `old_root` is a prefix of the tree of `proof.new_size` with `new_root`.
///
/// This is a standalone verifier — no [`Log`](crate::Log) is needed.
pub fn verify_consistency<H: TreeHasher>(
    hasher: &H,
    proof: &ConsistencyProof<H::Digest>,
    old_root: &H::Digest,
    new_root: &H::Digest,
) -> bool {
    if proof.old_size == 0 || proof.old_size >= proof.new_size {
        return false;
    }

    // Build the effective path: if old_size is a power of 2, prepend old_root.
    let path: Vec<&H::Digest> = if proof.old_size.is_power_of_two() {
        std::iter::once(old_root).chain(proof.path.iter()).collect()
    } else {
        proof.path.iter().collect()
    };

    if path.is_empty() {
        return false;
    }

    let mut fn_ = proof.old_size - 1;
    let mut sn = proof.new_size - 1;

    // Strip common prefix of trailing ones.
    while fn_ & 1 == 1 {
        fn_ >>= 1;
        sn >>= 1;
    }

    let mut fr = path[0].clone();
    let mut sr = path[0].clone();

    for c in &path[1..] {
        if sn == 0 {
            return false;
        }
        if (fn_ & 1 == 1) || fn_ == sn {
            fr = hasher.node(c, &fr);
            sr = hasher.node(c, &sr);
            while fn_ & 1 == 0 && fn_ != 0 {
                fn_ >>= 1;
                sn >>= 1;
            }
        } else {
            sr = hasher.node(&sr, c);
        }
        fn_ >>= 1;
        sn >>= 1;
    }

    sn == 0 && fr == *old_root && sr == *new_root
}