fix: rationalise witness for constant values (#984)

* rationalise witness for constant values

* Resolve issue 770

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

Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>

* Update crates/noirc_evaluator/src/ssa/context.rs

Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>

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

Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>

* Code review

* Code review

* Code review - styling

* clippy

---------

Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>

Author: guipublic

crates/noirc_evaluator/src/ssa/acir_gen.rs

@@ -67,7 +67,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)
@@ -89,7 +89,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)?
@@ -118,7 +118,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"),
+        }
+    }
 }

crates/noirc_evaluator/src/ssa/acir_gen/operations/binary.rs

@@ -1,8 +1,7 @@
 use crate::{
     ssa::{
         acir_gen::{
-            acir_mem::AcirMem, constraints, internal_var_cache::InternalVarCache, operations,
-            InternalVar,
+            constraints, internal_var_cache::InternalVarCache, operations, Acir, InternalVar,
         },
         context::SsaContext,
         node::{self, BinaryOp, Node, ObjectType},
@@ -29,8 +28,7 @@ fn get_predicate(
 pub(crate) fn evaluate(
     binary: &node::Binary,
     res_type: ObjectType,
-    var_cache: &mut InternalVarCache,
-    memory_map: &mut AcirMem,
+    acir_gen: &mut Acir,
     evaluator: &mut Evaluator,
     ctx: &SsaContext,
 ) -> Option<InternalVar> {
@@ -43,17 +41,17 @@ pub(crate) fn evaluate(
 
     let binary_output = match &binary.operator {
             BinaryOp::Add | BinaryOp::SafeAdd => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 InternalVar::from(constraints::add(
                     l_c.expression(),
                     FieldElement::one(),
                     r_c.expression(),
                 ))
             },
             BinaryOp::Sub { max_rhs_value } | BinaryOp::SafeSub { max_rhs_value } => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 if res_type == ObjectType::NativeField {
                     InternalVar::from(constraints::subtract(
                         l_c.expression(),
@@ -92,18 +90,18 @@ pub(crate) fn evaluate(
                 }
             }
             BinaryOp::Mul | BinaryOp::SafeMul => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 InternalVar::from(constraints::mul_with_witness(
                 evaluator,
                 l_c.expression(),
                 r_c.expression(),
             ))
             },
             BinaryOp::Udiv(_) => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
-                let predicate = get_predicate(var_cache,binary, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let predicate = get_predicate(&mut acir_gen.var_cache,binary, evaluator, ctx);
                 let (q_wit, _) = constraints::evaluate_udiv(
                     l_c.expression(),
                     r_c.expression(),
@@ -114,16 +112,16 @@ pub(crate) fn evaluate(
                 InternalVar::from(q_wit)
             }
             BinaryOp::Sdiv(_) => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 InternalVar::from(
                 constraints::evaluate_sdiv(l_c.expression(), r_c.expression(), evaluator).0,
             )
         },
             BinaryOp::Urem(_) => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
-                let predicate = get_predicate(var_cache,binary, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let predicate = get_predicate(&mut acir_gen.var_cache,binary, evaluator, ctx);
                 let (_, r_wit) = constraints::evaluate_udiv(
                     l_c.expression(),
                     r_c.expression(),
@@ -134,16 +132,16 @@ pub(crate) fn evaluate(
                 InternalVar::from(r_wit)
             }
             BinaryOp::Srem(_) => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 InternalVar::from(
                 // TODO: we should use variable naming here instead of .1
                 constraints::evaluate_sdiv(l_c.expression(), r_c.expression(), evaluator).1,
             )},
             BinaryOp::Div(_) => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let mut r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
-                let predicate = get_predicate(var_cache,binary, evaluator, ctx).expression().clone();
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let predicate = get_predicate(&mut acir_gen.var_cache,binary, evaluator, ctx).expression().clone();
                 if let Some(r_value) = r_c.to_const() {
                     if r_value.is_zero() {
                         panic!("Panic - division by zero");
@@ -152,7 +150,7 @@ pub(crate) fn evaluate(
                     }
                 } else {
                     //TODO avoid creating witnesses here.
-                    let x_witness = r_c.get_or_compute_witness(evaluator, false).expect("unexpected constant expression"); 
+                    let x_witness = acir_gen.var_cache.get_or_compute_witness(r_c, evaluator).expect("unexpected constant expression"); 
                     let inverse = Expression::from(&constraints::evaluate_inverse(
                         x_witness, &predicate, evaluator,
                     ));
@@ -164,20 +162,20 @@ pub(crate) fn evaluate(
                 }
             }
             BinaryOp::Eq => {
-                let l_c = var_cache.get_or_compute_internal_var(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var(binary.rhs, evaluator, ctx);
                 InternalVar::from(
-                operations::cmp::evaluate_eq(memory_map,binary.lhs, binary.rhs, l_c, r_c, ctx, evaluator),
+                operations::cmp::evaluate_eq(acir_gen,binary.lhs, binary.rhs, l_c, r_c, ctx, evaluator),
             )},
             BinaryOp::Ne => {
-                let l_c = var_cache.get_or_compute_internal_var(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var(binary.rhs, evaluator, ctx);
                 InternalVar::from(
-                operations::cmp::evaluate_neq(memory_map,binary.lhs, binary.rhs, l_c, r_c, ctx, evaluator),
+                operations::cmp::evaluate_neq(acir_gen,binary.lhs, binary.rhs, l_c, r_c, ctx, evaluator),
             )},
             BinaryOp::Ult => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 let size = ctx[binary.lhs].get_type().bits();
                 constraints::evaluate_cmp(
                     l_c.expression(),
@@ -189,8 +187,8 @@ pub(crate) fn evaluate(
                 .into()
             }
             BinaryOp::Ule => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 let size = ctx[binary.lhs].get_type().bits();
                 let e = constraints::evaluate_cmp(
                     r_c.expression(),
@@ -202,15 +200,15 @@ pub(crate) fn evaluate(
                 constraints::subtract(&Expression::one(), FieldElement::one(), &e).into()
             }
             BinaryOp::Slt => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 let s = ctx[binary.lhs].get_type().bits();
                 constraints::evaluate_cmp(l_c.expression(), r_c.expression(), s, true, evaluator)
                     .into()
             }
             BinaryOp::Sle => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 let s = ctx[binary.lhs].get_type().bits();
                 let e = constraints::evaluate_cmp(
                     r_c.expression(),
@@ -230,13 +228,13 @@ pub(crate) fn evaluate(
             )
             }
             BinaryOp::And | BinaryOp::Or | BinaryOp::Xor => {
-                let l_c = var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
-                let r_c = var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
+                let l_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.lhs, evaluator, ctx);
+                let r_c = acir_gen.var_cache.get_or_compute_internal_var_unwrap(binary.rhs, evaluator, ctx);
                 let bit_size = res_type.bits();
                 let opcode = binary.operator.clone();
                 let bitwise_result = match operations::bitwise::simplify_bitwise(&l_c, &r_c, bit_size, &opcode) {
                     Some(simplified_internal_var) => simplified_internal_var.expression().clone(),
-                    None => operations::bitwise::evaluate_bitwise(l_c, r_c, bit_size, evaluator, opcode),
+                    None => operations::bitwise::evaluate_bitwise(l_c, r_c, bit_size, evaluator, &mut acir_gen.var_cache, ctx, opcode),
                 };
                 InternalVar::from(bitwise_result)
             }

crates/noirc_evaluator/src/ssa/acir_gen/operations/bitwise.rs

@@ -8,7 +8,8 @@ use acvm::{
 
 use crate::{
     ssa::{
-        acir_gen::{constraints, InternalVar},
+        acir_gen::{constraints, internal_var_cache::InternalVarCache, InternalVar},
+        context::SsaContext,
         node::BinaryOp,
     },
     Evaluator,
@@ -76,10 +77,12 @@ pub(super) fn simplify_bitwise(
 }
 // Precondition: `lhs` and `rhs` do not represent constant expressions
 pub(super) fn evaluate_bitwise(
-    mut lhs: InternalVar,
-    mut rhs: InternalVar,
+    lhs: InternalVar,
+    rhs: InternalVar,
     bit_size: u32,
     evaluator: &mut Evaluator,
+    var_cache: &mut InternalVarCache,
+    ctx: &SsaContext,
     opcode: BinaryOp,
 ) -> Expression {
     // Check precondition
@@ -111,13 +114,8 @@ pub(super) fn evaluate_bitwise(
     // If the gate is implemented, it is expected to be better than going through bit decomposition, even if one of the operand is a constant
     // If the gate is not implemented, we rely on the ACIR simplification to remove these witnesses
     //
-
-    let mut a_witness = lhs
-        .get_or_compute_witness(evaluator, true)
-        .expect("infallible: `None` can only be returned when we disallow constant Expressions.");
-    let mut b_witness = rhs
-        .get_or_compute_witness(evaluator, true)
-        .expect("infallible: `None` can only be returned when we disallow constant Expressions.");
+    let mut a_witness = var_cache.get_or_compute_witness_unwrap(lhs, evaluator, ctx);
+    let mut b_witness = var_cache.get_or_compute_witness_unwrap(rhs, evaluator, ctx);
 
     let result = evaluator.add_witness_to_cs();
     let bit_size = if bit_size % 2 == 1 { bit_size + 1 } else { bit_size };
@@ -130,12 +128,12 @@ pub(super) fn evaluate_bitwise(
             a_witness = evaluator.create_intermediate_variable(constraints::subtract(
                 &Expression::from_field(max),
                 FieldElement::one(),
-                lhs.expression(),
+                &Expression::from(a_witness),
             ));
             b_witness = evaluator.create_intermediate_variable(constraints::subtract(
                 &Expression::from_field(max),
                 FieldElement::one(),
-                rhs.expression(),
+                &Expression::from(b_witness),
             ));
             // We do not have an OR gate yet, so we use the AND gate
             acvm::acir::BlackBoxFunc::AND

crates/noirc_evaluator/src/ssa/acir_gen/operations/cmp.rs

@@ -1,6 +1,6 @@
 use crate::{
     ssa::{
-        acir_gen::{acir_mem::AcirMem, constraints, InternalVar},
+        acir_gen::{acir_mem::AcirMem, constraints, Acir, InternalVar},
         context::SsaContext,
         mem::{MemArray, Memory},
         node::NodeId,
@@ -26,7 +26,7 @@ use iter_extended::vecmap;
 // so in reality, the NEQ instruction will be done on the fields
 // of the struct
 pub(super) fn evaluate_neq(
-    acir_mem: &mut AcirMem,
+    acir_gen: &mut Acir,
     lhs: NodeId,
     rhs: NodeId,
     l_c: Option<InternalVar>,
@@ -60,20 +60,22 @@ pub(super) fn evaluate_neq(
             )
         }
 
-        let mut x = InternalVar::from(array_eq(acir_mem, array_a, array_b, evaluator));
+        let x = InternalVar::from(array_eq(&mut acir_gen.memory, array_a, array_b, evaluator));
         // TODO we need a witness because of the directive, but we should use an expression
         // TODO if we change the Invert directive to take an `Expression`, then we
         // TODO can get rid of this extra gate.
-        let x_witness =
-            x.get_or_compute_witness(evaluator, false).expect("unexpected constant expression");
+        let x_witness = acir_gen
+            .var_cache
+            .get_or_compute_witness(x, evaluator)
+            .expect("unexpected constant expression");
 
         return Expression::from(&constraints::evaluate_zero_equality(x_witness, evaluator));
     }
 
     // Arriving here means that `lhs` and `rhs` are not Arrays
     let l_c = l_c.expect("ICE: unexpected array pointer");
     let r_c = r_c.expect("ICE: unexpected array pointer");
-    let mut x = InternalVar::from(constraints::subtract(
+    let x = InternalVar::from(constraints::subtract(
         l_c.expression(),
         FieldElement::one(),
         r_c.expression(),
@@ -89,22 +91,24 @@ pub(super) fn evaluate_neq(
         }
     } else {
         //todo we need a witness because of the directive, but we should use an expression
-        let x_witness =
-            x.get_or_compute_witness(evaluator, false).expect("unexpected constant expression");
+        let x_witness = acir_gen
+            .var_cache
+            .get_or_compute_witness(x, evaluator)
+            .expect("unexpected constant expression");
         Expression::from(&constraints::evaluate_zero_equality(x_witness, evaluator))
     }
 }
 
 pub(super) fn evaluate_eq(
-    memory_map: &mut AcirMem,
+    acir_gen: &mut Acir,
     lhs: NodeId,
     rhs: NodeId,
     l_c: Option<InternalVar>,
     r_c: Option<InternalVar>,
     ctx: &SsaContext,
     evaluator: &mut Evaluator,
 ) -> Expression {
-    let neq = evaluate_neq(memory_map, lhs, rhs, l_c, r_c, ctx, evaluator);
+    let neq = evaluate_neq(acir_gen, lhs, rhs, l_c, r_c, ctx, evaluator);
     constraints::subtract(&Expression::one(), FieldElement::one(), &neq)
 }
 

crates/noirc_evaluator/src/ssa/acir_gen/operations/intrinsics.rs

@@ -3,7 +3,7 @@ use crate::{
         acir_gen::{
             constraints::{bound_constraint_with_offset, to_radix_base},
             operations::sort::evaluate_permutation,
-            AcirMem, InternalVar, InternalVarCache,
+            Acir, AcirMem, InternalVar, InternalVarCache,
         },
         builtin,
         context::SsaContext,
@@ -33,8 +33,7 @@ pub(crate) fn evaluate(
     args: &[NodeId],
     instruction: &Instruction,
     opcode: builtin::Opcode,
-    var_cache: &mut InternalVarCache,
-    memory_map: &mut AcirMem,
+    acir_gen: &mut Acir,
     ctx: &SsaContext,
     evaluator: &mut Evaluator,
 ) -> Option<InternalVar> {
@@ -48,28 +47,30 @@ pub(crate) fn evaluate(
         Opcode::ToBits(endianess) => {
             // TODO: document where `0` and `1` are coming from, for args[0], args[1]
             let bit_size = ctx.get_as_constant(args[1]).unwrap().to_u128() as u32;
-            let l_c = var_cache.get_or_compute_internal_var_unwrap(args[0], evaluator, ctx);
+            let l_c =
+                acir_gen.var_cache.get_or_compute_internal_var_unwrap(args[0], evaluator, ctx);
             outputs = to_radix_base(l_c.expression(), 2, bit_size, endianess, evaluator);
             if let ObjectType::Pointer(a) = res_type {
-                memory_map.map_array(a, &outputs, ctx);
+                acir_gen.memory.map_array(a, &outputs, ctx);
             }
         }
         Opcode::ToRadix(endianess) => {
             // TODO: document where `0`, `1` and `2` are coming from, for args[0],args[1], args[2]
             let radix = ctx.get_as_constant(args[1]).unwrap().to_u128() as u32;
             let limb_size = ctx.get_as_constant(args[2]).unwrap().to_u128() as u32;
-            let l_c = var_cache.get_or_compute_internal_var_unwrap(args[0], evaluator, ctx);
+            let l_c =
+                acir_gen.var_cache.get_or_compute_internal_var_unwrap(args[0], evaluator, ctx);
             outputs = to_radix_base(l_c.expression(), radix, limb_size, endianess, evaluator);
             if let ObjectType::Pointer(a) = res_type {
-                memory_map.map_array(a, &outputs, ctx);
+                acir_gen.memory.map_array(a, &outputs, ctx);
             }
         }
         Opcode::Println(print_info) => {
             outputs = Vec::new(); // print statements do not output anything
             if print_info.show_output {
                 evaluate_println(
-                    var_cache,
-                    memory_map,
+                    &mut acir_gen.var_cache,
+                    &mut acir_gen.memory,
                     print_info.is_string_output,
                     args,
                     ctx,
@@ -78,9 +79,10 @@ pub(crate) fn evaluate(
             }
         }
         Opcode::LowLevel(op) => {
-            let inputs = prepare_inputs(var_cache, memory_map, args, ctx, evaluator);
+            let inputs = prepare_inputs(acir_gen, args, ctx, evaluator);
             let output_count = op.definition().output_size.0 as u32;
-            outputs = prepare_outputs(memory_map, instruction_id, output_count, ctx, evaluator);
+            outputs =
+                prepare_outputs(&mut acir_gen.memory, instruction_id, output_count, ctx, evaluator);
 
             let func_call = BlackBoxFuncCall {
                 name: op,
@@ -96,10 +98,15 @@ pub(crate) fn evaluate(
             let num_bits = array.element_type.bits();
             for i in 0..array.len {
                 in_expr.push(
-                    memory_map.load_array_element_constant_index(array, i).unwrap().to_expression(),
+                    acir_gen
+                        .memory
+                        .load_array_element_constant_index(array, i)
+                        .unwrap()
+                        .to_expression(),
                 );
             }
-            outputs = prepare_outputs(memory_map, instruction_id, array.len, ctx, evaluator);
+            outputs =
+                prepare_outputs(&mut acir_gen.memory, instruction_id, array.len, ctx, evaluator);
             let out_expr: Vec<Expression> = outputs.iter().map(|w| w.into()).collect();
             for i in 0..(out_expr.len() - 1) {
                 bound_constraint_with_offset(
@@ -119,7 +126,7 @@ pub(crate) fn evaluate(
                 sort_by: vec![0],
             }));
             if let node::ObjectType::Pointer(a) = res_type {
-                memory_map.map_array(a, &outputs, ctx);
+                acir_gen.memory.map_array(a, &outputs, ctx);
             } else {
                 unreachable!();
             }
@@ -134,24 +141,22 @@ pub(crate) fn evaluate(
 
 // Transform the arguments of intrinsic functions into witnesses
 fn prepare_inputs(
-    var_cache: &mut InternalVarCache,
-    memory_map: &mut AcirMem,
+    acir_gen: &mut Acir,
     arguments: &[NodeId],
     cfg: &SsaContext,
     evaluator: &mut Evaluator,
 ) -> Vec<FunctionInput> {
     let mut inputs: Vec<FunctionInput> = Vec::new();
 
     for argument in arguments {
-        inputs.extend(resolve_node_id(argument, var_cache, memory_map, cfg, evaluator))
+        inputs.extend(resolve_node_id(argument, acir_gen, cfg, evaluator))
     }
     inputs
 }
 
 fn resolve_node_id(
     node_id: &NodeId,
-    var_cache: &mut InternalVarCache,
-    memory_map: &mut AcirMem,
+    acir_gen: &mut Acir,
     cfg: &SsaContext,
     evaluator: &mut Evaluator,
 ) -> Vec<FunctionInput> {
@@ -162,7 +167,7 @@ fn resolve_node_id(
             match node_obj_type {
                 // If the `Variable` represents a Pointer
                 // Then we know that it is an `Array`
-                node::ObjectType::Pointer(a) => resolve_array(a, memory_map, cfg, evaluator),
+                node::ObjectType::Pointer(a) => resolve_array(a, acir_gen, cfg, evaluator),
                 // If it is not a pointer, we attempt to fetch the witness associated with it
                 _ => match v.witness {
                     Some(w) => {
@@ -176,24 +181,19 @@ fn resolve_node_id(
             // Upon the case that the `NodeObject` is not a `Variable`,
             // we attempt to fetch an associated `InternalVar`.
             // Otherwise, this is a internal compiler error.
-            let internal_var = var_cache.get(node_id);
-            match internal_var {
-                Some(var) => {
-                    let witness = var
-                        .clone()
-                        .get_or_compute_witness(evaluator, false)
-                        .expect("unexpected constant expression");
-                    vec![FunctionInput { witness, num_bits: node_object.size_in_bits() }]
-                }
-                None => unreachable!("invalid input: {:?}", node_object),
-            }
+            let internal_var = acir_gen.var_cache.get(node_id).expect("invalid input").clone();
+            let witness = acir_gen
+                .var_cache
+                .get_or_compute_witness(internal_var, evaluator)
+                .expect("unexpected constant expression");
+            vec![FunctionInput { witness, num_bits: node_object.size_in_bits() }]
         }
     }
 }
 
 fn resolve_array(
     array_id: ArrayId,
-    acir_mem: &mut AcirMem,
+    acir_gen: &mut Acir,
     cfg: &SsaContext,
     evaluator: &mut Evaluator,
 ) -> Vec<FunctionInput> {
@@ -202,16 +202,15 @@ fn resolve_array(
     let array = &cfg.mem[array_id];
     let num_bits = array.element_type.bits();
     for i in 0..array.len {
-        let mut arr_element = acir_mem
+        let mut arr_element = acir_gen
+            .memory
             .load_array_element_constant_index(array, i)
             .expect("array index out of bounds");
-
-        let witness = arr_element.get_or_compute_witness(evaluator, true).expect(
-            "infallible: `None` can only be returned when we disallow constant Expressions.",
-        );
+        let witness =
+            acir_gen.var_cache.get_or_compute_witness_unwrap(arr_element.clone(), evaluator, cfg);
         let func_input = FunctionInput { witness, num_bits };
-
-        acir_mem.insert(array.id, i, arr_element);
+        arr_element.set_witness(witness);
+        acir_gen.memory.insert(array.id, i, arr_element);
 
         inputs.push(func_input)
     }
@@ -287,7 +286,7 @@ fn evaluate_println(
                 log_string = format_field_string(field);
             }
             None => {
-                let mut var = var_cache
+                let var = var_cache
                     .get(&node_id)
                     .unwrap_or_else(|| {
                         panic!(
@@ -298,7 +297,7 @@ fn evaluate_println(
                     .clone();
                 if let Some(field) = var.to_const() {
                     log_string = format_field_string(field);
-                } else if let Some(w) = var.get_or_compute_witness(evaluator, false) {
+                } else if let Some(w) = var_cache.get_or_compute_witness(var, evaluator) {
                     // We check whether there has already been a cached witness for this node. If not, we generate a new witness and include it in the logs
                     // TODO we need a witness because of the directive, but we should use an expression
                     log_witnesses.push(w);

crates/noirc_evaluator/src/ssa/acir_gen/operations/return.rs

@@ -41,10 +41,8 @@ pub(crate) fn evaluate(
         };
 
         let mut witnesses: Vec<Witness> = Vec::new();
-        for mut object in objects {
-            let witness = object.get_or_compute_witness(evaluator, true).expect(
-                "infallible: `None` can only be returned when we disallow constant Expressions.",
-            );
+        for object in objects {
+            let witness = var_cache.get_or_compute_witness_unwrap(object, evaluator, ctx);
             // Before pushing to the public inputs, we need to check that
             // it was not a private ABI input
             if evaluator.is_private_abi_input(witness) {

crates/noirc_evaluator/src/ssa/context.rs

@@ -17,7 +17,6 @@ use iter_extended::vecmap;
 use noirc_errors::Location;
 use noirc_frontend::monomorphization::ast::{Definition, Expression, FuncId, Literal, Type};
 use num_bigint::BigUint;
-use num_traits::{One, Zero};
 use std::collections::{HashMap, HashSet};
 
 // This is a 'master' class for generating the SSA IR from the AST
@@ -35,6 +34,7 @@ pub(crate) struct SsaContext {
 
     pub(crate) functions: HashMap<FuncId, function::SsaFunction>,
     pub(crate) opcode_ids: HashMap<builtin::Opcode, NodeId>,
+    pub(crate) constants: HashMap<FieldElement, Vec<NodeId>>,
 
     //Adjacency Matrix of the call graph; list of rows where each row indicates the functions called by the function whose FuncIndex is the row number
     pub(crate) call_graph: Vec<Vec<u8>>,
@@ -62,6 +62,7 @@ impl Default for SsaContext {
             dummy_store: HashMap::new(),
             dummy_load: HashMap::new(),
             locations: HashMap::new(),
+            constants: HashMap::new(),
         };
         block::create_first_block(&mut pc);
         pc.one_with_type(node::ObjectType::Boolean);
@@ -72,11 +73,11 @@ impl Default for SsaContext {
 
 impl SsaContext {
     pub(crate) fn zero(&self) -> NodeId {
-        self.find_const_with_type(&BigUint::zero(), node::ObjectType::Boolean).unwrap()
+        self.find_const_with_type(&FieldElement::zero(), node::ObjectType::Boolean).unwrap()
     }
 
     pub(crate) fn one(&self) -> NodeId {
-        self.find_const_with_type(&BigUint::one(), node::ObjectType::Boolean).unwrap()
+        self.find_const_with_type(&FieldElement::one(), node::ObjectType::Boolean).unwrap()
     }
 
     pub(crate) fn zero_with_type(&mut self, obj_type: ObjectType) -> NodeId {
@@ -92,7 +93,7 @@ impl SsaContext {
             return false;
         }
         let typ = self.object_type(id);
-        if let Some(one) = self.find_const_with_type(&BigUint::one(), typ) {
+        if let Some(one) = self.find_const_with_type(&FieldElement::one(), typ) {
             id == one
         } else {
             false
@@ -104,7 +105,7 @@ impl SsaContext {
             return false;
         }
         let typ = self.object_type(id);
-        if let Some(zero) = self.find_const_with_type(&BigUint::zero(), typ) {
+        if let Some(zero) = self.find_const_with_type(&FieldElement::zero(), typ) {
             id == zero
         } else {
             false
@@ -571,14 +572,14 @@ impl SsaContext {
 
     pub(crate) fn find_const_with_type(
         &self,
-        value: &BigUint,
+        value: &FieldElement,
         e_type: node::ObjectType,
     ) -> Option<NodeId> {
-        //TODO We should map constant values to id
-        for (idx, o) in &self.nodes {
-            if let node::NodeObject::Const(c) = o {
-                if c.value == *value && c.get_type() == e_type {
-                    return Some(NodeId(idx));
+        // we look for the node in the constants map
+        if let Some(ids) = self.constants.get(value) {
+            for &id in ids {
+                if self[id].get_type() == e_type {
+                    return Some(id);
                 }
             }
         }
@@ -589,16 +590,20 @@ impl SsaContext {
     // If such object does not exist, we create one
     pub(crate) fn get_or_create_const(&mut self, x: FieldElement, t: node::ObjectType) -> NodeId {
         let value = BigUint::from_bytes_be(&x.to_be_bytes());
-        if let Some(prev_const) = self.find_const_with_type(&value, t) {
+        if let Some(prev_const) = self.find_const_with_type(&x, t) {
             return prev_const;
         }
 
-        self.add_const(node::Constant {
+        let id = self.add_const(node::Constant {
             id: NodeId::dummy(),
             value,
             value_str: String::new(),
             value_type: t,
-        })
+        });
+        // Adds the id into the constants map
+        let ids = self.constants.entry(x).or_default();
+        ids.push(id);
+        id
     }
 
     // Return the type of the operation result, based on the left hand type