Exposing vector_add to zok frontend

scripts/zx_tests/vector_add.zx

@@ -0,0 +1,7 @@
+from "EMBED" import vector_add
+
+def main() -> u32[5]:
+    u32[5] a = [1, 2, 3, 4, 5]
+    u32[5] b = [2, 3, 4, 5, 6]
+    assert(vector_add(a, b) == [3, 5, 7, 9, 11])
+    return vector_add(a,b)

scripts/zx_tests/vector_add.zxf

@@ -0,0 +1,6 @@
+from "EMBED" import vector_add
+
+def main() -> u32[5]:
+    u32[5] a = [1, 2, 3, 4, 5]
+    u32[5] b = []
+    return vector_add(a,b)

src/front/zsharp/mod.rs

@@ -268,7 +268,7 @@ impl<'ast> ZGen<'ast> {
             "bit_array_le" => {
                 if args.len() != 2 {
                     Err(format!(
-                        "Got {} args to EMBED/bit_array_le, expected 1",
+                        "Got {} args to EMBED/bit_array_le, expected 2",
                         args.len()
                     ))
                 } else if generics.len() != 1 {
@@ -305,6 +305,24 @@ impl<'ast> ZGen<'ast> {
                     Ok(uint_lit(DFL_T.modulus().significant_bits(), 32))
                 }
             }
+            "vector_add" => {
+                if args.len() != 2 {
+                    Err(format!(
+                        "Got {} args to EMBED/vector_add, expected 2",
+                        args.len()
+                    ))
+                } else if generics.len() != 1 {
+                    Err(format!(
+                        "Got {} generic args to EMBED/vector_add, expected 1",
+                        generics.len()
+                    ))
+                } else {
+                    assert!(args.iter().all(|t| matches!(t.type_(), Ty::Array(_, _))));
+                    let a = args.pop().unwrap();
+                    let b = args.pop().unwrap();
+                    vector_op(BV_ADD, a, b)
+                }
+            }
             _ => Err(format!("Unknown or unimplemented builtin '{}'", f_name)),
         }
     }
@@ -891,7 +909,6 @@ impl<'ast> ZGen<'ast> {
         } else {
             debug!("Expr: {}", e.span().as_str());
         }
-
         match e {
             ast::Expression::Ternary(u) => {
                 match self.expr_impl_::<true>(&u.first).ok().and_then(const_bool) {

src/front/zsharp/term.rs

@@ -908,6 +908,20 @@ pub fn bit_array_le(a: T, b: T, n: usize) -> Result<T, String> {
     ))
 }
 
+pub fn vector_op(op: Op, a: T, b: T) -> Result<T, String> {
+    match (&a.ty, &b.ty) {
+        (Ty::Array(a_s, a_ty), Ty::Array(b_s, b_ty)) => {
+            if a_s == b_s && a_ty == b_ty {
+                let t = term![Op::Map(Box::new(op)); a.term, b.term];
+                Ok(T::new(Ty::Array(*a_s, a_ty.clone()), t))
+            } else {
+                panic!("Mismatched array types (this is a bug: type checking should have caught this!)");
+            }
+        }
+        _ => Err("Cannot do vector_add on non-array types".to_string()),
+    }
+}
+
 pub struct ZSharp {
     values: Option<HashMap<String, Integer>>,
 }

src/ir/opt/cfold.rs

@@ -5,6 +5,7 @@ use crate::ir::term::*;
 use circ_fields::FieldV;
 use lazy_static::lazy_static;
 use rug::Integer;
+use std::collections::BTreeMap;
 use std::ops::DerefMut;
 use std::sync::RwLock;
 
@@ -244,6 +245,35 @@ pub fn fold_cache(node: &Term, cache: &mut TermCache<Term>) -> Term {
                     b.width() + w,
                 ))))
             }),
+            Op::Map(op) => match (get(0).as_array_opt(), get(1).as_array_opt()) {
+                (Some(a), Some(b)) => {
+                    assert!(a.size == b.size);
+                    let mut new_map: BTreeMap<Value, Value> = BTreeMap::new();
+                    for (a_ind, a_val) in &a.map {
+                        let b_val = &b.map[a_ind];
+                        let res = fold_cache(
+                            &term![*op.clone(); leaf_term(Op::Const(a_val.clone())), leaf_term(Op::Const(b_val.clone()))],
+                            cache,
+                        );
+                        match res.as_value_opt() {
+                            Some(r) => {
+                                new_map.insert(a_ind.clone(), r.clone());
+                            }
+                            None => {
+                                panic!("Unale to constant constant fold idx: {}", a_ind);
+                            }
+                        }
+                    }
+                    let new_arr = Array::new(
+                        a.key_sort.clone(),
+                        Box::new(a.key_sort.default_value()),
+                        new_map,
+                        a.size,
+                    );
+                    Some(leaf_term(Op::Const(Value::Array(new_arr))))
+                }
+                _ => None,
+            },
             _ => None,
         };
         let new_t = {

third_party/ZoKrates/zokrates_stdlib/stdlib/EMBED.zok

@@ -75,3 +75,7 @@ def u16_to_u32(u16 i) -> u32:
 
 def u8_to_u16(u8 i) -> u16:
     return 0u16
+
+// vector functions
+def vector_add<N>(u32[N] a, u32[N] b) -> u32[N]:
+    return [0; N]