feat: use constraint information to perform constant folding (#4060)

# Description

## Problem\*

Resolves #4070

## Summary\*

This PR uses `Instruction::Constrain` to perform simplifications in SSA
where one value is swapped for another simpler value due to the fact
that the circuit will be unsatisfiable should the two values not be
equivalent (because of the `Instruction::Constrain`.

A previous implementation of this resulted in the bug
#2645. This was because
constraints which were only active in certain areas of the code were
applied across all of the SSA. This implementation avoids this by using
a separate cache for the constrained values for each value of
`side_effects_enabled_var`. This means that passing a
`Instruction::EnableSideEffects` instruction will "forget" any
constraints which aren't currently active.

One trouble from this PR is that we now require multiple passes in order
to fully simplify the circuit. I've simply added another set of the
final 3 passes but we could in future perform extra passes until the SSA
stabilises.

## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

compiler/noirc_evaluator/src/ssa.rs

@@ -63,6 +63,10 @@ pub(crate) fn optimize_into_acir(
         // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
         .run_pass(Ssa::mem2reg, "After Mem2Reg:")
         .run_pass(Ssa::fold_constants, "After Constant Folding:")
+        .run_pass(
+            Ssa::fold_constants_using_constraints,
+            "After Constant Folding With Constraint Info:",
+        )
         .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:")
         .finish();
 

compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -22,6 +22,7 @@
 //! different blocks are merged, i.e. after the [`flatten_cfg`][super::flatten_cfg] pass.
 use std::collections::HashSet;
 
+use acvm::FieldElement;
 use iter_extended::vecmap;
 
 use crate::ssa::{
@@ -30,7 +31,8 @@ use crate::ssa::{
         dfg::{DataFlowGraph, InsertInstructionResult},
         function::Function,
         instruction::{Instruction, InstructionId},
-        value::ValueId,
+        types::Type,
+        value::{Value, ValueId},
     },
     ssa_gen::Ssa,
 };
@@ -43,16 +45,29 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn fold_constants(mut self) -> Ssa {
         for function in self.functions.values_mut() {
-            constant_fold(function);
+            constant_fold(function, false);
+        }
+        self
+    }
+
+    /// Performs constant folding on each instruction.
+    ///
+    /// Also uses constraint information to inform more optimizations.
+    ///
+    /// See [`constant_folding`][self] module for more information.
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub(crate) fn fold_constants_using_constraints(mut self) -> Ssa {
+        for function in self.functions.values_mut() {
+            constant_fold(function, true);
         }
         self
     }
 }
 
 /// The structure of this pass is simple:
 /// Go through each block and re-insert all instructions.
-fn constant_fold(function: &mut Function) {
-    let mut context = Context::default();
+fn constant_fold(function: &mut Function, use_constraint_info: bool) {
+    let mut context = Context { use_constraint_info, ..Default::default() };
     context.block_queue.push(function.entry_block());
 
     while let Some(block) = context.block_queue.pop() {
@@ -67,6 +82,7 @@ fn constant_fold(function: &mut Function) {
 
 #[derive(Default)]
 struct Context {
+    use_constraint_info: bool,
     /// Maps pre-folded ValueIds to the new ValueIds obtained by re-inserting the instruction.
     visited_blocks: HashSet<BasicBlockId>,
     block_queue: Vec<BasicBlockId>,
@@ -79,24 +95,43 @@ impl Context {
         // Cache of instructions without any side-effects along with their outputs.
         let mut cached_instruction_results: HashMap<Instruction, Vec<ValueId>> = HashMap::default();
 
+        // Contains sets of values which are constrained to be equivalent to each other.
+        //
+        // The mapping's structure is `side_effects_enabled_var => (constrained_value => simplified_value)`.
+        //
+        // We partition the maps of constrained values according to the side-effects flag at the point
+        // at which the values are constrained. This prevents constraints which are only sometimes enforced
+        // being used to modify the rest of the program.
+        let mut constraint_simplification_mappings: HashMap<ValueId, HashMap<ValueId, ValueId>> =
+            HashMap::default();
+        let mut side_effects_enabled_var =
+            function.dfg.make_constant(FieldElement::one(), Type::bool());
+
         for instruction_id in instructions {
-            Self::fold_constants_into_instruction(
+            self.fold_constants_into_instruction(
                 &mut function.dfg,
                 block,
                 instruction_id,
                 &mut cached_instruction_results,
+                &mut constraint_simplification_mappings,
+                &mut side_effects_enabled_var,
             );
         }
         self.block_queue.extend(function.dfg[block].successors());
     }
 
     fn fold_constants_into_instruction(
+        &self,
         dfg: &mut DataFlowGraph,
         block: BasicBlockId,
         id: InstructionId,
         instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        constraint_simplification_mappings: &mut HashMap<ValueId, HashMap<ValueId, ValueId>>,
+        side_effects_enabled_var: &mut ValueId,
     ) {
-        let instruction = Self::resolve_instruction(id, dfg);
+        let constraint_simplification_mapping =
+            constraint_simplification_mappings.entry(*side_effects_enabled_var).or_default();
+        let instruction = Self::resolve_instruction(id, dfg, constraint_simplification_mapping);
         let old_results = dfg.instruction_results(id).to_vec();
 
         // If a copy of this instruction exists earlier in the block, then reuse the previous results.
@@ -110,15 +145,49 @@ impl Context {
 
         Self::replace_result_ids(dfg, &old_results, &new_results);
 
-        Self::cache_instruction(instruction, new_results, dfg, instruction_result_cache);
+        self.cache_instruction(
+            instruction.clone(),
+            new_results,
+            dfg,
+            instruction_result_cache,
+            constraint_simplification_mapping,
+        );
+
+        // If we just inserted an `Instruction::EnableSideEffects`, we need to update `side_effects_enabled_var`
+        // so that we use the correct set of constrained values in future.
+        if let Instruction::EnableSideEffects { condition } = instruction {
+            *side_effects_enabled_var = condition;
+        };
     }
 
     /// Fetches an [`Instruction`] by its [`InstructionId`] and fully resolves its inputs.
-    fn resolve_instruction(instruction_id: InstructionId, dfg: &DataFlowGraph) -> Instruction {
+    fn resolve_instruction(
+        instruction_id: InstructionId,
+        dfg: &DataFlowGraph,
+        constraint_simplification_mapping: &HashMap<ValueId, ValueId>,
+    ) -> Instruction {
         let instruction = dfg[instruction_id].clone();
 
+        // Alternate between resolving `value_id` in the `dfg` and checking to see if the resolved value
+        // has been constrained to be equal to some simpler value in the current block.
+        //
+        // This allows us to reach a stable final `ValueId` for each instruction input as we add more
+        // constraints to the cache.
+        fn resolve_cache(
+            dfg: &DataFlowGraph,
+            cache: &HashMap<ValueId, ValueId>,
+            value_id: ValueId,
+        ) -> ValueId {
+            let resolved_id = dfg.resolve(value_id);
+            match cache.get(&resolved_id) {
+                Some(cached_value) => resolve_cache(dfg, cache, *cached_value),
+                None => resolved_id,
+            }
+        }
+
         // Resolve any inputs to ensure that we're comparing like-for-like instructions.
-        instruction.map_values(|value_id| dfg.resolve(value_id))
+        instruction
+            .map_values(|value_id| resolve_cache(dfg, constraint_simplification_mapping, value_id))
     }
 
     /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations
@@ -152,11 +221,42 @@ impl Context {
     }
 
     fn cache_instruction(
+        &self,
         instruction: Instruction,
         instruction_results: Vec<ValueId>,
         dfg: &DataFlowGraph,
         instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        constraint_simplification_mapping: &mut HashMap<ValueId, ValueId>,
     ) {
+        if self.use_constraint_info {
+            // If the instruction was a constraint, then create a link between the two `ValueId`s
+            // to map from the more complex to the simpler value.
+            if let Instruction::Constrain(lhs, rhs, _) = instruction {
+                // These `ValueId`s should be fully resolved now.
+                match (&dfg[lhs], &dfg[rhs]) {
+                    // Ignore trivial constraints
+                    (Value::NumericConstant { .. }, Value::NumericConstant { .. }) => (),
+
+                    // Prefer replacing with constants where possible.
+                    (Value::NumericConstant { .. }, _) => {
+                        constraint_simplification_mapping.insert(rhs, lhs);
+                    }
+                    (_, Value::NumericConstant { .. }) => {
+                        constraint_simplification_mapping.insert(lhs, rhs);
+                    }
+                    // Otherwise prefer block parameters over instruction results.
+                    // This is as block parameters are more likely to be a single witness rather than a full expression.
+                    (Value::Param { .. }, Value::Instruction { .. }) => {
+                        constraint_simplification_mapping.insert(rhs, lhs);
+                    }
+                    (Value::Instruction { .. }, Value::Param { .. }) => {
+                        constraint_simplification_mapping.insert(lhs, rhs);
+                    }
+                    (_, _) => (),
+                }
+            }
+        }
+
         // If the instruction doesn't have side-effects, cache the results so we can reuse them if
         // the same instruction appears again later in the block.
         if instruction.is_pure(dfg) {
@@ -336,9 +436,9 @@ mod test {
         //
         // fn main f0 {
         //   b0(v0: u16, Field 255: Field):
-        //      v5 = div v0, Field 255
-        //      v6 = truncate v5 to 8 bits, max_bit_size: 16
-        //      return v6
+        //      v6 = div v0, Field 255
+        //      v7 = truncate v6 to 8 bits, max_bit_size: 16
+        //      return v7
         // }
         main.dfg.set_value_from_id(v1, constant);
 
@@ -354,7 +454,7 @@ mod test {
         );
         assert_eq!(
             &main.dfg[instructions[1]],
-            &Instruction::Truncate { value: ValueId::test_new(5), bit_size: 8, max_bit_size: 16 }
+            &Instruction::Truncate { value: ValueId::test_new(6), bit_size: 8, max_bit_size: 16 }
         );
     }