apply suggestions (part 1)

changelogs/unreleased/1268-dark64

@@ -1 +1 @@
-Optimize assembly solver
+Reduce compiled program size by deduplicating assembly solvers

zokrates_ast/src/common/mod.rs

@@ -10,6 +10,6 @@ pub use self::embed::FlatEmbed;
 pub use self::error::RuntimeError;
 pub use self::metadata::SourceMetadata;
 pub use self::parameter::Parameter;
-pub use self::solvers::Solver;
+pub use self::solvers::{RefCall, Solver};
 pub use self::variable::Variable;
 pub use format_string::FormatString;

zokrates_ast/src/common/solvers.rs

@@ -2,6 +2,12 @@ use crate::zir::ZirFunction;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
+#[derive(Clone, PartialEq, Debug, Serialize, Deserialize, Hash, Eq)]
+pub struct RefCall {
+    pub index: usize,
+    pub argument_count: usize,
+}
+
 #[derive(Clone, PartialEq, Debug, Serialize, Deserialize, Hash, Eq)]
 pub enum Solver<'ast, T> {
     ConditionEq,
@@ -14,7 +20,7 @@ pub enum Solver<'ast, T> {
     EuclideanDiv,
     #[serde(borrow)]
     Zir(ZirFunction<'ast, T>),
-    IndexedCall(usize, usize),
+    Ref(RefCall),
     #[cfg(feature = "bellman")]
     Sha256Round,
     #[cfg(feature = "ark")]
@@ -33,7 +39,7 @@ impl<'ast, T> fmt::Display for Solver<'ast, T> {
             Solver::ShaCh => write!(f, "ShaCh"),
             Solver::EuclideanDiv => write!(f, "EuclideanDiv"),
             Solver::Zir(_) => write!(f, "Zir(..)"),
-            Solver::IndexedCall(index, argc) => write!(f, "IndexedCall@{}({})", index, argc),
+            Solver::Ref(call) => write!(f, "Ref@{}({})", call.index, call.argument_count),
             #[cfg(feature = "bellman")]
             Solver::Sha256Round => write!(f, "Sha256Round"),
             #[cfg(feature = "ark")]
@@ -54,7 +60,7 @@ impl<'ast, T> Solver<'ast, T> {
             Solver::ShaCh => (3, 1),
             Solver::EuclideanDiv => (2, 2),
             Solver::Zir(f) => (f.arguments.len(), 1),
-            Solver::IndexedCall(_, n) => (*n, 1),
+            Solver::Ref(c) => (c.argument_count, 1),
             #[cfg(feature = "bellman")]
             Solver::Sha256Round => (768, 26935),
             #[cfg(feature = "ark")]

zokrates_ast/src/ir/solver_indexer.rs

@@ -1,3 +1,4 @@
+use crate::common::RefCall;
 use crate::ir::folder::Folder;
 use crate::ir::Directive;
 use crate::ir::Solver;
@@ -20,14 +21,14 @@ fn hash<T: Field>(f: &ZirFunction<T>) -> Hash {
 #[derive(Debug, Default)]
 pub struct SolverIndexer<'ast, T> {
     pub solvers: Vec<Solver<'ast, T>>,
-    pub index_map: HashMap<u64, usize>,
+    pub index_map: HashMap<Hash, usize>,
 }
 
 impl<'ast, T: Field> Folder<'ast, T> for SolverIndexer<'ast, T> {
     fn fold_directive(&mut self, d: Directive<'ast, T>) -> Directive<'ast, T> {
         match d.solver {
             Solver::Zir(f) => {
-                let argc = f.arguments.len();
+                let argument_count = f.arguments.len();
                 let h = hash(&f);
                 let index = match self.index_map.entry(h) {
                     Entry::Occupied(v) => *v.get(),
@@ -41,7 +42,10 @@ impl<'ast, T: Field> Folder<'ast, T> for SolverIndexer<'ast, T> {
                 Directive {
                     inputs: d.inputs,
                     outputs: d.outputs,
-                    solver: Solver::IndexedCall(index, argc),
+                    solver: Solver::Ref(RefCall {
+                        index,
+                        argument_count,
+                    }),
                 }
             }
             _ => d,

zokrates_ast/src/zir/identifier.rs

@@ -31,6 +31,12 @@ impl<'ast> fmt::Display for SourceIdentifier<'ast> {
     }
 }
 
+impl<'ast> Identifier<'ast> {
+    pub fn internal<S: Into<String>>(name: S) -> Self {
+        Identifier::Internal(name.into())
+    }
+}
+
 impl<'ast> fmt::Display for Identifier<'ast> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
@@ -40,12 +46,6 @@ impl<'ast> fmt::Display for Identifier<'ast> {
     }
 }
 
-impl<'ast> From<String> for Identifier<'ast> {
-    fn from(id: String) -> Identifier<'ast> {
-        Identifier::Internal(id)
-    }
-}
-
 // this is only used in tests but somehow cfg(test) does not work
 impl<'ast> From<&'ast str> for Identifier<'ast> {
     fn from(id: &'ast str) -> Identifier<'ast> {

zokrates_ast/src/zir/substitution.rs

@@ -1,22 +1,19 @@
 use super::{Folder, Identifier};
-use std::{collections::HashMap, marker::PhantomData};
+use std::collections::HashMap;
 use zokrates_field::Field;
 
-pub struct ZirSubstitutor<'a, 'ast, T> {
-    substitution: &'a HashMap<Identifier<'ast>, Identifier<'ast>>,
-    _phantom: PhantomData<T>,
+#[derive(Default)]
+pub struct ZirSubstitutor<'ast> {
+    substitution: HashMap<Identifier<'ast>, Identifier<'ast>>,
 }
 
-impl<'a, 'ast, T: Field> ZirSubstitutor<'a, 'ast, T> {
-    pub fn new(substitution: &'a HashMap<Identifier<'ast>, Identifier<'ast>>) -> Self {
-        ZirSubstitutor {
-            substitution,
-            _phantom: PhantomData::default(),
-        }
+impl<'ast> ZirSubstitutor<'ast> {
+    pub fn new(substitution: HashMap<Identifier<'ast>, Identifier<'ast>>) -> Self {
+        Self { substitution }
     }
 }
 
-impl<'a, 'ast, T: Field> Folder<'ast, T> for ZirSubstitutor<'a, 'ast, T> {
+impl<'ast, T: Field> Folder<'ast, T> for ZirSubstitutor<'ast> {
     fn fold_name(&mut self, n: Identifier<'ast>) -> Identifier<'ast> {
         match self.substitution.get(&n) {
             Some(v) => v.clone(),

zokrates_codegen/src/lib.rs

@@ -2245,24 +2245,12 @@ impl<'ast, T: Field> Flattener<'ast, T> {
 
                 let mut substitution_map = HashMap::default();
                 for (index, p) in function.arguments.iter().enumerate() {
-                    let new_id = format!("i{}", index).into();
+                    let new_id = Identifier::internal(format!("i{}", index));
                     substitution_map.insert(p.id.id.clone(), new_id);
                 }
 
-                let mut substitutor = ZirSubstitutor::new(&substitution_map);
-                let function = ZirFunction {
-                    arguments: function
-                        .arguments
-                        .into_iter()
-                        .map(|p| substitutor.fold_parameter(p))
-                        .collect(),
-                    statements: function
-                        .statements
-                        .into_iter()
-                        .flat_map(|s| substitutor.fold_statement(s))
-                        .collect(),
-                    signature: function.signature,
-                };
+                let mut substitutor = ZirSubstitutor::new(substitution_map);
+                let function = substitutor.fold_function(function);
 
                 let solver = Solver::Zir(function);
                 let directive = FlatDirective::new(outputs, solver, inputs);

zokrates_interpreter/src/lib.rs

@@ -167,9 +167,9 @@ impl Interpreter {
         solvers: &[Solver<'ast, T>],
     ) -> Result<Vec<T>, String> {
         let solver = match solver {
-            Solver::IndexedCall(index, _) => solvers
-                .get(*index)
-                .ok_or_else(|| format!("Could not resolve solver at index {}", index))?,
+            Solver::Ref(call) => solvers
+                .get(call.index)
+                .ok_or_else(|| format!("Could not get solver at index {}", call.index))?,
             s => s,
         };