From 6ea437c83cb3804c5feb27bfbc67e7d066d780f5 Mon Sep 17 00:00:00 2001 From: dark64 Date: Wed, 31 May 2023 23:28:19 +0200 Subject: [PATCH 1/3] fix propagation in case of nested assembly --- .../src/flatten_complex_types.rs | 32 +++++++++++++++++++ zokrates_analysis/src/zir_propagation.rs | 21 ++++++++++++ zokrates_ast/src/zir/canonicalizer.rs | 13 +++++--- .../tests/assembly/nested_assignment.json | 16 ++++++++++ .../tests/assembly/nested_assignment.zok | 20 ++++++++++++ 5 files changed, 98 insertions(+), 4 deletions(-) create mode 100644 zokrates_core_test/tests/tests/assembly/nested_assignment.json create mode 100644 zokrates_core_test/tests/tests/assembly/nested_assignment.zok diff --git a/zokrates_analysis/src/flatten_complex_types.rs b/zokrates_analysis/src/flatten_complex_types.rs index 64705d1a..37b01c4b 100644 --- a/zokrates_analysis/src/flatten_complex_types.rs +++ b/zokrates_analysis/src/flatten_complex_types.rs @@ -510,6 +510,38 @@ pub struct ArgumentFinder<'ast, T> { } impl<'ast, T: Field> Folder<'ast, T> for ArgumentFinder<'ast, T> { + fn fold_assembly_block( + &mut self, + s: zir::AssemblyBlockStatement<'ast, T>, + ) -> Vec> { + let mut statements: Vec<_> = s + .inner + .into_iter() + .rev() + .flat_map(|s| self.fold_assembly_statement(s)) + .collect(); + statements.reverse(); + vec![zir::ZirStatement::Assembly( + zir::AssemblyBlockStatement::new(statements), + )] + } + + fn fold_assembly_assignment( + &mut self, + s: zir::AssemblyAssignment<'ast, T>, + ) -> Vec> { + let assignees: Vec> = s + .assignee + .into_iter() + .map(|v| self.fold_assignee(v)) + .collect(); + let expr = self.fold_function(s.expression); + for a in &assignees { + self.identifiers.remove(&a.id); + } + vec![zir::ZirAssemblyStatement::assignment(assignees, expr)] + } + fn fold_statement(&mut self, s: zir::ZirStatement<'ast, T>) -> Vec> { match s { zir::ZirStatement::Definition(s) => { diff --git a/zokrates_analysis/src/zir_propagation.rs b/zokrates_analysis/src/zir_propagation.rs index a677520c..9e7cbcfa 100644 --- a/zokrates_analysis/src/zir_propagation.rs +++ b/zokrates_analysis/src/zir_propagation.rs @@ -84,6 +84,27 @@ impl<'ast, T: Field> ResultFolder<'ast, T> for ZirPropagator<'ast, T> { }) } + fn fold_assembly_block( + &mut self, + s: zokrates_ast::zir::AssemblyBlockStatement<'ast, T>, + ) -> Result>, Self::Error> { + let block: Vec<_> = s + .inner + .into_iter() + .map(|s| self.fold_assembly_statement(s)) + .collect::, _>>()? + .into_iter() + .flatten() + .collect(); + + Ok(match block.is_empty() { + true => vec![], + false => vec![ZirStatement::Assembly( + zokrates_ast::zir::AssemblyBlockStatement::new(block), + )], + }) + } + fn fold_assembly_assignment( &mut self, s: zokrates_ast::zir::AssemblyAssignment<'ast, T>, diff --git a/zokrates_ast/src/zir/canonicalizer.rs b/zokrates_ast/src/zir/canonicalizer.rs index 61bf1970..0e4ffe56 100644 --- a/zokrates_ast/src/zir/canonicalizer.rs +++ b/zokrates_ast/src/zir/canonicalizer.rs @@ -9,11 +9,16 @@ pub struct ZirCanonicalizer<'ast> { impl<'ast, T: Field> Folder<'ast, T> for ZirCanonicalizer<'ast> { fn fold_parameter(&mut self, p: Parameter<'ast>) -> Parameter<'ast> { - let new_id = self.identifier_map.len(); - self.identifier_map.insert(p.id.id.clone(), new_id); - + let id = match self.identifier_map.get(&p.id.id) { + Some(v) => Identifier::internal(*v), + None => { + let new_id = self.identifier_map.len(); + self.identifier_map.insert(p.id.id.clone(), new_id); + Identifier::internal(new_id) + } + }; Parameter { - id: Variable::with_id_and_type(Identifier::internal(new_id), p.id.ty), + id: Variable::with_id_and_type(id, p.id.ty), ..p } } diff --git a/zokrates_core_test/tests/tests/assembly/nested_assignment.json b/zokrates_core_test/tests/tests/assembly/nested_assignment.json new file mode 100644 index 00000000..7d72a99f --- /dev/null +++ b/zokrates_core_test/tests/tests/assembly/nested_assignment.json @@ -0,0 +1,16 @@ +{ + "curves": ["Bn128"], + "max_constraint_count": 1, + "tests": [ + { + "input": { + "values": ["4"] + }, + "output": { + "Ok": { + "value": "32" + } + } + } + ] +} diff --git a/zokrates_core_test/tests/tests/assembly/nested_assignment.zok b/zokrates_core_test/tests/tests/assembly/nested_assignment.zok new file mode 100644 index 00000000..0abf3c5e --- /dev/null +++ b/zokrates_core_test/tests/tests/assembly/nested_assignment.zok @@ -0,0 +1,20 @@ +def bar(field a) -> field { + field tmp = a * a; + field mut b = 0; + asm { + b <-- tmp * 2; + } + return b; +} + +def foo(field a) -> field { + field mut b = 0; + asm { + b <-- bar(a); + } + return b; +} + +def main(field a) -> field { + return foo(a); +} \ No newline at end of file From 5c233ff3b49d460fc8a6912f77399d326c05a6fc Mon Sep 17 00:00:00 2001 From: dark64 Date: Mon, 12 Jun 2023 22:06:57 +0200 Subject: [PATCH 2/3] fix the format of assembly blocks in zir --- zokrates_ast/src/zir/mod.rs | 42 +++++++++++++++++++++++++++++-------- 1 file changed, 33 insertions(+), 9 deletions(-) diff --git a/zokrates_ast/src/zir/mod.rs b/zokrates_ast/src/zir/mod.rs index 5ae001bc..62025adf 100644 --- a/zokrates_ast/src/zir/mod.rs +++ b/zokrates_ast/src/zir/mod.rs @@ -63,8 +63,8 @@ pub struct ZirFunction<'ast, T> { pub type IdentifierOrExpression<'ast, T, E> = expressions::IdentifierOrExpression, E, >::Inner>; -impl<'ast, T: fmt::Display> fmt::Display for ZirFunction<'ast, T> { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { +impl<'ast, T: fmt::Display> ZirFunction<'ast, T> { + fn fmt_indented(&self, f: &mut fmt::Formatter, depth: usize) -> fmt::Result { writeln!( f, "({}) -> ({}) {{", @@ -82,11 +82,17 @@ impl<'ast, T: fmt::Display> fmt::Display for ZirFunction<'ast, T> { )?; for s in &self.statements { - s.fmt_indented(f, 1)?; + s.fmt_indented(f, depth + 1)?; writeln!(f)?; } - write!(f, "}}") + write!(f, "{}}}", "\t".repeat(depth)) + } +} + +impl<'ast, T: fmt::Display> fmt::Display for ZirFunction<'ast, T> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.fmt_indented(f, 0) } } @@ -174,19 +180,30 @@ impl<'ast, T> WithSpan for ZirAssemblyStatement<'ast, T> { } } -impl<'ast, T: fmt::Display> fmt::Display for ZirAssemblyStatement<'ast, T> { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { +impl<'ast, T: fmt::Display> ZirAssemblyStatement<'ast, T> { + fn fmt_indented(&self, f: &mut fmt::Formatter, depth: usize) -> fmt::Result { + write!(f, "{}", "\t".repeat(depth))?; match *self { ZirAssemblyStatement::Assignment(ref s) => { write!( f, - "{} <-- {};", + "{} <-- (", s.assignee .iter() .map(|a| a.to_string()) .collect::>() - .join(", "), + .join(", ") + )?; + s.expression.fmt_indented(f, depth)?; + write!( + f, + ")({});", s.expression + .arguments + .iter() + .map(|a| a.id.id.to_string()) + .collect::>() + .join(", ") ) } ZirAssemblyStatement::Constraint(ref s) => { @@ -196,6 +213,12 @@ impl<'ast, T: fmt::Display> fmt::Display for ZirAssemblyStatement<'ast, T> { } } +impl<'ast, T: fmt::Display> fmt::Display for ZirAssemblyStatement<'ast, T> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.fmt_indented(f, 0) + } +} + pub type DefinitionStatement<'ast, T> = common::expressions::DefinitionStatement, ZirExpression<'ast, T>>; pub type AssertionStatement<'ast, T> = @@ -462,7 +485,8 @@ impl<'ast, T: fmt::Display> ZirStatement<'ast, T> { ZirStatement::Assembly(s) => { writeln!(f, "asm {{")?; for s in &s.inner { - writeln!(f, "{}{}", "\t".repeat(depth + 1), s)?; + s.fmt_indented(f, depth + 1)?; + writeln!(f)?; } write!(f, "{}}}", "\t".repeat(depth)) } From 9c0c7a97253e7da8a76029cf8df8ab3837c77405 Mon Sep 17 00:00:00 2001 From: dark64 Date: Tue, 13 Jun 2023 12:22:09 +0200 Subject: [PATCH 3/3] add changelog --- changelogs/unreleased/1322-dark64 | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelogs/unreleased/1322-dark64 diff --git a/changelogs/unreleased/1322-dark64 b/changelogs/unreleased/1322-dark64 new file mode 100644 index 00000000..38595ad1 --- /dev/null +++ b/changelogs/unreleased/1322-dark64 @@ -0,0 +1 @@ +Fix panic in case of nested assembly blocks \ No newline at end of file