MainPod circuit skeleton (#151)

* feat: add MainPod circuit skeleton

* feat: use ValueTarget in mt, verify SignedPod type

* wip

* feat: match structure with mock

* apply feedback from @arnaucube

* add 2 operations

* fix test compilation

* Add missing todo

src/backends/plonky2/common.rs

@@ -0,0 +1,110 @@
+//! Common functionality to build Pod circuits with plonky2
+
+use crate::middleware::STATEMENT_ARG_F_LEN;
+use crate::middleware::{Params, Value, HASH_SIZE, VALUE_SIZE};
+use plonky2::field::extension::Extendable;
+use plonky2::field::types::PrimeField64;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::target::{BoolTarget, Target};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+
+#[derive(Copy, Clone)]
+pub struct ValueTarget {
+    pub elements: [Target; VALUE_SIZE],
+}
+
+#[derive(Clone)]
+pub struct StatementTarget {
+    pub code: [Target; HASH_SIZE + 2],
+    pub args: Vec<[Target; STATEMENT_ARG_F_LEN]>,
+}
+
+impl StatementTarget {
+    pub fn to_flattened(&self) -> Vec<Target> {
+        self.code
+            .iter()
+            .chain(self.args.iter().flatten())
+            .cloned()
+            .collect()
+    }
+}
+
+// TODO: Implement Operation::to_field to determine the size of each element
+#[derive(Clone)]
+pub struct OperationTarget {
+    pub code: [Target; 6],                        // TODO: Figure out the length
+    pub args: Vec<[Target; STATEMENT_ARG_F_LEN]>, // TODO: Figure out the length
+}
+
+pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
+    fn connect_values(&mut self, x: ValueTarget, y: ValueTarget);
+    fn connect_slice(&mut self, xs: &[Target], ys: &[Target]);
+    fn add_virtual_value(&mut self) -> ValueTarget;
+    fn add_virtual_statement(&mut self, params: &Params) -> StatementTarget;
+    fn add_virtual_operation(&mut self, params: &Params) -> OperationTarget;
+    fn select_value(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) -> ValueTarget;
+    fn select_bool(&mut self, b: BoolTarget, x: BoolTarget, y: BoolTarget) -> BoolTarget;
+    fn constant_value(&mut self, v: Value) -> ValueTarget;
+    fn is_equal_slice(&mut self, xs: &[Target], ys: &[Target]) -> BoolTarget;
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderPod<F, D>
+    for CircuitBuilder<F, D>
+{
+    fn connect_slice(&mut self, xs: &[Target], ys: &[Target]) {
+        assert_eq!(xs.len(), ys.len());
+        for (x, y) in xs.iter().zip(ys.iter()) {
+            self.connect(*x, *y);
+        }
+    }
+
+    fn connect_values(&mut self, x: ValueTarget, y: ValueTarget) {
+        self.connect_slice(&x.elements, &y.elements);
+    }
+
+    fn add_virtual_value(&mut self) -> ValueTarget {
+        ValueTarget {
+            elements: self.add_virtual_target_arr(),
+        }
+    }
+
+    fn add_virtual_statement(&mut self, params: &Params) -> StatementTarget {
+        StatementTarget {
+            code: self.add_virtual_target_arr::<6>(),
+            args: (0..params.max_statement_args)
+                .map(|_| self.add_virtual_target_arr::<STATEMENT_ARG_F_LEN>())
+                .collect(),
+        }
+    }
+
+    fn add_virtual_operation(&mut self, params: &Params) -> OperationTarget {
+        todo!()
+    }
+
+    fn select_value(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) -> ValueTarget {
+        ValueTarget {
+            elements: std::array::from_fn(|i| self.select(b, x.elements[i], y.elements[i])),
+        }
+    }
+
+    fn select_bool(&mut self, b: BoolTarget, x: BoolTarget, y: BoolTarget) -> BoolTarget {
+        BoolTarget::new_unsafe(self.select(b, x.target, y.target))
+    }
+
+    fn constant_value(&mut self, v: Value) -> ValueTarget {
+        ValueTarget {
+            elements: std::array::from_fn(|i| {
+                self.constant(F::from_noncanonical_u64(v.0[i].to_noncanonical_u64()))
+            }),
+        }
+    }
+
+    fn is_equal_slice(&mut self, xs: &[Target], ys: &[Target]) -> BoolTarget {
+        assert_eq!(xs.len(), ys.len());
+        let init = self._true();
+        xs.iter().zip(ys.iter()).fold(init, |ok, (x, y)| {
+            let is_eq = self.is_equal(*x, *y);
+            self.and(ok, is_eq)
+        })
+    }
+}

src/backends/plonky2/main.rs

@@ -0,0 +1,340 @@
+use crate::backends::plonky2::basetypes::{Hash, Value, D, EMPTY_HASH, EMPTY_VALUE, F, VALUE_SIZE};
+use crate::backends::plonky2::common::{
+    CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget,
+};
+use crate::backends::plonky2::primitives::merkletree::MerkleProofExistenceCircuit;
+use crate::backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree};
+use crate::middleware::{
+    hash_str, AnchoredKey, NativeOperation, NativePredicate, Operation, Params, PodType, Predicate,
+    Statement, StatementArg, ToFields, KEY_TYPE, SELF, STATEMENT_ARG_F_LEN,
+};
+use anyhow::Result;
+use itertools::Itertools;
+use plonky2::{
+    field::types::Field,
+    hash::{
+        hash_types::{HashOut, HashOutTarget},
+        poseidon::PoseidonHash,
+    },
+    iop::{
+        target::{BoolTarget, Target},
+        witness::{PartialWitness, WitnessWrite},
+    },
+    plonk::circuit_builder::CircuitBuilder,
+};
+use std::collections::HashMap;
+
+/// MerkleTree Max Depth
+const MD: usize = 32;
+
+//
+// SignedPod verification
+//
+
+struct SignedPodVerifyGate {
+    params: Params,
+}
+
+impl SignedPodVerifyGate {
+    fn eval(&self, builder: &mut CircuitBuilder<F, D>) -> Result<SignedPodVerifyTarget> {
+        // 2. Verify id
+        let id = builder.add_virtual_hash();
+        let mut mt_proofs = Vec::new();
+        for _ in 0..self.params.max_signed_pod_values {
+            let mt_proof = MerkleProofExistenceCircuit::<MD>::add_targets(builder)?;
+            builder.connect_hashes(id, mt_proof.root);
+            mt_proofs.push(mt_proof);
+        }
+
+        // 1. Verify type
+        let type_mt_proof = &mt_proofs[0];
+        let key_type = builder.constant_value(hash_str(KEY_TYPE).into());
+        builder.connect_values(type_mt_proof.key, key_type);
+        let value_type = builder.constant_value(Value::from(PodType::MockSigned));
+        builder.connect_values(type_mt_proof.value, value_type);
+
+        // 3. TODO: Verify signature
+
+        Ok(SignedPodVerifyTarget {
+            params: self.params.clone(),
+            id,
+            mt_proofs,
+        })
+    }
+}
+
+struct SignedPodVerifyTarget {
+    params: Params,
+    id: HashOutTarget,
+    // The KEY_TYPE entry must be the first one
+    // The KEY_SIGNER entry must be the second one
+    mt_proofs: Vec<MerkleProofExistenceCircuit<MD>>,
+}
+
+struct SignedPodVerifyInput {
+    kvs: HashMap<Value, Value>,
+}
+
+impl SignedPodVerifyTarget {
+    fn kvs(&self) -> Vec<(ValueTarget, ValueTarget)> {
+        let mut kvs = Vec::new();
+        for mt_proof in &self.mt_proofs {
+            kvs.push((mt_proof.key, mt_proof.value));
+        }
+        // TODO: when the slot is unused, do we force the kv to be (EMPTY, EMPTY), and then from
+        // it get a ValueOf((id, EMPTY), EMPTY)?  Or should we keep some boolean flags for unused
+        // slots and translate them to Statement::None instead?
+        kvs
+    }
+
+    fn pub_statements(&self) -> Vec<StatementTarget> {
+        // TODO: Here we need to use the self.id in the ValueOf statements
+        todo!()
+    }
+
+    fn set_targets(&self, pw: &mut PartialWitness<F>, input: &SignedPodVerifyInput) -> Result<()> {
+        assert!(input.kvs.len() <= self.params.max_signed_pod_values);
+        let tree = MerkleTree::new(MD, &input.kvs)?;
+        for (i, (k, v)) in input.kvs.iter().sorted_by_key(|kv| kv.0).enumerate() {
+            let (_, proof) = tree.prove(&k)?;
+            self.mt_proofs[i].set_targets(pw, tree.root(), proof, *k, *v)?;
+        }
+        // Padding
+        for i in input.kvs.len()..self.params.max_signed_pod_values {
+            // TODO: We need to disable the proofs for the unused slots.  We could add a flag
+            // "enable" to the MerkleTree proof circuit that skips the verification when false.
+            // self.mt_proofs[i].set_targets(pw, false, EMPTY_HASH, proof, *k, *v)?;
+        }
+        Ok(())
+    }
+}
+
+//
+// MainPod verification
+//
+
+struct OperationVerifyGate {
+    params: Params,
+}
+
+impl OperationVerifyGate {
+    fn eval(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        op: &OperationTarget,
+        prev_statements: &[StatementTarget],
+    ) -> Result<OperationVerifyTarget> {
+        // Verify that the operation `op` correctly generates the statement `st`.  The operation
+        // can reference any of the `prev_statements`.
+        // The verification may require aux data which needs to be stored in the
+        // `OperationVerifyTarget` so that we can set during witness generation.
+
+        // TODO: Figure out the right encoding of op.code
+        let op_none = builder.constant(F::from_canonical_u64(NativeOperation::None as u64));
+        let is_none = builder.is_equal(op.code[0], op_none);
+        let op_new_entry =
+            builder.constant(F::from_canonical_u64(NativeOperation::NewEntry as u64));
+        let is_new_entry = builder.is_equal(op.code[0], op_new_entry);
+        let op_copy_statement =
+            builder.constant(F::from_canonical_u64(NativeOperation::CopyStatement as u64));
+        let is_copy_statement = builder.is_equal(op.code[0], op_copy_statement);
+        let op_eq_from_entries = builder.constant(F::from_canonical_u64(
+            NativeOperation::EqualFromEntries as u64,
+        ));
+        let is_eq_from_entries = builder.is_equal(op.code[0], op_eq_from_entries);
+        let op_gt_from_entries =
+            builder.constant(F::from_canonical_u64(NativeOperation::GtFromEntries as u64));
+        let is_gt_from_entries = builder.is_equal(op.code[0], op_gt_from_entries);
+        let op_lt_from_entries =
+            builder.constant(F::from_canonical_u64(NativeOperation::LtFromEntries as u64));
+        let is_lt_from_entries = builder.is_equal(op.code[0], op_lt_from_entries);
+        let op_contains_from_entries = builder.constant(F::from_canonical_u64(
+            NativeOperation::ContainsFromEntries as u64,
+        ));
+        let is_contains_from_entries = builder.is_equal(op.code[0], op_contains_from_entries);
+
+        let ok = builder._true();
+        let none_ok = self.eval_none(builder, st, op);
+        let ok = builder.select_bool(is_none, none_ok, ok);
+        let new_entry_ok = self.eval_new_entry(builder, st, op);
+        let ok = builder.select_bool(is_new_entry, new_entry_ok, ok);
+
+        let _true = builder._true();
+        builder.connect(ok.target, _true.target);
+
+        todo!()
+    }
+
+    fn eval_none(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        _op: &OperationTarget,
+    ) -> BoolTarget {
+        let expected_statement_flattened =
+            builder.constants(&Statement::None.to_fields(&self.params));
+        builder.is_equal_slice(&st.to_flattened(), &expected_statement_flattened)
+    }
+
+    fn eval_new_entry(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        _op: &OperationTarget,
+    ) -> BoolTarget {
+        let value_of_st = &Statement::ValueOf(AnchoredKey(SELF, EMPTY_HASH), EMPTY_VALUE);
+        let expected_code =
+            builder.constants(&Predicate::Native(NativePredicate::ValueOf).to_fields(&self.params));
+        let code_ok = builder.is_equal_slice(&st.code, &expected_code);
+        let expected_arg_prefix = builder.constants(
+            &StatementArg::Key(AnchoredKey(SELF, EMPTY_HASH)).to_fields(&self.params)[..VALUE_SIZE],
+        );
+        let arg_prefix_ok = builder.is_equal_slice(&st.args[0][..VALUE_SIZE], &expected_arg_prefix);
+        builder.and(code_ok, arg_prefix_ok)
+    }
+}
+
+struct OperationVerifyTarget {
+    // TODO
+}
+
+struct OperationVerifyInputs {
+    // TODO
+}
+
+impl OperationVerifyTarget {
+    fn set_targets(&self, pw: &mut PartialWitness<F>, input: &OperationVerifyInputs) -> Result<()> {
+        todo!()
+    }
+}
+
+struct MainPodVerifyGate {
+    params: Params,
+}
+
+impl MainPodVerifyGate {
+    fn eval(&self, builder: &mut CircuitBuilder<F, D>) -> Result<MainPodVerifyTarget> {
+        let params = &self.params;
+        // 1. Verify all input signed pods
+        let mut signed_pods = Vec::new();
+        for _ in 0..params.max_input_signed_pods {
+            let signed_pod = SignedPodVerifyGate {
+                params: params.clone(),
+            }
+            .eval(builder)?;
+            signed_pods.push(signed_pod);
+        }
+
+        // Build the statement array
+        let mut statements = Vec::new();
+        for signed_pod in &signed_pods {
+            statements.extend_from_slice(signed_pod.pub_statements().as_slice());
+        }
+
+        // Add the input (private and public) statements and corresponding operations
+        let mut operations = Vec::new();
+        let input_statements_offset = statements.len();
+        for _ in 0..params.max_statements {
+            statements.push(builder.add_virtual_statement(params));
+            operations.push(builder.add_virtual_operation(params));
+        }
+
+        let input_statements = &statements[input_statements_offset..];
+        let pub_statements = &input_statements[statements.len() - params.max_public_statements..];
+
+        // 2. Calculate the Pod Id from the public statements
+        let pub_statements_flattened = pub_statements
+            .iter()
+            .map(|s| s.code.iter().chain(s.args.iter().flatten()))
+            .flatten()
+            .cloned()
+            .collect();
+        let id = builder.hash_n_to_hash_no_pad::<PoseidonHash>(pub_statements_flattened);
+
+        // 3. TODO check that all `input_statements` of type `ValueOf` with origin=SELF have unique
+        //    keys (no duplicates)
+
+        // 4. Verify type
+        let type_statement = &pub_statements[0];
+        // TODO: Store this hash in a global static with lazy init so that we don't have to
+        // compute it every time.
+        let key_type = hash_str(KEY_TYPE);
+        let expected_type_statement_flattened = builder.constants(
+            &Statement::ValueOf(AnchoredKey(SELF, key_type), Value::from(PodType::MockMain))
+                .to_fields(params),
+        );
+        builder.connect_slice(
+            &type_statement.to_flattened(),
+            &expected_type_statement_flattened,
+        );
+
+        // 5. Verify input statements
+        let mut op_verifications = Vec::new();
+        for (i, (st, op)) in input_statements.iter().zip(operations.iter()).enumerate() {
+            let prev_statements = &statements[..input_statements_offset + i - 1];
+            let op_verification = OperationVerifyGate {
+                params: params.clone(),
+            }
+            .eval(builder, st, op, prev_statements)?;
+            op_verifications.push(op_verification);
+        }
+
+        Ok(MainPodVerifyTarget {
+            params: params.clone(),
+            id,
+            signed_pods,
+            statements: input_statements.to_vec(),
+            operations,
+            op_verifications,
+        })
+    }
+}
+
+struct MainPodVerifyTarget {
+    params: Params,
+    id: HashOutTarget,
+    signed_pods: Vec<SignedPodVerifyTarget>,
+    // The KEY_TYPE statement must be the first public one
+    statements: Vec<StatementTarget>,
+    operations: Vec<OperationTarget>,
+    op_verifications: Vec<OperationVerifyTarget>,
+}
+
+struct MainPodVerifyInput {
+    signed_pods: Vec<SignedPodVerifyInput>,
+}
+
+impl MainPodVerifyTarget {
+    fn set_targets(&self, pw: &mut PartialWitness<F>, input: &MainPodVerifyInput) -> Result<()> {
+        assert!(input.signed_pods.len() <= self.params.max_input_signed_pods);
+        for (i, signed_pod) in input.signed_pods.iter().enumerate() {
+            self.signed_pods[i].set_targets(pw, signed_pod)?;
+        }
+        // Padding
+        for i in input.signed_pods.len()..self.params.max_input_signed_pods {
+            // TODO: We need to disable the verification for the unused slots.
+            // self.signed_pods[i].set_targets(pw, signed_pod)?;
+        }
+        // TODO: set_targets for:
+        // - statements
+        // - operations
+        // - op_verifications
+        Ok(())
+    }
+}
+
+struct MainPodVerifyCircuit {
+    params: Params,
+}
+
+impl MainPodVerifyCircuit {
+    fn eval(&self, builder: &mut CircuitBuilder<F, D>) -> Result<MainPodVerifyTarget> {
+        let main_pod = MainPodVerifyGate {
+            params: self.params.clone(),
+        }
+        .eval(builder)?;
+        builder.register_public_inputs(&main_pod.id.elements);
+        Ok(main_pod)
+    }
+}

src/backends/plonky2/mock_main/mod.rs

@@ -371,10 +371,12 @@ pub fn hash_statements(statements: &[Statement], _params: &Params) -> middleware
 
 impl Pod for MockMainPod {
     fn verify(&self) -> bool {
+        // 1. TODO: Verify input pods
+
         let input_statement_offset = self.offset_input_statements();
         // get the input_statements from the self.statements
         let input_statements = &self.statements[input_statement_offset..];
-        // get the id out of the public statements, and ensure it is equal to self.id
+        // 2. get the id out of the public statements, and ensure it is equal to self.id
         let ids_match = self.id == PodId(hash_statements(&self.public_statements, &self.params));
         // find a ValueOf statement from the public statements with key=KEY_TYPE and check that the
         // value is PodType::MockMainPod
@@ -391,7 +393,7 @@ impl Pod for MockMainPod {
                     }
             })
             .is_some();
-        // check that all `input_statements` of type `ValueOf` with origin=SELF have unique keys
+        // 3. check that all `input_statements` of type `ValueOf` with origin=SELF have unique keys
         // (no duplicates)
         // TODO: Instead of doing this, do a uniqueness check when verifying the output of a
         // `NewValue` operation.
@@ -421,7 +423,9 @@ impl Pod for MockMainPod {
                 .collect::<Vec<_>>();
             !(0..key_id_pairs.len() - 1).any(|i| key_id_pairs[i + 1..].contains(&key_id_pairs[i]))
         };
-        // verify that all `input_statements` are correctly generated
+        // 4. TODO: Verify type
+
+        // 5. verify that all `input_statements` are correctly generated
         // by `self.operations` (where each operation can only access previous statements)
         let statement_check = input_statements
             .iter()

src/backends/plonky2/mock_signed.rs

@@ -46,7 +46,7 @@ pub struct MockSignedPod {
 
 impl Pod for MockSignedPod {
     fn verify(&self) -> bool {
-        // Verify type
+        // 1. Verify type
         let value_at_type = match self.dict.get(&hash_str(KEY_TYPE).into()) {
             Ok(v) => v,
             Err(_) => return false,
@@ -55,7 +55,7 @@ impl Pod for MockSignedPod {
             return false;
         }
 
-        // Verify id
+        // 2. Verify id
         let mt = match MerkleTree::new(
             MAX_DEPTH,
             &self
@@ -72,7 +72,7 @@ impl Pod for MockSignedPod {
             return false;
         }
 
-        // Verify signature
+        // 3. Verify signature
         let pk_hash = match self.dict.get(&hash_str(KEY_SIGNER).into()) {
             Ok(v) => v,
             Err(_) => return false,

src/backends/plonky2/mod.rs

@@ -1,4 +1,6 @@
 pub mod basetypes;
+pub mod common;
+pub mod main;
 pub mod mock_main;
 pub mod mock_signed;
 pub mod primitives;

src/backends/plonky2/primitives/merkletree_circuit.rs

@@ -21,8 +21,12 @@ use plonky2::{
     },
     plonk::circuit_builder::CircuitBuilder,
 };
+use std::iter;
 
 use crate::backends::plonky2::basetypes::{Hash, Value, D, EMPTY_HASH, EMPTY_VALUE, F, VALUE_SIZE};
+use crate::backends::plonky2::common::{
+    CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget,
+};
 use crate::backends::plonky2::primitives::merkletree::MerkleProof;
 
 /// `MerkleProofCircuit` allows to verify both proofs of existence and proofs
@@ -30,30 +34,30 @@ use crate::backends::plonky2::primitives::merkletree::MerkleProof;
 /// If only proofs of existence are needed, use `MerkleProofExistenceCircuit`,
 /// which requires less amount of constraints.
 pub struct MerkleProofCircuit<const MAX_DEPTH: usize> {
-    root: HashOutTarget,
-    key: Vec<Target>,
-    value: Vec<Target>,
-    existence: BoolTarget,
-    siblings: Vec<HashOutTarget>,
-    case_ii_selector: BoolTarget, // for case ii)
-    other_key: Vec<Target>,
-    other_value: Vec<Target>,
+    pub root: HashOutTarget,
+    pub key: ValueTarget,
+    pub value: ValueTarget,
+    pub existence: BoolTarget,
+    pub siblings: Vec<HashOutTarget>,
+    pub case_ii_selector: BoolTarget, // for case ii)
+    pub other_key: ValueTarget,
+    pub other_value: ValueTarget,
 }
 
 impl<const MAX_DEPTH: usize> MerkleProofCircuit<MAX_DEPTH> {
     /// creates the targets and defines the logic of the circuit
     pub fn add_targets(builder: &mut CircuitBuilder<F, D>) -> Result<Self> {
         // create the targets
-        let key = builder.add_virtual_targets(VALUE_SIZE);
-        let value = builder.add_virtual_targets(VALUE_SIZE);
+        let key = builder.add_virtual_value();
+        let value = builder.add_virtual_value();
         // from proof struct:
         let existence = builder.add_virtual_bool_target_safe();
         // siblings are padded till MAX_DEPTH length
         let siblings = builder.add_virtual_hashes(MAX_DEPTH);
 
         let case_ii_selector = builder.add_virtual_bool_target_safe();
-        let other_key = builder.add_virtual_targets(VALUE_SIZE);
-        let other_value = builder.add_virtual_targets(VALUE_SIZE);
+        let other_key = builder.add_virtual_value();
+        let other_value = builder.add_virtual_value();
 
         // We have 3 cases for when computing the Leaf's hash:
         // - existence: leaf contains the given key & value
@@ -87,12 +91,8 @@ impl<const MAX_DEPTH: usize> MerkleProofCircuit<MAX_DEPTH> {
         // of existence or of non-existence, ie:
         // k = key * existence + other_key * (1-existence)
         // v = value * existence + other_value * (1-existence)
-        let k: Vec<Target> = (0..4)
-            .map(|j| builder.select(existence, key[j], other_key[j]))
-            .collect();
-        let v: Vec<Target> = (0..4)
-            .map(|j| builder.select(existence, value[j], other_value[j]))
-            .collect();
+        let k = builder.select_value(existence, key, other_key);
+        let v = builder.select_value(existence, value, other_value);
 
         // get leaf's hash for the selected k & v
         let h = kv_hash_target(builder, &k, &v);
@@ -139,8 +139,8 @@ impl<const MAX_DEPTH: usize> MerkleProofCircuit<MAX_DEPTH> {
         value: Value,
     ) -> Result<()> {
         pw.set_hash_target(self.root, HashOut::from_vec(root.0.to_vec()))?;
-        pw.set_target_arr(&self.key, &key.0)?;
-        pw.set_target_arr(&self.value, &value.0)?;
+        pw.set_target_arr(&self.key.elements, &key.0)?;
+        pw.set_target_arr(&self.value.elements, &value.0)?;
         pw.set_bool_target(self.existence, existence)?;
 
         // pad siblings with zeros to length MAX_DEPTH
@@ -156,14 +156,14 @@ impl<const MAX_DEPTH: usize> MerkleProofCircuit<MAX_DEPTH> {
             Some((k, v)) if !existence => {
                 // non-existence case ii) expected leaf does exist but it has a different key
                 pw.set_bool_target(self.case_ii_selector, true)?;
-                pw.set_target_arr(&self.other_key, &k.0)?;
-                pw.set_target_arr(&self.other_value, &v.0)?;
+                pw.set_target_arr(&self.other_key.elements, &k.0)?;
+                pw.set_target_arr(&self.other_value.elements, &v.0)?;
             }
             _ => {
                 // existence & non-existence case i) expected leaf does not exist
                 pw.set_bool_target(self.case_ii_selector, false)?;
-                pw.set_target_arr(&self.other_key, &EMPTY_VALUE.0)?;
-                pw.set_target_arr(&self.other_value, &EMPTY_VALUE.0)?;
+                pw.set_target_arr(&self.other_key.elements, &EMPTY_VALUE.0)?;
+                pw.set_target_arr(&self.other_value.elements, &EMPTY_VALUE.0)?;
             }
         }
 
@@ -174,18 +174,18 @@ impl<const MAX_DEPTH: usize> MerkleProofCircuit<MAX_DEPTH> {
 /// `MerkleProofExistenceCircuit` allows to verify proofs of existence only. If
 /// proofs of non-existence are needed, use `MerkleProofCircuit`.
 pub struct MerkleProofExistenceCircuit<const MAX_DEPTH: usize> {
-    root: HashOutTarget,
-    key: Vec<Target>,
-    value: Vec<Target>,
-    siblings: Vec<HashOutTarget>,
+    pub root: HashOutTarget,
+    pub key: ValueTarget,
+    pub value: ValueTarget,
+    pub siblings: Vec<HashOutTarget>,
 }
 
 impl<const MAX_DEPTH: usize> MerkleProofExistenceCircuit<MAX_DEPTH> {
     /// creates the targets and defines the logic of the circuit
     pub fn add_targets(builder: &mut CircuitBuilder<F, D>) -> Result<Self> {
         // create the targets
-        let key = builder.add_virtual_targets(VALUE_SIZE);
-        let value = builder.add_virtual_targets(VALUE_SIZE);
+        let key = builder.add_virtual_value();
+        let value = builder.add_virtual_value();
         // siblings are padded till MAX_DEPTH length
         let siblings = builder.add_virtual_hashes(MAX_DEPTH);
 
@@ -218,8 +218,8 @@ impl<const MAX_DEPTH: usize> MerkleProofExistenceCircuit<MAX_DEPTH> {
         value: Value,
     ) -> Result<()> {
         pw.set_hash_target(self.root, HashOut::from_vec(root.0.to_vec()))?;
-        pw.set_target_arr(&self.key, &key.0)?;
-        pw.set_target_arr(&self.value, &value.0)?;
+        pw.set_target_arr(&self.key.elements, &key.0)?;
+        pw.set_target_arr(&self.value.elements, &value.0)?;
 
         // pad siblings with zeros to length MAX_DEPTH
         let mut siblings = proof.siblings.clone();
@@ -297,21 +297,21 @@ fn compute_root_from_leaf<const MAX_DEPTH: usize>(
 // specially to be able to test it isolated.
 fn keypath_target<const MAX_DEPTH: usize>(
     builder: &mut CircuitBuilder<F, D>,
-    key: &Vec<Target>,
+    key: &ValueTarget,
 ) -> Vec<BoolTarget> {
-    assert_eq!(key.len(), VALUE_SIZE);
-
     let n_complete_field_elems: usize = MAX_DEPTH / F::BITS;
     let n_extra_bits: usize = MAX_DEPTH - n_complete_field_elems * F::BITS;
 
     let path: Vec<BoolTarget> = key
+        .elements
         .iter()
         .take(n_complete_field_elems)
         .flat_map(|e| builder.split_le(*e, F::BITS))
         .collect();
 
     let extra_bits = if n_extra_bits > 0 {
-        let extra_bits: Vec<BoolTarget> = builder.split_le(key[n_complete_field_elems], F::BITS);
+        let extra_bits: Vec<BoolTarget> =
+            builder.split_le(key.elements[n_complete_field_elems], F::BITS);
         extra_bits[..n_extra_bits].to_vec()
         // Note: ideally we would do:
         //     let extra_bits = builder.split_le(key[n_complete_field_elems], n_extra_bits);
@@ -326,10 +326,16 @@ fn keypath_target<const MAX_DEPTH: usize>(
 
 fn kv_hash_target(
     builder: &mut CircuitBuilder<F, D>,
-    key: &Vec<Target>,
-    value: &Vec<Target>,
+    key: &ValueTarget,
+    value: &ValueTarget,
 ) -> HashOutTarget {
-    let inputs: Vec<Target> = [key.clone(), value.clone(), vec![builder.one()]].concat();
+    let inputs = key
+        .elements
+        .iter()
+        .chain(value.elements.iter())
+        .cloned()
+        .chain(iter::once(builder.one()))
+        .collect();
     builder.hash_n_to_hash_no_pad::<PoseidonHash>(inputs)
 }
 
@@ -371,14 +377,14 @@ pub mod tests {
             let expected_path_targ: Vec<BoolTarget> = (0..MD)
                 .map(|_| builder.add_virtual_bool_target_safe())
                 .collect();
-            let key_targ = builder.add_virtual_targets(VALUE_SIZE);
+            let key_targ = builder.add_virtual_value();
             let computed_path_targ = keypath_target::<MD>(&mut builder, &key_targ);
             for i in 0..MD {
                 builder.connect(computed_path_targ[i].target, expected_path_targ[i].target);
             }
 
             // assign the input values to the targets
-            pw.set_target_arr(&key_targ, &key.0)?;
+            pw.set_target_arr(&key_targ.elements, &key.0)?;
             for i in 0..MD {
                 pw.set_bool_target(expected_path_targ[i], expected_path[i])?;
             }
@@ -404,15 +410,15 @@ pub mod tests {
             let mut pw = PartialWitness::<F>::new();
 
             let h_targ = builder.add_virtual_hash();
-            let key_targ = builder.add_virtual_targets(VALUE_SIZE);
-            let value_targ = builder.add_virtual_targets(VALUE_SIZE);
+            let key_targ = builder.add_virtual_value();
+            let value_targ = builder.add_virtual_value();
 
             let computed_h = kv_hash_target(&mut builder, &key_targ, &value_targ);
             builder.connect_hashes(computed_h, h_targ);
 
             // assign the input values to the targets
-            pw.set_target_arr(&key_targ, &key.0)?;
-            pw.set_target_arr(&value_targ, &value.0)?;
+            pw.set_target_arr(&key_targ.elements, &key.0)?;
+            pw.set_target_arr(&value_targ.elements, &value.0)?;
             pw.set_hash_target(h_targ, HashOut::from_vec(h.0.to_vec()))?;
 
             // generate & verify proof