cel_runtime/
raw_segment.rs

1use crate::raw_sequence::RawSequence;
2use crate::raw_stack::RawStack;
3use anyhow::Result;
4use std::cmp::max;
5
6type Operation = fn(&RawSequence, usize, &mut RawStack) -> Result<usize>;
7
8/// A segment represents a sequence of operations that can be executed.
9///
10/// Each operation is stored along with its data in the segment's storage,
11/// and can manipulate values on a stack during execution.
12pub struct RawSegment {
13    ops: Vec<Operation>,
14    storage: RawSequence,
15    dropper: Vec<fn(&mut RawSequence, usize) -> usize>,
16    base_alignment: usize,
17    frame: Option<(usize, usize)>,
18}
19
20impl Default for RawSegment {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl RawSegment {
27    /// Creates a new empty segment.
28    #[must_use]
29    pub fn new() -> Self {
30        RawSegment {
31            ops: Vec::new(),
32            storage: RawSequence::new(),
33            dropper: Vec::new(),
34            base_alignment: 0,
35            frame: None,
36        }
37    }
38
39    /// Returns the maximum alignment required by any value pushed onto the stack while executing this segment.
40    pub(crate) fn base_alignment(&self) -> usize {
41        self.base_alignment
42    }
43
44    /// Updates the stack base alignment to be at least `alignment`.
45    pub(crate) fn update_base_alignment(&mut self, alignment: usize) {
46        debug_assert!(alignment.is_power_of_two());
47        self.base_alignment = max(self.base_alignment, alignment);
48    }
49
50    /// Sets the caller-owned byte prefix and number of owned input values.
51    pub(crate) fn set_frame(&mut self, origin: usize, owned_inputs: usize) {
52        self.frame = Some((origin, owned_inputs));
53    }
54
55    /// Binds the initial frame unless an enclosing segment already supplied it.
56    pub(crate) fn bind_frame(&mut self, origin: usize, owned_inputs: usize) {
57        self.frame.get_or_insert((origin, owned_inputs));
58    }
59
60    /// Returns the byte length occupied by instruction closure state.
61    #[cfg(test)]
62    pub(crate) fn storage_size(&self) -> usize {
63        self.storage.len()
64    }
65
66    /// Pushes a value into the segment's storage and registers its dropper.
67    fn push_storage<T>(&mut self, value: T)
68    where
69        T: 'static,
70    {
71        self.storage.push(value);
72        self.dropper
73            .push(|storage, p| unsafe { storage.drop_in_place::<T>(p) });
74    }
75
76    /// Pushes a closure that manipulates the stack directly and returns `()`.
77    pub fn raw0_<F>(&mut self, op: F)
78    where
79        F: Fn(&mut RawStack) -> Result<()> + 'static,
80    {
81        self.push_storage(op);
82        self.ops.push(|storage, p, stack| {
83            let (f, r) = unsafe { storage.next::<F>(p) };
84            f(stack)?;
85            Ok(r)
86        });
87    }
88
89    /// Push a closure that manipulates the stack and returns a value `R`.
90    pub fn raw0<R, F>(&mut self, op: F)
91    where
92        F: Fn(&mut RawStack) -> Result<R> + 'static,
93        R: 'static,
94    {
95        self.push_storage(op);
96        self.ops.push(|storage, p, stack| {
97            let (f, r) = unsafe { storage.next::<F>(p) };
98            let result = f(stack)?;
99            stack.push(result);
100            Ok(r)
101        });
102        self.base_alignment = max(self.base_alignment, align_of::<R>());
103    }
104
105    /// Pushes a unary operation that takes one argument of type T and returns a value of type R.
106    pub fn push_op0<R, F>(&mut self, op: F)
107    where
108        F: Fn() -> R + 'static,
109        R: 'static,
110    {
111        self.push_storage(op);
112        self.ops.push(|storage, p, stack| {
113            let (f, r) = unsafe { storage.next::<F>(p) };
114            stack.push(f());
115            Ok(r)
116        });
117        self.base_alignment = max(self.base_alignment, align_of::<R>());
118    }
119
120    /// Pushes the op-dispatch closure for a unary infallible operation with compile-time padding.
121    fn push_op1_<const PADDING0: bool, T, R, F>(&mut self)
122    where
123        F: Fn(T) -> R + 'static,
124        T: 'static,
125        R: 'static,
126    {
127        self.ops.push(|storage, p, stack| {
128            let (f, r) = unsafe { storage.next::<F>(p) };
129            let x: T = unsafe { stack.pop(PADDING0) };
130            stack.push(f(x));
131            Ok(r)
132        });
133    }
134
135    /// Pushes the op-dispatch closure for a unary fallible operation with compile-time padding.
136    fn push_op1r_<const PADDING0: bool, T, R, F>(&mut self)
137    where
138        F: Fn(&mut RawStack, T) -> Result<R> + 'static,
139        T: 'static,
140        R: 'static,
141    {
142        self.ops.push(|storage, p, stack| {
143            let (f, r) = unsafe { storage.next::<F>(p) };
144            let x: T = unsafe { stack.pop(PADDING0) };
145            let result = f(stack, x)?;
146            stack.push(result);
147            Ok(r)
148        });
149    }
150
151    /// Pushes a unary operation that takes one argument of type T and returns a value of type R.
152    pub fn push_op1<T, R, F>(&mut self, op: F, padding0: bool)
153    where
154        F: Fn(T) -> R + 'static,
155        T: 'static,
156        R: 'static,
157    {
158        self.push_storage(op);
159        if padding0 {
160            self.push_op1_::<true, T, R, F>();
161        } else {
162            self.push_op1_::<false, T, R, F>();
163        }
164        self.base_alignment = max(self.base_alignment, align_of::<R>());
165    }
166
167    /// Push a fallible unary operation that can manipulate the stack.
168    pub fn raw1<T, R, F>(&mut self, op: F, padding0: bool)
169    where
170        F: Fn(&mut RawStack, T) -> Result<R> + 'static,
171        T: 'static,
172        R: 'static,
173    {
174        self.push_storage(op);
175        if padding0 {
176            self.push_op1r_::<true, T, R, F>();
177        } else {
178            self.push_op1r_::<false, T, R, F>();
179        }
180        self.base_alignment = max(self.base_alignment, align_of::<R>());
181    }
182
183    /// Pushes the op-dispatch closure for a unary drop operation with compile-time padding.
184    fn drop1_<const PADDING0: bool, T, F>(&mut self)
185    where
186        F: Fn(T) + 'static,
187        T: 'static,
188    {
189        self.ops.push(|storage, p, stack| {
190            let (f, r) = unsafe { storage.next::<F>(p) };
191            let x: T = unsafe { stack.pop(PADDING0) };
192            f(x); // drop the result
193            Ok(r)
194        });
195    }
196
197    /// Push an operation that consumes and drops one value produced by `op`.
198    pub fn drop1<T, F>(&mut self, op: F, padding0: bool)
199    where
200        F: Fn(T) + 'static,
201        T: 'static,
202    {
203        self.push_storage(op);
204        if padding0 {
205            self.drop1_::<true, T, F>();
206        } else {
207            self.drop1_::<false, T, F>();
208        }
209    }
210
211    /// Pushes the op-dispatch closure for a binary infallible operation with compile-time padding.
212    fn push_op2_<const PADDING0: bool, const PADDING1: bool, T, U, R, F>(&mut self)
213    where
214        F: Fn(T, U) -> R + 'static,
215        T: 'static,
216        U: 'static,
217        R: 'static,
218    {
219        self.ops.push(|storage, p, stack| {
220            let (f, r) = unsafe { storage.next::<F>(p) };
221            let y: U = unsafe { stack.pop(PADDING1) };
222            let x: T = unsafe { stack.pop(PADDING0) };
223            stack.push(f(x, y));
224            Ok(r)
225        });
226    }
227
228    /// Pushes a binary operation that takes two arguments of types T and U and returns a value of
229    /// type R.
230    pub fn push_op2<T, U, R, F>(&mut self, op: F, padding0: bool, padding1: bool)
231    where
232        F: Fn(T, U) -> R + 'static,
233        T: 'static,
234        U: 'static,
235        R: 'static,
236    {
237        self.push_storage(op);
238        match (padding0, padding1) {
239            (false, false) => self.push_op2_::<false, false, T, U, R, F>(),
240            (false, true) => self.push_op2_::<false, true, T, U, R, F>(),
241            (true, false) => self.push_op2_::<true, false, T, U, R, F>(),
242            (true, true) => self.push_op2_::<true, true, T, U, R, F>(),
243        }
244        self.base_alignment = max(self.base_alignment, align_of::<R>());
245    }
246
247    /// Pushes the op-dispatch closure for a binary fallible operation with compile-time padding.
248    fn push_op2r_<const PADDING0: bool, const PADDING1: bool, T, U, R, F>(&mut self)
249    where
250        F: Fn(&mut RawStack, T, U) -> Result<R> + 'static,
251        T: 'static,
252        U: 'static,
253        R: 'static,
254    {
255        self.ops.push(|storage, p, stack| {
256            let (f, r) = unsafe { storage.next::<F>(p) };
257            let y: U = unsafe { stack.pop(PADDING1) };
258            let x: T = unsafe { stack.pop(PADDING0) };
259            let result = f(stack, x, y)?;
260            stack.push(result);
261            Ok(r)
262        });
263    }
264
265    /// Push a fallible binary operation that can manipulate the stack.
266    pub fn raw2<T, U, R, F>(&mut self, op: F, padding0: bool, padding1: bool)
267    where
268        F: Fn(&mut RawStack, T, U) -> Result<R> + 'static,
269        T: 'static,
270        U: 'static,
271        R: 'static,
272    {
273        self.push_storage(op);
274        match (padding0, padding1) {
275            (false, false) => self.push_op2r_::<false, false, T, U, R, F>(),
276            (false, true) => self.push_op2r_::<false, true, T, U, R, F>(),
277            (true, false) => self.push_op2r_::<true, false, T, U, R, F>(),
278            (true, true) => self.push_op2r_::<true, true, T, U, R, F>(),
279        }
280        self.base_alignment = max(self.base_alignment, align_of::<R>());
281    }
282
283    /// Pushes a ternary operation that takes three arguments of types T, U, and V and returns a
284    /// value of type R.
285    #[expect(clippy::many_single_char_names, reason = "patterned code")]
286    fn push_op3_<const PADDING0: bool, const PADDING1: bool, const PADDING2: bool, T, U, V, R, F>(
287        &mut self,
288    ) where
289        F: Fn(T, U, V) -> R + 'static,
290        T: 'static,
291        U: 'static,
292        V: 'static,
293        R: 'static,
294    {
295        self.ops.push(|storage, p, stack| {
296            let (f, r) = unsafe { storage.next::<F>(p) };
297            let z: V = unsafe { stack.pop(PADDING2) };
298            let y: U = unsafe { stack.pop(PADDING1) };
299            let x: T = unsafe { stack.pop(PADDING0) };
300            stack.push(f(x, y, z));
301            Ok(r)
302        });
303    }
304
305    /// Pushes a ternary operation that takes three arguments of types T, U, and V and returns a value of type R.
306    pub fn push_op3<T, U, V, R, F>(&mut self, op: F, padding0: bool, padding1: bool, padding2: bool)
307    where
308        F: Fn(T, U, V) -> R + 'static,
309        T: 'static,
310        U: 'static,
311        V: 'static,
312        R: 'static,
313    {
314        self.push_storage(op);
315
316        match (padding0, padding1, padding2) {
317            (false, false, false) => self.push_op3_::<false, false, false, T, U, V, R, F>(),
318            (false, false, true) => self.push_op3_::<false, false, true, T, U, V, R, F>(),
319            (false, true, false) => self.push_op3_::<false, true, false, T, U, V, R, F>(),
320            (false, true, true) => self.push_op3_::<false, true, true, T, U, V, R, F>(),
321            (true, false, false) => self.push_op3_::<true, false, false, T, U, V, R, F>(),
322            (true, false, true) => self.push_op3_::<true, false, true, T, U, V, R, F>(),
323            (true, true, false) => self.push_op3_::<true, true, false, T, U, V, R, F>(),
324            (true, true, true) => self.push_op3_::<true, true, true, T, U, V, R, F>(),
325        }
326        self.base_alignment = max(self.base_alignment, align_of::<R>());
327    }
328
329    /// Pushes the op-dispatch closure for a ternary fallible operation with compile-time padding.
330    #[expect(clippy::many_single_char_names, reason = "patterned code")]
331    fn push_op3r_<const PADDING0: bool, const PADDING1: bool, const PADDING2: bool, T, U, V, R, F>(
332        &mut self,
333    ) where
334        F: Fn(&mut RawStack, T, U, V) -> Result<R> + 'static,
335        T: 'static,
336        U: 'static,
337        V: 'static,
338        R: 'static,
339    {
340        self.ops.push(|storage, p, stack| {
341            let (f, r) = unsafe { storage.next::<F>(p) };
342            let z: V = unsafe { stack.pop(PADDING2) };
343            let y: U = unsafe { stack.pop(PADDING1) };
344            let x: T = unsafe { stack.pop(PADDING0) };
345            let result = f(stack, x, y, z)?;
346            stack.push(result);
347            Ok(r)
348        });
349    }
350
351    /// Push a fallible ternary operation that can manipulate the stack.
352    pub fn raw3<T, U, V, R, F>(&mut self, op: F, padding0: bool, padding1: bool, padding2: bool)
353    where
354        F: Fn(&mut RawStack, T, U, V) -> Result<R> + 'static,
355        T: 'static,
356        U: 'static,
357        V: 'static,
358        R: 'static,
359    {
360        self.push_storage(op);
361        match (padding0, padding1, padding2) {
362            (false, false, false) => self.push_op3r_::<false, false, false, T, U, V, R, F>(),
363            (false, false, true) => self.push_op3r_::<false, false, true, T, U, V, R, F>(),
364            (false, true, false) => self.push_op3r_::<false, true, false, T, U, V, R, F>(),
365            (false, true, true) => self.push_op3r_::<false, true, true, T, U, V, R, F>(),
366            (true, false, false) => self.push_op3r_::<true, false, false, T, U, V, R, F>(),
367            (true, false, true) => self.push_op3r_::<true, false, true, T, U, V, R, F>(),
368            (true, true, false) => self.push_op3r_::<true, true, false, T, U, V, R, F>(),
369            (true, true, true) => self.push_op3r_::<true, true, true, T, U, V, R, F>(),
370        }
371        self.base_alignment = max(self.base_alignment, align_of::<R>());
372    }
373
374    /// Executes all operations in the segment on the supplied stack.
375    ///
376    /// # Errors
377    /// Halts execution and returns an error if any operation returns an error.
378    ///
379    /// # Safety
380    ///
381    /// Precondition: The stack must be in the correct state (base alignment, index, and values)
382    /// for the operations in the segment.
383    pub unsafe fn call0_stack(&self, stack: &mut RawStack) -> Result<()> {
384        let (origin, depth) = self.frame.map_or((0, 0), |(origin, inputs)| {
385            (
386                origin,
387                stack
388                    .depth()
389                    .checked_sub(inputs)
390                    .expect("valid input frame"),
391            )
392        });
393        stack.with_unwind(origin, depth, |stack| {
394            let mut p = 0;
395            for op in &self.ops {
396                p = op(&self.storage, p, stack)?;
397            }
398            Ok(())
399        })
400    }
401
402    /// Executes all operations in the segment and returns the final result.
403    ///
404    /// # Errors
405    /// Halts execution and returns an error if any operation returns an error.
406    ///
407    /// # Safety
408    /// This function is unsafe if the result type does not match the type returned by the
409    /// operations in the segment or if the operations expect any initial values on the stack.
410    pub unsafe fn call0<T>(&self) -> Result<T>
411    where
412        T: 'static,
413    {
414        let mut stack = RawStack::with_base_alignment(self.base_alignment);
415        unsafe {
416            self.call0_stack(&mut stack)?;
417        }
418        Ok(unsafe { stack.pop(false) })
419    }
420
421    /// Executes all operations in the segment with one argument of type A and returns the final
422    /// result.
423    ///
424    /// # Errors
425    /// Halts execution and returns an error if any operation returns an error.
426    ///
427    /// # Safety
428    /// This function is unsafe if the argument and result types do not match the types expected or
429    /// returned by the operations in the segment.
430    pub unsafe fn call1<A: 'static, T>(&self, arg: A) -> Result<T>
431    where
432        T: 'static,
433    {
434        let mut stack = RawStack::with_base_alignment(self.base_alignment);
435        stack.push(arg);
436        unsafe { self.call0_stack(&mut stack)? };
437        Ok(unsafe { stack.pop(false) })
438    }
439
440    /// Executes all operations in the segment with two arguments of types A and B and returns the
441    /// final result.
442    ///
443    /// # Errors
444    /// Halts execution and returns an error if any operation returns an error.
445    ///
446    /// # Safety
447    /// This function is unsafe if the arguments and result types do not match the types expected or
448    /// returned by the operations in the segment.
449    pub unsafe fn call2<A: 'static, B: 'static, T>(&self, arg: (A, B)) -> Result<T>
450    where
451        T: 'static,
452    {
453        let mut stack = RawStack::with_base_alignment(self.base_alignment);
454        stack.push(arg.0);
455        stack.push(arg.1);
456        unsafe { self.call0_stack(&mut stack)? };
457        Ok(unsafe { stack.pop(false) })
458    }
459}
460
461impl Drop for RawSegment {
462    fn drop(&mut self) {
463        let mut p = 0;
464        for e in &self.dropper {
465            p = e(&mut self.storage, p);
466        }
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    #[test]
475    fn nullary_operation() {
476        let mut segment = RawSegment::new();
477        segment.push_op0(|| 42);
478        unsafe {
479            assert_eq!(segment.call0::<i32>().unwrap(), 42);
480        }
481    }
482
483    #[test]
484    fn unary_operation() {
485        let mut segment = RawSegment::new();
486        segment.push_op0(|| 42);
487        segment.push_op1(|x: i32| x * 2, false);
488        unsafe {
489            assert_eq!(segment.call0::<i32>().unwrap(), 84);
490        }
491    }
492
493    #[test]
494    fn binary_operation() {
495        let mut segment = RawSegment::new();
496        segment.push_op0(|| 10);
497        segment.push_op0(|| 5);
498        segment.push_op2(|x: i32, y: i32| x + y, false, false);
499        unsafe {
500            assert_eq!(segment.call0::<i32>().unwrap(), 15);
501        }
502    }
503
504    #[test]
505    fn ternary_operation() {
506        let mut segment = RawSegment::new();
507        segment.push_op0(|| 2);
508        segment.push_op0(|| 3);
509        segment.push_op0(|| 4);
510        segment.push_op3(|x: i32, y: i32, z: i32| x + y + z, false, false, false);
511        unsafe {
512            assert_eq!(segment.call0::<i32>().unwrap(), 9);
513        }
514    }
515
516    #[test]
517    fn complex_chain() {
518        let mut segment = RawSegment::new();
519        segment.push_op0(|| 10);
520        segment.push_op1(|x: i32| x * 2, false);
521        segment.push_op0(|| 5);
522        segment.push_op2(|x: i32, y: i32| x + y, false, false);
523        unsafe {
524            assert_eq!(segment.call0::<i32>().unwrap(), 25);
525        }
526    }
527}
cel_runtime/raw_segment.rs

cel_runtime/
raw_segment.rs
