koto_runtime/core_lib/
iterator.rs

1//! The `iterator` core library module
2
3pub mod adaptors;
4pub mod generators;
5pub mod peekable;
6
7use crate::{KIteratorOutput as Output, Result, derive::*, prelude::*};
8
9static MODULE_NAME: &str = "core.iterator";
10
11/// Initializes the `iterator` core library module
12pub fn make_module() -> KMap {
13    let result = KMap::with_type(MODULE_NAME);
14
15    result.add_fn("advance", |ctx| {
16        let expected_error = "|Iterator, Number >= 0|";
17
18        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
19            (KValue::Iterator(iterator), [KValue::Number(n)]) if *n >= 0.0 => {
20                let mut iterator = iterator.clone();
21                let mut remaining = usize::from(n);
22
23                while remaining > 0 {
24                    match iterator.next() {
25                        Some(Output::Error(error)) => return Err(error),
26                        Some(_) => remaining -= 1,
27                        None => break,
28                    }
29                }
30
31                Ok(remaining.into())
32            }
33            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
34        }
35    });
36
37    result.add_fn("all", |ctx| {
38        let expected_error = "|Iterable, |Any| -> Bool|";
39
40        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
41            (iterable, [predicate]) if predicate.is_callable() => {
42                let iterable = iterable.clone();
43                let predicate = predicate.clone();
44
45                for output in ctx.vm.make_iterator(iterable)? {
46                    let predicate_result = match output {
47                        Output::Value(value) => ctx.vm.call_function(predicate.clone(), value),
48                        Output::ValuePair(a, b) => ctx
49                            .vm
50                            .call_function(predicate.clone(), CallArgs::AsTuple(&[a, b])),
51                        Output::Error(error) => return Err(error),
52                    };
53
54                    match predicate_result {
55                        Ok(KValue::Bool(result)) => {
56                            if !result {
57                                return Ok(false.into());
58                            }
59                        }
60                        Ok(unexpected) => {
61                            return unexpected_type(
62                                "a Bool to be returned from the predicate",
63                                &unexpected,
64                            );
65                        }
66                        error @ Err(_) => return error,
67                    }
68                }
69
70                Ok(true.into())
71            }
72            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
73        }
74    });
75
76    result.add_fn("any", |ctx| {
77        let expected_error = "|Iterable, |Any| -> Bool|";
78
79        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
80            (iterable, [predicate]) if predicate.is_callable() => {
81                let iterable = iterable.clone();
82                let predicate = predicate.clone();
83
84                for output in ctx.vm.make_iterator(iterable)? {
85                    let predicate_result = match output {
86                        Output::Value(value) => ctx.vm.call_function(predicate.clone(), value),
87                        Output::ValuePair(a, b) => ctx
88                            .vm
89                            .call_function(predicate.clone(), CallArgs::AsTuple(&[a, b])),
90                        Output::Error(error) => return Err(error),
91                    };
92
93                    match predicate_result {
94                        Ok(KValue::Bool(result)) => {
95                            if result {
96                                return Ok(true.into());
97                            }
98                        }
99                        Ok(unexpected) => {
100                            return unexpected_type(
101                                "a Bool to be returned from the predicate",
102                                &unexpected,
103                            );
104                        }
105                        Err(error) => return Err(error),
106                    }
107                }
108
109                Ok(false.into())
110            }
111            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
112        }
113    });
114
115    result.add_fn("chain", |ctx| {
116        let expected_error = "|Iterable, Iterable|";
117
118        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
119            (iterable_a, [iterable_b]) if iterable_b.is_iterable() => {
120                let iterable_a = iterable_a.clone();
121                let iterable_b = iterable_b.clone();
122                let result = KIterator::new(adaptors::Chain::new(
123                    ctx.vm.make_iterator(iterable_a)?,
124                    ctx.vm.make_iterator(iterable_b)?,
125                ));
126
127                Ok(KValue::Iterator(result))
128            }
129            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
130        }
131    });
132
133    result.add_fn("chunks", |ctx| {
134        let expected_error = "|Iterable, Number|";
135
136        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
137            (iterable, [KValue::Number(n)]) => {
138                let iterable = iterable.clone();
139                let n = *n;
140                match adaptors::Chunks::new(ctx.vm.make_iterator(iterable)?, n.into()) {
141                    Ok(result) => Ok(KIterator::new(result).into()),
142                    Err(e) => runtime_error!("iterator.chunks: {}", e),
143                }
144            }
145            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
146        }
147    });
148
149    result.add_fn("consume", |ctx| {
150        let expected_error = "|Iterable|, or |Iterable, |Any| -> Any|";
151
152        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
153            (iterable, []) => {
154                let iterable = iterable.clone();
155                for output in ctx.vm.make_iterator(iterable)? {
156                    if let Output::Error(error) = output {
157                        return Err(error);
158                    }
159                }
160                Ok(KValue::Null)
161            }
162            (iterable, [f]) if f.is_callable() => {
163                let iterable = iterable.clone();
164                let f = f.clone();
165                for output in ctx.vm.make_iterator(iterable)? {
166                    match output {
167                        Output::Value(value) => {
168                            ctx.vm.call_function(f.clone(), value)?;
169                        }
170                        Output::ValuePair(a, b) => {
171                            ctx.vm
172                                .call_function(f.clone(), CallArgs::AsTuple(&[a, b]))?;
173                        }
174                        Output::Error(error) => return Err(error),
175                    }
176                }
177                Ok(KValue::Null)
178            }
179            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
180        }
181    });
182
183    result.add_fn("count", |ctx| {
184        let expected_error = "|Iterable|";
185
186        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
187            (iterable, []) => {
188                let iterable = iterable.clone();
189                let mut result = 0;
190                for output in ctx.vm.make_iterator(iterable)? {
191                    if let Output::Error(error) = output {
192                        return Err(error);
193                    }
194                    result += 1;
195                }
196                Ok(KValue::Number(result.into()))
197            }
198            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
199        }
200    });
201
202    result.add_fn("cycle", |ctx| {
203        let expected_error = "|Iterable|";
204
205        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
206            (iterable, []) => {
207                let iterable = iterable.clone();
208                let result = adaptors::Cycle::new(ctx.vm.make_iterator(iterable)?);
209
210                Ok(KIterator::new(result).into())
211            }
212            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
213        }
214    });
215
216    result.add_fn("each", |ctx| {
217        let expected_error = "|Iterable, |Any| -> Any|";
218
219        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
220            (iterable, [f]) if f.is_callable() => {
221                let iterable = iterable.clone();
222                let f = f.clone();
223                let result = adaptors::Each::new(ctx.vm.make_iterator(iterable)?, f, ctx.vm);
224
225                Ok(KIterator::new(result).into())
226            }
227            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
228        }
229    });
230
231    result.add_fn("enumerate", |ctx| {
232        let expected_error = "|Iterable|";
233
234        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
235            (iterable, []) => {
236                let iterable = iterable.clone();
237                let result = adaptors::Enumerate::new(ctx.vm.make_iterator(iterable)?);
238                Ok(KIterator::new(result).into())
239            }
240            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
241        }
242    });
243
244    result.add_fn("find", |ctx| {
245        let expected_error = "|Iterable, |Any| -> Bool|";
246
247        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
248            (iterable, [predicate]) if predicate.is_callable() => {
249                let iterable = iterable.clone();
250                let predicate = predicate.clone();
251
252                for output in ctx.vm.make_iterator(iterable)?.map(collect_pair) {
253                    match output {
254                        Output::Value(value) => {
255                            match ctx.vm.call_function(predicate.clone(), value.clone()) {
256                                Ok(KValue::Bool(result)) => {
257                                    if result {
258                                        return Ok(value);
259                                    }
260                                }
261                                Ok(unexpected) => {
262                                    return unexpected_type(
263                                        "a Bool to be returned from the predicate",
264                                        &unexpected,
265                                    );
266                                }
267                                Err(error) => return Err(error),
268                            }
269                        }
270                        Output::Error(error) => return Err(error),
271                        _ => unreachable!(),
272                    }
273                }
274
275                Ok(KValue::Null)
276            }
277            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
278        }
279    });
280
281    result.add_fn("flatten", |ctx| {
282        let expected_error = "|Iterable|";
283
284        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
285            (iterable, []) => {
286                let iterable = iterable.clone();
287                let result = adaptors::Flatten::new(ctx.vm.make_iterator(iterable)?, ctx.vm);
288
289                Ok(KIterator::new(result).into())
290            }
291            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
292        }
293    });
294
295    result.add_fn("fold", |ctx| {
296        let expected_error = "|Iterable, Any, |Any, Any| -> Any|";
297
298        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
299            (iterable, [result, f]) if f.is_callable() => {
300                let iterable = iterable.clone();
301                let result = result.clone();
302                let f = f.clone();
303                let mut iter = ctx.vm.make_iterator(iterable)?;
304
305                match iter
306                    .borrow_internals(|iterator| {
307                        let mut fold_result = result.clone();
308                        for value in iterator.map(collect_pair) {
309                            match value {
310                                Output::Value(value) => {
311                                    match ctx.vm.call_function(f.clone(), &[fold_result, value]) {
312                                        Ok(result) => fold_result = result,
313                                        Err(error) => return Some(Output::Error(error)),
314                                    }
315                                }
316                                Output::Error(error) => return Some(Output::Error(error)),
317                                _ => unreachable!(),
318                            }
319                        }
320
321                        Some(Output::Value(fold_result))
322                    })
323                    // None is never returned from the closure
324                    .unwrap()
325                {
326                    Output::Value(result) => Ok(result),
327                    Output::Error(error) => Err(error),
328                    _ => unreachable!(),
329                }
330            }
331            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
332        }
333    });
334
335    result.add_fn("generate", |ctx| {
336        let instance = ctx.instance();
337        if !(matches!(instance, &KValue::Null) || instance.type_as_string() == MODULE_NAME) {
338            return runtime_error!("iterator.generate must be used as a standalone function");
339        }
340
341        match ctx.args() {
342            [f] if f.is_callable() => {
343                let result = generators::Generate::new(f.clone(), ctx.vm);
344                Ok(KIterator::new(result).into())
345            }
346            [f, KValue::Number(n)] if *n >= 0 && f.is_callable() => {
347                let result = generators::GenerateN::new(n.into(), f.clone(), ctx.vm);
348                Ok(KIterator::new(result).into())
349            }
350            unexpected => unexpected_args(
351                "|generator: || -> Any|, or |generator: || -> Any, n: Number >= 0|",
352                unexpected,
353            ),
354        }
355    });
356
357    result.add_fn("intersperse", |ctx| {
358        let expected_error = "|Iterable, Value|";
359
360        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
361            (iterable, [separator_fn]) if separator_fn.is_callable() => {
362                let iterable = iterable.clone();
363                let separator_fn = separator_fn.clone();
364                let result = adaptors::IntersperseWith::new(
365                    ctx.vm.make_iterator(iterable)?,
366                    separator_fn,
367                    ctx.vm,
368                );
369
370                Ok(KIterator::new(result).into())
371            }
372            (iterable, [separator]) => {
373                let iterable = iterable.clone();
374                let separator = separator.clone();
375                let result = adaptors::Intersperse::new(ctx.vm.make_iterator(iterable)?, separator);
376
377                Ok(KIterator::new(result).into())
378            }
379            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
380        }
381    });
382
383    result.add_fn("iter", |ctx| {
384        let expected_error = "|Iterable|";
385
386        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
387            (iterable, []) => {
388                let iterable = iterable.clone();
389                Ok(KValue::Iterator(ctx.vm.make_iterator(iterable)?))
390            }
391            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
392        }
393    });
394
395    result.add_fn("keep", |ctx| {
396        let expected_error = "|Iterable, |Any| -> Bool|";
397
398        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
399            (iterable, [predicate]) if predicate.is_callable() => {
400                let iterable = iterable.clone();
401                let predicate = predicate.clone();
402                let result =
403                    adaptors::Keep::new(ctx.vm.make_iterator(iterable)?, predicate, ctx.vm);
404                Ok(KIterator::new(result).into())
405            }
406            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
407        }
408    });
409
410    result.add_fn("last", |ctx| {
411        let expected_error = "|Iterable|";
412
413        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
414            (iterable, []) => {
415                let iterable = iterable.clone();
416                let mut result = KValue::Null;
417
418                let mut iter = ctx.vm.make_iterator(iterable)?.map(collect_pair);
419                for output in &mut iter {
420                    match output {
421                        Output::Value(value) => result = value,
422                        Output::Error(error) => return Err(error),
423                        _ => unreachable!(),
424                    }
425                }
426
427                Ok(result)
428            }
429            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
430        }
431    });
432
433    result.add_fn("max", |ctx| {
434        let expected_error = "|Iterable|, or |Iterable, |Any| -> Any|";
435
436        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
437            (iterable, []) => {
438                let iterable = iterable.clone();
439                run_iterator_comparison(ctx.vm, iterable, InvertResult::Yes)
440            }
441            (iterable, [key_fn]) if key_fn.is_callable() => {
442                let iterable = iterable.clone();
443                let key_fn = key_fn.clone();
444                run_iterator_comparison_by_key(ctx.vm, iterable, key_fn, InvertResult::Yes)
445            }
446            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
447        }
448    });
449
450    result.add_fn("min", |ctx| {
451        let expected_error = "|Iterable|, or |Iterable, |Any| -> Any|";
452
453        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
454            (iterable, []) => {
455                let iterable = iterable.clone();
456                run_iterator_comparison(ctx.vm, iterable, InvertResult::No)
457            }
458            (iterable, [key_fn]) if key_fn.is_callable() => {
459                let iterable = iterable.clone();
460                let key_fn = key_fn.clone();
461                run_iterator_comparison_by_key(ctx.vm, iterable, key_fn, InvertResult::No)
462            }
463            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
464        }
465    });
466
467    result.add_fn("min_max", |ctx| {
468        let expected_error = "|Iterable|, or |Iterable, |Any| -> Any|";
469
470        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
471            (iterable, []) => {
472                let iterable = iterable.clone();
473                let mut result = None;
474
475                for iter_output in ctx.vm.make_iterator(iterable)?.map(collect_pair) {
476                    match iter_output {
477                        Output::Value(value) => {
478                            result = Some(match result {
479                                Some((min, max)) => (
480                                    compare_values(ctx.vm, min, value.clone(), InvertResult::No)?,
481                                    compare_values(ctx.vm, max, value, InvertResult::Yes)?,
482                                ),
483                                None => (value.clone(), value),
484                            })
485                        }
486                        Output::Error(error) => return Err(error),
487                        _ => unreachable!(),
488                    }
489                }
490
491                Ok(result.map_or(KValue::Null, |(min, max)| {
492                    KValue::Tuple(vec![min, max].into())
493                }))
494            }
495            (iterable, [key_fn]) if key_fn.is_callable() => {
496                let iterable = iterable.clone();
497                let key_fn = key_fn.clone();
498                let mut result = None;
499
500                for iter_output in ctx.vm.make_iterator(iterable)?.map(collect_pair) {
501                    match iter_output {
502                        Output::Value(value) => {
503                            let key = ctx.vm.call_function(key_fn.clone(), value.clone())?;
504                            let value_and_key = (value, key);
505
506                            result = Some(match result {
507                                Some((min_and_key, max_and_key)) => (
508                                    compare_values_with_key(
509                                        ctx.vm,
510                                        min_and_key,
511                                        value_and_key.clone(),
512                                        InvertResult::No,
513                                    )?,
514                                    compare_values_with_key(
515                                        ctx.vm,
516                                        max_and_key,
517                                        value_and_key,
518                                        InvertResult::Yes,
519                                    )?,
520                                ),
521                                None => (value_and_key.clone(), value_and_key),
522                            })
523                        }
524                        Output::Error(error) => return Err(error),
525                        _ => unreachable!(), // value pairs have been collected in collect_pair
526                    }
527                }
528
529                Ok(result.map_or(KValue::Null, |((min, _), (max, _))| {
530                    KValue::Tuple(vec![min, max].into())
531                }))
532            }
533            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
534        }
535    });
536
537    result.add_fn("next", |ctx| {
538        let expected_error = "|Iterable|";
539
540        let mut iter = match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
541            (KValue::Iterator(i), []) => i.clone(),
542            (iterable, []) if iterable.is_iterable() => ctx.vm.make_iterator(iterable.clone())?,
543            (instance, args) => {
544                return unexpected_args_after_instance(expected_error, instance, args);
545            }
546        };
547
548        let output = match iter_output_to_result(iter.next())? {
549            None => KValue::Null,
550            Some(output) => IteratorOutput::from(output).into(),
551        };
552
553        Ok(output)
554    });
555
556    result.add_fn("next_back", |ctx| {
557        let expected_error = "|Iterable|";
558
559        let mut iter = match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
560            (KValue::Iterator(i), []) => i.clone(),
561            (iterable, []) if iterable.is_iterable() => ctx.vm.make_iterator(iterable.clone())?,
562            (instance, args) => {
563                return unexpected_args_after_instance(expected_error, instance, args);
564            }
565        };
566
567        let output = match iter_output_to_result(iter.next_back())? {
568            None => KValue::Null,
569            Some(output) => IteratorOutput::from(output).into(),
570        };
571
572        Ok(output)
573    });
574
575    result.add_fn("once", |ctx| {
576        let instance = ctx.instance();
577        if !(matches!(instance, &KValue::Null) || instance.type_as_string() == MODULE_NAME) {
578            return runtime_error!("iterator.once must be used as a standalone function");
579        }
580
581        match ctx.args() {
582            [value] => Ok(KIterator::new(generators::Once::new(value.clone())).into()),
583            unexpected => unexpected_args("|Any|", unexpected),
584        }
585    });
586
587    result.add_fn("peekable", |ctx| {
588        let expected_error = "|Iterable|";
589
590        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
591            (iterable, []) => {
592                let iterable = iterable.clone();
593                Ok(peekable::Peekable::make_value(
594                    ctx.vm.make_iterator(iterable)?,
595                ))
596            }
597            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
598        }
599    });
600
601    result.add_fn("position", |ctx| {
602        let expected_error = "|Iterable, |Any| -> Bool|";
603
604        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
605            (iterable, [predicate]) if predicate.is_callable() => {
606                let iterable = iterable.clone();
607                let predicate = predicate.clone();
608
609                for (i, output) in ctx.vm.make_iterator(iterable)?.enumerate() {
610                    let predicate_result = match output {
611                        Output::Value(value) => ctx.vm.call_function(predicate.clone(), value),
612                        Output::ValuePair(a, b) => ctx
613                            .vm
614                            .call_function(predicate.clone(), CallArgs::AsTuple(&[a, b])),
615                        Output::Error(error) => return Err(error),
616                    };
617
618                    match predicate_result {
619                        Ok(KValue::Bool(result)) => {
620                            if result {
621                                return Ok(i.into());
622                            }
623                        }
624                        Ok(unexpected) => {
625                            return unexpected_type(
626                                "a Bool to be returned from the predicate",
627                                &unexpected,
628                            );
629                        }
630                        Err(error) => return Err(error),
631                    }
632                }
633
634                Ok(KValue::Null)
635            }
636            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
637        }
638    });
639
640    result.add_fn("product", |ctx| {
641        let (iterable, initial_value) = {
642            let expected_error = "|Iterable|";
643
644            match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
645                (iterable, []) => (iterable.clone(), KValue::Number(1.into())),
646                (iterable, [initial_value]) => (iterable.clone(), initial_value.clone()),
647                (instance, args) => {
648                    return unexpected_args_after_instance(expected_error, instance, args);
649                }
650            }
651        };
652
653        fold_with_operator(ctx.vm, iterable, initial_value, BinaryOp::Multiply)
654    });
655
656    result.add_fn("repeat", |ctx| {
657        let instance = ctx.instance();
658        if !matches!(instance, &KValue::Null) && instance.type_as_string() != MODULE_NAME {
659            return runtime_error!("iterator.repeat must be used as a standalone function");
660        }
661
662        match ctx.args() {
663            [value] => {
664                let result = generators::Repeat::new(value.clone());
665                Ok(KIterator::new(result).into())
666            }
667            [value, KValue::Number(n)] if *n >= 0.0 => {
668                let result = generators::RepeatN::new(value.clone(), n.into());
669                Ok(KIterator::new(result).into())
670            }
671            unexpected => unexpected_args("|Any|, or |Any, Number >= 0|", unexpected),
672        }
673    });
674
675    result.add_fn("reversed", |ctx| {
676        let expected_error = "|Iterable|";
677
678        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
679            (iterable, []) => {
680                let iterable = iterable.clone();
681                match adaptors::Reversed::new(ctx.vm.make_iterator(iterable)?) {
682                    Ok(result) => Ok(KIterator::new(result).into()),
683                    Err(e) => runtime_error!("iterator.reversed: {}", e),
684                }
685            }
686            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
687        }
688    });
689
690    result.add_fn("skip", |ctx| {
691        let expected_error = "|Iterable, Number >= 0|";
692
693        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
694            (iterable, [KValue::Number(n)]) if *n >= 0.0 => {
695                let iterable = iterable.clone();
696                let n = *n;
697                let result = adaptors::Skip::new(ctx.vm.make_iterator(iterable)?, n.into());
698                Ok(KIterator::new(result).into())
699            }
700            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
701        }
702    });
703
704    result.add_fn("step", |ctx| {
705        let expected_error = "|Iterable, Number|";
706
707        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
708            (iterable, [KValue::Number(n)]) => {
709                if *n > 0 {
710                    let iterable = iterable.clone();
711                    let step_size = n.into();
712                    match adaptors::Step::new(ctx.vm.make_iterator(iterable)?, step_size) {
713                        Ok(result) => Ok(KIterator::new(result).into()),
714                        Err(e) => runtime_error!("iterator.step: {}", e),
715                    }
716                } else {
717                    runtime_error!("expected a non-negative number")
718                }
719            }
720            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
721        }
722    });
723
724    result.add_fn("sum", |ctx| {
725        let (iterable, initial_value) = {
726            let expected_error = "|Iterable|";
727
728            match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
729                (iterable, []) => (iterable.clone(), KValue::Number(0.into())),
730                (iterable, [initial_value]) => (iterable.clone(), initial_value.clone()),
731                (instance, args) => {
732                    return unexpected_args_after_instance(expected_error, instance, args);
733                }
734            }
735        };
736
737        fold_with_operator(ctx.vm, iterable, initial_value, BinaryOp::Add)
738    });
739
740    result.add_fn("take", |ctx| {
741        let expected_error = "|Iterable, Number >= 0|, or |Iterable, |Any| -> Bool|";
742
743        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
744            (iterable, [KValue::Number(n)]) if *n >= 0.0 => {
745                let iterable = iterable.clone();
746                let n = *n;
747                let result = adaptors::Take::new(ctx.vm.make_iterator(iterable)?, n.into());
748                Ok(KIterator::new(result).into())
749            }
750            (iterable, [predicate]) if predicate.is_callable() => {
751                let iterable = iterable.clone();
752                let predicate = predicate.clone();
753                let result =
754                    adaptors::TakeWhile::new(ctx.vm.make_iterator(iterable)?, predicate, ctx.vm);
755                Ok(KIterator::new(result).into())
756            }
757            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
758        }
759    });
760
761    result.add_fn("to_list", |ctx| {
762        let expected_error = "|Iterable|";
763
764        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
765            (iterable, []) => {
766                let iterable = iterable.clone();
767                let iterator = ctx.vm.make_iterator(iterable)?;
768                let (size_hint, _) = iterator.size_hint();
769                let mut result = ValueVec::with_capacity(size_hint);
770
771                for output in iterator.map(collect_pair) {
772                    match output {
773                        Output::Value(value) => result.push(value),
774                        Output::Error(error) => return Err(error),
775                        _ => unreachable!(),
776                    }
777                }
778
779                Ok(KValue::List(KList::with_data(result)))
780            }
781            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
782        }
783    });
784
785    result.add_fn("to_map", |ctx| {
786        let expected_error = "|Iterable|";
787
788        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
789            (iterable, []) => {
790                let iterable = iterable.clone();
791                let iterator = ctx.vm.make_iterator(iterable)?;
792                let (size_hint, _) = iterator.size_hint();
793                let mut result = ValueMap::with_capacity(size_hint);
794
795                for output in iterator {
796                    let (key, value) = match output {
797                        Output::ValuePair(key, value) => (key, value),
798                        Output::Value(KValue::Tuple(t)) if t.len() == 2 => {
799                            let key = t[0].clone();
800                            let value = t[1].clone();
801                            (key, value)
802                        }
803                        Output::Value(value) => (value, KValue::Null),
804                        Output::Error(error) => return Err(error),
805                    };
806
807                    result.insert(ValueKey::try_from(key)?, value);
808                }
809
810                Ok(KValue::Map(KMap::with_data(result)))
811            }
812            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
813        }
814    });
815
816    result.add_fn("to_string", |ctx| {
817        let expected_error = "|Iterable|";
818
819        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
820            (iterable, []) => {
821                let iterable = iterable.clone();
822                let iterator = ctx.vm.make_iterator(iterable)?;
823                let (size_hint, _) = iterator.size_hint();
824                let mut display_context = DisplayContext::with_vm_and_capacity(ctx.vm, size_hint);
825                for output in iterator.map(collect_pair) {
826                    match output {
827                        Output::Value(KValue::Str(s)) => display_context.append(s),
828                        Output::Value(value) => value.display(&mut display_context)?,
829                        Output::Error(error) => return Err(error),
830                        _ => unreachable!(),
831                    };
832                }
833
834                Ok(display_context.result().into())
835            }
836            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
837        }
838    });
839
840    result.add_fn("to_tuple", |ctx| {
841        let expected_error = "|Iterable|";
842
843        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
844            (iterable, []) => {
845                let iterable = iterable.clone();
846                let iterator = ctx.vm.make_iterator(iterable)?;
847                let (size_hint, _) = iterator.size_hint();
848                let mut result = Vec::with_capacity(size_hint);
849
850                for output in iterator.map(collect_pair) {
851                    match output {
852                        Output::Value(value) => result.push(value),
853                        Output::Error(error) => return Err(error),
854                        _ => unreachable!(),
855                    }
856                }
857
858                Ok(KValue::Tuple(result.into()))
859            }
860            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
861        }
862    });
863
864    result.add_fn("windows", |ctx| {
865        let expected_error = "|Iterable, Number|";
866
867        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
868            (iterable, [KValue::Number(n)]) => {
869                let iterable = iterable.clone();
870                let n = *n;
871                match adaptors::Windows::new(ctx.vm.make_iterator(iterable)?, n.into()) {
872                    Ok(result) => Ok(KIterator::new(result).into()),
873                    Err(e) => runtime_error!("iterator.windows: {}", e),
874                }
875            }
876            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
877        }
878    });
879
880    result.add_fn("zip", |ctx| {
881        let expected_error = "|Iterable|";
882
883        match ctx.instance_and_args(KValue::is_iterable, expected_error)? {
884            (iterable_a, [iterable_b]) if iterable_b.is_iterable() => {
885                let iterable_a = iterable_a.clone();
886                let iterable_b = iterable_b.clone();
887                let result = adaptors::Zip::new(
888                    ctx.vm.make_iterator(iterable_a)?,
889                    ctx.vm.make_iterator(iterable_b)?,
890                );
891                Ok(KIterator::new(result).into())
892            }
893            (instance, args) => unexpected_args_after_instance(expected_error, instance, args),
894        }
895    });
896
897    result
898}
899
900pub(crate) fn collect_pair(iterator_output: Output) -> Output {
901    match iterator_output {
902        Output::ValuePair(first, second) => {
903            Output::Value(KValue::Tuple(vec![first, second].into()))
904        }
905        _ => iterator_output,
906    }
907}
908
909pub(crate) fn iter_output_to_result(iterator_output: Option<Output>) -> Result<Option<KValue>> {
910    let output = match iterator_output {
911        Some(Output::Value(value)) => Some(value),
912        Some(Output::ValuePair(first, second)) => Some(KValue::Tuple(vec![first, second].into())),
913        Some(Output::Error(error)) => return Err(error),
914        None => None,
915    };
916
917    Ok(output)
918}
919
920/// The output type used by operations like `iterator.next()` and `next_back()`
921#[derive(Clone, KotoCopy, KotoType)]
922pub struct IteratorOutput(KValue);
923
924#[koto_impl(runtime = crate)]
925impl IteratorOutput {
926    /// Returns the wrapped output value
927    #[koto_method]
928    pub fn get(&self) -> KValue {
929        self.0.clone()
930    }
931}
932
933impl KotoObject for IteratorOutput {
934    fn display(&self, ctx: &mut DisplayContext) -> Result<()> {
935        ctx.append(Self::type_static());
936        ctx.append('(');
937
938        let mut wrapped_ctx = DisplayContext::default();
939        self.0.display(&mut wrapped_ctx)?;
940        ctx.append(wrapped_ctx.result());
941
942        ctx.append(')');
943        Ok(())
944    }
945}
946
947impl From<KValue> for IteratorOutput {
948    fn from(value: KValue) -> Self {
949        Self(value)
950    }
951}
952
953impl From<IteratorOutput> for KValue {
954    fn from(output: IteratorOutput) -> Self {
955        KObject::from(output).into()
956    }
957}
958
959fn fold_with_operator(
960    vm: &mut KotoVm,
961    iterable: KValue,
962    initial_value: KValue,
963    operator: BinaryOp,
964) -> Result<KValue> {
965    let mut result = initial_value;
966
967    for output in vm.make_iterator(iterable)?.map(collect_pair) {
968        match output {
969            Output::Value(rhs_value) => {
970                result = vm.run_binary_op(operator, result, rhs_value)?;
971            }
972            Output::Error(error) => return Err(error),
973            _ => unreachable!(),
974        }
975    }
976
977    Ok(result)
978}
979
980fn run_iterator_comparison(
981    vm: &mut KotoVm,
982    iterable: KValue,
983    invert_result: InvertResult,
984) -> Result<KValue> {
985    let mut result: Option<KValue> = None;
986
987    for iter_output in vm.make_iterator(iterable)?.map(collect_pair) {
988        match iter_output {
989            Output::Value(value) => {
990                result = Some(match result {
991                    Some(result) => {
992                        compare_values(vm, result.clone(), value.clone(), invert_result)?
993                    }
994                    None => value,
995                })
996            }
997            Output::Error(error) => return Err(error),
998            _ => unreachable!(),
999        }
1000    }
1001
1002    Ok(result.unwrap_or_default())
1003}
1004
1005fn run_iterator_comparison_by_key(
1006    vm: &mut KotoVm,
1007    iterable: KValue,
1008    key_fn: KValue,
1009    invert_result: InvertResult,
1010) -> Result<KValue> {
1011    let mut result_and_key: Option<(KValue, KValue)> = None;
1012
1013    for iter_output in vm.make_iterator(iterable)?.map(collect_pair) {
1014        match iter_output {
1015            Output::Value(value) => {
1016                let key = vm.call_function(key_fn.clone(), value.clone())?;
1017                let value_and_key = (value, key);
1018
1019                result_and_key = Some(match result_and_key {
1020                    Some(result_and_key) => {
1021                        compare_values_with_key(vm, result_and_key, value_and_key, invert_result)?
1022                    }
1023                    None => value_and_key,
1024                });
1025            }
1026            Output::Error(error) => return Err(error),
1027            _ => unreachable!(),
1028        }
1029    }
1030
1031    Ok(result_and_key.map_or(KValue::Null, |(value, _)| value))
1032}
1033
1034// Compares two values using BinaryOp::Less
1035//
1036// Returns the lesser of the two values, unless `invert_result` is set to Yes
1037fn compare_values(
1038    vm: &mut KotoVm,
1039    a: KValue,
1040    b: KValue,
1041    invert_result: InvertResult,
1042) -> Result<KValue> {
1043    use InvertResult::*;
1044    use KValue::Bool;
1045
1046    let comparison_result = vm.run_binary_op(BinaryOp::Less, a.clone(), b.clone())?;
1047
1048    match (comparison_result, invert_result) {
1049        (Bool(true), No) => Ok(a),
1050        (Bool(false), No) => Ok(b),
1051        (Bool(true), Yes) => Ok(b),
1052        (Bool(false), Yes) => Ok(a),
1053        (other, _) => runtime_error!(
1054            "Expected Bool from '<' comparison, found '{}'",
1055            other.type_as_string()
1056        ),
1057    }
1058}
1059
1060// Compares two values using BinaryOp::Less
1061//
1062// Returns the lesser of the two values, unless `invert_result` is set to Yes
1063fn compare_values_with_key(
1064    vm: &mut KotoVm,
1065    a_and_key: (KValue, KValue),
1066    b_and_key: (KValue, KValue),
1067    invert_result: InvertResult,
1068) -> Result<(KValue, KValue)> {
1069    use InvertResult::*;
1070    use KValue::Bool;
1071
1072    let comparison_result =
1073        vm.run_binary_op(BinaryOp::Less, a_and_key.1.clone(), b_and_key.1.clone())?;
1074
1075    match (comparison_result, invert_result) {
1076        (Bool(true), No) => Ok(a_and_key),
1077        (Bool(false), No) => Ok(b_and_key),
1078        (Bool(true), Yes) => Ok(b_and_key),
1079        (Bool(false), Yes) => Ok(a_and_key),
1080        (other, _) => runtime_error!(
1081            "Expected Bool from '<' comparison, found '{}'",
1082            other.type_as_string()
1083        ),
1084    }
1085}
1086
1087#[derive(Clone, Copy)]
1088enum InvertResult {
1089    Yes,
1090    No,
1091}