Skip to main content

cel_runtime/
segment.rs

1use crate::c_stack_list::{CNil, CStackList};
2use crate::dyn_segment::DynSegment;
3use crate::list_traits::{EmptyList, IntoList, List, ListTypeIteratorAdvance, TypeIdIterator};
4use crate::raw_segment::RawSegment;
5use crate::raw_stack::RawStack;
6use crate::{CStackListHeadLimit, CStackListHeadPadded, ReverseList};
7use anyhow::{Result, ensure};
8use std::any::TypeId;
9use std::result::Result::Ok;
10
11/// Extracts a typed segment's result or its empty-stack marker.
12///
13/// # Examples
14/// ```
15/// use cel_runtime::{Segment, SegmentCallable};
16/// assert!(Segment::<()>::new().call(()).is_ok());
17/// ```
18pub trait StackResult: List {
19    /// Moves the final result out, or returns the empty-stack marker.
20    ///
21    /// # Safety
22    /// The raw stack's live values must match this type-level stack exactly.
23    unsafe fn take_result(stack: &mut RawStack) -> Self::Head;
24}
25
26impl StackResult for CNil<()> {
27    /// Returns the empty-stack marker without reading a nonexistent slot.
28    unsafe fn take_result(_stack: &mut RawStack) -> crate::list_traits::Undefined {
29        crate::list_traits::Undefined
30    }
31}
32
33impl<H: 'static, T: List + CStackListHeadLimit> StackResult for CStackList<H, T> {
34    /// Transfers the concrete head out of its owned stack slot.
35    unsafe fn take_result(stack: &mut RawStack) -> H {
36        unsafe { stack.pop::<H>(Self::HEAD_PADDED) }
37    }
38}
39
40/// A type-safe segment that represents a sequence of operations.
41///
42/// The segment takes input arguments of type `Args` and maintains a type stack `Stack` that tracks
43/// the types of values produced by operations.
44///
45/// # Type Parameters
46/// - `Args`: The input argument types, must implement `IntoList`
47/// - `Stack`: The type stack tracking operation results, defaults to the reverse of `Args`
48///
49/// # Examples
50///
51/// ```rust
52/// use cel_runtime::segment::*;
53///
54/// assert_eq!(Segment::<(i32,)>::new() // create a new segment that takes an i32 argument
55///     .op1(|x| x * 2)                 // push a unary operation that multiplies the argument by 2
56///     .op0(|| 10)                     // push a nullary operation that returns 10
57///     .op2(|x, y| x + y)              // push a binary operation that adds two arguments
58///     .op1(|x| x.to_string())         // push a unary operation returning a string of the argument
59///     .call((42, )).unwrap(),         // call the segment with an argument 42
60///     "94");                          // the result is (42 * 2 + 10).to_string()
61/// ```
62pub struct Segment<
63    Args: IntoList + 'static,
64    Stack: List = ReverseList<<Args as IntoList>::Output<CNil<()>>>,
65> {
66    segment: RawSegment,
67    _phantom: std::marker::PhantomData<(Args, Stack)>,
68}
69
70impl<Args: IntoList + 'static> Default for Segment<Args> {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl<Args: IntoList + 'static> Segment<Args> {
77    /// Creates a new empty segment with no operations.
78    #[must_use]
79    pub fn new() -> Segment<Args> {
80        Segment {
81            segment: RawSegment::new(),
82            _phantom: std::marker::PhantomData,
83        }
84    }
85}
86
87impl<Args: IntoList + 'static, Stack: List + 'static> TryFrom<DynSegment> for Segment<Args, Stack>
88where
89    ReverseList<<Args as IntoList>::Output<CNil<()>>>: ListTypeIteratorAdvance<TypeId>,
90    Stack: ListTypeIteratorAdvance<TypeId>,
91{
92    type Error = anyhow::Error;
93    fn try_from(value: DynSegment) -> Result<Self, Self::Error> {
94        type ArgList<Args> = ReverseList<<Args as IntoList>::Output<CNil<()>>>;
95
96        ensure!(
97            ArgList::<Args>::LENGTH == value.argument_ids.len()
98                && TypeIdIterator::<ArgList::<Args>>::new().eq(value
99                    .argument_ids
100                    .iter()
101                    .rev()
102                    .copied()),
103            "argument type ids do not match"
104        );
105        for info in &value.stack_ids {
106            info.value_type
107                .validate_opaque_leaf(info.value_type.type_id())?;
108        }
109        ensure!(
110            value.argument_reads.is_empty() && value.base_stack_index == 0,
111            "typed conversion requires a direct preloaded frame without borrowed reads"
112        );
113        ensure!(
114            Stack::LENGTH == value.stack_ids.len()
115                && TypeIdIterator::<Stack>::new().eq(value
116                    .stack_ids
117                    .iter()
118                    .rev()
119                    .map(|info| info.value_type.type_id())),
120            "stack type ids do not match"
121        );
122        Ok(Segment {
123            segment: value.segment,
124            _phantom: std::marker::PhantomData,
125        })
126    }
127}
128
129impl<Args: IntoList + 'static, Stack> Segment<Args, Stack>
130where
131    Stack: StackResult + 'static,
132    Stack::Tail: CStackListHeadLimit,
133{
134    /// Private method to change the stack type.
135    fn into<NewStack: List + 'static>(self) -> Segment<Args, NewStack> {
136        Segment {
137            segment: self.segment,
138            _phantom: std::marker::PhantomData,
139        }
140    }
141
142    /*     pub fn indirect<N: 'static>(mut self) -> Segment<Args, Stack::PushFront<&Stack::Index<N>>> {
143        self.segment.indirect::<N, T>();
144        self.into()
145    } */
146
147    /// Pushes a nullary operation that takes no arguments and returns a value of type R.
148    pub fn op0<R, F>(mut self, op: F) -> Segment<Args, Stack::Push<R>>
149    where
150        F: Fn() -> R + 'static,
151        R: 'static,
152    {
153        self.segment.push_op0(op);
154        self.into()
155    }
156
157    /// Pushes a fallible nullary operation that returns `Result<R>`.
158    pub fn op0r<R, F>(mut self, op: F) -> Segment<Args, Stack::Push<R>>
159    where
160        F: Fn() -> Result<R> + 'static,
161        R: 'static,
162    {
163        self.segment.raw0(move |_| op());
164        self.into()
165    }
166
167    /// Pushes a unary operation that takes the current stack value and returns a new value.
168    pub fn op1<R, F>(mut self, op: F) -> Segment<Args, CStackList<R, Stack::Tail>>
169    where
170        Stack: CStackListHeadPadded,
171        F: Fn(Stack::Head) -> R + 'static,
172        R: 'static,
173    {
174        self.segment.push_op1(op, Stack::HEAD_PADDED);
175        self.into()
176    }
177
178    /// Pushes a fallible unary operation that takes the current stack value and returns `Result<R>`.
179    pub fn op1r<R, F>(mut self, op: F) -> Segment<Args, CStackList<R, Stack::Tail>>
180    where
181        Stack: CStackListHeadPadded,
182        F: Fn(Stack::Head) -> Result<R> + 'static,
183        R: 'static,
184    {
185        self.segment.raw1(move |_, x| op(x), Stack::HEAD_PADDED);
186        self.into()
187    }
188    /// Pushes a binary operation that takes the top two stack values and returns a new value.
189    pub fn op2<R, F>(
190        mut self,
191        op: F,
192    ) -> Segment<Args, <<Stack::Tail as List>::Tail as List>::Push<R>>
193    where
194        Stack: CStackListHeadPadded,
195        Stack::Tail: CStackListHeadPadded,
196        F: Fn(<Stack::Tail as List>::Head, Stack::Head) -> R + 'static,
197        R: 'static,
198    {
199        self.segment.push_op2(
200            op,
201            <Stack::Tail as CStackListHeadPadded>::HEAD_PADDED,
202            Stack::HEAD_PADDED,
203        );
204        self.into()
205    }
206
207    /// Executes all operations in the segment and returns the final result.
208    pub(crate) fn call0(&self) -> Result<Stack::Head> {
209        self.finish(RawStack::with_base_alignment(self.segment.base_alignment()))
210    }
211
212    /// Executes the owned frame and extracts its statically described result.
213    fn finish(&self, mut stack: RawStack) -> Result<Stack::Head> {
214        unsafe {
215            self.segment.call0_stack(&mut stack)?;
216            Ok(Stack::take_result(&mut stack))
217        }
218    }
219
220    /// Executes a one-argument segment and returns its statically described result.
221    pub(crate) fn call1<A: 'static>(&self, (a,): (A,)) -> Result<Stack::Head> {
222        let mut stack = RawStack::with_base_alignment(self.segment.base_alignment());
223        stack.push(a);
224        self.finish(stack)
225    }
226
227    /// Executes a two-argument segment and returns its statically described result.
228    pub(crate) fn call2<A: 'static, B: 'static>(&self, (a, b): (A, B)) -> Result<Stack::Head> {
229        let mut stack = RawStack::with_base_alignment(self.segment.base_alignment());
230        stack.push(a);
231        stack.push(b);
232        self.finish(stack)
233    }
234}
235
236/// Callable segment abstraction for uniform invocation across arities.
237///
238/// This trait invokes statically typed stack segments; [`crate::Callable`] is
239/// the separate shared native/interpreted value facade.
240///
241/// # Examples
242/// ```
243/// use cel_runtime::{Segment, SegmentCallable};
244/// let body = Segment::<()>::new().op0(|| 7i32);
245/// assert_eq!(body.call(()).unwrap(), 7);
246/// ```
247pub trait SegmentCallable<Args> {
248    /// The result type produced by calling with `Args`.
249    type Output;
250    /// Invokes the segment with the supplied arguments.
251    fn call(&self, args: Args) -> Self::Output;
252}
253
254impl<T: StackResult + 'static> SegmentCallable<()> for Segment<(), T>
255where
256    T::Tail: EmptyList + CStackListHeadLimit,
257{
258    type Output = Result<T::Head>;
259    fn call(&self, _args: ()) -> Self::Output {
260        self.call0()
261    }
262}
263
264impl<T: StackResult + 'static, A: 'static> SegmentCallable<(A,)> for Segment<(A,), T>
265where
266    T::Tail: EmptyList + CStackListHeadLimit,
267{
268    type Output = Result<T::Head>;
269    fn call(&self, args: (A,)) -> Self::Output {
270        self.call1(args)
271    }
272}
273
274impl<T: StackResult + 'static, A: 'static, B: 'static> SegmentCallable<(A, B)>
275    for Segment<(A, B), T>
276where
277    T::Tail: EmptyList + CStackListHeadLimit,
278{
279    type Output = Result<T::Head>;
280    fn call(&self, args: (A, B)) -> Self::Output {
281        self.call2(args)
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use anyhow::Result;
289    use std::sync::Arc;
290    use std::sync::atomic::AtomicUsize;
291    use std::sync::atomic::Ordering;
292
293    #[test]
294    fn segment_from_dyn_segment() -> Result<()> {
295        let mut dyn_segment = DynSegment::new::<(i32,)>();
296        dyn_segment.op0(|| 42);
297        dyn_segment.op2(|x: i32, y: i32| x + y)?;
298        let segment = Segment::<(i32,), CStackList<i32, CNil<()>>>::try_from(dyn_segment)?;
299
300        assert_eq!(segment.call((10,)).unwrap(), 52);
301        Ok(())
302    }
303
304    #[test]
305    fn unit_result() {
306        let segment = Segment::new();
307        let result = segment.call(());
308        assert!(result.is_ok());
309    }
310
311    /// Returns the empty-stack marker after consuming one or two input values.
312    #[test]
313    fn consumed_arguments_allow_empty_results() {
314        let mut one = DynSegment::new::<(String,)>();
315        one.drop_value().unwrap();
316        let one = Segment::<(String,), CNil<()>>::try_from(one).unwrap();
317        assert!(one.call((String::from("one"),)).is_ok());
318
319        let mut two = DynSegment::new::<(u8, String)>();
320        two.drop_value().unwrap();
321        two.drop_value().unwrap();
322        let two = Segment::<(u8, String), CNil<()>>::try_from(two).unwrap();
323        assert!(two.call((1, String::from("two"))).is_ok());
324    }
325
326    #[test]
327    fn drop_on_error() {
328        struct DropCounter(Arc<AtomicUsize>);
329
330        impl Drop for DropCounter {
331            fn drop(&mut self) {
332                self.0.fetch_add(1, Ordering::SeqCst);
333            }
334        }
335
336        impl Clone for DropCounter {
337            fn clone(&self) -> Self {
338                DropCounter(self.0.clone())
339            }
340        }
341
342        let drop_count = Arc::new(AtomicUsize::new(0));
343        let tracker = DropCounter(drop_count.clone());
344
345        let segment = Segment::new()
346            .op0(move || tracker.clone())
347            .op0r(|| -> Result<u32> { Err(anyhow::anyhow!("error")) })
348            .op2(|_: DropCounter, _: u32| 42u32);
349
350        assert_eq!(drop_count.load(Ordering::SeqCst), 0); // Nothing dropped yet
351        let result = segment.call(());
352        assert!(matches!(result, Err(e) if e.to_string() == "error"));
353        assert_eq!(drop_count.load(Ordering::SeqCst), 1); // The DropCounter from op0 was dropped
354    }
355
356    #[test]
357    fn op1r_error_does_not_panic() {
358        let parse_and_double = Segment::<(&str,)>::new()
359            .op1r(|s: &str| s.parse::<u32>().map_err(anyhow::Error::from))
360            .op1(|n| n * 2);
361
362        let result = parse_and_double.call(("nope",));
363        assert!(result.is_err());
364    }
365
366    #[test]
367    fn op1r_drops_remaining_stack_exactly_once_on_error() {
368        struct DropCounter(Arc<AtomicUsize>);
369
370        impl Drop for DropCounter {
371            fn drop(&mut self) {
372                self.0.fetch_add(1, Ordering::SeqCst);
373            }
374        }
375
376        impl Clone for DropCounter {
377            fn clone(&self) -> Self {
378                DropCounter(self.0.clone())
379            }
380        }
381
382        let drop_count = Arc::new(AtomicUsize::new(0));
383        let tracker = DropCounter(drop_count.clone());
384
385        let segment = Segment::new()
386            .op0(move || tracker.clone())
387            .op0(|| "nope")
388            .op1r(|s: &str| s.parse::<u32>().map_err(anyhow::Error::from))
389            .op2(|_: DropCounter, _: u32| 42u32);
390
391        assert_eq!(drop_count.load(Ordering::SeqCst), 0); // Nothing dropped yet
392        let result = segment.call(());
393        assert!(result.is_err());
394        assert_eq!(drop_count.load(Ordering::SeqCst), 1); // The DropCounter was dropped exactly once
395    }
396
397    #[test]
398    fn binary_operation_args_of_different_types() {
399        let result = Segment::new()
400            .op2(|x, y| format!("{} {}", x, y))
401            .call(("Hello", 12));
402
403        assert_eq!(result.unwrap(), "Hello 12");
404    }
405
406    #[test]
407    fn type_safe_operations() {
408        let result = Segment::new()
409            .op0(|| 42)
410            .op0(|| 10)
411            .op2(|x, y| x + y)
412            .op1(|x| x * 2)
413            .op1(|x| format!("{}", x))
414            .call(());
415
416        assert_eq!(result.unwrap(), "104");
417    }
418
419    #[test]
420    fn chain_operations() {
421        let result = Segment::<(&str,)>::new()
422            .op1(|s| s.len())
423            .op1(|n| n * 2)
424            .op1(|n| format!("Length * 2 = {}", n))
425            .call(("Hello",));
426
427        assert_eq!(result.unwrap(), "Length * 2 = 10");
428    }
429
430    #[test]
431    fn call_with_args() {
432        let result = Segment::new() //
433            .op1(|x: i32| x * 2) //
434            .call((21,));
435
436        assert_eq!(result.unwrap(), 42);
437    }
438}