Adding support to felt add mul and sub in constant propogation

Author: eytan-starkware

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

@@ -307,7 +307,7 @@ test_generated_function
 fn foo(a: u32) {
     let f = |a: felt252| {
         let mut b = @0;
-        if 1 == 2 {
+        if 1 == bar(2) {
             b = @a;
         } else {
             b = @a;
@@ -320,6 +320,10 @@ fn foo(a: u32) {
 foo
 
 //! > module_code
+#[inline(never)]
+fn bar(a: felt252) -> felt252 {
+    a + 1
+}
 
 //! > semantic_diagnostics
 
@@ -330,11 +334,11 @@ Main:
 Parameters: v0: core::integer::u32
 blk0 (root):
 Statements:
-  (v1: {[email protected]:2:13: 2:25}) <- struct_construct()
-  (v2: {[email protected]:2:13: 2:25}, v3: @{[email protected]:2:13: 2:25}) <- snapshot(v1)
+  (v1: {[email protected]:6:13: 6:25}) <- struct_construct()
+  (v2: {[email protected]:6:13: 6:25}, v3: @{[email protected]:6:13: 6:25}) <- snapshot(v1)
   (v4: core::felt252) <- 0
   (v5: (core::felt252,)) <- struct_construct(v4)
-  (v6: ()) <- Generated `core::ops::function::Fn::call` for {[email protected]:2:13: 2:25}(v3, v5)
+  (v6: ()) <- Generated `core::ops::function::Fn::call` for {[email protected]:6:13: 6:25}(v3, v5)
   (v7: ()) <- struct_construct()
 End:
   Return(v7)
@@ -344,13 +348,14 @@ 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)
+  (v1: core::felt252) <- 2
+  (v2: core::felt252) <- test::bar(v1)
+  (v3: core::felt252) <- 1
+  (v4: core::felt252) <- core::felt252_sub(v3, v2)
 End:
-  Match(match core::felt252_is_zero(v3) {
+  Match(match core::felt252_is_zero(v4) {
     IsZeroResult::Zero => blk1,
-    IsZeroResult::NonZero(v4) => blk2,
+    IsZeroResult::NonZero(v5) => blk2,
   })
 
 blk1:
@@ -368,7 +373,7 @@ Generated core::traits::Destruct::destruct lowering for source location:
     let f = |a: felt252| {
             ^^^^^^^^^^^^
 
-Parameters: v0: {[email protected]:2:13: 2:25}
+Parameters: v0: {[email protected]:6:13: 6:25}
 blk0 (root):
 Statements:
   () <- struct_destructure(v0)
@@ -378,7 +383,7 @@ End:
 
 
 Final lowering:
-Parameters: v0: {[email protected]:2:13: 2:25}
+Parameters: v0: {[email protected]:6:13: 6:25}
 blk0 (root):
 Statements:
 End:
@@ -389,55 +394,57 @@ Generated core::ops::function::Fn::call lowering for source location:
     let f = |a: felt252| {
             ^^^^^^^^^^^^
 
-Parameters: v0: @{[email protected]:2:13: 2:25}, v2: (core::felt252,)
+Parameters: v0: @{[email protected]:6:13: 6:25}, v2: (core::felt252,)
 blk0 (root):
 Statements:
-  (v1: {[email protected]:2:13: 2:25}) <- desnap(v0)
+  (v1: {[email protected]:6:13: 6:25}) <- desnap(v0)
   () <- struct_destructure(v1)
   (v3: core::felt252) <- struct_destructure(v2)
   (v4: core::felt252) <- 0
   (v5: core::felt252, v6: @core::felt252) <- snapshot(v4)
   (v7: core::felt252) <- 1
   (v8: core::felt252, v9: @core::felt252) <- snapshot(v7)
   (v10: core::felt252) <- 2
-  (v11: core::felt252, v12: @core::felt252) <- snapshot(v10)
-  (v13: core::bool) <- core::Felt252PartialEq::eq(v9, v12)
+  (v11: core::felt252) <- test::bar(v10)
+  (v12: core::felt252, v13: @core::felt252) <- snapshot(v11)
+  (v14: core::bool) <- core::Felt252PartialEq::eq(v9, v13)
 End:
-  Match(match_enum(v13) {
-    bool::False(v15) => blk2,
-    bool::True(v14) => blk1,
+  Match(match_enum(v14) {
+    bool::False(v16) => blk2,
+    bool::True(v15) => blk1,
   })
 
 blk1:
 Statements:
-  (v18: core::felt252, v19: @core::felt252) <- snapshot(v3)
+  (v19: core::felt252, v20: @core::felt252) <- snapshot(v3)
 End:
-  Goto(blk3, {v18 -> v20, v19 -> v21})
+  Goto(blk3, {v19 -> v21, v20 -> v22})
 
 blk2:
 Statements:
-  (v16: core::felt252, v17: @core::felt252) <- snapshot(v3)
+  (v17: core::felt252, v18: @core::felt252) <- snapshot(v3)
 End:
-  Goto(blk3, {v16 -> v20, v17 -> v21})
+  Goto(blk3, {v17 -> v21, v18 -> v22})
 
 blk3:
 Statements:
-  (v22: ()) <- struct_construct()
+  (v23: ()) <- struct_construct()
 End:
-  Return(v22)
+  Return(v23)
 
 
 Final lowering:
-Parameters: v0: @{[email protected]:2:13: 2:25}, v1: (core::felt252,)
+Parameters: v0: @{[email protected]:6:13: 6:25}, v1: (core::felt252,)
 blk0 (root):
 Statements:
-  (v2: core::felt252) <- 1
-  (v3: core::felt252) <- 2
-  (v4: core::felt252) <- core::felt252_sub(v2, v3)
+  (v2: core::felt252) <- 2
+  (v3: core::felt252) <- test::bar(v2)
+  (v4: core::felt252) <- 1
+  (v5: core::felt252) <- core::felt252_sub(v4, v3)
 End:
-  Match(match core::felt252_is_zero(v4) {
+  Match(match core::felt252_is_zero(v5) {
     IsZeroResult::Zero => blk1,
-    IsZeroResult::NonZero(v5) => blk2,
+    IsZeroResult::NonZero(v6) => blk2,
   })
 
 blk1:

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::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,39 @@ 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(false)
+                || rhs.map(|(v, _)| v.is_zero()).unwrap_or(false)
+            {
+                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 +1241,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 +1414,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"),