fix: rationalise witness for constant values

crates/noirc_evaluator/src/ssa/acir_gen.rs

@@ -46,7 +46,7 @@ impl Acir {
 
         let output = match &ins.operation {
             Operation::Binary(binary) => {
-                binary::evaluate(binary, ins.res_type, var_cache, acir_mem, evaluator, ctx)
+                binary::evaluate(binary, ins.res_type, self, evaluator, ctx)
             }
             Operation::Constrain(value, ..) => {
                 constrain::evaluate(value, var_cache, evaluator, ctx)
@@ -68,7 +68,7 @@ impl Acir {
                     }
                     _ => *opcode,
                 };
-                intrinsics::evaluate(args, ins, opcode, var_cache, acir_mem, ctx, evaluator)
+                intrinsics::evaluate(args, ins, opcode, self, ctx, evaluator)
             }
             Operation::Return(node_ids) => {
                 r#return::evaluate(node_ids, acir_mem, var_cache, evaluator, ctx)?
@@ -97,7 +97,7 @@ impl Acir {
         // then we add it to the `InternalVar` cache
         if let Some(mut output) = output {
             output.set_id(ins.id);
-            self.var_cache.update(ins.id, output);
+            self.var_cache.update(output);
         }
 
         Ok(())

crates/noirc_evaluator/src/ssa/acir_gen/internal_var.rs

@@ -1,4 +1,4 @@
-use crate::{ssa::acir_gen::expression_to_witness, ssa::node::NodeId, Evaluator};
+use crate::ssa::node::NodeId;
 use acvm::{
     acir::native_types::{Expression, Witness},
     FieldElement,
@@ -38,9 +38,16 @@ impl InternalVar {
     pub(crate) fn set_id(&mut self, id: NodeId) {
         self.id = Some(id)
     }
+    pub(crate) fn get_id(&self) -> Option<NodeId> {
+        self.id
+    }
     pub(crate) fn cached_witness(&self) -> &Option<Witness> {
         &self.cached_witness
     }
+    pub(crate) fn set_witness(&mut self, w: Witness) {
+        debug_assert!(self.cached_witness.is_none() || self.cached_witness == Some(w));
+        self.cached_witness = Some(w);
+    }
 
     pub(crate) fn to_expression(&self) -> Expression {
         if let Some(w) = self.cached_witness {
@@ -70,7 +77,7 @@ impl InternalVar {
     /// Example: f(x,y) = 10
     /// `f` is a constant expression because there are no
     /// bi-variate or uni-variate terms, just a constant.
-    fn is_const_expression(&self) -> bool {
+    pub(crate) fn is_const_expression(&self) -> bool {
         self.expression.is_const()
     }
 
@@ -101,32 +108,6 @@ impl InternalVar {
     pub(crate) fn from_constant(constant: FieldElement) -> InternalVar {
         InternalVar { expression: Expression::from_field(constant), cached_witness: None, id: None }
     }
-
-    /// Generates a `Witness` that is equal to the `expression`.
-    /// - If a `Witness` has previously been generated
-    /// we return that.
-    /// - If the Expression represents a constant, we return None.
-    pub(crate) fn get_or_compute_witness(
-        &mut self,
-        evaluator: &mut Evaluator,
-        create_witness_for_const: bool,
-    ) -> Option<Witness> {
-        // Check if we've already generated a `Witness` which is equal to
-        // the stored `Expression`
-        if let Some(witness) = self.cached_witness {
-            return Some(witness);
-        }
-
-        // There are cases where we need to convert a constant expression
-        // into a witness.
-        if !create_witness_for_const && self.is_const_expression() {
-            return None;
-        }
-
-        self.cached_witness = Some(expression_to_witness(self.expression.clone(), evaluator));
-
-        self.cached_witness
-    }
 }
 
 impl PartialEq for InternalVar {
@@ -184,21 +165,18 @@ impl From<FieldElement> for InternalVar {
 
 #[cfg(test)]
 mod tests {
-    use crate::{ssa::acir_gen::InternalVar, Evaluator};
-    use acvm::{acir::native_types::Witness, FieldElement};
+    use crate::ssa::acir_gen::InternalVar;
+    use acvm::FieldElement;
 
     #[test]
     fn internal_var_const_expression() {
-        let mut evaluator = Evaluator::default();
-
         let expected_constant = FieldElement::from(123456789u128);
 
         // Initialize an InternalVar with a FieldElement
-        let mut internal_var = InternalVar::from_constant(expected_constant);
+        let internal_var = InternalVar::from_constant(expected_constant);
 
         // We currently do not create witness when the InternalVar was created using a constant
-        let witness = internal_var.get_or_compute_witness(&mut evaluator, false);
-        assert!(witness.is_none());
+        assert!(internal_var.cached_witness().is_none());
 
         match internal_var.to_const() {
             Some(got_constant) => assert_eq!(got_constant, expected_constant),
@@ -207,19 +185,4 @@ mod tests {
             }
         }
     }
-    #[test]
-    fn internal_var_from_witness() {
-        let mut evaluator = Evaluator::default();
-
-        let expected_witness = Witness(1234);
-        // Initialize an InternalVar with a `Witness`
-        let mut internal_var = InternalVar::from_witness(expected_witness);
-
-        // We should get back the same `Witness`
-        let got_witness = internal_var.get_or_compute_witness(&mut evaluator, false);
-        match got_witness {
-            Some(got_witness) => assert_eq!(got_witness, expected_witness),
-            None => panic!("expected a `Witness` value"),
-        }
-    }
 }

crates/noirc_evaluator/src/ssa/acir_gen/internal_var_cache.rs

@@ -6,12 +6,22 @@ use crate::{
     },
     Evaluator,
 };
-use acvm::FieldElement;
+use acvm::{
+    acir::native_types::{Expression, Witness},
+    FieldElement,
+};
 use std::collections::HashMap;
 
+use super::expression_to_witness;
+
 #[derive(Default)]
 pub(crate) struct InternalVarCache {
+    /// Map node id to an InternalVar
     inner: HashMap<NodeId, InternalVar>,
+    /// Map field values to an InternalVar, which lazily gets a witness when required
+    /// This avoids us to re-create another witness for the same value
+    /// A witness for a field value should be avoided but can be needed as an opcode input, for example the logic opcode.
+    constants: HashMap<FieldElement, InternalVar>,
 }
 
 impl InternalVarCache {
@@ -31,8 +41,26 @@ impl InternalVarCache {
 
         let mut var = match ctx.try_get_node(id)? {
             NodeObject::Const(c) => {
+                // use the InternalVar from constants if exists
                 let field_value = FieldElement::from_be_bytes_reduce(&c.value.to_bytes_be());
-                InternalVar::from_constant(field_value)
+                let mut result = InternalVar::from_constant(field_value);
+                if let Some(c_var) = self.constants.get(&field_value) {
+                    result = c_var.clone();
+                }
+
+                //use witness from other nodes if exists
+                let new_vec = Vec::new();
+                let constant_ids = ctx.constants.get(&field_value).unwrap_or(&new_vec);
+                let cached_witness =
+                    constant_ids.iter().find_map(|c_id| match self.inner.get(c_id) {
+                        Some(i_var) => *i_var.cached_witness(),
+                        None => None,
+                    });
+                if let Some(w) = cached_witness {
+                    result.set_witness(w);
+                }
+
+                result
             }
             NodeObject::Variable(variable) => {
                 let variable_type = variable.get_type();
@@ -76,10 +104,184 @@ impl InternalVarCache {
             .expect("ICE: `NodeId` type cannot be converted into an `InternalVar`")
     }
 
-    pub(super) fn update(&mut self, id: NodeId, var: InternalVar) {
-        self.inner.insert(id, var);
-    }
     pub(super) fn get(&mut self, id: &NodeId) -> Option<&InternalVar> {
         self.inner.get(id)
     }
+
+    // Transform a field element into a witness
+    // It implements a 'get or create' pattern to ensure only one witness is created per element
+    fn const_to_witness(
+        &mut self,
+        value: FieldElement,
+        evaluator: &mut Evaluator,
+        ctx: &SsaContext,
+    ) -> Witness {
+        if let Some(ids) = ctx.constants.get(&value) {
+            //we have a constant node object for the value
+            if let Some(id) = ids.first() {
+                let var = self.get_or_compute_internal_var_unwrap(*id, evaluator, ctx);
+                if let Some(w) = var.cached_witness() {
+                    return *w;
+                }
+            }
+            // We generate a witness and assigns it
+            let w = evaluator.create_intermediate_variable(Expression::from(value));
+            for &id in ids {
+                let mut cached_var = self.get_or_compute_internal_var_unwrap(id, evaluator, ctx);
+                assert!(cached_var.cached_witness().is_none());
+                cached_var.set_witness(w);
+                self.update(cached_var);
+            }
+            w
+        } else {
+            //if not, we use the constants map
+            let var =
+                self.constants.entry(value).or_insert_with(|| InternalVar::from_constant(value));
+            Self::const_to_witness_helper(var, evaluator)
+        }
+    }
+
+    // Helper function which generates a witness for an InternalVar
+    // Do not call outside const_to_witness()
+    fn const_to_witness_helper(var: &mut InternalVar, evaluator: &mut Evaluator) -> Witness {
+        let w = Self::internal_get_or_compute_witness(var, evaluator);
+        if w.is_none() {
+            let witness = expression_to_witness(var.to_expression(), evaluator);
+            var.set_witness(witness);
+        }
+        var.cached_witness().expect("Infallible, the witness is computed before")
+    }
+
+    /// Get or compute a witness for an internal var
+    /// WARNING: It generates a witness even if the internal var is constant, so it should be used only if the var is an input
+    /// to some ACIR opcode which requires a witness
+    pub(crate) fn get_or_compute_witness_unwrap(
+        &mut self,
+        mut var: InternalVar,
+        evaluator: &mut Evaluator,
+        ctx: &SsaContext,
+    ) -> Witness {
+        if let Some(v) = var.to_const() {
+            self.const_to_witness(v, evaluator, ctx)
+        } else {
+            let w = Self::internal_get_or_compute_witness(&mut var, evaluator)
+                .expect("infallible non const expression");
+            var.set_witness(w);
+            self.update(var);
+            w
+        }
+    }
+
+    /// Get or compute a witness equating the internal var
+    /// It returns None when the variable is a constant instead of creating a witness
+    /// because we should not need a witness in that case
+    /// If you really need one, you can use get_or_compute_witness_unwrap()
+    pub(crate) fn get_or_compute_witness(
+        &mut self,
+        mut var: InternalVar,
+        evaluator: &mut Evaluator,
+    ) -> Option<Witness> {
+        let w = Self::internal_get_or_compute_witness(&mut var, evaluator);
+        if w.is_some() {
+            assert!(var.cached_witness().is_some());
+        } else {
+            return None;
+        };
+        self.update(var);
+
+        w
+    }
+
+    /// Generates a `Witness` that is equal to the `expression`.
+    /// - If a `Witness` has previously been generated, we return it.
+    /// - If the Expression represents a constant, we return None.
+    fn internal_get_or_compute_witness(
+        var: &mut InternalVar,
+        evaluator: &mut Evaluator,
+    ) -> Option<Witness> {
+        // Check if we've already generated a `Witness` which is equal to
+        // the stored `Expression`
+        if let Some(witness) = var.cached_witness() {
+            return Some(*witness);
+        }
+
+        // We do not generate a witness for constant values. It can only be done at the InternalVarCache level.
+        if var.is_const_expression() {
+            return None;
+        }
+
+        var.set_witness(expression_to_witness(var.expression().clone(), evaluator));
+
+        *var.cached_witness()
+    }
+
+    pub(super) fn update(&mut self, var: InternalVar) {
+        if let Some(id) = var.get_id() {
+            self.inner.insert(id, var);
+        } else if let Some(value) = var.to_const() {
+            self.constants.insert(value, var);
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use acvm::{acir::native_types::Witness, FieldElement};
+
+    use crate::{
+        ssa::{
+            acir_gen::internal_var_cache::{InternalVar, InternalVarCache},
+            context::SsaContext,
+        },
+        Evaluator,
+    };
+
+    // Check that only one witness is generated for const values
+    #[test]
+    fn test_const_witness() {
+        let mut eval = Evaluator::default();
+        let ctx = SsaContext::default();
+        let mut var_cache = InternalVarCache::default();
+        let v1 = var_cache.get_or_compute_internal_var_unwrap(ctx.one(), &mut eval, &ctx);
+        let v2 = var_cache.get_or_compute_internal_var_unwrap(ctx.zero(), &mut eval, &ctx);
+        let w1 = var_cache.get_or_compute_witness_unwrap(v1, &mut eval, &ctx);
+        let w2 = var_cache.get_or_compute_witness_unwrap(v2, &mut eval, &ctx);
+        let w11 = var_cache.get_or_compute_witness_unwrap(
+            InternalVar::from_constant(FieldElement::one()),
+            &mut eval,
+            &ctx,
+        );
+        let w21 = var_cache.get_or_compute_witness_unwrap(
+            InternalVar::from_constant(FieldElement::zero()),
+            &mut eval,
+            &ctx,
+        );
+        let two = FieldElement::one() + FieldElement::one();
+        assert!(var_cache.constants.is_empty());
+        assert_eq!(w1, w11);
+        assert_eq!(w2, w21);
+        var_cache.const_to_witness(two, &mut eval, &ctx);
+        assert!(var_cache.constants.len() == 1);
+        var_cache.const_to_witness(two, &mut eval, &ctx);
+        assert!(var_cache.constants.len() == 1);
+        var_cache.const_to_witness(FieldElement::one(), &mut eval, &ctx);
+        assert!(var_cache.constants.len() == 1);
+    }
+
+    #[test]
+    fn internal_var_from_witness() {
+        let mut evaluator = Evaluator::default();
+        let expected_witness = Witness(1234);
+        // Initialize an InternalVar with a `Witness`
+        let mut internal_var = InternalVar::from_witness(expected_witness);
+
+        // We should get back the same `Witness`
+        let got_witness =
+            InternalVarCache::internal_get_or_compute_witness(&mut internal_var, &mut evaluator);
+        match got_witness {
+            Some(got_witness) => assert_eq!(got_witness, expected_witness),
+            None => panic!("expected a `Witness` value"),
+        }
+    }
 }