[API] Refactor Merkle tree and its gadget

[alxiong]

Current API for MT and its gadgets need some refactoring, e.g. here's the code we use in ZPrice to generate a circuit of lots of Merkle Proof verification:

    // construct circuit constraining membership proof check
    let mut circuit = PlonkCircuit::new();
    // add root as a public input
    let root_var = circuit.create_public_variable(root)?;
    for (uid, proof) in merkle_proofs.iter().enumerate() {
        let leaf_var = circuit.create_variable(leaves[uid])?;
        let proof_var = AccMemberWitnessVar::new(&mut circuit, proof)?;
        let acc_elem_var = AccElemVars {
            uid: proof_var.uid,
            elem: leaf_var,
        };

        let claimed_root_var = circuit.compute_merkle_root(acc_elem_var, &proof_var.merkle_path)?;

        // enforce matching merkle root
        circuit.equal_gate(root_var, claimed_root_var)?;
    }

    // sanity check: the circuit must be satisfied.
    assert!(circuit.check_circuit_satisfiability(&[root]).is_ok());

• mixed style of Struct::new() and field filling to instantiate structs
• numerous MerklePath, MerkleNode etc are slightly confusing
• gadget APIs are not intuitive at all

[alxiong]

we should define something like trait AppendOnlyMerkleTree (or ForgettableMerkleTree or whatever, name subject to change), so that the interface of our MT is clear from the outset.

[alxiong]

I feel that the following is more ergonomic:

use ark_ff::PrimeField;
use jf_relation::{errors::CircuitError, BoolVar, PlonkCircuit};

use crate::{
    merkle_tree::{
        prelude::{RescueMerkleTree, RescueSparseMerkleTree},
        MerkleTreeScheme, UniversalMerkleTreeScheme,
    },
    rescue::RescueParameter,
};

/// Gadgets for rescue-based merkle tree
pub trait MerkleTreeGadget<F: PrimeField + RescueParameter> {
    /// given an element and its merkle proof in a `RescueMerkleTree`,
    /// return `BoolVar` indicating the correctness of its membership proof
    fn is_member(
        &mut self,
        elem: &<RescueMerkleTree<F> as MerkleTreeScheme>::Element,
        merkle_proof: &<RescueMerkleTree<F> as MerkleTreeScheme>::MembershipProof,
        merkle_root: &<RescueMerkleTree<F> as MerkleTreeScheme>::NodeValue,
    ) -> Result<BoolVar, CircuitError>;

    /// Enforce correct `merkle_proof` for the `elem` against
    /// `expected_merkle_root`.
    fn enforce_merkle_proof(
        &mut self,
        elem: &<RescueMerkleTree<F> as MerkleTreeScheme>::Element,
        merkle_proof: &<RescueMerkleTree<F> as MerkleTreeScheme>::MembershipProof,
        expected_merkle_root: &<RescueMerkleTree<F> as MerkleTreeScheme>::NodeValue,
    ) -> Result<(), CircuitError>;
}

/// Gadget for sparse merkle tree and key-value map tree
pub trait SparseMerkleTreeGadget<F: PrimeField + RescueParameter> {
    /// checking non-membership proof
    fn is_non_member(
        &mut self,
        elem: &<RescueSparseMerkleTree<F, F> as MerkleTreeScheme>::Element,
        nonmerkle_proof: &<RescueSparseMerkleTree<F, F> as UniversalMerkleTreeScheme>::NonMembershipProof,
        merkle_root: &<RescueSparseMerkleTree<F, F> as MerkleTreeScheme>::NodeValue,
    ) -> Result<BoolVar, CircuitError>;
}

impl<F> MerkleTreeGadget<F> for PlonkCircuit<F>
where
    F: PrimeField + RescueParameter,
{
    fn is_member(
        &mut self,
        _elem: &<RescueMerkleTree<F> as MerkleTreeScheme>::Element,
        _merkle_proof: &<RescueMerkleTree<F> as MerkleTreeScheme>::MembershipProof,
        _merkle_root: &<RescueMerkleTree<F> as MerkleTreeScheme>::NodeValue,
    ) -> Result<BoolVar, CircuitError> {
        todo!();
    }

    fn enforce_merkle_proof(
        &mut self,
        elem: &<RescueMerkleTree<F> as MerkleTreeScheme>::Element,
        merkle_proof: &<RescueMerkleTree<F> as MerkleTreeScheme>::MembershipProof,
        expected_merkle_root: &<RescueMerkleTree<F> as MerkleTreeScheme>::NodeValue,
    ) -> Result<(), CircuitError> {
        let is_member = self.is_member(elem, merkle_proof, expected_merkle_root)?;
        self.enforce_true(is_member.into())
    }
}

@mrain @tessico