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#[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 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#[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#[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#[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 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#[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 require(!lhs_range.lower.is_negative() && !rhs_range.lower.is_negative())
200 .ok_or(SpecializationError::UnsupportedGenericArg)?;
201 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
256pub enum BoundedIntDivRemAlgorithm {
258 KnownSmallRhs,
260 KnownSmallQuotient { q_upper_bound: BigInt },
262 KnownSmallLhs { lhs_upper_sqrt: BigInt },
266}
267impl BoundedIntDivRemAlgorithm {
268 pub fn try_new(lhs: &Range, rhs: &Range) -> Option<Self> {
273 let prime = Felt252::prime().to_bigint().unwrap();
274 let q_max = (&lhs.upper - 1) / std::cmp::max(&rhs.lower, &BigInt::one());
277 let u128_limit = BigInt::one().shl(128);
278 require(q_max < u128_limit)?;
280 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 if lhs_upper_sqrt.pow(2) != lhs.upper {
293 lhs_upper_sqrt += 1;
294 }
295 if &lhs_upper_sqrt * &u128_limit < prime {
296 require(lhs_upper_sqrt < u128_limit)?;
299 return Some(Self::KnownSmallLhs { lhs_upper_sqrt });
300 }
301 None
303 }
304}
305
306#[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#[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
463fn 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#[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 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 BranchSignature {
507 vars: vec![],
508 ap_change: SierraApChange::Known { new_vars_only: true },
509 },
510 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#[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 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
548pub 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}