From 684df1e05777fa0c19ad1147bddd07504bcab62a Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Tue, 10 Aug 2021 15:03:29 +0200
Subject: [PATCH] Pass cap index

---
 src/fri/recursive_verifier.rs |  8 +++-
 src/hash/merkle_proofs.rs     | 71 +++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/fri/recursive_verifier.rs b/src/fri/recursive_verifier.rs
index 5e801e9c..5d62d26c 100644
--- a/src/fri/recursive_verifier.rs
+++ b/src/fri/recursive_verifier.rs
@@ -284,6 +284,11 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         // TODO: Do we need to range check `x_index` to a target smaller than `p`?
         let x_index = challenger.get_challenge(self);
         let mut x_index_bits = self.low_bits(x_index, n_log, 64);
+        let cap_index = self.le_sum(
+            x_index_bits[x_index_bits.len() - common_data.config.fri_config.cap_height..]
+                .to_vec()
+                .into_iter(),
+        );
         let mut domain_size = n;
         with_context!(
             self,
@@ -346,9 +351,10 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             with_context!(
                 self,
                 "verify FRI round Merkle proof.",
-                self.verify_merkle_proof(
+                self.verify_merkle_proof_with_cap_index(
                     flatten_target(&evals),
                     &high_x_index_bits,
+                    cap_index,
                     &proof.commit_phase_merkle_roots[i],
                     &round_proof.steps[i].merkle_proof,
                 )
diff --git a/src/hash/merkle_proofs.rs b/src/hash/merkle_proofs.rs
index fb303775..3a9e7a31 100644
--- a/src/hash/merkle_proofs.rs
+++ b/src/hash/merkle_proofs.rs
@@ -51,6 +51,7 @@ pub(crate) fn verify_merkle_proof<F: Field>(
             compress(current_digest, sibling_digest)
         }
     }
+    dbg!(index);
     ensure!(
         current_digest == merkle_cap.0[index],
         "Invalid Merkle proof."
@@ -132,6 +133,76 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         // self.named_assert_hashes_equal(state, merkle_root, "check Merkle root".into())
     }
 
+    pub(crate) fn verify_merkle_proof_with_cap_index(
+        &mut self,
+        leaf_data: Vec<Target>,
+        leaf_index_bits: &[Target],
+        cap_index: Target,
+        merkle_root: &MerkleCapTarget,
+        proof: &MerkleProofTarget,
+    ) {
+        let zero = self.zero();
+
+        let mut state: HashOutTarget = self.hash_or_noop(leaf_data);
+
+        for (&bit, &sibling) in leaf_index_bits.iter().zip(&proof.siblings) {
+            let gate_type = GMiMCGate::<F, D, GMIMC_ROUNDS>::new_automatic_constants();
+            let gate = self.add_gate(gate_type, vec![]);
+
+            let swap_wire = GMiMCGate::<F, D, GMIMC_ROUNDS>::WIRE_SWAP;
+            let swap_wire = Target::Wire(Wire {
+                gate,
+                input: swap_wire,
+            });
+            self.generate_copy(bit, swap_wire);
+
+            let input_wires = (0..12)
+                .map(|i| {
+                    Target::Wire(Wire {
+                        gate,
+                        input: GMiMCGate::<F, D, GMIMC_ROUNDS>::wire_input(i),
+                    })
+                })
+                .collect::<Vec<_>>();
+
+            for i in 0..4 {
+                self.route(state.elements[i], input_wires[i]);
+                self.route(sibling.elements[i], input_wires[4 + i]);
+                self.route(zero, input_wires[8 + i]);
+            }
+
+            state = HashOutTarget::from_vec(
+                (0..4)
+                    .map(|i| {
+                        Target::Wire(Wire {
+                            gate,
+                            input: GMiMCGate::<F, D, GMIMC_ROUNDS>::wire_output(i),
+                        })
+                    })
+                    .collect(),
+            )
+        }
+
+        let mut state_ext = [zero; D];
+        for i in 0..D {
+            state_ext[i] = state.elements[i];
+        }
+        let state_ext = ExtensionTarget(state_ext);
+        let cap_ext = merkle_root
+            .0
+            .iter()
+            .map(|h| {
+                let mut tmp = [zero; D];
+                for i in 0..D {
+                    tmp[i] = h.elements[i];
+                }
+                ExtensionTarget(tmp)
+            })
+            .collect();
+        self.random_access(cap_index, state_ext, cap_ext);
+        // self.named_assert_hashes_equal(state, merkle_root, "check Merkle root".into())
+    }
+
     pub(crate) fn assert_hashes_equal(&mut self, x: HashOutTarget, y: HashOutTarget) {
         for i in 0..4 {
             self.assert_equal(x.elements[i], y.elements[i]);