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![OutputVarInfo::new_builtin(range_check_type, 0), OutputVarInfo {
123                        ty: range_ty,
124                        ref_info: OutputVarReferenceInfo::SimpleDerefs,
125                    }],
126                    ap_change: SierraApChange::Known { new_vars_only: false },
127                },
128            ],
129            fallthrough: Some(0),
130        })
131    }
132}
133
134/// Libfunc that takes the range `[x, y)` and if `x < y`, returns the range `[x + 1, y)` and the
135/// value `x`.
136#[derive(Default)]
137pub struct IntRangePopFrontLibfunc {}
138impl SignatureOnlyGenericLibfunc for IntRangePopFrontLibfunc {
139    const STR_ID: &'static str = "int_range_pop_front";
140
141    fn specialize_signature(
142        &self,
143        context: &dyn SignatureSpecializationContext,
144        args: &[GenericArg],
145    ) -> Result<LibfuncSignature, SpecializationError> {
146        let ty = args_as_single_type(args)?;
147        let range_ty = context.get_wrapped_concrete_type(IntRangeType::id(), ty.clone())?;
148
149        Ok(LibfuncSignature {
150            param_signatures: vec![ParamSignature::new(range_ty.clone())],
151            branch_signatures: vec![
152                // Failure.
153                BranchSignature {
154                    vars: vec![],
155                    ap_change: SierraApChange::Known { new_vars_only: false },
156                },
157                // Success.
158                BranchSignature {
159                    vars: vec![
160                        OutputVarInfo {
161                            ty: range_ty,
162                            ref_info: OutputVarReferenceInfo::Deferred(
163                                DeferredOutputKind::AddConst { param_idx: 0 },
164                            ),
165                        },
166                        OutputVarInfo {
167                            ty,
168                            ref_info: OutputVarReferenceInfo::PartialParam { param_idx: 0 },
169                        },
170                    ],
171                    ap_change: SierraApChange::Known { new_vars_only: false },
172                },
173            ],
174            fallthrough: Some(0),
175        })
176    }
177}