Skip to main content

rune/runtime/
generator.rs

1use core::fmt;
2use core::iter;
3
4use crate as rune;
5use crate::alloc::clone::TryClone;
6use crate::runtime::{
7    Dismantle, GeneratorState, Handover, Value, Vm, VmError, VmErrorKind, VmExecution, VmHaltInfo,
8    VmOutcome,
9};
10use crate::Any;
11
12/// A generator produced by a generator function.
13///
14/// Generator are functions or closures which contain the `yield` expressions.
15///
16/// # Examples
17///
18/// ```rune
19/// use std::ops::generator::Generator;
20///
21/// let f = |n| {
22///     yield n;
23///     yield n + 1;
24/// };
25///
26/// let g = f(10);
27///
28/// assert!(g is Generator);
29/// ```
30#[derive(Any)]
31#[rune(crate, item = ::std::ops::generator, dismantle)]
32pub struct Generator {
33    execution: Option<VmExecution<Vm>>,
34}
35
36impl Generator {
37    /// Construct a generator from a virtual machine.
38    pub(crate) fn new(vm: Vm) -> Self {
39        Self {
40            execution: Some(VmExecution::new(vm)),
41        }
42    }
43
44    /// Get the next value produced by this stream.
45    pub fn next(&mut self) -> Result<Option<Value>, VmError> {
46        let Some(execution) = self.execution.as_mut() else {
47            return Ok(None);
48        };
49
50        let state = execution.resume().complete()?;
51
52        match state {
53            VmOutcome::Complete(_) => {
54                self.execution = None;
55                Ok(None)
56            }
57            VmOutcome::Yielded(value) => Ok(Some(value)),
58            VmOutcome::Limited => Err(VmError::from(VmErrorKind::Halted {
59                halt: VmHaltInfo::Limited,
60            })),
61        }
62    }
63
64    /// Resume the generator with a value and get the next [`GeneratorState`].
65    pub fn resume(&mut self, value: Value) -> Result<GeneratorState, VmError> {
66        let execution = self
67            .execution
68            .as_mut()
69            .ok_or(VmErrorKind::GeneratorComplete)?;
70
71        let outcome = execution.resume().with_value(value).complete()?;
72
73        match outcome {
74            VmOutcome::Complete(value) => {
75                self.execution = None;
76                Ok(GeneratorState::Complete(value))
77            }
78            VmOutcome::Yielded(value) => Ok(GeneratorState::Yielded(value)),
79            VmOutcome::Limited => Err(VmError::from(VmErrorKind::Halted {
80                halt: VmHaltInfo::Limited,
81            })),
82        }
83    }
84}
85
86impl Generator {
87    /// Convert into iterator
88    pub fn rune_iter(self) -> Iter {
89        self.into_iter()
90    }
91}
92
93/// A suspended generator holds every value its machine was working over, so a
94/// generator which captured another one nests just like a container does.
95impl Dismantle for Generator {
96    fn dismantle(&mut self, out: &mut Handover<'_>) {
97        let Some(execution) = self.execution.as_mut() else {
98            return;
99        };
100
101        out.consume(execution.vm_mut().stack_mut());
102    }
103}
104
105impl IntoIterator for Generator {
106    type Item = Result<Value, VmError>;
107    type IntoIter = Iter;
108
109    #[inline]
110    fn into_iter(self) -> Self::IntoIter {
111        Iter { generator: self }
112    }
113}
114
115/// An iterator holds the generator it drives, and a script can nest one inside
116/// another, so it hands the generator over rather than being dropped in place.
117#[derive(Any)]
118#[rune(item = ::std::ops::generator)]
119pub struct Iter {
120    #[rune(dismantle)]
121    generator: Generator,
122}
123
124impl Iter {
125    #[rune::function(instance, keep, protocol = NEXT)]
126    pub(crate) fn next(&mut self) -> Result<Option<Value>, VmError> {
127        self.generator.next()
128    }
129}
130
131impl iter::Iterator for Iter {
132    type Item = Result<Value, VmError>;
133
134    #[inline]
135    fn next(&mut self) -> Option<Result<Value, VmError>> {
136        match Iter::next(self) {
137            Ok(Some(value)) => Some(Ok(value)),
138            Ok(None) => None,
139            Err(error) => Some(Err(error)),
140        }
141    }
142}
143
144impl fmt::Debug for Generator {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        f.debug_struct("Generator")
147            .field("completed", &self.execution.is_none())
148            .finish()
149    }
150}
151
152impl TryClone for Generator {
153    #[inline]
154    fn try_clone(&self) -> Result<Self, rune_alloc::Error> {
155        Ok(Self {
156            execution: self.execution.try_clone()?,
157        })
158    }
159}