cairo_lang_sierra_to_casm/environment/
frame_state.rs

1use thiserror::Error;
2
3use super::{ApTracking, ApTrackingBase};
4
5#[derive(Error, Debug, Eq, PartialEq)]
6pub enum FrameStateError {
7    #[error("InvalidTransition")]
8    InvalidTransition,
9    #[error("alloc_local is not allowed at this point.")]
10    InvalidAllocLocal(FrameState),
11    #[error("finalize_locals is not allowed at this point.")]
12    InvalidFinalizeLocals(FrameState),
13    #[error("locals were allocated but finalize_locals was not called.")]
14    FinalizeLocalsMissing(FrameState),
15}
16
17/// The frame state of the current function.
18/// This state keeps track of how many locals have been allocated and whether the
19/// frame has been finalized.
20#[derive(Clone, Debug, Eq, PartialEq)]
21pub enum FrameState {
22    /// Before any locals were allocated (`alloc_local` wasn't called).
23    BeforeAllocation,
24    /// `finalize_locals` wasn't called yet.
25    /// `allocated` is the number of stack slot that were already allocated for local variables.
26    /// `locals_start_ap_offset` is the ap change between the first `alloc_local` and the beginning
27    /// of the function. It is used to validate that there were no ap changes between the
28    /// allocations and the call to `handle_finalize_locals`.
29    Allocating { allocated: usize, locals_start_ap_offset: usize },
30    /// finalize_locals was called and the frame has been finalized.
31    Finalized { allocated: usize },
32}
33
34/// Returns the number of slots that were allocated for locals and the new frame state.
35pub fn handle_finalize_locals(
36    frame_state: FrameState,
37    ap_tracking: ApTracking,
38) -> Result<(usize, FrameState), FrameStateError> {
39    match frame_state {
40        FrameState::BeforeAllocation => {
41            // TODO(ilya, 10/10/2022): Do we want to support allocating 0 locals?
42            if matches!(
43                ap_tracking,
44                ApTracking::Enabled { ap_change: _, base: ApTrackingBase::FunctionStart }
45            ) {
46                Ok((0, FrameState::Finalized { allocated: 0 }))
47            } else {
48                Err(FrameStateError::InvalidFinalizeLocals(frame_state))
49            }
50        }
51        FrameState::Allocating { allocated, locals_start_ap_offset } => {
52            if matches!(
53                ap_tracking,
54                ApTracking::Enabled { ap_change, base: ApTrackingBase::FunctionStart }
55                if ap_change == locals_start_ap_offset
56            ) {
57                Ok((allocated, FrameState::Finalized { allocated }))
58            } else {
59                Err(FrameStateError::InvalidFinalizeLocals(frame_state))
60            }
61        }
62        FrameState::Finalized { .. } => Err(FrameStateError::InvalidFinalizeLocals(frame_state)),
63    }
64}
65
66/// Returns the offset of the newly allocated variable and the new frame state.
67pub fn handle_alloc_local(
68    frame_state: FrameState,
69    ap_tracking: ApTracking,
70    allocation_size: usize,
71) -> Result<(usize, FrameState), FrameStateError> {
72    match frame_state {
73        FrameState::BeforeAllocation => {
74            if let ApTracking::Enabled { ap_change, base: ApTrackingBase::FunctionStart } =
75                ap_tracking
76            {
77                Ok((
78                    ap_change,
79                    FrameState::Allocating {
80                        allocated: allocation_size,
81                        locals_start_ap_offset: ap_change,
82                    },
83                ))
84            } else {
85                Err(FrameStateError::InvalidAllocLocal(frame_state))
86            }
87        }
88        FrameState::Allocating { allocated, locals_start_ap_offset } => {
89            if matches!(
90                ap_tracking,
91                ApTracking::Enabled { ap_change, base: ApTrackingBase::FunctionStart }
92                if ap_change == locals_start_ap_offset
93            ) {
94                Ok((
95                    locals_start_ap_offset + allocated,
96                    FrameState::Allocating {
97                        allocated: allocated + allocation_size,
98                        locals_start_ap_offset,
99                    },
100                ))
101            } else {
102                Err(FrameStateError::InvalidAllocLocal(frame_state))
103            }
104        }
105        FrameState::Finalized { .. } => Err(FrameStateError::InvalidAllocLocal(frame_state)),
106    }
107}
108
109/// Validates that the state at the end of a function is valid.
110pub fn validate_final_frame_state(frame_state: &FrameState) -> Result<(), FrameStateError> {
111    if matches!(frame_state, FrameState::Allocating { .. }) {
112        Err(FrameStateError::FinalizeLocalsMissing(frame_state.clone()))
113    } else {
114        Ok(())
115    }
116}