Merge pull request #227 from koto-lang/peekable-iterator

Introduce iterator.peekable()

CHANGELOG.md

@@ -43,7 +43,9 @@ The Koto project adheres to
       # true
       ```
   - `is_inclusive` and `intersection` have been added.
-- `iterator.next_back` has been added.
+- `iterator` additions:
+  - `iterator.next_back`
+  - `iterator.peekable`
 
 #### Internals
 

core/bytecode/src/compiler.rs

@@ -2684,10 +2684,10 @@ impl Compiler {
         // Do we need to modify the accessed value?
         if access_assignment {
             let Some(rhs) = rhs else {
-                return compiler_error!(self, "compile_lookup: Missing rhs")
+                return compiler_error!(self, "compile_lookup: Missing rhs");
             };
             let Some(rhs_op) = rhs_op else {
-                return compiler_error!(self, "compile_lookup: Missing rhs_op")
+                return compiler_error!(self, "compile_lookup: Missing rhs_op");
             };
 
             self.push_op(rhs_op, &[access_register, rhs]);

core/runtime/src/core/iterator.rs

@@ -2,6 +2,7 @@
 
 pub mod adaptors;
 pub mod generators;
+pub mod peekable;
 
 use crate::{prelude::*, ValueIteratorOutput as Output};
 
@@ -468,12 +469,7 @@ pub fn make_module() -> ValueMap {
             }
         };
 
-        match iter.next().map(collect_pair) {
-            Some(Output::Value(value)) => Ok(value),
-            Some(Output::Error(error)) => Err(error),
-            None => Ok(Value::Null),
-            _ => unreachable!(),
-        }
+        iter_output_to_result(iter.next())
     });
 
     result.add_fn("next_back", |vm, args| {
@@ -485,12 +481,15 @@ pub fn make_module() -> ValueMap {
             }
         };
 
-        match iter.next_back().map(collect_pair) {
-            Some(Output::Value(value)) => Ok(value),
-            Some(Output::Error(error)) => Err(error),
-            None => Ok(Value::Null),
-            _ => unreachable!(),
+        iter_output_to_result(iter.next_back())
+    });
+
+    result.add_fn("peekable", |vm, args| match vm.get_args(args) {
+        [iterable] if iterable.is_iterable() => {
+            let iterable = iterable.clone();
+            Ok(peekable::Peekable::make_value(vm.make_iterator(iterable)?))
         }
+        unexpected => type_error_with_slice("an iterable value as argument", unexpected),
     });
 
     result.add_fn("position", |vm, args| match vm.get_args(args) {
@@ -757,6 +756,15 @@ pub(crate) fn collect_pair(iterator_output: Output) -> Output {
     }
 }
 
+pub(crate) fn iter_output_to_result(iterator_output: Option<Output>) -> RuntimeResult {
+    match iterator_output.map(collect_pair) {
+        Some(Output::Value(value)) => Ok(value),
+        Some(Output::ValuePair(first, second)) => Ok(Value::Tuple(vec![first, second].into())),
+        Some(Output::Error(error)) => Err(error),
+        None => Ok(Value::Null),
+    }
+}
+
 fn fold_with_operator(
     vm: &mut Vm,
     iterable: Value,

core/runtime/src/core/iterator/peekable.rs

@@ -0,0 +1,104 @@
+//! A double-ended peekable iterator for Koto
+
+use {super::iter_output_to_result, crate::prelude::*};
+
+/// A double-ended peekable iterator for Koto
+#[derive(Clone, Debug)]
+pub struct Peekable {
+    iter: ValueIterator,
+    peeked_front: Option<Value>,
+    peeked_back: Option<Value>,
+}
+
+impl ExternalData for Peekable {
+    fn data_type(&self) -> ValueString {
+        PEEKABLE_TYPE_STRING.with(|x| x.clone())
+    }
+
+    fn make_copy(&self) -> PtrMut<dyn ExternalData> {
+        make_data_ptr(self.clone())
+    }
+}
+
+impl Peekable {
+    /// Initializes a Peekable that wraps the given iterator
+    pub fn new(iter: ValueIterator) -> Self {
+        Self {
+            iter,
+            peeked_front: None,
+            peeked_back: None,
+        }
+    }
+
+    /// Makes an instance of Peekable along with a meta map that allows it be used as a Koto Value
+    pub fn make_value(iter: ValueIterator) -> Value {
+        Value::External(External::with_shared_meta_map(
+            Self::new(iter),
+            PEEKABLE_META.with(|meta| meta.clone()),
+        ))
+    }
+
+    fn peek(&mut self) -> RuntimeResult {
+        match self.peeked_front.clone() {
+            Some(peeked) => Ok(peeked),
+            None => match self.next()? {
+                Value::Null => Ok(Value::Null),
+                peeked => {
+                    self.peeked_front = Some(peeked.clone());
+                    Ok(peeked)
+                }
+            },
+        }
+    }
+
+    fn peek_back(&mut self) -> RuntimeResult {
+        match self.peeked_back.clone() {
+            Some(peeked) => Ok(peeked),
+            None => match self.next_back()? {
+                Value::Null => Ok(Value::Null),
+                peeked => {
+                    self.peeked_back = Some(peeked.clone());
+                    Ok(peeked)
+                }
+            },
+        }
+    }
+
+    fn next(&mut self) -> RuntimeResult {
+        match self.peeked_front.take() {
+            Some(value) => Ok(value),
+            None => match iter_output_to_result(self.iter.next())? {
+                Value::Null => Ok(self.peeked_back.take().unwrap_or(Value::Null)),
+                other => Ok(other),
+            },
+        }
+    }
+
+    fn next_back(&mut self) -> RuntimeResult {
+        match self.peeked_back.take() {
+            Some(value) => Ok(value),
+            None => match iter_output_to_result(self.iter.next_back())? {
+                Value::Null => Ok(self.peeked_front.take().unwrap_or(Value::Null)),
+                other => Ok(other),
+            },
+        }
+    }
+}
+
+fn make_peekable_meta_map() -> PtrMut<MetaMap> {
+    use UnaryOp::*;
+
+    MetaMapBuilder::<Peekable>::new(PEEKABLE_TYPE)
+        .function("peek", |context| context.data_mut()?.peek())
+        .function("peek_back", |context| context.data_mut()?.peek_back())
+        .function(Next, |context| context.data_mut()?.next())
+        .function(NextBack, |context| context.data_mut()?.next_back())
+        .build()
+}
+
+const PEEKABLE_TYPE: &str = "Peekable";
+
+thread_local! {
+    static PEEKABLE_META: PtrMut<MetaMap> = make_peekable_meta_map();
+    static PEEKABLE_TYPE_STRING: ValueString = PEEKABLE_TYPE.into();
+}

core/runtime/src/value.rs

@@ -133,7 +133,12 @@ impl Value {
         use Value::*;
         match self {
             Range(_) | List(_) | Tuple(_) | Map(_) | Str(_) | Iterator(_) => true,
-            External(v) if v.contains_meta_key(&UnaryOp::Iterator.into()) => true,
+            External(v)
+                if v.contains_meta_key(&UnaryOp::Iterator.into())
+                    || v.contains_meta_key(&UnaryOp::Next.into()) =>
+            {
+                true
+            }
             _ => false,
         }
     }

core/runtime/src/vm.rs

@@ -2440,7 +2440,9 @@ impl Vm {
                 Some(value) => self.set_register(result_register, value),
                 None => {
                     // Iterator fallback?
-                    if e.contains_meta_key(&UnaryOp::Iterator.into()) {
+                    if e.contains_meta_key(&UnaryOp::Iterator.into())
+                        || e.contains_meta_key(&UnaryOp::Next.into())
+                    {
                         let iterator_op = self.get_core_op(
                             &key,
                             &self.context.core_lib.iterator,

core/runtime/tests/iterator_tests.rs

@@ -1,6 +1,6 @@
 mod runtime_test_utils;
 
-use crate::runtime_test_utils::*;
+use {crate::runtime_test_utils::*, koto_runtime::Value};
 
 mod iterator {
     use super::*;
@@ -166,6 +166,101 @@ y.next()
         }
     }
 
+    mod peekable {
+        use super::*;
+        use Value::Null;
+
+        #[test]
+        fn peek() {
+            use Value::Null;
+
+            let script = "
+i = (1, 2, 3).peekable()
+result = []
+result.push i.peek() # 1
+result.push i.next() # 1
+result.push i.next() # 2
+result.push i.peek() # 3
+result.push i.next() # 3
+result.push i.peek() # null
+result.push i.next() # null
+result
+";
+            test_script(
+                script,
+                value_list(&[1.into(), 1.into(), 2.into(), 3.into(), 3.into(), Null, Null]),
+            );
+        }
+
+        #[test]
+        fn peek_back_forwards() {
+            let script = "
+i = (1, 2, 3).peekable()
+result = []
+result.push i.peek() # 1
+result.push i.peek_back() # 3
+result.push i.peek_back() # 3
+result.push i.next() # 1
+result.push i.next() # 2
+result.push i.peek() # 3
+result.push i.next() # 3
+result.push i.next() # null
+result.push i.next_back() # null
+result.push i.peek_back() # null
+result
+";
+            test_script(
+                script,
+                value_list(&[
+                    1.into(),
+                    3.into(),
+                    3.into(),
+                    1.into(),
+                    2.into(),
+                    3.into(),
+                    3.into(),
+                    Null,
+                    Null,
+                    Null,
+                ]),
+            );
+        }
+
+        #[test]
+        fn peek_back_backwards() {
+            let script = "
+i = (1, 2, 3).peekable()
+result = []
+result.push i.peek() # 1
+result.push i.peek_back() # 3
+result.push i.peek_back() # 3
+result.push i.next_back() # 3
+result.push i.next_back() # 2
+result.push i.peek_back() # 1
+result.push i.next_back() # 1
+result.push i.peek_back() # null
+result.push i.next_back() # null
+result.push i.next() # null
+result
+";
+            test_script(
+                script,
+                value_list(&[
+                    1.into(),
+                    3.into(),
+                    3.into(),
+                    3.into(),
+                    2.into(),
+                    1.into(),
+                    1.into(),
+                    Null,
+                    Null,
+                    Null,
+                ]),
+            );
+        }
+    }
+
     mod skip {
         use super::*;
 

docs/core_lib/iterator.md

@@ -575,6 +575,64 @@ check! null
 - [`iterator.next`](#next)
 - [`iterator.reversed`](#reversed)
 
+## peekable
+
+```kototype
+|Iterable| -> Peekable
+```
+
+Wraps the given iterable value in a peekable iterator.
+
+### Peekable.peek
+
+Returns the next value from the iterator without advancing it. 
+The peeked value is cached until the iterator is advanced.
+
+#### Example
+
+```koto
+x = 'abcdef'.peekable()
+print! x.peek()
+check! a
+print! x.peek()
+check! a
+print! x.next()
+check! a
+print! x.peek()
+check! b
+```
+
+#### See Also
+
+- [`iterator.next`](#next)
+
+### Peekable.peek_back
+
+Returns the next value from the end of the iterator without advancing it. 
+The peeked value is cached until the iterator is advanced.
+
+#### Example
+
+```koto
+x = 'abcdef'.peekable()
+print! x.peek_back()
+check! f
+print! x.next_back()
+check! f
+print! x.peek()
+check! a
+print! x.peek_back()
+check! e
+print! x.next_back()
+check! e
+print! x.next()
+check! a
+```
+
+#### See Also
+
+- [`iterator.peek_back`](#peek_back)
+
 ## position
 
 ```kototype