openvm_circuit/arch/execution_mode/
metered_cost.rs

1use std::num::NonZero;
2
3use getset::WithSetters;
4use openvm_instructions::riscv::RV32_IMM_AS;
5
6use crate::{
7    arch::{
8        execution_mode::metered::segment_ctx::DEFAULT_MAX_CELLS as DEFAULT_SEGMENT_MAX_CELLS,
9        ExecutionCtxTrait, MeteredExecutionCtxTrait, SystemConfig, VmExecState,
10    },
11    system::memory::online::GuestMemory,
12};
13
14const DEFAULT_MAX_SEGMENTS: u64 = 100;
15pub const DEFAULT_MAX_COST: u64 = DEFAULT_MAX_SEGMENTS * DEFAULT_SEGMENT_MAX_CELLS as u64;
16
17#[derive(Clone, Debug)]
18pub struct AccessAdapterCtx {
19    min_block_size_bits: Vec<u8>,
20    idx_offset: usize,
21}
22
23impl AccessAdapterCtx {
24    pub fn new(config: &SystemConfig) -> Self {
25        Self {
26            min_block_size_bits: config.memory_config.min_block_size_bits(),
27            idx_offset: config.access_adapter_air_id_offset(),
28        }
29    }
30
31    #[inline(always)]
32    pub fn update_cells(
33        &self,
34        cost: &mut u64,
35        address_space: u32,
36        size_bits: u32,
37        widths: &[usize],
38    ) {
39        debug_assert!((address_space as usize) < self.min_block_size_bits.len());
40
41        // SAFETY: address_space passed is usually a hardcoded constant or derived from an
42        // Instruction where it is bounds checked before passing
43        let align_bits = unsafe {
44            *self
45                .min_block_size_bits
46                .get_unchecked(address_space as usize)
47        };
48        debug_assert!(
49            align_bits as u32 <= size_bits,
50            "align_bits ({}) must be <= size_bits ({})",
51            align_bits,
52            size_bits
53        );
54
55        for adapter_bits in (align_bits as u32 + 1..=size_bits).rev() {
56            let adapter_idx = self.idx_offset + adapter_bits as usize - 1;
57            debug_assert!(adapter_idx < widths.len());
58            // SAFETY: widths is initialized taking access adapters into account
59            let width = unsafe { *widths.get_unchecked(adapter_idx) };
60            let height_delta = 1 << (size_bits - adapter_bits + 1);
61            *cost += (height_delta as u64) * (width as u64);
62        }
63    }
64}
65
66#[derive(Clone, Debug, WithSetters)]
67pub struct MeteredCostCtx {
68    pub widths: Vec<usize>,
69    pub access_adapter_ctx: AccessAdapterCtx,
70    #[getset(set_with = "pub")]
71    pub max_execution_cost: u64,
72    // Cost is number of trace cells (height * width)
73    pub cost: u64,
74}
75
76impl MeteredCostCtx {
77    pub fn new(widths: Vec<usize>, config: &SystemConfig) -> Self {
78        let access_adapter_ctx = AccessAdapterCtx::new(config);
79        Self {
80            widths,
81            access_adapter_ctx,
82            max_execution_cost: DEFAULT_MAX_COST,
83            cost: 0,
84        }
85    }
86
87    #[cold]
88    fn panic_cost_exceeded(&self) -> ! {
89        panic!(
90            "Execution cost {} exceeded maximum allowed cost of {}",
91            self.cost,
92            2 * DEFAULT_MAX_COST
93        );
94    }
95}
96
97impl ExecutionCtxTrait for MeteredCostCtx {
98    #[inline(always)]
99    fn on_memory_operation(&mut self, address_space: u32, _ptr: u32, size: u32) {
100        debug_assert!(
101            address_space != RV32_IMM_AS,
102            "address space must not be immediate"
103        );
104        debug_assert!(size > 0, "size must be greater than 0, got {}", size);
105        debug_assert!(
106            size.is_power_of_two(),
107            "size must be a power of 2, got {}",
108            size
109        );
110        // Prevent unbounded memory accesses per instruction
111        if self.cost > 2 * std::cmp::max(self.max_execution_cost, DEFAULT_MAX_COST) {
112            self.panic_cost_exceeded();
113        }
114
115        // Handle access adapter updates
116        // SAFETY: size passed is always a non-zero power of 2
117        let size_bits = unsafe { NonZero::new_unchecked(size).ilog2() };
118        self.access_adapter_ctx.update_cells(
119            &mut self.cost,
120            address_space,
121            size_bits,
122            &self.widths,
123        );
124    }
125
126    fn should_suspend<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
127        vm_state.ctx.cost > vm_state.ctx.max_execution_cost
128    }
129}
130
131impl MeteredExecutionCtxTrait for MeteredCostCtx {
132    #[inline(always)]
133    fn on_height_change(&mut self, chip_idx: usize, height_delta: u32) {
134        debug_assert!(chip_idx < self.widths.len(), "chip_idx out of bounds");
135        // SAFETY: chip_idx is created in executor_idx_to_air_idx and is always within bounds
136        let width = unsafe { *self.widths.get_unchecked(chip_idx) };
137        self.cost += (height_delta as u64) * (width as u64);
138    }
139}