revm_context_interface/
local.rs1use core::{
3 cell::{Ref, RefCell},
4 ops::Range,
5};
6use std::{rc::Rc, vec::Vec};
7
8#[derive(Debug, Clone)]
10pub struct FrameStack<T> {
11 stack: Vec<T>,
12 index: Option<usize>,
13}
14
15impl<T> Default for FrameStack<T> {
16 fn default() -> Self {
17 Self::new()
18 }
19}
20
21impl<T> FrameStack<T> {
22 pub fn new() -> Self {
24 Self {
25 stack: Vec::with_capacity(4),
26 index: None,
27 }
28 }
29
30 #[inline]
32 pub fn start_init(&mut self) -> OutFrame<'_, T> {
33 self.index = None;
34 if self.stack.is_empty() {
35 self.stack.reserve(1);
36 }
37 self.out_frame_at(0)
38 }
39
40 #[inline]
42 pub fn end_init(&mut self, token: FrameToken) {
43 token.assert();
44 if self.stack.is_empty() {
45 unsafe { self.stack.set_len(1) };
46 }
47 self.index = Some(0);
48 }
49
50 #[inline]
52 pub fn index(&self) -> Option<usize> {
53 self.index
54 }
55
56 #[inline]
58 pub fn push(&mut self, token: FrameToken) {
59 token.assert();
60 let index = self.index.as_mut().unwrap();
61 if *index + 1 == self.stack.len() {
62 unsafe { self.stack.set_len(self.stack.len() + 1) };
63 self.stack.reserve(1);
64 }
65 *index += 1;
66 }
67
68 #[inline]
71 pub fn clear(&mut self) {
72 self.index = None;
73 }
74
75 #[inline]
77 pub fn pop(&mut self) {
78 self.index = self.index.unwrap_or(0).checked_sub(1);
79 }
80
81 #[inline]
83 pub fn get(&mut self) -> &mut T {
84 debug_assert!(self.stack.capacity() > self.index.unwrap() + 1);
85 unsafe { &mut *self.stack.as_mut_ptr().add(self.index.unwrap()) }
86 }
87
88 #[inline]
90 pub fn get_next(&mut self) -> OutFrame<'_, T> {
91 self.out_frame_at(self.index.unwrap() + 1)
92 }
93
94 fn out_frame_at(&mut self, idx: usize) -> OutFrame<'_, T> {
95 unsafe {
96 OutFrame::new_maybe_uninit(self.stack.as_mut_ptr().add(idx), idx < self.stack.len())
97 }
98 }
99}
100
101#[allow(missing_debug_implementations)]
103pub struct OutFrame<'a, T> {
104 ptr: *mut T,
105 init: bool,
106 lt: core::marker::PhantomData<&'a mut T>,
107}
108
109impl<'a, T> OutFrame<'a, T> {
110 pub fn new_init(slot: &'a mut T) -> Self {
112 unsafe { Self::new_maybe_uninit(slot, true) }
113 }
114
115 pub fn new_uninit(slot: &'a mut core::mem::MaybeUninit<T>) -> Self {
117 unsafe { Self::new_maybe_uninit(slot.as_mut_ptr(), false) }
118 }
119
120 pub unsafe fn new_maybe_uninit(ptr: *mut T, init: bool) -> Self {
128 Self {
129 ptr,
130 init,
131 lt: Default::default(),
132 }
133 }
134
135 pub fn get(&mut self, f: impl FnOnce() -> T) -> &mut T {
137 if !self.init {
138 self.do_init(f);
139 }
140 unsafe { &mut *self.ptr }
141 }
142
143 #[inline(never)]
144 #[cold]
145 fn do_init(&mut self, f: impl FnOnce() -> T) {
146 unsafe {
147 self.init = true;
148 self.ptr.write(f());
149 }
150 }
151
152 pub unsafe fn get_unchecked(&mut self) -> &mut T {
158 debug_assert!(self.init, "OutFrame must be initialized before use");
159 unsafe { &mut *self.ptr }
160 }
161
162 pub fn consume(self) -> FrameToken {
164 FrameToken(self.init)
165 }
166}
167
168#[allow(missing_debug_implementations)]
170pub struct FrameToken(bool);
171
172impl FrameToken {
173 #[cfg_attr(debug_assertions, track_caller)]
175 pub fn assert(self) {
176 assert!(self.0, "FrameToken must be initialized before use");
177 }
178}
179
180pub trait LocalContextTr {
182 fn shared_memory_buffer(&self) -> &Rc<RefCell<Vec<u8>>>;
184
185 fn shared_memory_buffer_slice(&self, range: Range<usize>) -> Option<Ref<'_, [u8]>> {
187 let buffer = self.shared_memory_buffer();
188 buffer.borrow().get(range.clone())?;
189 Some(Ref::map(buffer.borrow(), |b| {
190 b.get(range).unwrap_or_default()
191 }))
192 }
193
194 fn clear(&mut self);
196}
197
198#[cfg(test)]
199mod tests {
200 use super::*;
201
202 #[test]
203 fn frame_stack() {
204 let mut stack = FrameStack::new();
205 let mut frame = stack.start_init();
206 frame.get(|| 1);
207 let token = frame.consume();
208 stack.end_init(token);
209
210 assert_eq!(stack.index(), Some(0));
211 assert_eq!(stack.stack.len(), 1);
212
213 let a = stack.get();
214 assert_eq!(a, &mut 1);
215 let mut b = stack.get_next();
216 assert!(!b.init);
217 assert_eq!(b.get(|| 2), &mut 2);
218 let token = b.consume(); stack.push(token);
220
221 assert_eq!(stack.index(), Some(1));
222 assert_eq!(stack.stack.len(), 2);
223 let a = stack.get();
224 assert_eq!(a, &mut 2);
225 let b = stack.get_next();
226 assert!(!b.init);
227
228 stack.pop();
229
230 assert_eq!(stack.index(), Some(0));
231 assert_eq!(stack.stack.len(), 2);
232 let a = stack.get();
233 assert_eq!(a, &mut 1);
234 let mut b = stack.get_next();
235 assert!(b.init);
236 assert_eq!(unsafe { b.get_unchecked() }, &mut 2);
237 }
238}