Use Montgomery representation for Felts

Cargo.toml

@@ -134,6 +134,7 @@ cairo-lang-sierra-type-size.workspace = true
 cairo-lang-utils.workspace = true
 educe.workspace = true
 itertools.workspace = true
+lambdaworks-math.workspace = true
 lazy_static.workspace = true
 libc.workspace = true
 libloading.workspace = true

benches/benches.rs

@@ -88,6 +88,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     compare(c, "programs/benches/dict_snapshot.cairo");
     compare(c, "programs/benches/dict_insert.cairo");
     compare(c, "programs/benches/factorial_2M.cairo");
+    compare(c, "programs/benches/factorial_2M_inv.cairo");
     compare(c, "programs/benches/fib_2M.cairo");
     compare(c, "programs/benches/linear_search.cairo");
     compare(c, "programs/benches/logistic_map.cairo");

programs/benches/factorial_2M_inv.cairo

@@ -0,0 +1,15 @@
+fn factorial_inv(value: felt252, n: felt252) -> felt252 {
+    if (n == 1) {
+        value
+    } else {
+        factorial_inv(felt252_div(value, n.try_into().unwrap()), n - 1)
+    }
+}
+
+fn main() {
+    let result = factorial_inv(0x4d6e41de886ac83938da3456ccf1481182687989ead34d9d35236f0864575a0, 2_000_000);
+    assert(
+        result == 1,
+        'invalid result'
+    );
+}

src/arch/aarch64.rs

@@ -10,7 +10,11 @@
 #![cfg(target_arch = "aarch64")]
 
 use super::AbiArgument;
-use crate::{error::Error, starknet::U256, utils::get_integer_layout};
+use crate::{
+    error::Error,
+    starknet::U256,
+    utils::{get_integer_layout, montgomery::MontyBytes},
+};
 use cairo_lang_sierra::ids::ConcreteTypeId;
 use num_traits::ToBytes;
 use starknet_types_core::felt::Felt;
@@ -197,7 +201,8 @@ impl AbiArgument for Felt {
         if buffer.len() >= 56 {
             align_to(buffer, get_integer_layout(252).align());
         }
-        buffer.extend_from_slice(&self.to_bytes_le());
+
+        buffer.extend_from_slice(&self.to_monty_bytes_le());
         Ok(())
     }
 }
@@ -489,8 +494,10 @@ mod test {
             buffer,
             [0; 40]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
 
@@ -504,8 +511,10 @@ mod test {
             buffer,
             [0; 48]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
 
@@ -519,8 +528,10 @@ mod test {
             buffer,
             [0; 64]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
     }

src/arch/x86_64.rs

@@ -10,7 +10,11 @@
 #![cfg(target_arch = "x86_64")]
 
 use super::AbiArgument;
-use crate::{error::Error, starknet::U256, utils::get_integer_layout};
+use crate::{
+    error::Error,
+    starknet::U256,
+    utils::{get_integer_layout, montgomery::MontyBytes},
+};
 use cairo_lang_sierra::ids::ConcreteTypeId;
 use num_traits::ToBytes;
 use starknet_types_core::felt::Felt;
@@ -159,7 +163,7 @@ impl AbiArgument for Felt {
             align_to(buffer, get_integer_layout(252).align());
         }
 
-        buffer.extend_from_slice(&self.to_bytes_le());
+        buffer.extend_from_slice(&self.to_monty_bytes_le());
         Ok(())
     }
 }
@@ -240,8 +244,10 @@ mod test {
             buffer,
             [0; 24]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
 
@@ -255,8 +261,10 @@ mod test {
             buffer,
             [0; 32]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
 
@@ -270,8 +278,10 @@ mod test {
             buffer,
             [0; 48]
                 .into_iter()
-                .chain([0xFF; 31])
-                .chain([0x00])
+                .chain([
+                    37, 0, 126, 115, 253, 255, 255, 255, 255, 15, 51, 1, 0, 0, 0, 0, 128, 111, 255,
+                    255, 255, 255, 255, 255, 184, 2, 94, 171, 212, 255, 255, 7
+                ])
                 .collect::<Vec<_>>()
         );
     }

src/executor.rs

@@ -454,11 +454,17 @@ fn parse_result(
 
                 #[cfg(target_arch = "aarch64")]
                 Ok(Value::Felt252({
+                    use crate::utils::montgomery;
+
                     let data = unsafe {
                         std::mem::transmute::<&mut [u64; 4], &mut [u8; 32]>(&mut ret_registers)
                     };
                     data[31] &= 0x0F; // Filter out first 4 bits (they're outside an i252).
-                    starknet_types_core::felt::Felt::from_bytes_le(data)
+
+                    // Felts are represented in Montgomery form. Due to this, we
+                    // need to convert them back to their original representation.
+                    montgomery::felt_from_monty_bytes(&data)
+                        .to_native_assert_error("Couldn't create felt from Montgomery bytes")?
                 }))
             }
         },

src/executor/contract.rs

@@ -51,7 +51,9 @@ use crate::{
     types::TypeBuilder,
     utils::{
         decode_error_message, generate_function_name, get_integer_layout, get_types_total_size,
-        libc_free, libc_malloc, BuiltinCosts,
+        libc_free, libc_malloc,
+        montgomery::{self, MontyBytes},
+        BuiltinCosts,
     },
     OptLevel,
 };
@@ -501,7 +503,8 @@ impl AotContractExecutor {
         };
 
         for (idx, elem) in args.iter().enumerate() {
-            let f = elem.to_bytes_le();
+            let f = elem.to_monty_bytes_le();
+
             unsafe {
                 std::ptr::copy_nonoverlapping(
                     f.as_ptr().cast::<u8>(),
@@ -667,9 +670,15 @@ impl AotContractExecutor {
                 let cur_elem_ptr = unsafe { array_ptr.byte_add(elem_stride * i as usize) };
 
                 let mut data = unsafe { cur_elem_ptr.cast::<[u8; 32]>().read() };
+
                 data[31] &= 0x0F; // Filter out first 4 bits (they're outside an i252).
 
-                array_value.push(Felt::from_bytes_le(&data));
+                // Felts are represented in Montgomery form. Due to this, we
+                // need to convert them back to their original representation.
+                array_value.push(
+                    montgomery::felt_from_monty_bytes(&data)
+                        .to_native_assert_error("Couldn't create felt from Montgomery bytes")?,
+                );
             }
 
             unsafe {

src/libfuncs/bool.rs

@@ -5,7 +5,7 @@ use crate::{
     error::{panic::ToNativeAssertError, Result},
     metadata::MetadataStorage,
     types::TypeBuilder,
-    utils::ProgramRegistryExt,
+    utils::{montgomery, ProgramRegistryExt},
 };
 use cairo_lang_sierra::{
     extensions::{
@@ -200,9 +200,10 @@ pub fn build_bool_to_felt252<'ctx, 'this>(
     let value = entry.arg(0)?;
     let tag_value = entry.extract_value(context, location, value, tag_ty, 0)?;
 
-    let result = entry.extui(tag_value, felt252_ty, location)?;
+    // Convert into Montgomery representation.
+    let felt = montgomery::mlir::monty_transform(context, entry, tag_value, felt252_ty, location)?;
 
-    helper.br(entry, 0, &[result], location)
+    helper.br(entry, 0, &[felt], location)
 }
 
 #[cfg(test)]

src/libfuncs/bounded_int.rs

@@ -7,7 +7,7 @@ use crate::{
     metadata::MetadataStorage,
     native_assert,
     types::TypeBuilder,
-    utils::RangeExt,
+    utils::{montgomery, RangeExt},
 };
 use cairo_lang_sierra::{
     extensions::{
@@ -149,6 +149,17 @@ fn build_add<'ctx, 'this>(
         rhs_value
     };
 
+    let lhs_value = if lhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, lhs_value, lhs_value.r#type(), location)?
+    } else {
+        lhs_value
+    };
+    let rhs_value = if rhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, rhs_value, rhs_value.r#type(), location)?
+    } else {
+        rhs_value
+    };
+
     // Addition and get the result value on the desired range
     let res_value = entry.addi(lhs_value, rhs_value, location)?;
     let res_value = if compute_width > dst_width {
@@ -248,6 +259,17 @@ fn build_sub<'ctx, 'this>(
         rhs_value
     };
 
+    let lhs_value = if lhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, lhs_value, lhs_value.r#type(), location)?
+    } else {
+        lhs_value
+    };
+    let rhs_value = if rhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, rhs_value, rhs_value.r#type(), location)?
+    } else {
+        rhs_value
+    };
+
     let compile_time_val =
         entry.const_int_from_type(context, location, compile_time_val, compute_ty)?;
     // First we do -> intermediate_res = Ad - Bd
@@ -357,13 +379,17 @@ fn build_mul<'ctx, 'this>(
         let lhs_offset =
             entry.const_int_from_type(context, location, lhs_range.lower, compute_ty)?;
         entry.addi(lhs_value, lhs_offset, location)?
+    } else if lhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, lhs_value, lhs_value.r#type(), location)?
     } else {
         lhs_value
     };
     let rhs_value = if rhs_ty.is_bounded_int(registry)? && rhs_range.lower != BigInt::ZERO {
         let rhs_offset =
             entry.const_int_from_type(context, location, rhs_range.lower, compute_ty)?;
         entry.addi(rhs_value, rhs_offset, location)?
+    } else if rhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, rhs_value, rhs_value.r#type(), location)?
     } else {
         rhs_value
     };
@@ -489,13 +515,17 @@ fn build_div_rem<'ctx, 'this>(
         let lhs_offset =
             entry.const_int_from_type(context, location, lhs_range.lower, compute_ty)?;
         entry.addi(lhs_value, lhs_offset, location)?
+    } else if lhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, lhs_value, lhs_value.r#type(), location)?
     } else {
         lhs_value
     };
     let rhs_value = if rhs_ty.is_bounded_int(registry)? && rhs_range.lower != BigInt::ZERO {
         let rhs_offset =
             entry.const_int_from_type(context, location, rhs_range.lower, compute_ty)?;
         entry.addi(rhs_value, rhs_offset, location)?
+    } else if rhs_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(context, entry, rhs_value, rhs_value.r#type(), location)?
     } else {
         rhs_value
     };
@@ -578,7 +608,6 @@ fn build_constrain<'ctx, 'this>(
     info: &BoundedIntConstrainConcreteLibfunc,
 ) -> Result<()> {
     let range_check = super::increment_builtin_counter(context, entry, location, entry.arg(0)?)?;
-    let src_value: Value = entry.arg(1)?;
 
     let src_ty = registry.get_type(&info.param_signatures()[1].ty)?;
     let src_range = src_ty.integer_range(registry)?;
@@ -589,6 +618,18 @@ fn build_constrain<'ctx, 'this>(
         src_range.zero_based_bit_width()
     };
 
+    let src_value: Value = if src_ty.is_felt252(registry)? {
+        montgomery::mlir::monty_reduce(
+            context,
+            entry,
+            entry.arg(1)?,
+            IntegerType::new(context, 252).into(),
+            location,
+        )?
+    } else {
+        entry.arg(1)?
+    };
+
     let lower_range = registry
         .get_type(&info.branch_signatures()[0].vars[1].ty)?
         .integer_range(registry)?;
@@ -783,6 +824,7 @@ fn build_is_zero<'ctx, 'this>(
     } else {
         entry.const_int_from_type(context, location, 0, src_value.r#type())?
     };
+
     let src_is_zero = entry.cmpi(context, CmpiPredicate::Eq, src_value, k0, location)?;
 
     helper.cond_br(

src/libfuncs/bytes31.rs

@@ -4,7 +4,7 @@ use super::LibfuncHelper;
 use crate::{
     error::{Error, Result},
     metadata::MetadataStorage,
-    utils::ProgramRegistryExt,
+    utils::{montgomery, ProgramRegistryExt},
 };
 use cairo_lang_sierra::{
     extensions::{
@@ -22,7 +22,7 @@ use melior::{
         cf,
     },
     helpers::{ArithBlockExt, BuiltinBlockExt},
-    ir::{Attribute, Block, BlockLike, Location, Value},
+    ir::{r#type::IntegerType, Attribute, Block, BlockLike, Location, Value},
     Context,
 };
 use num_bigint::BigUint;
@@ -96,7 +96,11 @@ pub fn build_to_felt252<'ctx, 'this>(
     )?;
     let value: Value = entry.arg(0)?;
 
-    let result = entry.extui(value, felt252_ty, location)?;
+    let result = {
+        let result = entry.extui(value, felt252_ty, location)?;
+
+        montgomery::mlir::monty_transform(context, entry, result, felt252_ty, location)?
+    };
 
     helper.br(entry, 0, &[result], location)
 }
@@ -116,7 +120,13 @@ pub fn build_from_felt252<'ctx, 'this>(
     let range_check: Value =
         super::increment_builtin_counter_by(context, entry, location, entry.arg(0)?, 3)?;
 
-    let value: Value = entry.arg(1)?;
+    let value: Value = montgomery::mlir::monty_reduce(
+        context,
+        entry,
+        entry.arg(1)?,
+        IntegerType::new(context, 252).into(),
+        location,
+    )?;
 
     let felt252_ty =
         registry.build_type(context, helper, metadata, &info.param_signatures()[1].ty)?;