handle CStore in linearization pass

Author: alex-ozdemir

circ_hc/src/rc/example_u8.rs

@@ -14,6 +14,7 @@ struct NodeData {
     cs: Box<[Node]>,
 }
 
+#[allow(dead_code)]
 struct NodeDataRef<'a, Q: Borrow<[Node]>>(&'a u8, &'a Q);
 
 #[derive(Clone)]

circ_hc/src/rc/macro_.rs

@@ -17,6 +17,7 @@ macro_rules! generate_hashcons_rc {
             cs: Box<[Node]>,
         }
 
+        #[allow(dead_code)]
         struct NodeDataRef<'a, Q: Borrow<[Node]>>(&'a $Op, &'a Q);
 
         #[derive(Clone)]

circ_hc/src/rc/template.rs

@@ -14,6 +14,7 @@ struct NodeData {
     cs: Box<[Node]>,
 }
 
+#[allow(dead_code)]
 struct NodeDataRef<'a, Q: Borrow<[Node]>>(&'a TemplateOp, &'a Q);
 
 #[derive(Clone)]

examples/ZoKrates/pf/2024_05_24_benny_bug.zok

@@ -0,0 +1,7 @@
+def main(field x) -> field:
+    field[25] A = [0; 25]
+    for field counter in 0..30 do
+        cond_store(A, counter, x, counter < x)
+    endfor
+
+    return A[x]

examples/ZoKrates/pf/2024_05_24_benny_bug_tr.zok

@@ -0,0 +1,7 @@
+def main(field x) -> field:
+    transcript field[25] A = [0; 25]
+    for field counter in 0..30 do
+        cond_store(A, counter, x, counter < x)
+    endfor
+
+    return A[x]

scripts/zokrates_test.zsh

@@ -86,6 +86,8 @@ r1cs_test ./third_party/ZoKrates/zokrates_stdlib/stdlib/utils/casts/bool_128_to_
 r1cs_test ./third_party/ZoKrates/zokrates_stdlib/stdlib/ecc/edwardsScalarMult.zok
 r1cs_test ./third_party/ZoKrates/zokrates_stdlib/stdlib/hashes/mimc7/mimc7R20.zok
 r1cs_test ./third_party/ZoKrates/zokrates_stdlib/stdlib/hashes/pedersen/512bit.zok
+r1cs_test ./examples/ZoKrates/pf/2024_05_24_benny_bug.zok
+r1cs_test ./examples/ZoKrates/pf/2024_05_24_benny_bug_tr.zok
 
 pf_test_only_pf sha_temp1
 pf_test_only_pf sha_rot

src/circify/mod.rs

@@ -110,12 +110,6 @@ struct ScopeIdx {
     lex: usize,
 }
 
-#[derive(Hash, PartialEq, Eq)]
-struct VerVar {
-    name: VarName,
-    ver: Version,
-}
-
 #[derive(Debug)]
 struct LexEntry<Ty> {
     ver: Version,

src/ir/opt/mem/lin.rs

@@ -52,11 +52,15 @@ impl RewritePass for Linearizer {
                 let cs = rewritten_children();
                 let idx = &cs[1];
                 let tup = &cs[0];
-                if let Sort::Array(key_sort, _, size) = check(&orig.cs()[0]) {
+                if let Sort::Array(key_sort, val_sort, size) = check(&orig.cs()[0]) {
                     assert!(size > 0);
                     if idx.is_const() {
                         let idx_usize = extras::as_uint_constant(idx).unwrap().to_usize().unwrap();
-                        Some(term![Op::Field(idx_usize); tup.clone()])
+                        if idx_usize < size {
+                            Some(term![Op::Field(idx_usize); tup.clone()])
+                        } else {
+                            Some(val_sort.default_term())
+                        }
                     } else {
                         let mut fields = (0..size).map(|idx| term![Op::Field(idx); tup.clone()]);
                         let first = fields.next().unwrap();
@@ -77,7 +81,11 @@ impl RewritePass for Linearizer {
                     assert!(size > 0);
                     if idx.is_const() {
                         let idx_usize = extras::as_uint_constant(idx).unwrap().to_usize().unwrap();
-                        Some(term![Op::Update(idx_usize); tup.clone(), val.clone()])
+                        if idx_usize < size {
+                            Some(term![Op::Update(idx_usize); tup.clone(), val.clone()])
+                        } else {
+                            Some(tup.clone())
+                        }
                     } else {
                         let mut updates =
                             (0..size).map(|idx| term![Op::Update(idx); tup.clone(), val.clone()]);
@@ -90,6 +98,35 @@ impl RewritePass for Linearizer {
                     unreachable!()
                 }
             }
+            Op::CStore => {
+                let cs = rewritten_children();
+                let tup = &cs[0];
+                let idx = &cs[1];
+                let val = &cs[2];
+                let cond = &cs[3];
+                if let Sort::Array(key_sort, _, size) = check(&orig.cs()[0]) {
+                    assert!(size > 0);
+                    if idx.is_const() {
+                        let idx_usize = extras::as_uint_constant(idx).unwrap().to_usize().unwrap();
+                        if idx_usize < size {
+                            Some(
+                                term![Op::Ite; cond.clone(), term![Op::Update(idx_usize); tup.clone(), val.clone()], tup.clone()],
+                            )
+                        } else {
+                            Some(tup.clone())
+                        }
+                    } else {
+                        let mut updates =
+                            (0..size).map(|idx| term![Op::Update(idx); tup.clone(), val.clone()]);
+                        let first = updates.next().unwrap();
+                        Some(key_sort.elems_iter().take(size).skip(1).zip(updates).fold(first, |acc, (idx_c, update)| {
+                        term![Op::Ite; term![AND; term![Op::Eq; idx.clone(), idx_c], cond.clone()], update, acc]
+                    }))
+                    }
+                } else {
+                    unreachable!()
+                }
+            }
             // ITEs and EQs are correctly rewritten by default.
             _ => None,
         }

src/ir/opt/mem/ram/volatile.rs

@@ -275,7 +275,7 @@ impl RewritePass for Extactor {
                 "Bad ram term: {}",
                 t
             );
-            debug_assert!(self.graph.children.get(t).is_some());
+            debug_assert!(self.graph.children.contains_key(t));
             debug_assert_eq!(1, self.graph.children.get(t).unwrap().len());
             // Get dependency's RAM
             let child = self.graph.children.get(t).unwrap()[0].clone();

src/ir/opt/visit.rs

@@ -102,6 +102,7 @@ pub trait RewritePass {
 
 /// An analysis pass that repeated sweeps all terms, visiting them, untill a pass makes no more
 /// progress.
+#[allow(dead_code)]
 pub trait ProgressAnalysisPass {
     /// The visit function. Returns whether progress was made.
     fn visit(&mut self, term: &Term) -> bool;

src/ir/term/dist.rs

@@ -1,5 +1,5 @@
 //! Distributions over terms (useful for fuzz testing)
-
+#![allow(dead_code)]
 use super::*;
 
 use circ_fields::{FieldT, FieldV};