chore(ssa): enable cse for assert (#1350)

enable cse for assert
noir-lang · May 16, 2023 · a830ce5 · a830ce5
1 parent 63d84a3
commit a830ce5
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/crates/noirc_evaluator/src/ssa/node.rs b/crates/noirc_evaluator/src/ssa/node.rs
@@ -506,6 +506,36 @@ impl Instruction {
             }
         }
     }
+
+    pub(crate) fn get_location(&self) -> Option<Location> {
+        match &self.operation {
+            Operation::Binary(bin) => match bin.operator {
+                BinaryOp::Udiv(location)
+                | BinaryOp::Sdiv(location)
+                | BinaryOp::Urem(location)
+                | BinaryOp::Srem(location)
+                | BinaryOp::Div(location)
+                | BinaryOp::Shr(location) => Some(location),
+                _ => None,
+            },
+            Operation::Call { location, .. } => Some(*location),
+            Operation::Load { location, .. }
+            | Operation::Store { location, .. }
+            | Operation::Constrain(_, location) => *location,
+            Operation::Cast(_)
+            | Operation::Truncate { .. }
+            | Operation::Not(_)
+            | Operation::Jne(_, _)
+            | Operation::Jeq(_, _)
+            | Operation::Jmp(_)
+            | Operation::Phi { .. }
+            | Operation::Return(_)
+            | Operation::Result { .. }
+            | Operation::Cond { .. }
+            | Operation::Intrinsic(_, _)
+            | Operation::Nop => None,
+        }
+    }
 }
 
 //adapted from LLVM IR

diff --git a/crates/noirc_evaluator/src/ssa/optimizations.rs b/crates/noirc_evaluator/src/ssa/optimizations.rs
@@ -507,6 +507,22 @@ fn cse_block_with_anchor(
                         new_list.push(*ins_id);
                     }
                 }
+                Operation::Constrain(condition, location) => {
+                    if let Some(similar) = anchor.find_similar_instruction(&operator) {
+                        assert_ne!(similar, ins.id);
+                        *modified = true;
+                        let similar_ins = ctx
+                            .try_get_mut_instruction(similar)
+                            .expect("Similar instructions are instructions");
+                        if location.is_some() && similar_ins.get_location().is_none() {
+                            similar_ins.operation = Operation::Constrain(*condition, *location);
+                        }
+                        new_mark = Mark::ReplaceWith(similar);
+                    } else {
+                        new_list.push(*ins_id);
+                        anchor.push_front(&ins.operation, *ins_id);
+                    }
+                }
                 _ => {
                     //TODO: checks we do not need to propagate res arguments
                     new_list.push(*ins_id);