Adding support to felt add mul and sub in constant propogation

Author: eytan-starkware

crates/cairo-lang-lowering/src/lower/test_data/closure

@@ -343,22 +343,6 @@ End:
 Final lowering:
 Parameters: v0: core::integer::u32
 blk0 (root):
-Statements:
-  (v1: core::felt252) <- 1
-  (v2: core::felt252) <- 2
-  (v3: core::felt252) <- core::felt252_sub(v1, v2)
-End:
-  Match(match core::felt252_is_zero(v3) {
-    IsZeroResult::Zero => blk1,
-    IsZeroResult::NonZero(v4) => blk2,
-  })
-
-blk1:
-Statements:
-End:
-  Return()
-
-blk2:
 Statements:
 End:
   Return()
@@ -430,22 +414,6 @@ End:
 Final lowering:
 Parameters: v0: @{[email protected]:2:13: 2:25}, v1: (core::felt252,)
 blk0 (root):
-Statements:
-  (v2: core::felt252) <- 1
-  (v3: core::felt252) <- 2
-  (v4: core::felt252) <- core::felt252_sub(v2, v3)
-End:
-  Match(match core::felt252_is_zero(v4) {
-    IsZeroResult::Zero => blk1,
-    IsZeroResult::NonZero(v5) => blk2,
-  })
-
-blk1:
-Statements:
-End:
-  Return()
-
-blk2:
 Statements:
 End:
   Return()

crates/cairo-lang-lowering/src/lower/test_data/for

@@ -55,7 +55,7 @@ Parameters: v0: core::RangeCheck, v1: core::gas::GasBuiltin
 blk0 (root):
 Statements:
   (v2: core::felt252) <- 1
-  (v3: core::felt252) <- core::felt252_add(v2, v2)
+  (v3: core::felt252) <- 2
   (v4: core::array::Array::<core::felt252>) <- core::array::array_new::<core::felt252>()
   (v5: core::felt252) <- 10
   (v6: core::array::Array::<core::felt252>) <- core::array::array_append::<core::felt252>(v4, v5)

crates/cairo-lang-lowering/src/optimizations/const_folding.rs

@@ -18,6 +18,7 @@ use cairo_lang_semantic::types::{TypeSizeInformation, TypesSemantic};
 use cairo_lang_semantic::{
     ConcreteTypeId, GenericArgumentId, MatchArmSelector, TypeId, TypeLongId, corelib,
 };
+use cairo_lang_utils::bigint::{CAIRO_PRIME_BIGINT, felt252_mod};
 use cairo_lang_utils::byte_array::BYTE_ARRAY_MAGIC;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use cairo_lang_utils::ordered_hash_set::OrderedHashSet;
@@ -493,10 +494,37 @@ impl<'db, 'mt> ConstFoldingContext<'db, 'mt> {
         }
         let (id, _generic_args) = stmt.function.get_extern(db)?;
         if id == self.felt_sub {
-            // (a - 0) can be replaced by a.
-            let val = self.as_int(stmt.inputs[1].var_id)?;
-            if val.is_zero() {
-                self.var_info.insert(stmt.outputs[0], VarInfo::Var(stmt.inputs[0]));
+            if let Some(rhs) = self.as_int(stmt.inputs[1].var_id) {
+                if rhs.is_zero() {
+                    self.var_info.insert(stmt.outputs[0], VarInfo::Var(stmt.inputs[0]));
+                    return None;
+                }
+            }
+            if let (Some(lhs), Some(rhs)) =
+                (self.as_int(stmt.inputs[0].var_id), self.as_int(stmt.inputs[1].var_id))
+            {
+                let value = felt252_mod(&(lhs - rhs));
+                return Some(self.propagate_const_and_get_statement(value, stmt.outputs[0], false));
+            }
+            None
+        } else if id == self.felt_add {
+            if let (Some(lhs), Some(rhs)) =
+                (self.as_int(stmt.inputs[0].var_id), self.as_int(stmt.inputs[1].var_id))
+            {
+                let value = felt252_mod(&(lhs + rhs));
+                return Some(self.propagate_const_and_get_statement(value, stmt.outputs[0], false));
+            }
+            None
+        } else if id == self.felt_mul {
+            let lhs = self.as_int_ex(stmt.inputs[0].var_id);
+            let rhs = self.as_int_ex(stmt.inputs[1].var_id);
+            if lhs.map(|(v, _)| v.is_zero()).unwrap_or_default() || rhs.map(|(v, _)| v.is_zero()).unwrap_or_default() {
+                return Some(self.propagate_zero_and_get_statement(stmt.outputs[0]));
+            }
+            if let (Some((lhs_val, lhs_nz)), Some((rhs_val, rhs_nz))) = (lhs, rhs) {
+                let value = felt252_mod(&(lhs_val * rhs_val));
+                let nz_ty = lhs_nz && rhs_nz;
+                return Some(self.propagate_const_and_get_statement(value, stmt.outputs[0], nz_ty));
             }
             None
         } else if self.wide_mul_fns.contains(&id) {
@@ -1211,6 +1239,10 @@ fn priv_const_folding_info<'db>(
 pub struct ConstFoldingLibfuncInfo<'db> {
     /// The `felt252_sub` libfunc.
     felt_sub: ExternFunctionId<'db>,
+    /// The `felt252_add` libfunc.
+    felt_add: ExternFunctionId<'db>,
+    /// The `felt252_mul` libfunc.
+    felt_mul: ExternFunctionId<'db>,
     /// The `into_box` libfunc.
     into_box: ExternFunctionId<'db>,
     /// The `unbox` libfunc.
@@ -1380,6 +1412,8 @@ impl<'db> ConstFoldingLibfuncInfo<'db> {
             );
         Self {
             felt_sub: core.extern_function_id("felt252_sub"),
+            felt_add: core.extern_function_id("felt252_add"),
+            felt_mul: core.extern_function_id("felt252_mul"),
             into_box: box_module.extern_function_id("into_box"),
             unbox: box_module.extern_function_id("unbox"),
             box_forward_snapshot: box_module.generic_function_id("box_forward_snapshot"),

crates/cairo-lang-lowering/src/optimizations/test_data/const_folding

@@ -6150,3 +6150,175 @@ End:
   Return(v7)
 
 //! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 add const fold.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> felt252 {
+    1 + 2
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 1
+  (v1: core::felt252) <- 2
+  (v2: core::felt252) <- core::felt252_add(v0, v1)
+End:
+  Return(v2)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 1
+  (v1: core::felt252) <- 2
+  (v2: core::felt252) <- 3
+End:
+  Return(v2)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 sub const fold.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> felt252 {
+    5 - 3
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 5
+  (v1: core::felt252) <- 3
+  (v2: core::felt252) <- core::felt252_sub(v0, v1)
+End:
+  Return(v2)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 5
+  (v1: core::felt252) <- 3
+  (v2: core::felt252) <- 2
+End:
+  Return(v2)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 mul const fold.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> felt252 {
+    2 * 3
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 2
+  (v1: core::felt252) <- 3
+  (v2: core::felt252) <- core::felt252_mul(v0, v1)
+End:
+  Return(v2)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 2
+  (v1: core::felt252) <- 3
+  (v2: core::felt252) <- 6
+End:
+  Return(v2)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 combined operations const fold.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> felt252 {
+    1 + 2 * 3 - 4
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 1
+  (v1: core::felt252) <- 2
+  (v2: core::felt252) <- 3
+  (v3: core::felt252) <- core::felt252_mul(v1, v2)
+  (v4: core::felt252) <- core::felt252_add(v0, v3)
+  (v5: core::felt252) <- 4
+  (v6: core::felt252) <- core::felt252_sub(v4, v5)
+End:
+  Return(v6)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 1
+  (v1: core::felt252) <- 2
+  (v2: core::felt252) <- 3
+  (v3: core::felt252) <- 6
+  (v4: core::felt252) <- 7
+  (v5: core::felt252) <- 4
+  (v6: core::felt252) <- 3
+End:
+  Return(v6)
+
+//! > lowering_diagnostics

crates/cairo-lang-lowering/src/optimizations/test_data/reorder_statements

@@ -74,7 +74,7 @@ End:
 
 blk4:
 Statements:
-  (v12: core::felt252) <- core::felt252_add(v7, v8)
+  (v12: core::felt252) <- 3
 End:
   Goto(blk9, {v12 -> v13})
 
@@ -89,7 +89,7 @@ End:
 
 blk6:
 Statements:
-  (v17: core::felt252) <- core::felt252_add(v9, v14)
+  (v17: core::felt252) <- 7
 End:
   Goto(blk8, {v17 -> v18})
 
@@ -133,7 +133,6 @@ End:
 
 blk3:
 Statements:
-  (v7: core::felt252) <- 1
 End:
   Match(match_enum(v2) {
     Option::Some(v10) => blk4,
@@ -142,8 +141,7 @@ End:
 
 blk4:
 Statements:
-  (v8: core::felt252) <- 2
-  (v12: core::felt252) <- core::felt252_add(v7, v8)
+  (v12: core::felt252) <- 3
 End:
   Goto(blk9, {v12 -> v13})
 
@@ -157,14 +155,13 @@ End:
 
 blk6:
 Statements:
-  (v9: core::felt252) <- 3
-  (v14: core::felt252) <- 4
-  (v17: core::felt252) <- core::felt252_add(v9, v14)
+  (v17: core::felt252) <- 7
 End:
   Goto(blk8, {v17 -> v18})
 
 blk7:
 Statements:
+  (v7: core::felt252) <- 1
 End:
   Goto(blk8, {v7 -> v18})
 
@@ -356,16 +353,14 @@ Statements:
   (v2: core::felt252) <- 5
   (v3: ()) <- struct_construct()
   (v4: core::felt252) <- 3
-  (v5: core::felt252) <- core::felt252_add(v2, v4)
+  (v5: core::felt252) <- 8
 End:
   Return(v5)
 
 //! > after
 Parameters:
 blk0 (root):
 Statements:
-  (v2: core::felt252) <- 5
-  (v4: core::felt252) <- 3
-  (v5: core::felt252) <- core::felt252_add(v2, v4)
+  (v5: core::felt252) <- 8
 End:
   Return(v5)