Added wrap non-zero method.

commit-id:6c35f136
starkware-libs · Jun 4, 2024 · 1660595 · 1660595
1 parent 293c82e
commit 1660595
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 2 deletions.
diff --git a/crates/cairo-lang-sierra-ap-change/src/core_libfunc_ap_change.rs b/crates/cairo-lang-sierra-ap-change/src/core_libfunc_ap_change.rs
@@ -388,6 +388,9 @@ pub fn core_libfunc_ap_change<InfoProvider: InvocationApChangeInfoProvider>(
                 ]
             }
             BoundedIntConcreteLibfunc::IsZero(_) => vec![ApChange::Known(0), ApChange::Known(0)],
+            BoundedIntConcreteLibfunc::WrapNonZero(_) => {
+                vec![ApChange::Known(0)]
+            }
         },
         Circuit(CircuitConcreteLibfunc::U384IsZero(_)) => {
             vec![ApChange::Known(0), ApChange::Known(0)]

diff --git a/crates/cairo-lang-sierra-gas/src/core_libfunc_cost_base.rs b/crates/cairo-lang-sierra-gas/src/core_libfunc_cost_base.rs
@@ -491,6 +491,9 @@ pub fn core_libfunc_cost(
             BoundedIntConcreteLibfunc::IsZero(_) => {
                 vec![ConstCost::steps(1).into(), ConstCost::steps(1).into()]
             }
+            BoundedIntConcreteLibfunc::WrapNonZero(_) => {
+                vec![ConstCost::steps(0).into()]
+            }
         },
         Circuit(CircuitConcreteLibfunc::FillInput(_)) => {
             vec![ConstCost::steps(7).into(), ConstCost::steps(7).into()]

diff --git a/crates/cairo-lang-sierra-to-casm/src/invocations/int/bounded.rs b/crates/cairo-lang-sierra-to-casm/src/invocations/int/bounded.rs
@@ -11,7 +11,7 @@ use num_bigint::BigInt;
 use num_traits::One;
 
 use crate::invocations::felt252::build_felt252_op_with_var;
-use crate::invocations::misc::build_is_zero;
+use crate::invocations::misc::{build_identity, build_is_zero};
 use crate::invocations::{
     add_input_variables, get_non_fallthrough_statement_id, CompiledInvocation,
     CompiledInvocationBuilder, CostValidationInfo, InvocationError,
@@ -39,6 +39,7 @@ pub fn build(
             build_constrain(builder, &libfunc.boundary)
         }
         BoundedIntConcreteLibfunc::IsZero(_) => build_is_zero(builder),
+        BoundedIntConcreteLibfunc::WrapNonZero(_) => build_identity(builder),
     }
 }
 

diff --git a/crates/cairo-lang-sierra/src/extensions/modules/bounded_int.rs b/crates/cairo-lang-sierra/src/extensions/modules/bounded_int.rs
@@ -8,7 +8,7 @@ use starknet_types_core::felt::Felt as Felt252;
 
 use super::non_zero::nonzero_ty;
 use super::range_check::RangeCheckType;
-use super::utils::Range;
+use super::utils::{reinterpret_cast_signature, Range};
 use crate::define_libfunc_hierarchy;
 use crate::extensions::lib_func::{
     BranchSignature, DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature,
@@ -80,6 +80,7 @@ define_libfunc_hierarchy! {
         DivRem(BoundedIntDivRemLibfunc),
         Constrain(BoundedIntConstrainLibfunc),
         IsZero(BoundedIntIsZeroLibfunc),
+        WrapNonZero(BoundedIntWrapNonZeroLibfunc),
     }, BoundedIntConcreteLibfunc
 }
 
@@ -404,6 +405,30 @@ impl SignatureOnlyGenericLibfunc for BoundedIntIsZeroLibfunc {
     }
 }
 
+/// Libfunc for wrapping a given bounded int with non-zero, given 0 is not in the range.
+#[derive(Default)]
+pub struct BoundedIntWrapNonZeroLibfunc {}
+impl SignatureOnlyGenericLibfunc for BoundedIntWrapNonZeroLibfunc {
+    const STR_ID: &'static str = "bounded_int_wrap_non_zero";
+
+    fn specialize_signature(
+        &self,
+        context: &dyn SignatureSpecializationContext,
+        args: &[GenericArg],
+    ) -> Result<LibfuncSignature, SpecializationError> {
+        let ty = args_as_single_type(args)?;
+        let range = Range::from_type(context, ty.clone())?;
+        // Making sure 0 is not in the given range.
+        require(range.lower.is_positive() || !range.upper.is_positive())
+            .ok_or(SpecializationError::UnsupportedGenericArg)?;
+        let prime: BigInt = Felt252::prime().to_bigint().unwrap();
+        require(range.upper <= prime && range.lower > -prime)
+            .ok_or(SpecializationError::UnsupportedGenericArg)?;
+        let nz_ty = nonzero_ty(context, &ty)?;
+        Ok(reinterpret_cast_signature(ty, nz_ty))
+    }
+}
+
 /// Returns the concrete type for a BoundedInt<min, max>.
 pub fn bounded_int_ty(
     context: &dyn SignatureSpecializationContext,

diff --git a/crates/cairo-lang-starknet-classes/src/allowed_libfuncs_lists/all.json b/crates/cairo-lang-starknet-classes/src/allowed_libfuncs_lists/all.json
@@ -24,6 +24,7 @@
         "bounded_int_is_zero",
         "bounded_int_mul",
         "bounded_int_sub",
+        "bounded_int_wrap_non_zero",
         "box_forward_snapshot",
         "branch_align",
         "bytes31_const",

diff --git a/tests/e2e_test_data/libfuncs/bounded_int b/tests/e2e_test_data/libfuncs/bounded_int
@@ -641,3 +641,73 @@ store_temp<core::zeroable::IsZeroResult::<test::BoundedInt::<0, 6805647338418769
 return([4]); // 9
 
 test::foo@0([0]: BoundedInt<0, 680564733841876926926749214863536422911>) -> (core::zeroable::IsZeroResult::<test::BoundedInt::<0, 680564733841876926926749214863536422911>>);
+
+//! > ==========================================================================
+
+//! > bounded_int_wrap_non_zero libfunc positive.
+
+//! > test_runner_name
+SmallE2ETestRunner
+
+//! > cairo
+extern type BoundedInt<const MIN: felt252, const MAX: felt252>;
+extern fn bounded_int_wrap_non_zero<T>(value: T) -> NonZero<T> implicits() nopanic;
+
+fn foo(value: BoundedInt<1, 0x1ffffffffffffffffffffffffffffffff>) -> NonZero<BoundedInt<1, 0x1ffffffffffffffffffffffffffffffff>> {
+    bounded_int_wrap_non_zero(value)
+}
+
+//! > casm
+[ap + 0] = [fp + -3], ap++;
+ret;
+
+//! > function_costs
+test::foo: OrderedHashMap({Const: 100})
+
+//! > sierra_code
+type BoundedInt<1, 680564733841876926926749214863536422911> = BoundedInt<1, 680564733841876926926749214863536422911> [storable: true, drop: true, dup: true, zero_sized: false];
+type NonZero<BoundedInt<1, 680564733841876926926749214863536422911>> = NonZero<BoundedInt<1, 680564733841876926926749214863536422911>> [storable: true, drop: true, dup: true, zero_sized: false];
+
+libfunc bounded_int_wrap_non_zero<BoundedInt<1, 680564733841876926926749214863536422911>> = bounded_int_wrap_non_zero<BoundedInt<1, 680564733841876926926749214863536422911>>;
+libfunc store_temp<NonZero<BoundedInt<1, 680564733841876926926749214863536422911>>> = store_temp<NonZero<BoundedInt<1, 680564733841876926926749214863536422911>>>;
+
+bounded_int_wrap_non_zero<BoundedInt<1, 680564733841876926926749214863536422911>>([0]) -> ([1]); // 0
+store_temp<NonZero<BoundedInt<1, 680564733841876926926749214863536422911>>>([1]) -> ([1]); // 1
+return([1]); // 2
+
+test::foo@0([0]: BoundedInt<1, 680564733841876926926749214863536422911>) -> (NonZero<BoundedInt<1, 680564733841876926926749214863536422911>>);
+
+//! > ==========================================================================
+
+//! > bounded_int_wrap_non_zero libfunc negative.
+
+//! > test_runner_name
+SmallE2ETestRunner
+
+//! > cairo
+extern type BoundedInt<const MIN: felt252, const MAX: felt252>;
+extern fn bounded_int_wrap_non_zero<T>(value: T) -> NonZero<T> implicits() nopanic;
+
+fn foo(value: BoundedInt<-0x1ffffffffffffffffffffffffffffffff, -1>) -> NonZero<BoundedInt<-0x1ffffffffffffffffffffffffffffffff, -1>> {
+    bounded_int_wrap_non_zero(value)
+}
+
+//! > casm
+[ap + 0] = [fp + -3], ap++;
+ret;
+
+//! > function_costs
+test::foo: OrderedHashMap({Const: 100})
+
+//! > sierra_code
+type BoundedInt<-680564733841876926926749214863536422911, -1> = BoundedInt<-680564733841876926926749214863536422911, -1> [storable: true, drop: true, dup: true, zero_sized: false];
+type NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>> = NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>> [storable: true, drop: true, dup: true, zero_sized: false];
+
+libfunc bounded_int_wrap_non_zero<BoundedInt<-680564733841876926926749214863536422911, -1>> = bounded_int_wrap_non_zero<BoundedInt<-680564733841876926926749214863536422911, -1>>;
+libfunc store_temp<NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>>> = store_temp<NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>>>;
+
+bounded_int_wrap_non_zero<BoundedInt<-680564733841876926926749214863536422911, -1>>([0]) -> ([1]); // 0
+store_temp<NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>>>([1]) -> ([1]); // 1
+return([1]); // 2
+
+test::foo@0([0]: BoundedInt<-680564733841876926926749214863536422911, -1>) -> (NonZero<BoundedInt<-680564733841876926926749214863536422911, -1>>);