From a35cd98b03b9533fa7946a58327a5bae3bc9990b Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Mon, 18 Oct 2021 15:45:52 +0200
Subject: [PATCH] New random access gadget

---
 src/gadgets/random_access.rs | 28 ++++++++++++++++++++++++++--
 src/gates/random_access.rs   |  3 ---
 src/hash/merkle_proofs.rs    | 29 +++++++----------------------
 3 files changed, 33 insertions(+), 27 deletions(-)

diff --git a/src/gadgets/random_access.rs b/src/gadgets/random_access.rs
index bae7c9cd..f4396466 100644
--- a/src/gadgets/random_access.rs
+++ b/src/gadgets/random_access.rs
@@ -6,8 +6,32 @@ use crate::iop::target::Target;
 use crate::plonk::circuit_builder::CircuitBuilder;
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
-    /// Checks that a `Target` matches a vector at a non-deterministic index.
-    /// Note: `index` is not range-checked.
+    /// Checks that an `ExtensionTarget` matches a vector at a non-deterministic index.
+    /// Note: `access_index` is not range-checked.
+    pub fn random_access(&mut self, access_index: Target, claimed_element: Target, v: Vec<Target>) {
+        debug_assert!(!v.is_empty());
+        if v.len() == 1 {
+            return self.connect(claimed_element, v[0]);
+        }
+        let gate = RandomAccessGate::new(1, v.len());
+        let gate_index = self.add_gate(gate.clone(), vec![]);
+
+        let copy = 0;
+        v.iter().enumerate().for_each(|(i, &val)| {
+            self.connect(val, Target::wire(gate_index, gate.wire_list_item(i, copy)));
+        });
+        self.connect(
+            access_index,
+            Target::wire(gate_index, gate.wire_access_index(copy)),
+        );
+        self.connect(
+            claimed_element,
+            Target::wire(gate_index, gate.wire_claimed_element(copy)),
+        );
+    }
+
+    /// Checks that an `ExtensionTarget` matches a vector at a non-deterministic index.
+    /// Note: `access_index` is not range-checked.
     pub fn random_access_extension(
         &mut self,
         access_index: Target,
diff --git a/src/gates/random_access.rs b/src/gates/random_access.rs
index 8e1618fc..e11a33d0 100644
--- a/src/gates/random_access.rs
+++ b/src/gates/random_access.rs
@@ -106,7 +106,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for RandomAccessGa
                 let index_matches = vars.local_wires[self.wire_index_matches_for_index(i, copy)];
 
                 // The two index equality constraints.
-                dbg!(difference, equality_dummy, index_matches);
                 constraints.push(difference * equality_dummy - (F::Extension::ONE - index_matches));
                 constraints.push(index_matches * difference);
                 // Value equality constraint.
@@ -382,7 +381,6 @@ mod tests {
             .zip(&access_indices)
             .map(|(l, &i)| l[i])
             .collect();
-        dbg!(&lists, &access_indices, &good_claimed_elements);
         let good_vars = EvaluationVars {
             local_constants: &[],
             local_wires: &get_wires(lists.clone(), access_indices.clone(), good_claimed_elements),
@@ -395,7 +393,6 @@ mod tests {
             public_inputs_hash: &HashOut::rand(),
         };
 
-        dbg!(gate.eval_unfiltered(good_vars));
         assert!(
             gate.eval_unfiltered(good_vars).iter().all(|x| x.is_zero()),
             "Gate constraints are not satisfied."
diff --git a/src/hash/merkle_proofs.rs b/src/hash/merkle_proofs.rs
index 062ff521..3e265979 100644
--- a/src/hash/merkle_proofs.rs
+++ b/src/hash/merkle_proofs.rs
@@ -75,17 +75,6 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
 
         let index = self.le_sum(leaf_index_bits[proof.siblings.len()..].to_vec().into_iter());
-        let state_ext = state.elements[..].try_into().expect("requires D = 4");
-        let state_ext = ExtensionTarget(state_ext);
-        let cap_ext = merkle_cap
-            .0
-            .iter()
-            .map(|h| {
-                let tmp = h.elements[..].try_into().expect("requires D = 4");
-                ExtensionTarget(tmp)
-            })
-            .collect();
-        self.random_access_extension(index, state_ext, cap_ext);
     }
 
     /// Same a `verify_merkle_proof` but with the final "cap index" as extra parameter.
@@ -112,17 +101,13 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             };
         }
 
-        let state_ext = state.elements[..].try_into().expect("requires D = 4");
-        let state_ext = ExtensionTarget(state_ext);
-        let cap_ext = merkle_cap
-            .0
-            .iter()
-            .map(|h| {
-                let tmp = h.elements[..].try_into().expect("requires D = 4");
-                ExtensionTarget(tmp)
-            })
-            .collect();
-        self.random_access_extension(cap_index, state_ext, cap_ext);
+        for i in 0..4 {
+            self.random_access(
+                cap_index,
+                state.elements[i],
+                merkle_cap.0.iter().map(|h| h.elements[i]).collect(),
+            );
+        }
     }
 
     pub fn assert_hashes_equal(&mut self, x: HashOutTarget, y: HashOutTarget) {