rune/modules/ops/
generator.rs

1//! Overloadable operators and associated types.
2
3use crate as rune;
4use crate::alloc;
5use crate::alloc::fmt::TryWrite;
6use crate::alloc::prelude::*;
7use crate::runtime::generator::Iter;
8use crate::runtime::{EnvProtocolCaller, Formatter, Generator, GeneratorState, Value, VmError};
9use crate::{docstring, ContextError, Module};
10
11/// Types related to generators.
12#[rune::module(::std::ops::generator)]
13pub fn module() -> Result<Module, ContextError> {
14    let mut m = Module::from_meta(self::module__meta)?;
15
16    {
17        m.ty::<Generator>()?;
18        m.function_meta(generator_next__meta)?;
19        m.function_meta(generator_resume__meta)?;
20        m.function_meta(generator_iter__meta)?;
21        m.function_meta(generator_into_iter__meta)?;
22        m.function_meta(generator_debug__meta)?;
23        m.function_meta(generator_clone__meta)?;
24        m.implement_trait::<Generator>(rune::item!(::std::clone::Clone))?;
25
26        m.ty::<Iter>()?.docs(docstring! {
27            /// An iterator over a generator.
28        })?;
29        m.function_meta(Iter::next__meta)?;
30        m.implement_trait::<Iter>(rune::item!(::std::iter::Iterator))?;
31    }
32
33    {
34        m.ty::<GeneratorState>()?.docs(docstring! {
35            /// Enum indicating the state of a generator.
36        })?;
37
38        m.function_meta(generator_state_partial_eq__meta)?;
39        m.implement_trait::<GeneratorState>(rune::item!(::std::cmp::PartialEq))?;
40        m.function_meta(generator_state_eq__meta)?;
41        m.implement_trait::<GeneratorState>(rune::item!(::std::cmp::Eq))?;
42        m.function_meta(generator_state_debug__meta)?;
43        m.function_meta(generator_state_clone__meta)?;
44        m.implement_trait::<GeneratorState>(rune::item!(::std::clone::Clone))?;
45    }
46
47    Ok(m)
48}
49
50/// Advance a generator producing the next value yielded.
51///
52/// Unlike [`Generator::resume`], this can only consume the yielded values.
53///
54/// # Examples
55///
56/// ```rune
57/// use std::ops::{Generator, GeneratorState};
58///
59/// fn generate() {
60///     yield 1;
61///     yield 2;
62/// }
63///
64/// let g = generate();
65///
66/// assert_eq!(g.next(), Some(1));
67/// assert_eq!(g.next(), Some(2));
68/// assert_eq!(g.next(), None);
69/// ``
70#[rune::function(keep, instance, path = next)]
71fn generator_next(this: &mut Generator) -> Result<Option<Value>, VmError> {
72    this.next()
73}
74
75/// Resumes the execution of this generator.
76///
77/// This function will resume execution of the generator or start execution if
78/// it hasn't already. This call will return back into the generator's last
79/// suspension point, resuming execution from the latest `yield`. The generator
80/// will continue executing until it either yields or returns, at which point
81/// this function will return.
82///
83/// # Return value
84///
85/// The `GeneratorState` enum returned from this function indicates what state
86/// the generator is in upon returning. If the `Yielded` variant is returned
87/// then the generator has reached a suspension point and a value has been
88/// yielded out. Generators in this state are available for resumption at a
89/// later point.
90///
91/// If `Complete` is returned then the generator has completely finished with
92/// the value provided. It is invalid for the generator to be resumed again.
93///
94/// # Panics
95///
96/// This function may panic if it is called after the `Complete` variant has
97/// been returned previously. While generator literals in the language are
98/// guaranteed to panic on resuming after `Complete`, this is not guaranteed for
99/// all implementations of the `Generator`.
100///
101/// # Examples
102///
103/// ```rune
104/// use std::ops::{Generator, GeneratorState};
105///
106/// fn generate() {
107///     let n = yield 1;
108///     yield 2 + n;
109/// }
110///
111/// let g = generate();
112///
113/// assert_eq!(g.resume(()), GeneratorState::Yielded(1));
114/// assert_eq!(g.resume(1), GeneratorState::Yielded(3));
115/// assert_eq!(g.resume(()), GeneratorState::Complete(()));
116/// ``
117#[rune::function(keep, instance, path = resume)]
118fn generator_resume(this: &mut Generator, value: Value) -> Result<GeneratorState, VmError> {
119    this.resume(value)
120}
121
122/// Convert a generator into an iterator.
123///
124/// # Examples
125///
126/// ```rune
127/// fn count_numbers(limit) {
128///     for n in 0..limit.unwrap_or(10) {
129///         yield n;
130///     }
131/// }
132///
133/// assert_eq!(count_numbers(None).iter().collect::<Vec>(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
134/// assert_eq!(count_numbers(Some(2)).iter().collect::<Vec>(), [0, 1]);
135/// ```
136#[rune::function(keep, instance, path = iter)]
137#[inline]
138fn generator_iter(this: Generator) -> Iter {
139    this.rune_iter()
140}
141
142/// Construct an iterator over a generator.
143///
144/// # Examples
145///
146/// ```rune
147/// fn count_numbers(limit) {
148///     for n in 0..limit {
149///         yield n;
150///     }
151/// }
152///
153/// let result = 0;
154///
155/// for n in count_numbers(3) {
156///     result += n;
157/// }
158///
159/// assert_eq!(result, 3);
160/// ```
161#[rune::function(keep, instance, protocol = INTO_ITER)]
162#[inline]
163fn generator_into_iter(this: Generator) -> Iter {
164    this.rune_iter()
165}
166
167/// Debug print this generator.
168///
169/// # Examples
170///
171/// ```rune
172/// use std::ops::GeneratorState;
173///
174/// fn generate() {
175///     let n = yield 1;
176///     yield 2 + n;
177/// }
178///
179/// let a = generate();
180///
181/// println!("{a:?}");
182/// ``
183#[rune::function(keep, instance, protocol = DEBUG_FMT)]
184fn generator_debug(this: &Generator, f: &mut Formatter) -> alloc::Result<()> {
185    write!(f, "{this:?}")
186}
187
188/// Clone a generator.
189///
190/// This clones the state of the generator too, allowing it to be resumed
191/// independently.
192///
193/// # Examples
194///
195/// ```rune
196/// use std::ops::GeneratorState;
197///
198/// fn generate() {
199///     let n = yield 1;
200///     yield 2 + n;
201/// }
202///
203/// let a = generate();
204///
205/// assert_eq!(a.resume(()), GeneratorState::Yielded(1));
206/// let b = a.clone();
207/// assert_eq!(a.resume(2), GeneratorState::Yielded(4));
208/// assert_eq!(b.resume(3), GeneratorState::Yielded(5));
209///
210/// assert_eq!(a.resume(()), GeneratorState::Complete(()));
211/// assert_eq!(b.resume(()), GeneratorState::Complete(()));
212/// ``
213#[rune::function(keep, instance, protocol = CLONE)]
214fn generator_clone(this: &Generator) -> alloc::Result<Generator> {
215    this.try_clone()
216}
217
218/// Test for partial equality over a generator state.
219///
220/// # Examples
221///
222/// ```rune
223/// use std::ops::{Generator, GeneratorState};
224///
225/// fn generate() {
226///     let n = yield 1;
227///     yield 2 + n;
228/// }
229///
230/// let g = generate();
231///
232/// assert_eq!(g.resume(()), GeneratorState::Yielded(1));
233/// assert_eq!(g.resume(1), GeneratorState::Yielded(3));
234/// assert_eq!(g.resume(()), GeneratorState::Complete(()));
235/// ``
236#[rune::function(keep, instance, protocol = PARTIAL_EQ)]
237fn generator_state_partial_eq(
238    this: &GeneratorState,
239    other: &GeneratorState,
240) -> Result<bool, VmError> {
241    this.partial_eq_with(other, &mut EnvProtocolCaller)
242}
243
244/// Test for total equality over a generator state.
245///
246/// # Examples
247///
248/// ```rune
249/// use std::ops::{Generator, GeneratorState};
250/// use std::ops::eq;
251///
252/// fn generate() {
253///     let n = yield 1;
254///     yield 2 + n;
255/// }
256///
257/// let g = generate();
258///
259/// assert!(eq(g.resume(()), GeneratorState::Yielded(1)));
260/// assert!(eq(g.resume(1), GeneratorState::Yielded(3)));
261/// assert!(eq(g.resume(()), GeneratorState::Complete(())));
262/// ``
263#[rune::function(keep, instance, protocol = EQ)]
264fn generator_state_eq(this: &GeneratorState, other: &GeneratorState) -> Result<bool, VmError> {
265    this.eq_with(other, &mut EnvProtocolCaller)
266}
267
268/// Debug print this generator state.
269///
270/// # Examples
271///
272/// ```rune
273/// use std::ops::GeneratorState;
274///
275/// let a = GeneratorState::Yielded(1);
276/// let b = GeneratorState::Complete(());
277///
278/// println!("{a:?}");
279/// println!("{b:?}");
280/// ``
281#[rune::function(keep, instance, protocol = DEBUG_FMT)]
282fn generator_state_debug(this: &GeneratorState, f: &mut Formatter) -> Result<(), VmError> {
283    match this {
284        GeneratorState::Yielded(value) => {
285            write!(f, "Yielded(")?;
286            value.debug_fmt_with(f, &mut EnvProtocolCaller)?;
287            write!(f, ")")?;
288        }
289        GeneratorState::Complete(value) => {
290            write!(f, "Complete(")?;
291            value.debug_fmt_with(f, &mut EnvProtocolCaller)?;
292            write!(f, ")")?;
293        }
294    }
295
296    Ok(())
297}
298
299/// Clone a generator state.
300///
301/// # Examples
302///
303/// ```rune
304/// use std::ops::GeneratorState;
305///
306/// let a = GeneratorState::Yielded(1);
307/// let b = a.clone();
308///
309/// assert_eq!(a, b);
310/// ``
311#[rune::function(keep, instance, protocol = CLONE)]
312fn generator_state_clone(this: &GeneratorState) -> alloc::Result<GeneratorState> {
313    this.try_clone()
314}