Add iterators to the new AST rows

core/src/bytecode/ast/typ.rs

@@ -1,7 +1,8 @@
 //! Representation of Nickel types in the AST.
 
 use super::{Ast, AstAlloc, TermPos};
-use crate::{traverse::*, typ as mainline_typ};
+use crate::{identifier::Ident, traverse::*, typ as mainline_typ};
+use iter::*;
 pub use mainline_typ::{EnumRowF, EnumRowsF, RecordRowF, RecordRowsF, TypeF};
 
 use std::fmt;
@@ -218,6 +219,79 @@ impl<'ast> TraverseAlloc<'ast, Type<'ast>> for RecordRows<'ast> {
     }
 }
 
+impl<'ast> RecordRows<'ast> {
+    /// Find a nested binding in a record row type. The nested field is given as a list of
+    /// successive fields, that is, as a path. Return `None` if there is no such binding.
+    ///
+    /// # Example
+    ///
+    /// - self: `{a : {b : Number }}`
+    /// - path: `["a", "b"]`
+    /// - result: `Some(Number)`
+    pub fn find_path<'a>(&'a self, path: &[Ident]) -> Option<&'a RecordRow<'ast>> {
+        if path.is_empty() {
+            return None;
+        }
+
+        // While going through the record rows, we use this helper for recursion instead of
+        // `find_path`, to avoid cloning a lot of intermediate rows, and rather only clone the
+        // final one to return.
+        fn find_path_ref<'a, 'ast>(
+            rrows: &'a RecordRows<'ast>,
+            path: &[Ident],
+        ) -> Option<&'a RecordRow<'ast>> {
+            let next = rrows.iter().find_map(|item| match item {
+                RecordRowsItem::Row(row) if row.id.ident() == path[0] => Some(row),
+                _ => None,
+            });
+
+            if path.len() == 1 {
+                next
+            } else {
+                match next.map(|row| &row.typ.typ) {
+                    Some(TypeF::Record(rrows)) => find_path_ref(rrows, &path[1..]),
+                    _ => None,
+                }
+            }
+        }
+
+        find_path_ref(self, path)
+    }
+
+    /// Find the row with the given identifier in the record type. Return `None` if there is no such
+    /// row.
+    ///
+    /// Equivalent to `find_path(&[id])`.
+    pub fn find_row(&self, id: Ident) -> Option<&RecordRow<'ast>> {
+        self.find_path(&[id])
+    }
+
+    pub fn iter(&self) -> RecordRowsIter<Type<'ast>, RecordRows<'ast>> {
+        RecordRowsIter {
+            rrows: Some(self),
+            ty: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<'ast> EnumRows<'ast> {
+    /// Find the row with the given identifier in the enum type. Return `None` if there is no such
+    /// row.
+    pub fn find_row<'a>(&'a self, id: Ident) -> Option<&'a EnumRow<'ast>> {
+        self.iter().find_map(|row_item| match row_item {
+            EnumRowsItem::Row(row) if row.id.ident() == id => Some(row),
+            _ => None,
+        })
+    }
+
+    pub fn iter(&self) -> EnumRowsIter<Type<'ast>, EnumRows<'ast>> {
+        EnumRowsIter {
+            erows: Some(self),
+            ty: std::marker::PhantomData,
+        }
+    }
+}
+
 //TODO: get rid of this expensive implementation once we migrate pretty::*.
 impl fmt::Display for Type<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -227,3 +301,77 @@ impl fmt::Display for Type<'_> {
         write!(f, "{}", typ::Type::from_ast(self))
     }
 }
+
+pub mod iter {
+    use super::*;
+    use crate::identifier::LocIdent;
+
+    /// An iterator over the rows of a record type.
+    pub struct RecordRowsIter<'a, Ty, RRows> {
+        pub(crate) rrows: Option<&'a RRows>,
+        pub(crate) ty: std::marker::PhantomData<Ty>,
+    }
+
+    /// The item produced by an iterator over record rows.
+    pub enum RecordRowsItem<'a, Ty> {
+        TailDyn,
+        TailVar(&'a LocIdent),
+        Row(&'a RecordRowF<Ty>),
+    }
+
+    impl<'a, 'ast> Iterator for RecordRowsIter<'a, Type<'ast>, RecordRows<'ast>> {
+        type Item = RecordRowsItem<'a, &'ast Type<'ast>>;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.rrows.and_then(|next| match next.0 {
+                RecordRowsF::Empty => {
+                    self.rrows = None;
+                    None
+                }
+                RecordRowsF::TailDyn => {
+                    self.rrows = None;
+                    Some(RecordRowsItem::TailDyn)
+                }
+                RecordRowsF::TailVar(ref id) => {
+                    self.rrows = None;
+                    Some(RecordRowsItem::TailVar(id))
+                }
+                RecordRowsF::Extend { ref row, ref tail } => {
+                    self.rrows = Some(tail);
+                    Some(RecordRowsItem::Row(row))
+                }
+            })
+        }
+    }
+
+    pub struct EnumRowsIter<'a, Ty, ERows> {
+        pub(crate) erows: Option<&'a ERows>,
+        pub(crate) ty: std::marker::PhantomData<Ty>,
+    }
+
+    pub enum EnumRowsItem<'a, Ty> {
+        TailVar(&'a LocIdent),
+        Row(&'a EnumRowF<Ty>),
+    }
+
+    impl<'a, 'ast> Iterator for EnumRowsIter<'a, Type<'ast>, EnumRows<'ast>> {
+        type Item = EnumRowsItem<'a, &'ast Type<'ast>>;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.erows.and_then(|next| match next.0 {
+                EnumRowsF::Empty => {
+                    self.erows = None;
+                    None
+                }
+                EnumRowsF::TailVar(ref id) => {
+                    self.erows = None;
+                    Some(EnumRowsItem::TailVar(id))
+                }
+                EnumRowsF::Extend { ref row, ref tail } => {
+                    self.erows = Some(tail);
+                    Some(EnumRowsItem::Row(row))
+                }
+            })
+        }
+    }
+}