Skip to main content

rquickjs_core/value/
iterable.rs

1//! JavaScript iterator and iterable types.
2//!
3//! - [`Iterable`] — wraps a Rust closure or iterator as a JS iterator
4//! - [`JsIterator`] — consumes a JS iterator from Rust
5
6use crate::js_lifetime::JsLifetime;
7use crate::object::Property;
8use crate::safe_ref::Mut;
9use crate::{
10    atom::PredefinedAtom,
11    function::{MutFn, This},
12    Ctx, Error, FromJs, Function, IntoJs, Object, Result, Value,
13};
14use core::{iter::FusedIterator, marker::PhantomData};
15
16struct IteratorPrototypeCache<'js>(Object<'js>);
17
18unsafe impl<'js> JsLifetime<'js> for IteratorPrototypeCache<'js> {
19    type Changed<'to> = IteratorPrototypeCache<'to>;
20}
21
22fn get_iterator_prototype<'js>(ctx: &Ctx<'js>) -> Result<Object<'js>> {
23    if let Some(guard) = ctx.userdata::<IteratorPrototypeCache>() {
24        return Ok(guard.0.clone());
25    }
26
27    let iterator_ctor: Object = ctx.globals().get(PredefinedAtom::Iterator)?;
28    let proto: Object = iterator_ctor.get(PredefinedAtom::Prototype)?;
29
30    let _ = ctx.store_userdata(IteratorPrototypeCache(proto.clone()));
31
32    Ok(proto)
33}
34
35fn build_iterator_object<'js, F>(ctx: &Ctx<'js>, next_fn: F) -> Result<Object<'js>>
36where
37    F: FnMut(&Ctx<'js>) -> Option<Result<Value<'js>>> + 'js,
38{
39    let iter_proto = get_iterator_prototype(ctx)?;
40
41    let proto = Object::new(ctx.clone())?;
42    proto.set_prototype(Some(&iter_proto))?;
43
44    let state = Mut::new(next_fn);
45    let next = Function::new(
46        ctx.clone(),
47        MutFn::new(move |ctx: Ctx<'js>| -> Result<Object<'js>> {
48            let result = Object::new(ctx.clone())?;
49            match (state.lock())(&ctx) {
50                Some(Ok(value)) => {
51                    result.set(PredefinedAtom::Value, value)?;
52                    result.set(PredefinedAtom::Done, false)?;
53                }
54                Some(Err(e)) => return Err(e),
55                None => {
56                    result.set(PredefinedAtom::Done, true)?;
57                }
58            }
59            Ok(result)
60        }),
61    )?;
62
63    proto.prop(
64        PredefinedAtom::Next,
65        Property::from(next).enumerable().writable().configurable(),
66    )?;
67
68    let iter_obj = Object::new(ctx.clone())?;
69    iter_obj.set_prototype(Some(&proto))?;
70
71    Ok(iter_obj)
72}
73
74/// Creates a JavaScript iterator from a Rust closure or iterator.
75///
76/// The returned JS object has `next()` and inherits
77/// `[Symbol.iterator]() { return this }` from `%IteratorPrototype%`,
78/// so it works with `for...of`, spread, and destructuring.
79///
80/// # From a closure
81/// ```
82/// # use rquickjs::{Runtime, Context, Result, Iterable};
83/// # let rt = Runtime::new().unwrap();
84/// # let ctx = Context::full(&rt).unwrap();
85/// # ctx.with(|ctx| -> Result<()> {
86/// let mut i = 0;
87/// let iter = Iterable::from_fn(move || {
88///     i += 1;
89///     if i <= 3 { Some(i) } else { None }
90/// });
91/// ctx.globals().set("myIter", iter)?;
92/// let result: Vec<i32> = ctx.eval("[...myIter]")?;
93/// assert_eq!(result, vec![1, 2, 3]);
94/// # Ok(())
95/// # }).unwrap();
96/// ```
97///
98/// # From an iterator
99/// ```
100/// # use rquickjs::{Runtime, Context, Result, Iterable};
101/// # let rt = Runtime::new().unwrap();
102/// # let ctx = Context::full(&rt).unwrap();
103/// # ctx.with(|ctx| -> Result<()> {
104/// let iter = Iterable::from(vec![1, 2, 3]);
105/// ctx.globals().set("myIter", iter)?;
106/// let result: Vec<i32> = ctx.eval("[...myIter]")?;
107/// assert_eq!(result, vec![1, 2, 3]);
108/// # Ok(())
109/// # }).unwrap();
110/// ```
111pub struct Iterable<F>(pub F);
112
113impl<F> Iterable<ClosureWrapper<F>> {
114    /// Create from a `FnMut() -> Option<T>` closure.
115    pub fn from_fn(f: F) -> Self {
116        Iterable(ClosureWrapper(f))
117    }
118}
119
120impl<I: IntoIterator> From<I> for Iterable<IteratorWrapper<I::IntoIter>> {
121    fn from(iter: I) -> Self {
122        Iterable(IteratorWrapper(Some(iter.into_iter())))
123    }
124}
125
126/// Trait for types that can produce optional values, used by [`Iterable`].
127///
128/// Implemented for `FnMut() -> Option<T>` and `Iterator<Item = T>`,
129/// allowing [`Iterable`] to wrap both closures and iterators uniformly.
130/// External users can implement this trait for custom iteration sources.
131pub trait IterableFn {
132    type Item;
133    fn call(&mut self) -> Option<Self::Item>;
134}
135
136/// Wrapper that adapts a `FnMut() -> Option<T>` closure into an [`IterableFn`].
137pub struct ClosureWrapper<F>(F);
138
139impl<T, F: FnMut() -> Option<T>> IterableFn for ClosureWrapper<F> {
140    type Item = T;
141    fn call(&mut self) -> Option<T> {
142        (self.0)()
143    }
144}
145
146/// Wrapper that adapts a Rust `Iterator` into an [`IterableFn`].
147pub struct IteratorWrapper<I>(Option<I>);
148
149impl<I: Iterator> IterableFn for IteratorWrapper<I> {
150    type Item = I::Item;
151    fn call(&mut self) -> Option<I::Item> {
152        self.0.as_mut()?.next()
153    }
154}
155
156impl<'js, F> IntoJs<'js> for Iterable<F>
157where
158    F: IterableFn + 'js,
159    F::Item: IntoJs<'js> + 'js,
160{
161    fn into_js(self, ctx: &Ctx<'js>) -> Result<Value<'js>> {
162        let mut f = self.0;
163        let iter_obj = build_iterator_object(ctx, move |ctx| f.call().map(|v| v.into_js(ctx)))?;
164        Ok(iter_obj.into_value())
165    }
166}
167
168/// Consumes a JavaScript iterator from Rust.
169///
170/// Wraps a JS iterator object and implements Rust's [`Iterator`] trait.
171/// Can be created from any JS iterable (arrays, maps, sets, generators, etc.).
172///
173/// # Example
174/// ```
175/// # use rquickjs::{Runtime, Context, Result, JsIterator};
176/// # let rt = Runtime::new().unwrap();
177/// # let ctx = Context::full(&rt).unwrap();
178/// # ctx.with(|ctx| -> Result<()> {
179/// let iter: JsIterator<i32> = ctx.eval("[1, 2, 3]")?;
180/// let values: Vec<i32> = iter.filter_map(|r| r.ok()).collect();
181/// assert_eq!(values, vec![1, 2, 3]);
182/// # Ok(())
183/// # }).unwrap();
184/// ```
185pub struct JsIterator<'js, T = Value<'js>> {
186    iterator: Object<'js>,
187    done: bool,
188    _marker: PhantomData<T>,
189}
190
191impl<'js, T> JsIterator<'js, T> {
192    /// Returns the underlying JS iterator object.
193    pub fn into_inner(self) -> Object<'js> {
194        self.iterator
195    }
196
197    /// Maps this iterator to yield a different type.
198    ///
199    /// This is useful when you have a `JsIterator<Value>` and want to convert
200    /// values to a specific type.
201    pub fn typed<U: FromJs<'js>>(self) -> JsIterator<'js, U> {
202        JsIterator {
203            iterator: self.iterator,
204            done: self.done,
205            _marker: PhantomData,
206        }
207    }
208}
209
210impl<'js, T: FromJs<'js>> Iterator for JsIterator<'js, T> {
211    type Item = Result<T>;
212
213    fn next(&mut self) -> Option<Self::Item> {
214        if self.done {
215            return None;
216        }
217
218        let next_fn: Function<'js> = match self.iterator.get(PredefinedAtom::Next) {
219            Ok(f) => f,
220            Err(e) => return Some(Err(e)),
221        };
222
223        let result: Object<'js> = match next_fn.call((This(self.iterator.clone()),)) {
224            Ok(r) => r,
225            Err(e) => return Some(Err(e)),
226        };
227
228        let done: bool = match result.get(PredefinedAtom::Done) {
229            Ok(d) => d,
230            Err(e) => return Some(Err(e)),
231        };
232
233        if done {
234            self.done = true;
235            return None;
236        }
237
238        let value: Value<'js> = match result.get(PredefinedAtom::Value) {
239            Ok(v) => v,
240            Err(e) => return Some(Err(e)),
241        };
242
243        Some(T::from_js(self.iterator.ctx(), value))
244    }
245}
246
247impl<'js, T: FromJs<'js>> FusedIterator for JsIterator<'js, T> {}
248
249impl<'js, T: FromJs<'js>> FromJs<'js> for JsIterator<'js, T> {
250    fn from_js(_ctx: &Ctx<'js>, value: Value<'js>) -> Result<Self> {
251        let obj = Object::from_value(value)?;
252
253        // Try Symbol.iterator first (for iterables like arrays)
254        if let Ok(iter_fn) = obj.get::<_, Function<'js>>(PredefinedAtom::SymbolIterator) {
255            let iterator: Object<'js> = iter_fn.call((This(obj),))?;
256            return Ok(JsIterator {
257                iterator,
258                done: false,
259                _marker: PhantomData,
260            });
261        }
262
263        // Fall back to treating it as an iterator (has `next` method)
264        if obj.contains_key(PredefinedAtom::Next)? {
265            return Ok(JsIterator {
266                iterator: obj,
267                done: false,
268                _marker: PhantomData,
269            });
270        }
271
272        Err(Error::new_from_js(
273            "value",
274            "iterable (object with Symbol.iterator or next)",
275        ))
276    }
277}
278
279#[cfg(test)]
280mod test {
281    use super::*;
282    use crate::*;
283
284    #[test]
285    fn from_vec() {
286        test_with(|ctx| {
287            let iter = Iterable::from(vec![1i32, 2, 3]);
288            ctx.globals().set("myIter", iter).unwrap();
289            let result: Vec<i32> = ctx.eval("[...myIter]").unwrap();
290            assert_eq!(result, vec![1, 2, 3]);
291        });
292    }
293
294    #[test]
295    fn from_range() {
296        test_with(|ctx| {
297            let iter = Iterable::from(0..5);
298            ctx.globals().set("myIter", iter).unwrap();
299            let result: Vec<i32> = ctx.eval("[...myIter]").unwrap();
300            assert_eq!(result, vec![0, 1, 2, 3, 4]);
301        });
302    }
303
304    #[test]
305    fn from_closure() {
306        test_with(|ctx| {
307            let mut i = 0i32;
308            let iter = Iterable::from_fn(move || {
309                i += 1;
310                if i <= 3 {
311                    Some(i)
312                } else {
313                    None
314                }
315            });
316            ctx.globals().set("myIter", iter).unwrap();
317            let result: Vec<i32> = ctx.eval("[...myIter]").unwrap();
318            assert_eq!(result, vec![1, 2, 3]);
319        });
320    }
321
322    #[test]
323    fn for_of() {
324        test_with(|ctx| {
325            let iter = Iterable::from(vec!["a", "b", "c"]);
326            ctx.globals().set("myIter", iter).unwrap();
327            let result: alloc::string::String = ctx
328                .eval(r#"let s = ""; for (const x of myIter) { s += x; } s"#)
329                .unwrap();
330            assert_eq!(result, "abc");
331        });
332    }
333
334    #[test]
335    fn symbol_iterator_returns_this() {
336        test_with(|ctx| {
337            let iter = Iterable::from(vec![1i32]);
338            ctx.globals().set("myIter", iter).unwrap();
339            let ok: bool = ctx.eval("myIter[Symbol.iterator]() === myIter").unwrap();
340            assert!(ok);
341        });
342    }
343
344    #[test]
345    fn prototype_chain() {
346        test_with(|ctx| {
347            let iter = Iterable::from(vec![1i32]);
348            ctx.globals().set("myIter", iter).unwrap();
349            let ok: bool = ctx
350                .eval(
351                    r#"
352                    const iterProto = Object.getPrototypeOf(
353                        Object.getPrototypeOf([][Symbol.iterator]())
354                    );
355                    Object.getPrototypeOf(Object.getPrototypeOf(myIter)) === iterProto
356                    "#,
357                )
358                .unwrap();
359            assert!(ok);
360        });
361    }
362
363    #[test]
364    fn next_descriptors() {
365        test_with(|ctx| {
366            let iter = Iterable::from(vec![1i32]);
367            ctx.globals().set("myIter", iter).unwrap();
368            let ok: bool = ctx
369                .eval(
370                    r#"
371                    const proto = Object.getPrototypeOf(myIter);
372                    const desc = Object.getOwnPropertyDescriptor(proto, "next");
373                    desc.enumerable && desc.writable && desc.configurable
374                    "#,
375                )
376                .unwrap();
377            assert!(ok);
378        });
379    }
380
381    #[test]
382    fn js_iter_from_array() {
383        test_with(|ctx| {
384            let iter: JsIterator<i32> = ctx.eval("[1, 2, 3][Symbol.iterator]()").unwrap();
385            let values: Vec<i32> = iter.filter_map(|r| r.ok()).collect();
386            assert_eq!(values, vec![1, 2, 3]);
387        });
388    }
389
390    #[test]
391    fn js_iter_from_iterable() {
392        test_with(|ctx| {
393            let iter: JsIterator<i32> = ctx.eval("[4, 5, 6]").unwrap();
394            let values: Vec<i32> = iter.filter_map(|r| r.ok()).collect();
395            assert_eq!(values, vec![4, 5, 6]);
396        });
397    }
398
399    #[test]
400    fn js_iter_from_generator() {
401        test_with(|ctx| {
402            let iter: JsIterator<i32> = ctx
403                .eval("(function*() { yield 10; yield 20; yield 30; })()")
404                .unwrap();
405            let values: Vec<i32> = iter.filter_map(|r| r.ok()).collect();
406            assert_eq!(values, vec![10, 20, 30]);
407        });
408    }
409
410    #[test]
411    fn js_iter_roundtrip() {
412        test_with(|ctx| {
413            let rust_iter = Iterable::from(vec![100i32, 200, 300]);
414            ctx.globals().set("myIter", rust_iter).unwrap();
415            let js_iter: JsIterator<i32> = ctx.eval("myIter").unwrap();
416            let values: Vec<i32> = js_iter.filter_map(|r| r.ok()).collect();
417            assert_eq!(values, vec![100, 200, 300]);
418        });
419    }
420
421    #[test]
422    fn js_iter_raw_values() {
423        test_with(|ctx| {
424            let iter: JsIterator<Value> = ctx.eval("[1, 'two', 3]").unwrap();
425            let values: Vec<Value> = iter.filter_map(|r| r.ok()).collect();
426            assert_eq!(values.len(), 3);
427        });
428    }
429
430    #[test]
431    fn js_iter_typed() {
432        test_with(|ctx| {
433            let iter: JsIterator<Value> = ctx.eval("[1, 2, 3]").unwrap();
434            let values: Vec<i32> = iter.typed::<i32>().filter_map(|r| r.ok()).collect();
435            assert_eq!(values, vec![1, 2, 3]);
436        });
437    }
438
439    #[test]
440    fn js_iter_map_entries() {
441        test_with(|ctx| {
442            let iter: JsIterator<Array> =
443                ctx.eval("new Map([['a', 1], ['b', 2]]).entries()").unwrap();
444            let entries: Vec<Array> = iter.filter_map(|r| r.ok()).collect();
445            assert_eq!(entries.len(), 2);
446            assert_eq!(entries[0].get::<alloc::string::String>(0).unwrap(), "a");
447            assert_eq!(entries[0].get::<i32>(1).unwrap(), 1);
448        });
449    }
450
451    #[test]
452    fn js_iter_set() {
453        test_with(|ctx| {
454            let iter: JsIterator<i32> = ctx.eval("new Set([1, 2, 3])").unwrap();
455            let values: Vec<i32> = iter.filter_map(|r| r.ok()).collect();
456            assert_eq!(values, vec![1, 2, 3]);
457        });
458    }
459
460    #[test]
461    fn custom_iterable_fn() {
462        struct Counter(i32);
463
464        impl IterableFn for Counter {
465            type Item = i32;
466            fn call(&mut self) -> Option<i32> {
467                self.0 += 1;
468                (self.0 <= 3).then_some(self.0)
469            }
470        }
471
472        test_with(|ctx| {
473            let iter = Iterable(Counter(0));
474            ctx.globals().set("myIter", iter).unwrap();
475            let result: Vec<i32> = ctx.eval("[...myIter]").unwrap();
476            assert_eq!(result, vec![1, 2, 3]);
477        });
478    }
479
480    #[test]
481    fn custom_iterator_as_iterable_fn() {
482        struct Doubles(i32);
483
484        impl Iterator for Doubles {
485            type Item = i32;
486            fn next(&mut self) -> Option<i32> {
487                self.0 += 1;
488                (self.0 <= 3).then_some(self.0 * 2)
489            }
490        }
491
492        test_with(|ctx| {
493            let iter = Iterable::from(Doubles(0));
494            ctx.globals().set("myIter", iter).unwrap();
495            let result: Vec<i32> = ctx.eval("[...myIter]").unwrap();
496            assert_eq!(result, vec![2, 4, 6]);
497        });
498    }
499}