cairo_lang_sierra/extensions/modules/
bounded_int.rs

1use std::ops::Shl;
2
3use cairo_lang_utils::require;
4use itertools::Itertools;
5use num_bigint::{BigInt, ToBigInt};
6use num_traits::{One, Signed, Zero};
7use starknet_types_core::felt::Felt as Felt252;
8
9use super::non_zero::{NonZeroType, nonzero_ty};
10use super::range_check::RangeCheckType;
11use super::utils::{Range, reinterpret_cast_signature};
12use crate::define_libfunc_hierarchy;
13use crate::extensions::lib_func::{
14    BranchSignature, DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature,
15    SierraApChange, SignatureOnlyGenericLibfunc, SignatureSpecializationContext,
16    SpecializationContext,
17};
18use crate::extensions::type_specialization_context::TypeSpecializationContext;
19use crate::extensions::types::TypeInfo;
20use crate::extensions::{
21    ConcreteType, NamedLibfunc, NamedType, OutputVarReferenceInfo, SignatureBasedConcreteLibfunc,
22    SpecializationError, args_as_single_type, args_as_two_types,
23};
24use crate::ids::{ConcreteTypeId, GenericTypeId};
25use crate::program::GenericArg;
26
27/// Type for BoundedInt.
28/// The native type of the Cairo architecture.
29#[derive(Default)]
30pub struct BoundedIntType {}
31impl NamedType for BoundedIntType {
32    type Concrete = BoundedIntConcreteType;
33
34    const ID: GenericTypeId = GenericTypeId::new_inline("BoundedInt");
35    fn specialize(
36        &self,
37        _context: &dyn TypeSpecializationContext,
38        args: &[GenericArg],
39    ) -> Result<Self::Concrete, SpecializationError> {
40        let (min, max) = match args {
41            [GenericArg::Value(min), GenericArg::Value(max)] => (min.clone(), max.clone()),
42            [_, _] => return Err(SpecializationError::UnsupportedGenericArg),
43            _ => return Err(SpecializationError::WrongNumberOfGenericArgs),
44        };
45
46        let prime: BigInt = Felt252::prime().into();
47        if !(-&prime < min && min <= max && max < prime && &max - &min < prime) {
48            return Err(SpecializationError::UnsupportedGenericArg);
49        }
50
51        let long_id = Self::concrete_type_long_id(args);
52        let ty_info = TypeInfo {
53            long_id,
54            zero_sized: false,
55            storable: true,
56            droppable: true,
57            duplicatable: true,
58        };
59
60        Ok(Self::Concrete { info: ty_info, range: Range::closed(min, max) })
61    }
62}
63
64pub struct BoundedIntConcreteType {
65    pub info: TypeInfo,
66    /// The range bounds for a value of this type.
67    pub range: Range,
68}
69impl ConcreteType for BoundedIntConcreteType {
70    fn info(&self) -> &TypeInfo {
71        &self.info
72    }
73}
74
75define_libfunc_hierarchy! {
76    pub enum BoundedIntLibfunc {
77        Add(BoundedIntAddLibfunc),
78        Sub(BoundedIntSubLibfunc),
79        Mul(BoundedIntMulLibfunc),
80        DivRem(BoundedIntDivRemLibfunc),
81        Constrain(BoundedIntConstrainLibfunc),
82        TrimMin(BoundedIntTrimLibfunc<false>),
83        TrimMax(BoundedIntTrimLibfunc<true>),
84        IsZero(BoundedIntIsZeroLibfunc),
85        WrapNonZero(BoundedIntWrapNonZeroLibfunc),
86    }, BoundedIntConcreteLibfunc
87}
88
89/// Libfunc for adding two BoundedInts.
90/// The result is a BoundedInt.
91#[derive(Default)]
92pub struct BoundedIntAddLibfunc {}
93impl SignatureOnlyGenericLibfunc for BoundedIntAddLibfunc {
94    const STR_ID: &'static str = "bounded_int_add";
95
96    fn specialize_signature(
97        &self,
98        context: &dyn SignatureSpecializationContext,
99        args: &[GenericArg],
100    ) -> Result<LibfuncSignature, SpecializationError> {
101        specialize_helper(context, args, |lhs_range, rhs_range| {
102            (lhs_range.lower + rhs_range.lower, (lhs_range.upper - 1) + (rhs_range.upper - 1))
103        })
104    }
105}
106
107/// Libfunc for subtracting two BoundedInts.
108/// The result is a BoundedInt.
109#[derive(Default)]
110pub struct BoundedIntSubLibfunc {}
111impl SignatureOnlyGenericLibfunc for BoundedIntSubLibfunc {
112    const STR_ID: &'static str = "bounded_int_sub";
113
114    fn specialize_signature(
115        &self,
116        context: &dyn SignatureSpecializationContext,
117        args: &[GenericArg],
118    ) -> Result<LibfuncSignature, SpecializationError> {
119        specialize_helper(context, args, |lhs_range, rhs_range| {
120            (lhs_range.lower - (rhs_range.upper - 1), (lhs_range.upper - 1) - rhs_range.lower)
121        })
122    }
123}
124
125/// Libfunc for multiplying two BoundedInts.
126/// The result is a BoundedInt.
127#[derive(Default)]
128pub struct BoundedIntMulLibfunc {}
129impl SignatureOnlyGenericLibfunc for BoundedIntMulLibfunc {
130    const STR_ID: &'static str = "bounded_int_mul";
131
132    fn specialize_signature(
133        &self,
134        context: &dyn SignatureSpecializationContext,
135        args: &[GenericArg],
136    ) -> Result<LibfuncSignature, SpecializationError> {
137        let (lhs, rhs) = args_as_two_types(args)?;
138        let lhs_info = context.get_type_info(lhs.clone())?;
139        let rhs_info = context.get_type_info(rhs.clone())?;
140        let is_nz = {
141            let lhs_is_nz = lhs_info.long_id.generic_id == NonZeroType::ID;
142            let rhs_is_nz = rhs_info.long_id.generic_id == NonZeroType::ID;
143            if lhs_is_nz != rhs_is_nz {
144                return Err(SpecializationError::UnsupportedGenericArg);
145            }
146            lhs_is_nz
147        };
148
149        let [lhs_range, rhs_range] = if is_nz {
150            [lhs_info, rhs_info].map(|info| {
151                let inner_ty = args_as_single_type(&info.long_id.generic_args)?;
152                Range::from_type(context, inner_ty)
153            })
154        } else {
155            [lhs_info, rhs_info].map(|info| Range::from_type_info(&info))
156        };
157        let (lhs_range, rhs_range) = (lhs_range?, rhs_range?);
158        // The result is the minimum and maximum of the four possible extremes.
159        // Done to properly handle multiplication by negative values.
160        let extremes = [
161            &lhs_range.lower * &rhs_range.lower,
162            lhs_range.lower * (&rhs_range.upper - 1),
163            (&lhs_range.upper - 1) * rhs_range.lower,
164            (lhs_range.upper - 1) * (rhs_range.upper - 1),
165        ];
166        let (min_result, max_result) = extremes.into_iter().minmax().into_option().unwrap();
167        let res_ty = bounded_int_ty(context, min_result, max_result)?;
168        Ok(LibfuncSignature::new_non_branch_ex(
169            vec![ParamSignature::new(lhs), ParamSignature::new(rhs).with_allow_const()],
170            vec![OutputVarInfo {
171                ty: if is_nz { nonzero_ty(context, &res_ty)? } else { res_ty },
172                ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::Generic),
173            }],
174            SierraApChange::Known { new_vars_only: true },
175        ))
176    }
177}
178
179/// Libfunc for dividing two non negative BoundedInts and getting the quotient and remainder.
180#[derive(Default)]
181pub struct BoundedIntDivRemLibfunc {}
182impl NamedLibfunc for BoundedIntDivRemLibfunc {
183    type Concrete = BoundedIntDivRemConcreteLibfunc;
184
185    const STR_ID: &'static str = "bounded_int_div_rem";
186
187    fn specialize_signature(
188        &self,
189        context: &dyn SignatureSpecializationContext,
190        args: &[GenericArg],
191    ) -> Result<LibfuncSignature, SpecializationError> {
192        let (lhs, rhs) = args_as_two_types(args)?;
193        let lhs_range = Range::from_type(context, lhs.clone())?;
194        let rhs_range = Range::from_type(context, rhs.clone())?;
195        // Supporting only division of a non-negative number by a positive number (non zero and non
196        // negative).
197        // TODO(orizi): Consider relaxing the constraint, and defining the
198        // div_rem of negatives.
199        require(!lhs_range.lower.is_negative() && !rhs_range.lower.is_negative())
200            .ok_or(SpecializationError::UnsupportedGenericArg)?;
201        // Making sure the algorithm is runnable.
202        BoundedIntDivRemAlgorithm::try_new(&lhs_range, &rhs_range)
203            .ok_or(SpecializationError::UnsupportedGenericArg)?;
204        require(rhs_range.upper >= BigInt::from(2))
205            .ok_or(SpecializationError::UnsupportedGenericArg)?;
206        let quotient_min = lhs_range.lower / (&rhs_range.upper - 1);
207        let quotient_max = (&lhs_range.upper - 1) / std::cmp::max(rhs_range.lower, BigInt::one());
208        let range_check_type = context.get_concrete_type(RangeCheckType::id(), &[])?;
209        Ok(LibfuncSignature::new_non_branch_ex(
210            vec![
211                ParamSignature::new(range_check_type.clone()).with_allow_add_const(),
212                ParamSignature::new(lhs.clone()),
213                ParamSignature::new(nonzero_ty(context, &rhs)?),
214            ],
215            vec![
216                OutputVarInfo::new_builtin(range_check_type.clone(), 0),
217                OutputVarInfo {
218                    ty: bounded_int_ty(context, quotient_min, quotient_max)?,
219                    ref_info: OutputVarReferenceInfo::SimpleDerefs,
220                },
221                OutputVarInfo {
222                    ty: bounded_int_ty(context, 0.into(), rhs_range.upper - 2)?,
223                    ref_info: OutputVarReferenceInfo::SimpleDerefs,
224                },
225            ],
226            SierraApChange::Known { new_vars_only: false },
227        ))
228    }
229
230    fn specialize(
231        &self,
232        context: &dyn SpecializationContext,
233        args: &[GenericArg],
234    ) -> Result<Self::Concrete, SpecializationError> {
235        let (lhs, rhs) = args_as_two_types(args)?;
236        let context = context.upcast();
237        Ok(Self::Concrete {
238            lhs: Range::from_type(context, lhs.clone())?,
239            rhs: Range::from_type(context, rhs.clone())?,
240            signature: self.specialize_signature(context, args)?,
241        })
242    }
243}
244
245pub struct BoundedIntDivRemConcreteLibfunc {
246    pub lhs: Range,
247    pub rhs: Range,
248    signature: LibfuncSignature,
249}
250impl SignatureBasedConcreteLibfunc for BoundedIntDivRemConcreteLibfunc {
251    fn signature(&self) -> &LibfuncSignature {
252        &self.signature
253    }
254}
255
256/// The algorithm to use for division and remainder of bounded integers.
257pub enum BoundedIntDivRemAlgorithm {
258    /// The rhs is small enough to be multiplied by `2**128` without wraparound.
259    KnownSmallRhs,
260    /// The quotient is small enough to be multiplied by `2**128` without wraparound.
261    KnownSmallQuotient { q_upper_bound: BigInt },
262    /// The lhs is small enough so that its square root plus 1 can be multiplied by `2**128`
263    /// without wraparound.
264    /// `lhs_upper_sqrt` is the square root of the upper bound of the lhs, rounded up.
265    KnownSmallLhs { lhs_upper_sqrt: BigInt },
266}
267impl BoundedIntDivRemAlgorithm {
268    /// Returns the algorithm to use for division and remainder of bounded integers.
269    /// Fails if the div_rem of the ranges is not supported yet.
270    ///
271    /// Assumption: `lhs` is non-negative and `rhs` is positive.
272    pub fn try_new(lhs: &Range, rhs: &Range) -> Option<Self> {
273        let prime = Felt252::prime().to_bigint().unwrap();
274        // Note that `rhs.lower` may be 0 - but since it is a non-zero type, it is guaranteed to be
275        // at least 1.
276        let q_max = (&lhs.upper - 1) / std::cmp::max(&rhs.lower, &BigInt::one());
277        let u128_limit = BigInt::one().shl(128);
278        // `q` is range checked in all algorithm variants, so `q_max` must be smaller than `2**128`.
279        require(q_max < u128_limit)?;
280        // `r` is range checked in all algorithm variants, so `rhs.upper` must be at most
281        // `2**128 + 1`.
282        require(rhs.upper <= &u128_limit + 1)?;
283        if &rhs.upper * &u128_limit < prime {
284            return Some(Self::KnownSmallRhs);
285        }
286        let q_upper_bound = q_max + 1;
287        if &q_upper_bound * &u128_limit < prime {
288            return Some(Self::KnownSmallQuotient { q_upper_bound });
289        }
290        let mut lhs_upper_sqrt = lhs.upper.sqrt();
291        // Round lhs_upper_sqrt up.
292        if lhs_upper_sqrt.pow(2) != lhs.upper {
293            lhs_upper_sqrt += 1;
294        }
295        if &lhs_upper_sqrt * &u128_limit < prime {
296            // Make sure `lhs_upper_sqrt < 2**128`, since the value bounded by root is range
297            // checked.
298            require(lhs_upper_sqrt < u128_limit)?;
299            return Some(Self::KnownSmallLhs { lhs_upper_sqrt });
300        }
301        // No algorithm found.
302        None
303    }
304}
305
306/// Libfunc for constraining a BoundedInt<Min, Max> to one of two non-empty ranges: [Min, Boundary)
307/// or [Boundary, Max]. The libfunc is also applicable for standard types such as u* and i*.
308#[derive(Default)]
309pub struct BoundedIntConstrainLibfunc {}
310impl NamedLibfunc for BoundedIntConstrainLibfunc {
311    type Concrete = BoundedIntConstrainConcreteLibfunc;
312
313    const STR_ID: &'static str = "bounded_int_constrain";
314
315    fn specialize_signature(
316        &self,
317        context: &dyn SignatureSpecializationContext,
318        args: &[GenericArg],
319    ) -> Result<LibfuncSignature, SpecializationError> {
320        let (ty, boundary) = match args {
321            [GenericArg::Type(ty), GenericArg::Value(boundary)] => Ok((ty, boundary)),
322            [_, _] => Err(SpecializationError::UnsupportedGenericArg),
323            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
324        }?;
325        let ty_info = context.get_type_info(ty.clone())?;
326        let is_nz = ty_info.long_id.generic_id == NonZeroType::ID;
327        let range = if is_nz {
328            let inner_ty = args_as_single_type(&ty_info.long_id.generic_args)?;
329            Range::from_type(context, inner_ty)?
330        } else {
331            Range::from_type_info(&ty_info)?
332        };
333        require(&range.lower < boundary && boundary < &range.upper)
334            .ok_or(SpecializationError::UnsupportedGenericArg)?;
335        let low_range = Range::half_open(range.lower, boundary.clone());
336        let high_range = Range::half_open(boundary.clone(), range.upper);
337        require(low_range.is_small_range() && high_range.is_small_range())
338            .ok_or(SpecializationError::UnsupportedGenericArg)?;
339        let range_check_type = context.get_concrete_type(RangeCheckType::id(), &[])?;
340        let branch_signature = |rng: Range| {
341            let inner_res_ty = bounded_int_ty(context, rng.lower, rng.upper - 1)?;
342            let res_ty = if is_nz { nonzero_ty(context, &inner_res_ty)? } else { inner_res_ty };
343            Ok(BranchSignature {
344                vars: vec![
345                    OutputVarInfo::new_builtin(range_check_type.clone(), 0),
346                    OutputVarInfo {
347                        ty: res_ty,
348                        ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 1 },
349                    },
350                ],
351                ap_change: SierraApChange::Known { new_vars_only: false },
352            })
353        };
354        Ok(LibfuncSignature {
355            param_signatures: vec![
356                ParamSignature::new(range_check_type.clone()).with_allow_add_const(),
357                ParamSignature::new(ty.clone()),
358            ],
359            branch_signatures: vec![branch_signature(low_range)?, branch_signature(high_range)?],
360            fallthrough: Some(0),
361        })
362    }
363
364    fn specialize(
365        &self,
366        context: &dyn SpecializationContext,
367        args: &[GenericArg],
368    ) -> Result<Self::Concrete, SpecializationError> {
369        let boundary = match args {
370            [GenericArg::Type(_), GenericArg::Value(boundary)] => Ok(boundary.clone()),
371            [_, _] => Err(SpecializationError::UnsupportedGenericArg),
372            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
373        }?;
374        let context = context.upcast();
375        Ok(Self::Concrete { boundary, signature: self.specialize_signature(context, args)? })
376    }
377}
378
379pub struct BoundedIntConstrainConcreteLibfunc {
380    pub boundary: BigInt,
381    signature: LibfuncSignature,
382}
383impl SignatureBasedConcreteLibfunc for BoundedIntConstrainConcreteLibfunc {
384    fn signature(&self) -> &LibfuncSignature {
385        &self.signature
386    }
387}
388
389/// Libfunc for trimming a BoundedInt<Min, Max> by removing `Min` or `Max` from the range.
390/// The libfunc is also applicable for standard types such as u* and i*.
391#[derive(Default)]
392pub struct BoundedIntTrimLibfunc<const IS_MAX: bool> {}
393impl<const IS_MAX: bool> NamedLibfunc for BoundedIntTrimLibfunc<IS_MAX> {
394    type Concrete = BoundedIntTrimConcreteLibfunc;
395
396    const STR_ID: &'static str =
397        if IS_MAX { "bounded_int_trim_max" } else { "bounded_int_trim_min" };
398
399    fn specialize_signature(
400        &self,
401        context: &dyn SignatureSpecializationContext,
402        args: &[GenericArg],
403    ) -> Result<LibfuncSignature, SpecializationError> {
404        Ok(Self::Concrete::new::<IS_MAX>(context, args)?.signature)
405    }
406
407    fn specialize(
408        &self,
409        context: &dyn SpecializationContext,
410        args: &[GenericArg],
411    ) -> Result<Self::Concrete, SpecializationError> {
412        Self::Concrete::new::<IS_MAX>(context.upcast(), args)
413    }
414}
415
416pub struct BoundedIntTrimConcreteLibfunc {
417    pub trimmed_value: BigInt,
418    signature: LibfuncSignature,
419}
420impl BoundedIntTrimConcreteLibfunc {
421    fn new<const IS_MAX: bool>(
422        context: &dyn SignatureSpecializationContext,
423        args: &[GenericArg],
424    ) -> Result<Self, SpecializationError> {
425        let ty = args_as_single_type(args)?;
426        let ty_info = context.get_type_info(ty.clone())?;
427        let range = Range::from_type_info(&ty_info)?;
428        let (res_ty, trimmed_value) = if IS_MAX {
429            (
430                bounded_int_ty(context, range.lower.clone(), range.upper.clone() - 2)?,
431                range.upper - 1,
432            )
433        } else {
434            (
435                bounded_int_ty(context, range.lower.clone() + 1, range.upper.clone() - 1)?,
436                range.lower,
437            )
438        };
439        let ap_change = SierraApChange::Known { new_vars_only: trimmed_value.is_zero() };
440        let signature = LibfuncSignature {
441            param_signatures: vec![ParamSignature::new(ty.clone())],
442            branch_signatures: vec![
443                BranchSignature { vars: vec![], ap_change: ap_change.clone() },
444                BranchSignature {
445                    vars: vec![OutputVarInfo {
446                        ty: res_ty,
447                        ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 0 },
448                    }],
449                    ap_change,
450                },
451            ],
452            fallthrough: Some(0),
453        };
454        Ok(Self { trimmed_value, signature })
455    }
456}
457impl SignatureBasedConcreteLibfunc for BoundedIntTrimConcreteLibfunc {
458    fn signature(&self) -> &LibfuncSignature {
459        &self.signature
460    }
461}
462
463/// Helper function for specializing the signature of a simple bounded int operation libfunc.
464fn specialize_helper(
465    context: &dyn SignatureSpecializationContext,
466    args: &[GenericArg],
467    result_range: impl Fn(Range, Range) -> (BigInt, BigInt),
468) -> Result<LibfuncSignature, SpecializationError> {
469    let (lhs, rhs) = args_as_two_types(args)?;
470
471    let (min_result, max_result) = result_range(
472        Range::from_type(context, lhs.clone())?,
473        Range::from_type(context, rhs.clone())?,
474    );
475    Ok(LibfuncSignature::new_non_branch_ex(
476        vec![ParamSignature::new(lhs), ParamSignature::new(rhs).with_allow_const()],
477        vec![OutputVarInfo {
478            ty: bounded_int_ty(context, min_result, max_result)?,
479            ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::Generic),
480        }],
481        SierraApChange::Known { new_vars_only: true },
482    ))
483}
484
485/// Libfunc for checking whether the given bounded int is zero or not, and returning a non-zero
486/// wrapped value in case of success.
487#[derive(Default)]
488pub struct BoundedIntIsZeroLibfunc;
489impl SignatureOnlyGenericLibfunc for BoundedIntIsZeroLibfunc {
490    const STR_ID: &'static str = "bounded_int_is_zero";
491
492    fn specialize_signature(
493        &self,
494        context: &dyn SignatureSpecializationContext,
495        args: &[GenericArg],
496    ) -> Result<LibfuncSignature, SpecializationError> {
497        let ty = args_as_single_type(args)?;
498        let range = Range::from_type(context, ty.clone())?;
499        // Make sure 0 is actually in the given range.
500        require(!range.lower.is_positive() && range.upper.is_positive())
501            .ok_or(SpecializationError::UnsupportedGenericArg)?;
502        Ok(LibfuncSignature {
503            param_signatures: vec![ParamSignature::new(ty.clone())],
504            branch_signatures: vec![
505                // Zero.
506                BranchSignature {
507                    vars: vec![],
508                    ap_change: SierraApChange::Known { new_vars_only: true },
509                },
510                // NonZero.
511                BranchSignature {
512                    vars: vec![OutputVarInfo {
513                        ty: nonzero_ty(context, &ty)?,
514                        ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 0 },
515                    }],
516                    ap_change: SierraApChange::Known { new_vars_only: true },
517                },
518            ],
519            fallthrough: Some(0),
520        })
521    }
522}
523
524/// Libfunc for wrapping a given bounded int with non-zero, given 0 is not in the range.
525#[derive(Default)]
526pub struct BoundedIntWrapNonZeroLibfunc {}
527impl SignatureOnlyGenericLibfunc for BoundedIntWrapNonZeroLibfunc {
528    const STR_ID: &'static str = "bounded_int_wrap_non_zero";
529
530    fn specialize_signature(
531        &self,
532        context: &dyn SignatureSpecializationContext,
533        args: &[GenericArg],
534    ) -> Result<LibfuncSignature, SpecializationError> {
535        let ty = args_as_single_type(args)?;
536        let range = Range::from_type(context, ty.clone())?;
537        // Make sure 0 is not in the given range.
538        require(range.lower.is_positive() || !range.upper.is_positive())
539            .ok_or(SpecializationError::UnsupportedGenericArg)?;
540        let prime: BigInt = Felt252::prime().to_bigint().unwrap();
541        require(range.upper <= prime && range.lower > -prime)
542            .ok_or(SpecializationError::UnsupportedGenericArg)?;
543        let nz_ty = nonzero_ty(context, &ty)?;
544        Ok(reinterpret_cast_signature(ty, nz_ty))
545    }
546}
547
548/// Returns the concrete type for a BoundedInt<min, max>.
549pub fn bounded_int_ty(
550    context: &dyn SignatureSpecializationContext,
551    min: BigInt,
552    max: BigInt,
553) -> Result<ConcreteTypeId, SpecializationError> {
554    context.get_concrete_type(BoundedIntType::ID, &[GenericArg::Value(min), GenericArg::Value(max)])
555}