cairo_lang_sierra/extensions/modules/
function_call.rs

1use cairo_lang_utils::extract_matches;
2use itertools::{Itertools, chain};
3use num_traits::Zero;
4
5use super::coupon::coupon_ty;
6use crate::extensions::lib_func::{
7    LibfuncSignature, OutputVarInfo, SierraApChange, SignatureBasedConcreteLibfunc,
8    SignatureSpecializationContext, SpecializationContext,
9};
10use crate::extensions::{
11    NamedLibfunc, OutputVarReferenceInfo, SpecializationError, args_as_single_user_func,
12};
13use crate::program::{Function, FunctionSignature, GenericArg};
14
15/// Returns the [OutputVarInfo] instances for the return types of the given function signature.
16fn get_output_var_infos(
17    context: &dyn SignatureSpecializationContext,
18    signature: FunctionSignature,
19) -> Result<Vec<OutputVarInfo>, SpecializationError> {
20    let mut curr_stack_idx = 0;
21    let mut get_stack_idx = || {
22        let idx = curr_stack_idx;
23        curr_stack_idx += 1;
24        idx
25    };
26
27    signature
28        .ret_types
29        .iter()
30        .map(|ty| {
31            Ok(OutputVarInfo {
32                ty: ty.clone(),
33                ref_info: if context.get_type_info(ty)?.zero_sized {
34                    OutputVarReferenceInfo::ZeroSized
35                } else {
36                    OutputVarReferenceInfo::NewTempVar { idx: get_stack_idx() }
37                },
38            })
39        })
40        .collect::<Result<Vec<_>, _>>()
41}
42
43/// Libfunc used to call user functions.
44#[derive(Default)]
45pub struct FunctionCallLibfunc {}
46impl NamedLibfunc for FunctionCallLibfunc {
47    type Concrete = SignatureAndFunctionConcreteLibfunc;
48
49    const STR_ID: &'static str = "function_call";
50
51    fn specialize_signature(
52        &self,
53        context: &dyn SignatureSpecializationContext,
54        args: &[GenericArg],
55    ) -> Result<LibfuncSignature, SpecializationError> {
56        let function_id = args_as_single_user_func(args)?;
57
58        let signature = context.get_function_signature(function_id)?;
59        let ap_change = context.get_function_ap_change(function_id)?;
60        Ok(LibfuncSignature::new_non_branch(
61            signature.param_types.clone(),
62            get_output_var_infos(context, signature)?,
63            ap_change,
64        ))
65    }
66
67    fn specialize(
68        &self,
69        context: &dyn SpecializationContext,
70        args: &[GenericArg],
71    ) -> Result<Self::Concrete, SpecializationError> {
72        let function_id = args_as_single_user_func(args)?;
73        Ok(Self::Concrete {
74            function: context.get_function(function_id)?,
75            signature: self.specialize_signature(context, args)?,
76        })
77    }
78}
79
80pub struct SignatureAndFunctionConcreteLibfunc {
81    pub function: Function,
82    pub signature: LibfuncSignature,
83}
84impl SignatureBasedConcreteLibfunc for SignatureAndFunctionConcreteLibfunc {
85    fn signature(&self) -> &LibfuncSignature {
86        &self.signature
87    }
88}
89
90/// Libfunc used to call user functions.
91#[derive(Default)]
92pub struct CouponCallLibfunc {}
93impl NamedLibfunc for CouponCallLibfunc {
94    type Concrete = SignatureAndFunctionConcreteLibfunc;
95    const STR_ID: &'static str = "coupon_call";
96
97    fn specialize_signature(
98        &self,
99        context: &dyn SignatureSpecializationContext,
100        args: &[GenericArg],
101    ) -> Result<LibfuncSignature, SpecializationError> {
102        let function_id = args_as_single_user_func(args)?;
103
104        let signature = context.get_function_signature(function_id)?;
105        let ap_change = context.get_function_ap_change(function_id)?;
106
107        let coupon_ty = coupon_ty(context, function_id.clone())?;
108        Ok(LibfuncSignature::new_non_branch(
109            chain!(signature.param_types.iter().cloned(), [coupon_ty]).collect(),
110            get_output_var_infos(context, signature)?,
111            ap_change,
112        ))
113    }
114
115    fn specialize(
116        &self,
117        context: &dyn SpecializationContext,
118        args: &[GenericArg],
119    ) -> Result<Self::Concrete, SpecializationError> {
120        let function_id = args_as_single_user_func(args)?;
121
122        Ok(Self::Concrete {
123            function: context.get_function(function_id)?,
124            signature: self.specialize_signature(context, args)?,
125        })
126    }
127}
128
129/// Libfunc to replace function calls in size estimation.
130/// The libfunc consumes the function arguments and returns the function return values.
131/// The ap_change and signature are encoded in the generic arguments, see
132/// `try_extract_dummy_func_info` for details.
133#[derive(Default)]
134pub struct DummyFunctionCallLibfunc {}
135impl NamedLibfunc for DummyFunctionCallLibfunc {
136    type Concrete = SignatureAndFunctionConcreteLibfunc;
137
138    const STR_ID: &'static str = "dummy_function_call";
139
140    fn specialize_signature(
141        &self,
142        context: &dyn SignatureSpecializationContext,
143        args: &[GenericArg],
144    ) -> Result<LibfuncSignature, SpecializationError> {
145        let Some((signature, ap_change)) = try_extract_dummy_func_info(args.iter()) else {
146            return Err(SpecializationError::UnsupportedGenericArg);
147        };
148
149        Ok(LibfuncSignature::new_non_branch(
150            signature.param_types.clone(),
151            get_output_var_infos(context, signature).unwrap(),
152            ap_change,
153        ))
154    }
155
156    fn specialize(
157        &self,
158        context: &dyn SpecializationContext,
159        args: &[GenericArg],
160    ) -> Result<Self::Concrete, SpecializationError> {
161        let function_id = args_as_single_user_func(&args[..1])?;
162
163        Ok(Self::Concrete {
164            function: context.get_function(function_id)?,
165            signature: self.specialize_signature(context, args)?,
166        })
167    }
168}
169
170/// Given the generic arguments of a dummy function call, returns the function signature and
171/// ap_change.
172fn try_extract_dummy_func_info<'a>(
173    mut args: impl Iterator<Item = &'a GenericArg>,
174) -> Option<(FunctionSignature, SierraApChange)> {
175    let [_user_func_id, GenericArg::Value(ap_change), GenericArg::Value(n_params)] =
176        args.next_array()?
177    else {
178        return None;
179    };
180    let ap_change = if ap_change.is_zero() {
181        SierraApChange::Known { new_vars_only: false }
182    } else {
183        SierraApChange::Unknown
184    };
185    let extract_ty = |garg: &GenericArg| extract_matches!(garg, GenericArg::Type).clone();
186    let param_types = args.by_ref().take(n_params.try_into().ok()?).map(extract_ty).collect();
187
188    let GenericArg::Value(n_ret) = args.next()? else {
189        return None;
190    };
191    let ret_types = args.by_ref().take(n_ret.try_into().ok()?).map(extract_ty).collect();
192    Some((FunctionSignature { param_types, ret_types }, ap_change))
193}
cairo_lang_sierra/extensions/modules/function_call.rs

cairo_lang_sierra/extensions/modules/
function_call.rs
