chore(ssa refactor): Add basic instruction simplification (#1329)

* Add basic instruction simplification

* Cargo fmt

* Add comments

crates/noirc_evaluator/src/ssa_refactor/ir/dfg.rs

@@ -100,15 +100,21 @@ impl DataFlowGraph {
         id
     }
 
-    /// Replace an instruction id with another.
-    ///
-    /// This function should generally be avoided if possible in favor of inserting new
-    /// instructions since it does not check whether the instruction results of the removed
-    /// instruction are still in use. Users of this function thus need to ensure the old
-    /// instruction's results are no longer in use or are otherwise compatible with the
-    /// new instruction's result count and types.
-    pub(crate) fn replace_instruction(&mut self, id: Id<Instruction>, instruction: Instruction) {
-        self.instructions[id] = instruction;
+    /// Inserts a new instruction at the end of the given block and returns its results
+    pub(crate) fn insert_instruction(
+        &mut self,
+        instruction: Instruction,
+        block: BasicBlockId,
+        ctrl_typevars: Option<Vec<Type>>,
+    ) -> InsertInstructionResult {
+        match instruction.simplify(self) {
+            Some(simplification) => InsertInstructionResult::SimplifiedTo(simplification),
+            None => {
+                let id = self.make_instruction(instruction, ctrl_typevars);
+                self.insert_instruction_in_block(block, id);
+                InsertInstructionResult::Results(self.instruction_results(id))
+            }
+        }
     }
 
     /// Insert a value into the dfg's storage and return an id to reference it.
@@ -300,6 +306,52 @@ impl std::ops::IndexMut<BasicBlockId> for DataFlowGraph {
     }
 }
 
+// The result of calling DataFlowGraph::insert_instruction can
+// be a list of results or a single ValueId if the instruction was simplified
+// to an existing value.
+pub(crate) enum InsertInstructionResult<'dfg> {
+    Results(&'dfg [ValueId]),
+    SimplifiedTo(ValueId),
+    InstructionRemoved,
+}
+
+impl<'dfg> InsertInstructionResult<'dfg> {
+    /// Retrieve the first (and expected to be the only) result.
+    pub(crate) fn first(&self) -> ValueId {
+        match self {
+            InsertInstructionResult::SimplifiedTo(value) => *value,
+            InsertInstructionResult::Results(results) => results[0],
+            InsertInstructionResult::InstructionRemoved => {
+                panic!("Instruction was removed, no results")
+            }
+        }
+    }
+
+    /// Return all the results contained in the internal results array.
+    /// This is used for instructions returning multiple results that were
+    /// not simplified - like function calls.
+    pub(crate) fn results(&self) -> &'dfg [ValueId] {
+        match self {
+            InsertInstructionResult::Results(results) => results,
+            InsertInstructionResult::SimplifiedTo(_) => {
+                panic!("InsertInstructionResult::results called on a simplified instruction")
+            }
+            InsertInstructionResult::InstructionRemoved => {
+                panic!("InsertInstructionResult::results called on a removed instruction")
+            }
+        }
+    }
+
+    /// Returns the amount of ValueIds contained
+    pub(crate) fn len(&self) -> usize {
+        match self {
+            InsertInstructionResult::SimplifiedTo(_) => 1,
+            InsertInstructionResult::Results(results) => results.len(),
+            InsertInstructionResult::InstructionRemoved => 0,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::DataFlowGraph;

crates/noirc_evaluator/src/ssa_refactor/ir/instruction.rs

@@ -1,7 +1,13 @@
-use acvm::acir::BlackBoxFunc;
+use acvm::{acir::BlackBoxFunc, FieldElement};
 use iter_extended::vecmap;
 
-use super::{basic_block::BasicBlockId, map::Id, types::Type, value::ValueId};
+use super::{
+    basic_block::BasicBlockId,
+    dfg::DataFlowGraph,
+    map::Id,
+    types::Type,
+    value::{Value, ValueId},
+};
 
 /// Reference to an instruction
 ///
@@ -151,6 +157,48 @@ impl Instruction {
             }
         }
     }
+
+    /// Try to simplify this instruction. If the instruction can be simplified to a known value,
+    /// that value is returned. Otherwise None is returned.
+    pub(crate) fn simplify(&self, dfg: &mut DataFlowGraph) -> Option<ValueId> {
+        match self {
+            Instruction::Binary(binary) => binary.simplify(dfg),
+            Instruction::Cast(value, typ) => (*typ == dfg.type_of_value(*value)).then_some(*value),
+            Instruction::Not(value) => {
+                match &dfg[*value] {
+                    // Limit optimizing ! on constants to only booleans. If we tried it on fields,
+                    // there is no Not on FieldElement, so we'd need to convert between u128. This
+                    // would be incorrect however since the extra bits on the field would not be flipped.
+                    Value::NumericConstant { constant, typ } if *typ == Type::bool() => {
+                        let value = dfg[*constant].value().is_zero() as u128;
+                        Some(dfg.make_constant(value.into(), Type::bool()))
+                    }
+                    Value::Instruction { instruction, .. } => {
+                        // !!v => v
+                        match &dfg[*instruction] {
+                            Instruction::Not(value) => Some(*value),
+                            _ => None,
+                        }
+                    }
+                    _ => None,
+                }
+            }
+            Instruction::Constrain(value) => {
+                if let Some(constant) = dfg.get_numeric_constant(*value) {
+                    if constant.is_one() {
+                        // "simplify" to a unit literal that will just be thrown away anyway
+                        return Some(dfg.make_constant(0u128.into(), Type::Unit));
+                    }
+                }
+                None
+            }
+            Instruction::Truncate { .. } => None,
+            Instruction::Call { .. } => None,
+            Instruction::Allocate { .. } => None,
+            Instruction::Load { .. } => None,
+            Instruction::Store { .. } => None,
+        }
+    }
 }
 
 /// The possible return values for Instruction::return_types
@@ -219,6 +267,156 @@ impl Binary {
             _ => InstructionResultType::Operand(self.lhs),
         }
     }
+
+    /// Try to simplify this binary instruction, returning the new value if possible.
+    fn simplify(&self, dfg: &mut DataFlowGraph) -> Option<ValueId> {
+        let lhs = dfg.get_numeric_constant(self.lhs);
+        let rhs = dfg.get_numeric_constant(self.rhs);
+        let operand_type = dfg.type_of_value(self.lhs);
+
+        if let (Some(lhs), Some(rhs)) = (lhs, rhs) {
+            return self.eval_constants(dfg, lhs, rhs, operand_type);
+        }
+
+        let lhs_is_zero = lhs.map_or(false, |lhs| lhs.is_zero());
+        let rhs_is_zero = rhs.map_or(false, |rhs| rhs.is_zero());
+
+        let lhs_is_one = lhs.map_or(false, |lhs| lhs.is_one());
+        let rhs_is_one = rhs.map_or(false, |rhs| rhs.is_one());
+
+        match self.operator {
+            BinaryOp::Add => {
+                if lhs_is_zero {
+                    return Some(self.rhs);
+                }
+                if rhs_is_zero {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Sub => {
+                if rhs_is_zero {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Mul => {
+                if lhs_is_one {
+                    return Some(self.rhs);
+                }
+                if rhs_is_one {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Div => {
+                if rhs_is_one {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Mod => {
+                if rhs_is_one {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Eq => {
+                if self.lhs == self.rhs {
+                    return Some(dfg.make_constant(FieldElement::one(), Type::bool()));
+                }
+            }
+            BinaryOp::Lt => {
+                if self.lhs == self.rhs {
+                    return Some(dfg.make_constant(FieldElement::zero(), Type::bool()));
+                }
+            }
+            BinaryOp::And => {
+                if lhs_is_zero || rhs_is_zero {
+                    return Some(dfg.make_constant(FieldElement::zero(), operand_type));
+                }
+            }
+            BinaryOp::Or => {
+                if lhs_is_zero {
+                    return Some(self.rhs);
+                }
+                if rhs_is_zero {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Xor => (),
+            BinaryOp::Shl => {
+                if rhs_is_zero {
+                    return Some(self.lhs);
+                }
+            }
+            BinaryOp::Shr => {
+                if rhs_is_zero {
+                    return Some(self.lhs);
+                }
+            }
+        }
+        None
+    }
+
+    /// Evaluate the two constants with the operation specified by self.operator.
+    /// Pushes the resulting value to the given DataFlowGraph's constants and returns it.
+    fn eval_constants(
+        &self,
+        dfg: &mut DataFlowGraph,
+        lhs: FieldElement,
+        rhs: FieldElement,
+        operand_type: Type,
+    ) -> Option<Id<Value>> {
+        let value = match self.operator {
+            BinaryOp::Add => lhs + rhs,
+            BinaryOp::Sub => lhs - rhs,
+            BinaryOp::Mul => lhs * rhs,
+            BinaryOp::Div => lhs / rhs,
+            BinaryOp::Eq => (lhs == rhs).into(),
+            BinaryOp::Lt => (lhs < rhs).into(),
+
+            // The rest of the operators we must try to convert to u128 first
+            BinaryOp::Mod => self.eval_constant_u128_operations(lhs, rhs)?,
+            BinaryOp::And => self.eval_constant_u128_operations(lhs, rhs)?,
+            BinaryOp::Or => self.eval_constant_u128_operations(lhs, rhs)?,
+            BinaryOp::Xor => self.eval_constant_u128_operations(lhs, rhs)?,
+            BinaryOp::Shl => self.eval_constant_u128_operations(lhs, rhs)?,
+            BinaryOp::Shr => self.eval_constant_u128_operations(lhs, rhs)?,
+        };
+        // TODO: Keep original type of constant
+        Some(dfg.make_constant(value, operand_type))
+    }
+
+    /// Try to evaluate the given operands as u128s for operators that are only valid on u128s,
+    /// like the bitwise operators and modulus.
+    fn eval_constant_u128_operations(
+        &self,
+        lhs: FieldElement,
+        rhs: FieldElement,
+    ) -> Option<FieldElement> {
+        let lhs = lhs.try_into_u128()?;
+        let rhs = rhs.try_into_u128()?;
+        match self.operator {
+            BinaryOp::Mod => Some((lhs % rhs).into()),
+            BinaryOp::And => Some((lhs & rhs).into()),
+            BinaryOp::Or => Some((lhs | rhs).into()),
+            BinaryOp::Shr => Some((lhs >> rhs).into()),
+            // Check for overflow and return None if anything does overflow
+            BinaryOp::Shl => {
+                let rhs = rhs.try_into().ok()?;
+                lhs.checked_shl(rhs).map(Into::into)
+            }
+
+            // Converting a field xor to a u128 xor would be incorrect since we wouldn't have the
+            // extra bits of the field. So we don't optimize it here.
+            BinaryOp::Xor => None,
+
+            op @ (BinaryOp::Add
+            | BinaryOp::Sub
+            | BinaryOp::Mul
+            | BinaryOp::Div
+            | BinaryOp::Eq
+            | BinaryOp::Lt) => panic!(
+                "eval_constant_u128_operations invalid for {op:?} use eval_constants instead"
+            ),
+        }
+    }
 }
 
 /// Binary Operations allowed in the IR.

crates/noirc_evaluator/src/ssa_refactor/opt/inlining.rs

@@ -9,6 +9,7 @@ use iter_extended::vecmap;
 use crate::ssa_refactor::{
     ir::{
         basic_block::BasicBlockId,
+        dfg::InsertInstructionResult,
         function::{Function, FunctionId},
         instruction::{Instruction, InstructionId, TerminatorInstruction},
         value::{Value, ValueId},
@@ -343,7 +344,8 @@ impl<'function> PerFunctionContext<'function> {
         let old_results = self.source_function.dfg.instruction_results(call_id);
         let arguments = vecmap(arguments, |arg| self.translate_value(*arg));
         let new_results = self.context.inline_function(ssa, function, &arguments);
-        Self::insert_new_instruction_results(&mut self.values, old_results, &new_results);
+        let new_results = InsertInstructionResult::Results(&new_results);
+        Self::insert_new_instruction_results(&mut self.values, old_results, new_results);
     }
 
     /// Push the given instruction from the source_function into the current block of the
@@ -365,11 +367,20 @@ impl<'function> PerFunctionContext<'function> {
     fn insert_new_instruction_results(
         values: &mut HashMap<ValueId, ValueId>,
         old_results: &[ValueId],
-        new_results: &[ValueId],
+        new_results: InsertInstructionResult,
     ) {
         assert_eq!(old_results.len(), new_results.len());
-        for (old_result, new_result) in old_results.iter().zip(new_results) {
-            values.insert(*old_result, *new_result);
+
+        match new_results {
+            InsertInstructionResult::SimplifiedTo(new_result) => {
+                values.insert(old_results[0], new_result);
+            }
+            InsertInstructionResult::Results(new_results) => {
+                for (old_result, new_result) in old_results.iter().zip(new_results) {
+                    values.insert(*old_result, *new_result);
+                }
+            }
+            InsertInstructionResult::InstructionRemoved => (),
         }
     }