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
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use cairo_lang_sierra::extensions::gas::CostTokenType;
use cairo_lang_sierra::ids::ConcreteTypeId;
use cairo_lang_sierra::program::Function;
use cairo_lang_utils::collection_arithmetics::{add_maps, sub_maps};
use cairo_lang_utils::ordered_hash_map::OrderedHashMap;

/// Represents constant cost.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct ConstCost {
    pub steps: i32,
    pub holes: i32,
    pub range_checks: i32,
}
impl ConstCost {
    pub const fn cost(&self) -> i32 {
        self.steps * 100 + self.holes * 10 + self.range_checks * 70
    }
    pub const fn steps(value: i32) -> Self {
        Self { steps: value, holes: 0, range_checks: 0 }
    }
    pub const fn holes(value: i32) -> Self {
        Self { holes: value, steps: 0, range_checks: 0 }
    }
    pub const fn range_checks(value: i32) -> Self {
        Self { range_checks: value, steps: 0, holes: 0 }
    }
}

/// Adds two [ConstCost] instances.
impl ConstCost {
    // Note: this is necessary because `impl Add` does not support `const fn`.
    pub const fn add(self, rhs: Self) -> Self {
        Self {
            steps: self.steps + rhs.steps,
            holes: self.holes + rhs.holes,
            range_checks: self.range_checks + rhs.range_checks,
        }
    }
}

/// Adds two [ConstCost] instances.
impl std::ops::Add for ConstCost {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        self.add(rhs)
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct PreCost(pub OrderedHashMap<CostTokenType, i32>);
impl PreCost {
    pub fn builtin(token_type: CostTokenType) -> Self {
        Self(OrderedHashMap::from_iter(([(token_type, 1)]).into_iter()))
    }
}

/// Adds two [ConstCost] instances.
impl std::ops::Add for PreCost {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        PreCost(add_maps(self.0, rhs.0))
    }
}

/// Subtracts two [ConstCost] instances.
impl std::ops::Sub for PreCost {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        PreCost(sub_maps(self.0, rhs.0))
    }
}

/// The cost of executing a libfunc for a specific output branch.
#[derive(Clone)]
pub enum BranchCost {
    /// The cost of the statement is independent on other statements.
    Regular { const_cost: ConstCost, pre_cost: PreCost },
    /// A cost of a function call.
    FunctionCall { const_cost: ConstCost, function: Function },
    /// The cost of the `branch_align` libfunc.
    BranchAlign,
    /// The cost of `withdraw_gas` and `withdraw_gas_all` libfuncs.
    WithdrawGas { const_cost: ConstCost, success: bool, with_builtin_costs: bool },
    /// The cost of the `redeposit_gas` libfunc.
    RedepositGas,
}

/// Converts [ConstCost] into [BranchCost].
impl From<ConstCost> for BranchCost {
    fn from(value: ConstCost) -> Self {
        BranchCost::Regular { const_cost: value, pre_cost: PreCost::default() }
    }
}

/// Trait for providing extra information required for calculating costs for a specific libfunc
/// invocation.
pub trait CostInfoProvider {
    /// Provides the sizes of types.
    fn type_size(&self, ty: &ConcreteTypeId) -> usize;
}