cairo_lang_sierra_ap_change/
lib.rs

1//! Sierra AP change model.
2use ap_change_info::ApChangeInfo;
3use cairo_lang_sierra::extensions::core::{CoreLibfunc, CoreType};
4use cairo_lang_sierra::extensions::gas::CostTokenType;
5use cairo_lang_sierra::ids::{ConcreteTypeId, FunctionId};
6use cairo_lang_sierra::program::{Program, StatementIdx};
7use cairo_lang_sierra::program_registry::{ProgramRegistry, ProgramRegistryError};
8use cairo_lang_sierra_type_size::{TypeSizeMap, get_type_size_map};
9use cairo_lang_utils::casts::IntoOrPanic;
10use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
11use core_libfunc_ap_change::InvocationApChangeInfoProvider;
12use generate_equations::{Effects, Var};
13use itertools::Itertools;
14use thiserror::Error;
15
16pub mod ap_change_info;
17/// Direct linear computation of AP-Changes instead of equation solver.
18pub mod compute;
19pub mod core_libfunc_ap_change;
20mod generate_equations;
21
22/// Describes the effect on the `ap` register in a given libfunc branch.
23#[derive(Clone, Debug, Eq, PartialEq)]
24pub enum ApChange {
25    /// The libfunc changes `ap` in an unknown way.
26    Unknown,
27    /// The libfunc changes `ap` by a known size.
28    Known(usize),
29    /// The libfunc changes `ap` by a known size, provided in the metadata. Currently this only
30    /// includes `branch_align` libfunc.
31    FromMetadata,
32    /// The libfunc changes `ap` by a known size at locals finalization stage.
33    AtLocalsFinalization(usize),
34    /// The libfunc is a function call - it changes according to the given function and call cost.
35    FunctionCall(FunctionId),
36    /// The libfunc allocates locals, the `ap` change depends on the environment.
37    FinalizeLocals,
38    /// The libfunc is the ap tracking enabler.
39    EnableApTracking,
40    /// The libfunc is the ap tracking disabler.
41    DisableApTracking,
42}
43
44/// Error occurring while calculating the costing of a program's variables.
45#[derive(Error, Debug, Eq, PartialEq)]
46pub enum ApChangeError {
47    #[error("error from the program registry")]
48    ProgramRegistryError(#[from] Box<ProgramRegistryError>),
49    #[error("got a statement out of order during ap change calculations")]
50    StatementOutOfOrder(StatementIdx),
51    #[error("Wrong number of libfunc branches in ap-change information")]
52    WrongNumApChangeBranches(StatementIdx),
53    #[error("Attempted to merge branches with different number of allocated locals")]
54    BadMergeAllocatedLocalsMismatch(StatementIdx),
55    #[error("Attempted to merge branches with different bases to align")]
56    BadMergeBaseMismatch(StatementIdx),
57    #[error("failed solving the ap changes")]
58    SolvingApChangeEquationFailed,
59}
60
61/// Helper to implement the `InvocationApChangeInfoProvider` for the equation generation.
62struct InvocationApChangeInfoProviderForEqGen<'a, TokenUsages: Fn(CostTokenType) -> usize> {
63    /// Registry for providing the sizes of the types.
64    type_sizes: &'a TypeSizeMap,
65    /// Closure providing the token usages for the invocation.
66    token_usages: TokenUsages,
67}
68
69impl<TokenUsages: Fn(CostTokenType) -> usize> InvocationApChangeInfoProvider
70    for InvocationApChangeInfoProviderForEqGen<'_, TokenUsages>
71{
72    fn type_size(&self, ty: &ConcreteTypeId) -> usize {
73        self.type_sizes[ty].into_or_panic()
74    }
75
76    fn token_usages(&self, token_type: CostTokenType) -> usize {
77        (self.token_usages)(token_type)
78    }
79}
80
81/// Calculates gas information for a given program.
82pub fn calc_ap_changes<TokenUsages: Fn(StatementIdx, CostTokenType) -> usize>(
83    program: &Program,
84    token_usages: TokenUsages,
85) -> Result<ApChangeInfo, ApChangeError> {
86    let registry = ProgramRegistry::<CoreType, CoreLibfunc>::new(program)?;
87    let type_sizes = get_type_size_map(program, &registry).unwrap();
88    let equations = generate_equations::generate_equations(program, |idx, libfunc_id| {
89        let libfunc = registry.get_libfunc(libfunc_id)?;
90        core_libfunc_ap_change::core_libfunc_ap_change(
91            libfunc,
92            &InvocationApChangeInfoProviderForEqGen {
93                type_sizes: &type_sizes,
94                token_usages: |token_type| token_usages(idx, token_type),
95            },
96        )
97        .into_iter()
98        .map(|ap_change| {
99            Ok(match ap_change {
100                ApChange::AtLocalsFinalization(known) => {
101                    Effects { ap_change: ApChange::Known(0), locals: known }
102                }
103                _ => Effects { ap_change, locals: 0 },
104            })
105        })
106        .collect::<Result<Vec<_>, _>>()
107    })?;
108    let minimization_vars =
109        equations.iter().flat_map(|eq| eq.var_to_coef.keys()).unique().cloned().collect();
110    let solution = cairo_lang_eq_solver::try_solve_equations(equations, vec![minimization_vars])
111        .ok_or(ApChangeError::SolvingApChangeEquationFailed)?;
112
113    let mut variable_values = OrderedHashMap::default();
114    let mut function_ap_change = OrderedHashMap::default();
115    for (var, value) in solution {
116        match var {
117            Var::LibfuncImplicitApChangeVariable(idx) => {
118                variable_values.insert(idx, value as usize)
119            }
120            Var::FunctionApChange(func_id) => function_ap_change.insert(func_id, value as usize),
121        };
122    }
123    Ok(ApChangeInfo { variable_values, function_ap_change })
124}
cairo_lang_sierra_ap_change/lib.rs

cairo_lang_sierra_ap_change/
lib.rs
