(optimization): felt252_div constant propogation support

Author: eytan-starkware

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

@@ -550,6 +550,37 @@ impl<'db, 'mt> ConstFoldingContext<'db, 'mt> {
             } else {
                 None
             }
+        } else if id == self.felt_div {
+            // There is no need to check if divisor is 0 and throw error, because casting to NonZero
+            // will add a panic that will always be thrown. TODO(eytan-starkware): Make
+            // sure that NonZero on a const zero leads to a compile error.
+            if let Some(rhs) = self.as_int(stmt.inputs[1].var_id)
+                && rhs.is_one()
+            {
+                // Check if dividing by 1 (returns the original value)
+                self.var_info.insert(stmt.outputs[0], VarInfo::Var(stmt.inputs[0]));
+                None
+            } else if let Some(lhs) = self.as_int(stmt.inputs[0].var_id)
+                && lhs.is_zero()
+            {
+                // Check if 0 is being divided (returns 0)
+                Some(self.propagate_zero_and_get_statement(stmt.outputs[0]))
+            } else if let (Some(lhs), Some(rhs)) =
+                (self.as_int(stmt.inputs[0].var_id), self.as_int(stmt.inputs[1].var_id))
+                && !rhs.is_zero()
+            {
+                // Constant fold when both operands are constants
+
+                // Use field_div for Felt252 division (requires non-zero divisor)
+                let lhs_felt = Felt252::from(lhs);
+                let rhs_felt = Felt252::from(rhs);
+                // For non-zero divisor, use field_div; the libfunc should handle zero checks
+                let rhs_nonzero = rhs_felt.try_into().expect("Non-zero divisor");
+                let value = lhs_felt.field_div(&rhs_nonzero).to_bigint();
+                Some(self.propagate_const_and_get_statement(value, stmt.outputs[0], false))
+            } else {
+                None
+            }
         } else if self.wide_mul_fns.contains(&id) {
             let lhs = self.as_int_ex(stmt.inputs[0].var_id);
             let rhs = self.as_int(stmt.inputs[1].var_id);
@@ -1266,6 +1297,8 @@ pub struct ConstFoldingLibfuncInfo<'db> {
     felt_add: ExternFunctionId<'db>,
     /// The `felt252_mul` libfunc.
     felt_mul: ExternFunctionId<'db>,
+    /// The `felt252_div` libfunc.
+    felt_div: ExternFunctionId<'db>,
     /// The `into_box` libfunc.
     into_box: ExternFunctionId<'db>,
     /// The `unbox` libfunc.
@@ -1437,6 +1470,7 @@ impl<'db> ConstFoldingLibfuncInfo<'db> {
             felt_sub: core.extern_function_id("felt252_sub"),
             felt_add: core.extern_function_id("felt252_add"),
             felt_mul: core.extern_function_id("felt252_mul"),
+            felt_div: core.extern_function_id("felt252_div"),
             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/const_folding_test.rs

@@ -33,8 +33,16 @@ fn test_match_optimizer(
     .split();
     let function_id =
         ConcreteFunctionWithBodyId::from_semantic(db, test_function.concrete_function_id);
-
-    let mut before = db.lowered_body(function_id, LoweringStage::PreOptimizations).unwrap().clone();
+    let mut before = db
+        .lowered_body(function_id, LoweringStage::PreOptimizations)
+        .unwrap_or_else(|_| {
+            panic!(
+                "Failed to get lowered body for function {:?}. Diagnostics: {:?}",
+                function_id,
+                db.module_lowering_diagnostics(test_function.module_id)
+            )
+        })
+        .clone();
     OptimizationPhase::ApplyInlining { enable_const_folding: false }
         .apply(db, function_id, &mut before)
         .unwrap();

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

@@ -6493,3 +6493,130 @@ End:
   Return(v6)
 
 //! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 div const fold.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> felt252 {
+    let x = core::felt252_div(6, 30);
+    x * 5
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::felt252;
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 6
+  (v1: core::zeroable::NonZero::<core::felt252>) <- NonZero(30)
+  (v2: core::felt252) <- core::felt252_div(v0, v1)
+  (v3: core::felt252) <- 5
+  (v4: core::felt252) <- core::felt252_mul(v2, v3)
+End:
+  Return(v4)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::felt252) <- 6
+  (v1: core::zeroable::NonZero::<core::felt252>) <- NonZero(30)
+  (v2: core::felt252) <- 2894802230932904970957858226476056084498485772265277359978473644908697616385
+  (v3: core::felt252) <- 5
+  (v4: core::felt252) <- 1
+End:
+  Return(v4)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 div by one.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo(a: felt252) -> felt252 {
+    core::felt252_div(a, 1)
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::felt252;
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters: v0: core::felt252
+blk0 (root):
+Statements:
+  (v1: core::zeroable::NonZero::<core::felt252>) <- NonZero(1)
+  (v2: core::felt252) <- core::felt252_div(v0, v1)
+End:
+  Return(v2)
+
+//! > after
+Parameters: v0: core::felt252
+blk0 (root):
+Statements:
+  (v1: core::zeroable::NonZero::<core::felt252>) <- NonZero(1)
+  (v2: core::felt252) <- core::felt252_div(v0, v1)
+End:
+  Return(v0)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > Felt252 zero div.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo(a: NonZero<felt252>) -> felt252 {
+    core::felt252_div(0, a)
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::felt252;
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters: v0: core::zeroable::NonZero::<core::felt252>
+blk0 (root):
+Statements:
+  (v1: core::felt252) <- 0
+  (v2: core::felt252) <- core::felt252_div(v1, v0)
+End:
+  Return(v2)
+
+//! > after
+Parameters: v0: core::zeroable::NonZero::<core::felt252>
+blk0 (root):
+Statements:
+  (v1: core::felt252) <- 0
+  (v2: core::felt252) <- 0
+End:
+  Return(v2)
+
+//! > lowering_diagnostics