1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

//! Sierra AP change model.
use std::collections::HashMap;

use ap_change_info::ApChangeInfo;
use cairo_lang_sierra::extensions::core::{CoreLibfunc, CoreType};
use cairo_lang_sierra::extensions::ConcreteType;
use cairo_lang_sierra::ids::{ConcreteTypeId, FunctionId};
use cairo_lang_sierra::program::{Program, StatementIdx};
use cairo_lang_sierra::program_registry::{ProgramRegistry, ProgramRegistryError};
use generate_equations::{Effects, Var};
use thiserror::Error;

pub mod ap_change_info;
pub mod core_libfunc_ap_change;
mod generate_equations;

/// Describes the effect on the `ap` register in a given libfunc branch.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ApChange {
    /// The libfunc changes `ap` in an unknown way.
    Unknown,
    /// The libfunc changes `ap` by a known size.
    Known(usize),
    /// The libfunc changes `ap` by a known size, provided in the metadata. Currently this only
    /// includes `branch_align` libfunc.
    FromMetadata,
    /// The libfunc changes `ap` by a known size, which is the size of the given type at locals
    /// finalization stage.
    AtLocalsFinalizationByTypeSize(ConcreteTypeId),
    /// The libfunc changes `ap` by a known size, which is the size of the given type.
    KnownByTypeSize(ConcreteTypeId),
    /// The libfunc is a function call - it changes according to the given function and call cost.
    FunctionCall(FunctionId),
    // The libfunc allocates locals, the `ap` change depends on the environment.
    FinalizeLocals,
}

/// Error occurring while calculating the costing of a program's variables.
#[derive(Error, Debug, Eq, PartialEq)]
pub enum ApChangeError {
    #[error("error from the program registry")]
    ProgramRegistryError(#[from] Box<ProgramRegistryError>),
    #[error("found an illegal statement index during ap change calculations")]
    StatementOutOfBounds(StatementIdx),
    #[error("found an illegal statement index during ap change calculations")]
    StatementOutOfOrder(StatementIdx),
    #[error("found an illegal invocation during cost calculations")]
    IllegalInvocation(StatementIdx),
    #[error("failed solving the ap changes")]
    SolvingApChangeEquationFailed,
}

/// Calculates gas information for a given program.
pub fn calc_ap_changes(program: &Program) -> Result<ApChangeInfo, ApChangeError> {
    let registry = ProgramRegistry::<CoreType, CoreLibfunc>::new(program)?;
    let equations = generate_equations::generate_equations(program, |libfunc_id| {
        let libfunc = registry.get_libfunc(libfunc_id)?;
        core_libfunc_ap_change::core_libfunc_ap_change(libfunc)
            .into_iter()
            .map(|ap_change| {
                Ok(match ap_change {
                    ApChange::KnownByTypeSize(ty) => Effects {
                        ap_change: ApChange::Known(registry.get_type(&ty)?.info().size as usize),
                        locals: 0,
                    },
                    ApChange::AtLocalsFinalizationByTypeSize(ty) => Effects {
                        ap_change: ApChange::Known(0),
                        locals: registry.get_type(&ty)?.info().size as usize,
                    },
                    _ => Effects { ap_change, locals: 0 },
                })
            })
            .collect::<Result<Vec<_>, _>>()
    })?;
    let solution = cairo_lang_eq_solver::try_solve_equations(equations)
        .ok_or(ApChangeError::SolvingApChangeEquationFailed)?;

    let mut variable_values = HashMap::<StatementIdx, usize>::default();
    let mut function_ap_change = HashMap::<cairo_lang_sierra::ids::FunctionId, usize>::default();
    for (var, value) in solution {
        match var {
            Var::LibfuncImplicitApChangeVariable(idx) => {
                variable_values.insert(idx, value as usize)
            }
            Var::FunctionApChange(func_id) => function_ap_change.insert(func_id, value as usize),
        };
    }
    Ok(ApChangeInfo { variable_values, function_ap_change })
}