cairo_lang_sierra/extensions/modules/
non_zero.rs

1use crate::extensions::lib_func::{
2    LibfuncSignature, OutputVarInfo, SierraApChange, SignatureOnlyGenericLibfunc,
3    SignatureSpecializationContext,
4};
5use crate::extensions::type_specialization_context::TypeSpecializationContext;
6use crate::extensions::types::{
7    GenericTypeArgGenericType, GenericTypeArgGenericTypeWrapper, TypeInfo,
8};
9use crate::extensions::{
10    NamedType, OutputVarReferenceInfo, SpecializationError, args_as_single_type,
11};
12use crate::ids::{ConcreteTypeId, GenericTypeId};
13use crate::program::GenericArg;
14
15/// Type wrapping a value as non zero.
16#[derive(Default)]
17pub struct NonZeroTypeWrapped {}
18impl GenericTypeArgGenericType for NonZeroTypeWrapped {
19    const ID: GenericTypeId = GenericTypeId::new_inline("NonZero");
20
21    fn calc_info(
22        &self,
23        _context: &dyn TypeSpecializationContext,
24        long_id: crate::program::ConcreteTypeLongId,
25        TypeInfo { zero_sized, storable, droppable, duplicatable, .. }: TypeInfo,
26    ) -> Result<TypeInfo, SpecializationError> {
27        if storable {
28            Ok(TypeInfo { long_id, zero_sized, storable, droppable, duplicatable })
29        } else {
30            Err(SpecializationError::UnsupportedGenericArg)
31        }
32    }
33}
34pub type NonZeroType = GenericTypeArgGenericTypeWrapper<NonZeroTypeWrapped>;
35
36/// Returns the type `NonZero<T>` for a given type `T`.
37pub fn nonzero_ty(
38    context: &dyn SignatureSpecializationContext,
39    ty: &ConcreteTypeId,
40) -> Result<ConcreteTypeId, SpecializationError> {
41    context.get_wrapped_concrete_type(NonZeroType::id(), ty.clone())
42}
43
44/// Libfunc for unwrapping a `NonZero<T>` back into a T.
45#[derive(Default)]
46pub struct UnwrapNonZeroLibfunc {}
47impl SignatureOnlyGenericLibfunc for UnwrapNonZeroLibfunc {
48    const STR_ID: &'static str = "unwrap_non_zero";
49
50    fn specialize_signature(
51        &self,
52        context: &dyn SignatureSpecializationContext,
53        args: &[GenericArg],
54    ) -> Result<LibfuncSignature, SpecializationError> {
55        let ty = args_as_single_type(args)?;
56        Ok(LibfuncSignature::new_non_branch(
57            vec![nonzero_ty(context, &ty)?],
58            vec![OutputVarInfo {
59                ty,
60                ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 0 },
61            }],
62            SierraApChange::Known { new_vars_only: true },
63        ))
64    }
65}