Tidy up AST

1 files changed, 133 insertions, 116 deletions

diff --git a/src/syntax/mod.rs b/src/syntax/mod.rs
index 022b51de..fc05ad50 100644
--- a/src/syntax/mod.rs
+++ b/src/syntax/mod.rs
@@ -5,7 +5,6 @@ mod pretty;
 mod span;
 
 use std::fmt::{self, Debug, Display, Formatter};
-use std::mem;
 use std::rc::Rc;
 
 pub use pretty::*;
@@ -40,14 +39,6 @@ impl Green {
         self.data().kind()
     }
 
-    /// Set the type of the node.
-    pub fn set_kind(&mut self, kind: NodeKind) {
-        match self {
-            Self::Node(node) => Rc::make_mut(node).data.set_kind(kind),
-            Self::Token(data) => data.set_kind(kind),
-        }
-    }
-
     /// The length of the node.
     pub fn len(&self) -> usize {
         self.data().len()
@@ -68,6 +59,18 @@ impl Green {
             Green::Token(_) => &[],
         }
     }
+
+    /// Change the type of the node.
+    pub fn convert(&mut self, kind: NodeKind) {
+        match self {
+            Self::Node(node) => {
+                let node = Rc::make_mut(node);
+                node.erroneous |= kind.is_error();
+                node.data.kind = kind;
+            }
+            Self::Token(data) => data.kind = kind,
+        }
+    }
 }
 
 impl Default for Green {
@@ -161,11 +164,6 @@ impl GreenData {
         &self.kind
     }
 
-    /// Set the type of the node.
-    pub fn set_kind(&mut self, kind: NodeKind) {
-        self.kind = kind;
-    }
-
     /// The length of the node.
     pub fn len(&self) -> usize {
         self.len
@@ -178,7 +176,82 @@ impl From<GreenData> for Green {
     }
 }
 
-/// A borrowed wrapper for a [`GreenNode`] with span information.
+/// A owned wrapper for a green node with span information.
+///
+/// Owned variant of [`RedRef`]. Can be [cast](Self::cast) to an AST nodes.
+#[derive(Clone, PartialEq)]
+pub struct RedNode {
+    id: SourceId,
+    offset: usize,
+    green: Green,
+}
+
+impl RedNode {
+    /// Create a new root node from a [`GreenNode`].
+    pub fn new_root(root: Rc<GreenNode>, id: SourceId) -> Self {
+        Self { id, offset: 0, green: root.into() }
+    }
+
+    /// Convert to a borrowed representation.
+    pub fn as_ref(&self) -> RedRef<'_> {
+        RedRef {
+            id: self.id,
+            offset: self.offset,
+            green: &self.green,
+        }
+    }
+
+    /// The type of the node.
+    pub fn kind(&self) -> &NodeKind {
+        self.as_ref().kind()
+    }
+
+    /// The length of the node.
+    pub fn len(&self) -> usize {
+        self.as_ref().len()
+    }
+
+    /// The span of the node.
+    pub fn span(&self) -> Span {
+        self.as_ref().span()
+    }
+
+    /// The error messages for this node and its descendants.
+    pub fn errors(&self) -> Vec<Error> {
+        self.as_ref().errors()
+    }
+
+    /// Convert the node to a typed AST node.
+    pub fn cast<T>(self) -> Option<T>
+    where
+        T: TypedNode,
+    {
+        self.as_ref().cast()
+    }
+
+    /// The children of the node.
+    pub fn children(&self) -> Children<'_> {
+        self.as_ref().children()
+    }
+
+    /// Get the first child that can cast to some AST type.
+    pub fn cast_first_child<T: TypedNode>(&self) -> Option<T> {
+        self.as_ref().cast_first_child()
+    }
+
+    /// Get the last child that can cast to some AST type.
+    pub fn cast_last_child<T: TypedNode>(&self) -> Option<T> {
+        self.as_ref().cast_last_child()
+    }
+}
+
+impl Debug for RedNode {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        self.as_ref().fmt(f)
+    }
+}
+
+/// A borrowed wrapper for a green node with span information.
 ///
 /// Borrowed variant of [`RedNode`]. Can be [cast](Self::cast) to an AST node.
 #[derive(Copy, Clone, PartialEq)]
@@ -213,30 +286,25 @@ impl<'a> RedRef<'a> {
         Span::new(self.id, self.offset, self.offset + self.green.len())
     }
 
-    /// Whether the node or its children contain an error.
-    pub fn erroneous(self) -> bool {
-        self.green.erroneous()
-    }
-
     /// The error messages for this node and its descendants.
     pub fn errors(self) -> Vec<Error> {
-        if !self.erroneous() {
+        if !self.green.erroneous() {
             return vec![];
         }
 
         match self.kind() {
             NodeKind::Error(pos, msg) => {
                 let span = match pos {
-                    ErrorPosition::Start => self.span().at_start(),
-                    ErrorPosition::Full => self.span(),
-                    ErrorPosition::End => self.span().at_end(),
+                    ErrorPos::Start => self.span().at_start(),
+                    ErrorPos::Full => self.span(),
+                    ErrorPos::End => self.span().at_end(),
                 };
 
                 vec![Error::new(span, msg.to_string())]
             }
             _ => self
                 .children()
-                .filter(|red| red.erroneous())
+                .filter(|red| red.green.erroneous())
                 .flat_map(|red| red.errors())
                 .collect(),
         }
@@ -251,34 +319,28 @@ impl<'a> RedRef<'a> {
     }
 
     /// The node's children.
-    pub fn children(self) -> impl Iterator<Item = RedRef<'a>> {
+    pub fn children(self) -> Children<'a> {
         let children = match &self.green {
             Green::Node(node) => node.children(),
             Green::Token(_) => &[],
         };
 
-        let mut cursor = self.offset;
-        children.iter().map(move |green| {
-            let offset = cursor;
-            cursor += green.len();
-            RedRef { id: self.id, offset, green }
-        })
-    }
-
-    /// Get the first child of some type.
-    pub(crate) fn typed_child(self, kind: &NodeKind) -> Option<RedRef<'a>> {
-        self.children()
-            .find(|x| mem::discriminant(x.kind()) == mem::discriminant(kind))
+        Children {
+            id: self.id,
+            iter: children.iter(),
+            front: self.offset,
+            back: self.offset + self.len(),
+        }
     }
 
     /// Get the first child that can cast to some AST type.
-    pub(crate) fn cast_first_child<T: TypedNode>(self) -> Option<T> {
+    pub fn cast_first_child<T: TypedNode>(self) -> Option<T> {
         self.children().find_map(RedRef::cast)
     }
 
     /// Get the last child that can cast to some AST type.
-    pub(crate) fn cast_last_child<T: TypedNode>(self) -> Option<T> {
-        self.children().filter_map(RedRef::cast).last()
+    pub fn cast_last_child<T: TypedNode>(self) -> Option<T> {
+        self.children().rev().find_map(RedRef::cast)
     }
 }
 
@@ -294,86 +356,41 @@ impl Debug for RedRef<'_> {
     }
 }
 
-/// A owned wrapper for a [`GreenNode`] with span information.
-///
-/// Owned variant of [`RedRef`]. Can be [cast](Self::cast) to an AST nodes.
-#[derive(Clone, PartialEq)]
-pub struct RedNode {
+/// An iterator over the children of a red node.
+pub struct Children<'a> {
     id: SourceId,
-    offset: usize,
-    green: Green,
+    iter: std::slice::Iter<'a, Green>,
+    front: usize,
+    back: usize,
 }
 
-impl RedNode {
-    /// Create a new root node from a [`GreenNode`].
-    pub fn new_root(root: Rc<GreenNode>, id: SourceId) -> Self {
-        Self { id, offset: 0, green: root.into() }
-    }
-
-    /// Convert to a borrowed representation.
-    pub fn as_ref(&self) -> RedRef<'_> {
-        RedRef {
-            id: self.id,
-            offset: self.offset,
-            green: &self.green,
-        }
-    }
-
-    /// The type of the node.
-    pub fn kind(&self) -> &NodeKind {
-        self.as_ref().kind()
-    }
-
-    /// The length of the node.
-    pub fn len(&self) -> usize {
-        self.as_ref().len()
-    }
-
-    /// The span of the node.
-    pub fn span(&self) -> Span {
-        self.as_ref().span()
-    }
-
-    /// The error messages for this node and its descendants.
-    pub fn errors(&self) -> Vec<Error> {
-        self.as_ref().errors()
-    }
+impl<'a> Iterator for Children<'a> {
+    type Item = RedRef<'a>;
 
-    /// Convert the node to a typed AST node.
-    pub fn cast<T>(self) -> Option<T>
-    where
-        T: TypedNode,
-    {
-        self.as_ref().cast()
-    }
-
-    /// The children of the node.
-    pub fn children(&self) -> impl Iterator<Item = RedRef<'_>> {
-        self.as_ref().children()
-    }
-
-    /// Get the first child of some type.
-    pub(crate) fn typed_child(&self, kind: &NodeKind) -> Option<RedNode> {
-        self.as_ref().typed_child(kind).map(RedRef::own)
-    }
-
-    /// Get the first child that can cast to some AST type.
-    pub(crate) fn cast_first_child<T: TypedNode>(&self) -> Option<T> {
-        self.as_ref().cast_first_child()
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next().map(|green| {
+            let offset = self.front;
+            self.front += green.len();
+            RedRef { id: self.id, offset, green }
+        })
     }
 
-    /// Get the last child that can cast to some AST type.
-    pub(crate) fn cast_last_child<T: TypedNode>(&self) -> Option<T> {
-        self.as_ref().cast_last_child()
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
     }
 }
 
-impl Debug for RedNode {
-    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-        self.as_ref().fmt(f)
+impl DoubleEndedIterator for Children<'_> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back().map(|green| {
+            self.back -= green.len();
+            RedRef { id: self.id, offset: self.back, green }
+        })
     }
 }
 
+impl ExactSizeIterator for Children<'_> {}
+
 /// All syntactical building blocks that can be part of a Typst document.
 ///
 /// Can be emitted as a token by the tokenizer or as part of a green node by
@@ -533,7 +550,7 @@ pub enum NodeKind {
     Array,
     /// A dictionary expression: `(thickness: 3pt, pattern: dashed)`.
     Dict,
-    /// A named argument: `thickness: 3pt`.
+    /// A named pair: `thickness: 3pt`.
     Named,
     /// A grouped expression: `(1 + 2)`.
     Group,
@@ -582,12 +599,12 @@ pub enum NodeKind {
     /// The comment can contain nested block comments.
     BlockComment,
     /// Tokens that appear in the wrong place.
-    Error(ErrorPosition, EcoString),
+    Error(ErrorPos, EcoString),
     /// Unknown character sequences.
     Unknown(EcoString),
 }
 
-/// Payload of a raw block: `` `...` ``.
+/// Payload of a raw block node.
 #[derive(Debug, Clone, PartialEq)]
 pub struct RawData {
     /// The raw text in the block.
@@ -600,19 +617,19 @@ pub struct RawData {
     pub block: bool,
 }
 
-/// Payload of a math formula: `$2pi + x$` or `$[f'(x) = x^2]$`.
+/// Payload of a math formula node.
 #[derive(Debug, Clone, PartialEq)]
 pub struct MathData {
-    /// The formula between the dollars.
+    /// The formula between the dollars / brackets.
     pub formula: EcoString,
     /// Whether the formula is display-level, that is, it is surrounded by
-    /// `$[..]`.
+    /// `$[..]$`.
     pub display: bool,
 }
 
 /// Where in a node an error should be annotated.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub enum ErrorPosition {
+pub enum ErrorPos {
     /// At the start of the node.
     Start,
     /// Over the full width of the node.