cairo_lang_sierra/extensions/modules/
qm31.rs

1use num_bigint::BigInt;
2use num_traits::Zero;
3use num_traits::cast::ToPrimitive;
4
5use super::bounded_int::bounded_int_ty;
6use super::is_zero::{IsZeroLibfunc, IsZeroTraits};
7use super::non_zero::nonzero_ty;
8use super::range_check::RangeCheckType;
9use super::utils::reinterpret_cast_signature;
10use crate::define_libfunc_hierarchy;
11use crate::extensions::lib_func::{
12    DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature, SierraApChange,
13    SignatureSpecializationContext, SpecializationContext,
14};
15use crate::extensions::{
16    GenericLibfunc, NamedLibfunc, NamedType, NoGenericArgsGenericLibfunc, NoGenericArgsGenericType,
17    OutputVarReferenceInfo, SignatureBasedConcreteLibfunc, SpecializationError,
18};
19use crate::ids::{ConcreteTypeId, GenericLibfuncId, GenericTypeId};
20use crate::program::GenericArg;
21
22/// Type for qm31.
23/// The native type of the Cairo architecture.
24#[derive(Default)]
25pub struct QM31Type {}
26impl NoGenericArgsGenericType for QM31Type {
27    const ID: GenericTypeId = GenericTypeId::new_inline("qm31");
28    const STORABLE: bool = true;
29    const DUPLICATABLE: bool = true;
30    const DROPPABLE: bool = true;
31    const ZERO_SIZED: bool = false;
32}
33
34define_libfunc_hierarchy! {
35    pub enum QM31Libfunc {
36        BinaryOperation(QM31BinaryOperationLibfunc),
37        Const(QM31ConstLibfunc),
38        IsZero(QM31JumpNotZeroLibfunc),
39        Pack(QM31PackLibfunc),
40        Unpack(QM31UnpackLibfunc),
41        FromM31(QM31FromM31Libfunc),
42    }, QM31Concrete
43}
44
45const M31_BOUND: i32 = i32::MAX - 1;
46
47#[derive(Default)]
48pub struct QM31Traits {}
49impl IsZeroTraits for QM31Traits {
50    const IS_ZERO: &'static str = "qm31_is_zero";
51    const GENERIC_TYPE_ID: GenericTypeId = <QM31Type as NamedType>::ID;
52}
53pub type QM31JumpNotZeroLibfunc = IsZeroLibfunc<QM31Traits>;
54
55/// QM31 binary operators.
56#[derive(Copy, Clone, Debug, PartialEq, Eq)]
57pub enum QM31BinaryOperator {
58    Add,
59    Sub,
60    Mul,
61    Div,
62}
63
64/// M31 value type.
65#[derive(Copy, Clone, Debug, PartialEq, Eq)]
66pub enum M31ValueType {
67    /// A single M31 value.
68    Single,
69    /// A full QM31 type.
70    Quad,
71}
72
73/// Libfunc for qm31 binary operations.
74pub struct QM31BinaryOperationLibfunc {
75    pub operator: QM31BinaryOperator,
76    pub value_type: M31ValueType,
77}
78impl QM31BinaryOperationLibfunc {
79    fn new(operator: QM31BinaryOperator, value_type: M31ValueType) -> Self {
80        Self { operator, value_type }
81    }
82    const QM31_ADD: &'static str = "qm31_add";
83    const QM31_SUB: &'static str = "qm31_sub";
84    const QM31_MUL: &'static str = "qm31_mul";
85    const QM31_DIV: &'static str = "qm31_div";
86    const M31_ADD: &'static str = "m31_add";
87    const M31_SUB: &'static str = "m31_sub";
88    const M31_MUL: &'static str = "m31_mul";
89    const M31_DIV: &'static str = "m31_div";
90}
91impl GenericLibfunc for QM31BinaryOperationLibfunc {
92    type Concrete = QM31BinaryOpConcreteLibfunc;
93
94    fn supported_ids() -> Vec<GenericLibfuncId> {
95        vec![
96            GenericLibfuncId::from(Self::QM31_ADD),
97            GenericLibfuncId::from(Self::QM31_SUB),
98            GenericLibfuncId::from(Self::QM31_MUL),
99            GenericLibfuncId::from(Self::QM31_DIV),
100            GenericLibfuncId::from(Self::M31_ADD),
101            GenericLibfuncId::from(Self::M31_SUB),
102            GenericLibfuncId::from(Self::M31_MUL),
103            GenericLibfuncId::from(Self::M31_DIV),
104        ]
105    }
106
107    fn by_id(id: &GenericLibfuncId) -> Option<Self> {
108        match id.0.as_str() {
109            Self::QM31_ADD => Some(Self::new(QM31BinaryOperator::Add, M31ValueType::Quad)),
110            Self::QM31_SUB => Some(Self::new(QM31BinaryOperator::Sub, M31ValueType::Quad)),
111            Self::QM31_MUL => Some(Self::new(QM31BinaryOperator::Mul, M31ValueType::Quad)),
112            Self::QM31_DIV => Some(Self::new(QM31BinaryOperator::Div, M31ValueType::Quad)),
113            Self::M31_ADD => Some(Self::new(QM31BinaryOperator::Add, M31ValueType::Single)),
114            Self::M31_SUB => Some(Self::new(QM31BinaryOperator::Sub, M31ValueType::Single)),
115            Self::M31_MUL => Some(Self::new(QM31BinaryOperator::Mul, M31ValueType::Single)),
116            Self::M31_DIV => Some(Self::new(QM31BinaryOperator::Div, M31ValueType::Single)),
117            _ => None,
118        }
119    }
120
121    fn specialize_signature(
122        &self,
123        context: &dyn SignatureSpecializationContext,
124        args: &[GenericArg],
125    ) -> Result<LibfuncSignature, SpecializationError> {
126        let ty = match self.value_type {
127            M31ValueType::Quad => context.get_concrete_type(QM31Type::id(), &[]),
128            M31ValueType::Single => m31_ty(context),
129        }?;
130        let second_param_type = if matches!(self.operator, QM31BinaryOperator::Div) {
131            nonzero_ty(context, &ty)?
132        } else {
133            ty.clone()
134        };
135        match args {
136            [] => Ok(LibfuncSignature::new_non_branch_ex(
137                vec![
138                    ParamSignature::new(ty.clone()),
139                    ParamSignature::new(second_param_type).with_allow_const(),
140                ],
141                vec![OutputVarInfo { ty, ref_info: OutputVarReferenceInfo::NewTempVar { idx: 0 } }],
142                SierraApChange::Known { new_vars_only: true },
143            )),
144            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
145        }
146    }
147
148    fn specialize(
149        &self,
150        context: &dyn SpecializationContext,
151        args: &[GenericArg],
152    ) -> Result<Self::Concrete, SpecializationError> {
153        match args {
154            [] => Ok({
155                QM31BinaryOpConcreteLibfunc {
156                    operator: self.operator,
157                    signature: self.specialize_signature(context.upcast(), args)?,
158                }
159            }),
160            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
161        }
162    }
163}
164
165pub struct QM31BinaryOpConcreteLibfunc {
166    pub operator: QM31BinaryOperator,
167    pub signature: LibfuncSignature,
168}
169impl SignatureBasedConcreteLibfunc for QM31BinaryOpConcreteLibfunc {
170    fn signature(&self) -> &LibfuncSignature {
171        &self.signature
172    }
173}
174
175/// Libfunc for creating a constant qm31.
176#[derive(Default)]
177pub struct QM31ConstLibfunc {}
178impl NamedLibfunc for QM31ConstLibfunc {
179    type Concrete = QM31ConstConcreteLibfunc;
180    const STR_ID: &'static str = "qm31_const";
181
182    fn specialize_signature(
183        &self,
184        context: &dyn SignatureSpecializationContext,
185        _args: &[GenericArg],
186    ) -> Result<LibfuncSignature, SpecializationError> {
187        Ok(LibfuncSignature::new_non_branch(
188            vec![],
189            vec![OutputVarInfo {
190                ty: context.get_concrete_type(QM31Type::id(), &[])?,
191                ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::Const),
192            }],
193            SierraApChange::Known { new_vars_only: true },
194        ))
195    }
196
197    fn specialize(
198        &self,
199        context: &dyn SpecializationContext,
200        args: &[GenericArg],
201    ) -> Result<Self::Concrete, SpecializationError> {
202        let to_m31 = |arg: &GenericArg| -> Result<u32, SpecializationError> {
203            match arg {
204                GenericArg::Value(val) if *val <= M31_BOUND.into() => {
205                    val.to_u32().ok_or(SpecializationError::UnsupportedGenericArg)
206                }
207                _ => Err(SpecializationError::UnsupportedGenericArg),
208            }
209        };
210        match args {
211            [w0, w1, w2, w3] => Ok(QM31ConstConcreteLibfunc {
212                w0: to_m31(w0)?,
213                w1: to_m31(w1)?,
214                w2: to_m31(w2)?,
215                w3: to_m31(w3)?,
216                signature: <Self as NamedLibfunc>::specialize_signature(
217                    self,
218                    context.upcast(),
219                    args,
220                )?,
221            }),
222            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
223        }
224    }
225}
226
227pub struct QM31ConstConcreteLibfunc {
228    pub w0: u32,
229    pub w1: u32,
230    pub w2: u32,
231    pub w3: u32,
232    pub signature: LibfuncSignature,
233}
234impl SignatureBasedConcreteLibfunc for QM31ConstConcreteLibfunc {
235    fn signature(&self) -> &LibfuncSignature {
236        &self.signature
237    }
238}
239
240/// Libfunc for packing 4 `m31`s into a `qm31`.
241#[derive(Default)]
242pub struct QM31PackLibfunc {}
243impl NoGenericArgsGenericLibfunc for QM31PackLibfunc {
244    const STR_ID: &'static str = "qm31_pack";
245
246    fn specialize_signature(
247        &self,
248        context: &dyn SignatureSpecializationContext,
249    ) -> Result<LibfuncSignature, SpecializationError> {
250        Ok(LibfuncSignature::new_non_branch_ex(
251            vec![ParamSignature::new(m31_ty(context)?); 4],
252            vec![OutputVarInfo {
253                ty: context.get_concrete_type(QM31Type::id(), &[])?,
254                ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::Generic),
255            }],
256            SierraApChange::Known { new_vars_only: false },
257        ))
258    }
259}
260
261/// Libfunc for unpacking a `qm31` into 4 `m31`s.
262#[derive(Default)]
263pub struct QM31UnpackLibfunc {}
264impl NoGenericArgsGenericLibfunc for QM31UnpackLibfunc {
265    const STR_ID: &'static str = "qm31_unpack";
266
267    fn specialize_signature(
268        &self,
269        context: &dyn SignatureSpecializationContext,
270    ) -> Result<LibfuncSignature, SpecializationError> {
271        let range_check_ty = context.get_concrete_type(RangeCheckType::id(), &[])?;
272        let output_var_info =
273            OutputVarInfo { ty: m31_ty(context)?, ref_info: OutputVarReferenceInfo::SimpleDerefs };
274        Ok(LibfuncSignature::new_non_branch_ex(
275            vec![
276                ParamSignature::new(range_check_ty.clone()).with_allow_add_const(),
277                ParamSignature::new(context.get_concrete_type(QM31Type::id(), &[])?),
278            ],
279            vec![
280                OutputVarInfo::new_builtin(range_check_ty, 0),
281                output_var_info.clone(),
282                output_var_info.clone(),
283                output_var_info.clone(),
284                output_var_info,
285            ],
286            SierraApChange::Known { new_vars_only: false },
287        ))
288    }
289}
290
291/// Libfunc for casting from an m31 to qm31.
292#[derive(Default)]
293pub struct QM31FromM31Libfunc {}
294impl NoGenericArgsGenericLibfunc for QM31FromM31Libfunc {
295    const STR_ID: &'static str = "qm31_from_m31";
296
297    fn specialize_signature(
298        &self,
299        context: &dyn SignatureSpecializationContext,
300    ) -> Result<LibfuncSignature, SpecializationError> {
301        Ok(reinterpret_cast_signature(
302            m31_ty(context)?,
303            context.get_concrete_type(QM31Type::id(), &[])?,
304        ))
305    }
306}
307
308fn m31_ty(
309    context: &dyn SignatureSpecializationContext,
310) -> Result<ConcreteTypeId, SpecializationError> {
311    bounded_int_ty(context, BigInt::zero(), M31_BOUND.into())
312}
cairo_lang_sierra/extensions/modules/qm31.rs

cairo_lang_sierra/extensions/modules/
qm31.rs
