rune/runtime/
generator.rs

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