Merge pull request #1258 from Zokrates/fix-assembly-propagation

Fix zir assembly propagation

zokrates_analysis/src/zir_propagation.rs

@@ -7,7 +7,7 @@ use zokrates_ast::zir::types::UBitwidth;
 use zokrates_ast::zir::{
     result_folder::*, Conditional, ConditionalExpression, ConditionalOrExpression, Constant, Expr,
     Id, IdentifierExpression, IdentifierOrExpression, SelectExpression, SelectOrExpression,
-    ZirAssemblyStatement,
+    ZirAssemblyStatement, ZirFunction,
 };
 use zokrates_ast::zir::{
     BooleanExpression, FieldElementExpression, Identifier, RuntimeError, UExpression,
@@ -57,6 +57,28 @@ impl<'ast, T: Field> ZirPropagator<'ast, T> {
 impl<'ast, T: Field> ResultFolder<'ast, T> for ZirPropagator<'ast, T> {
     type Error = Error;
 
+    fn fold_function(
+        &mut self,
+        f: ZirFunction<'ast, T>,
+    ) -> Result<ZirFunction<'ast, T>, Self::Error> {
+        Ok(ZirFunction {
+            arguments: f
+                .arguments
+                .into_iter()
+                .filter(|p| !self.constants.contains_key(&p.id.id))
+                .collect(),
+            statements: f
+                .statements
+                .into_iter()
+                .map(|s| self.fold_statement(s))
+                .collect::<Result<Vec<_>, _>>()?
+                .into_iter()
+                .flatten()
+                .collect(),
+            ..f
+        })
+    }
+
     fn fold_assembly_statement(
         &mut self,
         s: ZirAssemblyStatement<'ast, T>,

zokrates_core_test/tests/tests/assembly/propagation/array_propagation.json

@@ -0,0 +1,16 @@
+{
+  "curves": ["Bn128"],
+  "max_constraint_count": 1,
+  "tests": [
+    {
+      "input": {
+        "values": ["2"]
+      },
+      "output": {
+        "Ok": {
+          "value": "4"
+        }
+      }
+    }
+  ]
+}

zokrates_core_test/tests/tests/assembly/propagation/array_propagation.zok

@@ -0,0 +1,12 @@
+def mul(field a, field b) -> field {
+    field mut res = 0;
+    asm {
+        res <== a * b;
+    }
+    return res;
+}
+
+def main(field a) -> field {
+    field[2] b = [2, a]; // this will get propagated in zir
+    return mul(b[0], b[1]);
+}