cairo_lang_sierra/extensions/modules/
range.rs

1use super::bounded_int::BoundedIntType;
2use super::int::signed::{Sint8Type, Sint16Type, Sint32Type, Sint64Type};
3use super::int::signed128::Sint128Type;
4use super::int::unsigned::{Uint8Type, Uint16Type, Uint32Type, Uint64Type};
5use super::int::unsigned128::Uint128Type;
6use super::range_check::RangeCheckType;
7use super::utils::Range;
8use crate::define_libfunc_hierarchy;
9use crate::extensions::lib_func::{
10    BranchSignature, DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature,
11    SierraApChange, SignatureOnlyGenericLibfunc, SignatureSpecializationContext,
12};
13use crate::extensions::type_specialization_context::TypeSpecializationContext;
14use crate::extensions::types::{
15    GenericTypeArgGenericType, GenericTypeArgGenericTypeWrapper, TypeInfo,
16};
17use crate::extensions::{
18    NamedType, OutputVarReferenceInfo, SpecializationError, args_as_single_type,
19};
20use crate::ids::GenericTypeId;
21use crate::program::GenericArg;
22
23fn check_inner_type(ty_info: &TypeInfo) -> Result<(), SpecializationError> {
24    // Note: the implementation assumes the following types are of size 1.
25    match (&ty_info.long_id.generic_id, &ty_info.long_id.generic_args[..]) {
26        (id, []) if *id == Uint8Type::id() => (),
27        (id, []) if *id == Uint16Type::id() => (),
28        (id, []) if *id == Uint32Type::id() => (),
29        (id, []) if *id == Uint64Type::id() => (),
30        (id, []) if *id == Uint128Type::id() => (),
31        (id, []) if *id == Sint8Type::id() => (),
32        (id, []) if *id == Sint16Type::id() => (),
33        (id, []) if *id == Sint32Type::id() => (),
34        (id, []) if *id == Sint64Type::id() => (),
35        (id, []) if *id == Sint128Type::id() => (),
36        (id, [GenericArg::Value(_), GenericArg::Value(_)]) if *id == BoundedIntType::id() => (),
37        _ => return Err(SpecializationError::UnsupportedGenericArg),
38    };
39    Ok(())
40}
41
42/// Type for `IntRange(x, y)` where `x <= y`.
43#[derive(Default)]
44pub struct IntRangeTypeWrapped {}
45impl GenericTypeArgGenericType for IntRangeTypeWrapped {
46    const ID: GenericTypeId = GenericTypeId::new_inline("IntRange");
47
48    fn calc_info(
49        &self,
50        _context: &dyn TypeSpecializationContext,
51        long_id: crate::program::ConcreteTypeLongId,
52        wrapped_info: TypeInfo,
53    ) -> Result<TypeInfo, SpecializationError> {
54        check_inner_type(&wrapped_info)?;
55
56        // The following assert is a sanity check. It should follow from the fact that
57        // `check_inner_type` passed.
58        assert!(
59            wrapped_info.storable
60                && wrapped_info.duplicatable
61                && wrapped_info.droppable
62                && !wrapped_info.zero_sized
63        );
64        Ok(TypeInfo {
65            long_id,
66            duplicatable: true,
67            droppable: true,
68            storable: true,
69            zero_sized: false,
70        })
71    }
72}
73pub type IntRangeType = GenericTypeArgGenericTypeWrapper<IntRangeTypeWrapped>;
74
75define_libfunc_hierarchy! {
76    pub enum IntRangeLibfunc {
77        TryNew(IntRangeTryNewLibfunc),
78        PopFront(IntRangePopFrontLibfunc),
79    }, IntRangeConcreteLibfunc
80}
81
82/// Libfunc that constructs the range `[x, y)` if `x <= y`.
83/// Otherwise, returns the empty range `[y, y)`.
84#[derive(Default)]
85pub struct IntRangeTryNewLibfunc {}
86impl SignatureOnlyGenericLibfunc for IntRangeTryNewLibfunc {
87    const STR_ID: &'static str = "int_range_try_new";
88
89    fn specialize_signature(
90        &self,
91        context: &dyn SignatureSpecializationContext,
92        args: &[GenericArg],
93    ) -> Result<LibfuncSignature, SpecializationError> {
94        let ty = args_as_single_type(args)?;
95        let range_ty = context.get_wrapped_concrete_type(IntRangeType::id(), ty.clone())?;
96        let range_check_type = context.get_concrete_type(RangeCheckType::id(), &[])?;
97
98        if !Range::from_type(context, ty.clone())?.is_small_range() {
99            return Err(SpecializationError::UnsupportedGenericArg);
100        }
101
102        Ok(LibfuncSignature {
103            param_signatures: vec![
104                ParamSignature::new(range_check_type.clone()).with_allow_add_const(),
105                ParamSignature::new(ty.clone()),
106                ParamSignature::new(ty.clone()),
107            ],
108            branch_signatures: vec![
109                // Success.
110                BranchSignature {
111                    vars: vec![
112                        OutputVarInfo::new_builtin(range_check_type.clone(), 0),
113                        OutputVarInfo {
114                            ty: range_ty.clone(),
115                            ref_info: OutputVarReferenceInfo::SimpleDerefs,
116                        },
117                    ],
118                    ap_change: SierraApChange::Known { new_vars_only: false },
119                },
120                // Failure.
121                BranchSignature {
122                    vars: vec![
123                        OutputVarInfo::new_builtin(range_check_type, 0),
124                        OutputVarInfo {
125                            ty: range_ty,
126                            ref_info: OutputVarReferenceInfo::SimpleDerefs,
127                        },
128                    ],
129                    ap_change: SierraApChange::Known { new_vars_only: false },
130                },
131            ],
132            fallthrough: Some(0),
133        })
134    }
135}
136
137/// Libfunc that takes the range `[x, y)` and if `x < y`, returns the range `[x + 1, y)` and the
138/// value `x`.
139#[derive(Default)]
140pub struct IntRangePopFrontLibfunc {}
141impl SignatureOnlyGenericLibfunc for IntRangePopFrontLibfunc {
142    const STR_ID: &'static str = "int_range_pop_front";
143
144    fn specialize_signature(
145        &self,
146        context: &dyn SignatureSpecializationContext,
147        args: &[GenericArg],
148    ) -> Result<LibfuncSignature, SpecializationError> {
149        let ty = args_as_single_type(args)?;
150        let range_ty = context.get_wrapped_concrete_type(IntRangeType::id(), ty.clone())?;
151
152        Ok(LibfuncSignature {
153            param_signatures: vec![ParamSignature::new(range_ty.clone())],
154            branch_signatures: vec![
155                // Failure.
156                BranchSignature {
157                    vars: vec![],
158                    ap_change: SierraApChange::Known { new_vars_only: false },
159                },
160                // Success.
161                BranchSignature {
162                    vars: vec![
163                        OutputVarInfo {
164                            ty: range_ty,
165                            ref_info: OutputVarReferenceInfo::Deferred(
166                                DeferredOutputKind::AddConst { param_idx: 0 },
167                            ),
168                        },
169                        OutputVarInfo {
170                            ty,
171                            ref_info: OutputVarReferenceInfo::PartialParam { param_idx: 0 },
172                        },
173                    ],
174                    ap_change: SierraApChange::Known { new_vars_only: false },
175                },
176            ],
177            fallthrough: Some(0),
178        })
179    }
180}