Final touches

3 files changed, 59 insertions, 75 deletions

diff --git a/src/syntax/ast.rs b/src/syntax/ast.rs
index 067bd6da..288c749a 100644
--- a/src/syntax/ast.rs
+++ b/src/syntax/ast.rs
@@ -68,11 +68,8 @@ impl Markup {
             NodeKind::EnDash => Some(MarkupNode::Text("\u{2013}".into())),
             NodeKind::EmDash => Some(MarkupNode::Text("\u{2014}".into())),
             NodeKind::NonBreakingSpace => Some(MarkupNode::Text("\u{00A0}".into())),
-            NodeKind::Raw(raw) => Some(MarkupNode::Raw(RawNode {
-                block: raw.block,
-                lang: raw.lang.clone(),
-                text: raw.text.clone(),
-            })),
+            NodeKind::Math(math) => Some(MarkupNode::Math(math.as_ref().clone())),
+            NodeKind::Raw(raw) => Some(MarkupNode::Raw(raw.as_ref().clone())),
             NodeKind::Heading => node.cast().map(MarkupNode::Heading),
             NodeKind::List => node.cast().map(MarkupNode::List),
             NodeKind::Enum => node.cast().map(MarkupNode::Enum),
@@ -98,6 +95,8 @@ pub enum MarkupNode {
     Text(EcoString),
     /// A raw block with optional syntax highlighting: `` `...` ``.
     Raw(RawNode),
+    /// A math formula: `$a^2 = b^2 + c^2$`.
+    Math(MathNode),
     /// A section heading: `= Introduction`.
     Heading(HeadingNode),
     /// An item in an unordered list: `- ...`.
@@ -121,6 +120,16 @@ pub struct RawNode {
     pub block: bool,
 }
 
+/// A math formula: `$a^2 + b^2 = c^2$`.
+#[derive(Debug, Clone, PartialEq)]
+pub struct MathNode {
+    /// The formula between the dollars / brackets.
+    pub formula: EcoString,
+    /// Whether the formula is display-level, that is, it is surrounded by
+    /// `$[..]$`.
+    pub display: bool,
+}
+
 node! {
     /// A section heading: `= Introduction`.
     HeadingNode: Heading
@@ -133,12 +142,8 @@ impl HeadingNode {
     }
 
     /// The section depth (numer of equals signs).
-    pub fn level(&self) -> u8 {
-        self.0
-            .children()
-            .filter(|n| n.kind() == &NodeKind::Eq)
-            .count()
-            .min(u8::MAX.into()) as u8
+    pub fn level(&self) -> usize {
+        self.0.children().filter(|n| n.kind() == &NodeKind::Eq).count()
     }
 }
 
diff --git a/src/syntax/mod.rs b/src/syntax/mod.rs
index 0660d57b..ca6ed243 100644
--- a/src/syntax/mod.rs
+++ b/src/syntax/mod.rs
@@ -10,7 +10,7 @@ use std::rc::Rc;
 pub use pretty::*;
 pub use span::*;
 
-use self::ast::TypedNode;
+use self::ast::{MathNode, RawNode, TypedNode};
 use crate::diag::Error;
 use crate::geom::{AngularUnit, LengthUnit};
 use crate::source::SourceId;
@@ -178,7 +178,7 @@ impl From<GreenData> for Green {
 
 /// A owned wrapper for a green node with span information.
 ///
-/// Owned variant of [`RedRef`]. Can be [cast](Self::cast) to an AST nodes.
+/// Owned variant of [`RedRef`]. Can be [cast](Self::cast) to an AST node.
 #[derive(Clone, PartialEq)]
 pub struct RedNode {
     id: SourceId,
@@ -192,15 +192,6 @@ impl RedNode {
         Self { id, offset: 0, green: root.into() }
     }
 
-    /// Create a new red node from a node kind and a span.
-    pub fn from_data(kind: NodeKind, span: Span) -> Self {
-        Self {
-            id: span.source,
-            offset: span.start,
-            green: Green::Token(GreenData { kind, len: span.len() }),
-        }
-    }
-
     /// Convert to a borrowed representation.
     pub fn as_ref(&self) -> RedRef<'_> {
         RedRef {
@@ -527,13 +518,11 @@ pub enum NodeKind {
     EnumNumbering(Option<usize>),
     /// An item in an unordered list: `- ...`.
     List,
-    /// The bullet character of an item in an unordered list: `-`.
-    ListBullet,
     /// An arbitrary number of backticks followed by inner contents, terminated
     /// with the same number of backticks: `` `...` ``.
-    Raw(Rc<RawData>),
+    Raw(Rc<RawNode>),
     /// Dollar signs surrounding inner contents.
-    Math(Rc<MathData>),
+    Math(Rc<MathNode>),
     /// An identifier: `center`.
     Ident(EcoString),
     /// A boolean: `true`, `false`.
@@ -613,29 +602,6 @@ pub enum NodeKind {
     Unknown(EcoString),
 }
 
-/// Payload of a raw block node.
-#[derive(Debug, Clone, PartialEq)]
-pub struct RawData {
-    /// The raw text in the block.
-    pub text: EcoString,
-    /// The programming language of the raw text.
-    pub lang: Option<EcoString>,
-    /// The number of opening backticks.
-    pub backticks: u8,
-    /// Whether to display this as a block.
-    pub block: bool,
-}
-
-/// Payload of a math formula node.
-#[derive(Debug, Clone, PartialEq)]
-pub struct MathData {
-    /// 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 ErrorPos {
@@ -730,7 +696,6 @@ impl NodeKind {
             Self::Enum => "enumeration item",
             Self::EnumNumbering(_) => "enumeration item numbering",
             Self::List => "list item",
-            Self::ListBullet => "list bullet",
             Self::Raw(_) => "raw block",
             Self::Math(_) => "math formula",
             Self::Ident(_) => "identifier",
diff --git a/src/syntax/pretty.rs b/src/syntax/pretty.rs
index 9e4510b6..c453fb56 100644
--- a/src/syntax/pretty.rs
+++ b/src/syntax/pretty.rs
@@ -63,7 +63,6 @@ impl Printer {
             write_item(item, self);
             count += 1;
         }
-
         count
     }
 
@@ -99,6 +98,7 @@ impl Pretty for MarkupNode {
             Self::Emph => p.push('_'),
             Self::Text(text) => p.push_str(text),
             Self::Raw(raw) => raw.pretty(p),
+            Self::Math(math) => math.pretty(p),
             Self::Heading(heading) => heading.pretty(p),
             Self::List(list) => list.pretty(p),
             Self::Enum(enum_) => enum_.pretty(p),
@@ -168,6 +168,20 @@ impl Pretty for RawNode {
     }
 }
 
+impl Pretty for MathNode {
+    fn pretty(&self, p: &mut Printer) {
+        p.push('$');
+        if self.display {
+            p.push('[');
+        }
+        p.push_str(&self.formula);
+        if self.display {
+            p.push(']');
+        }
+        p.push('$');
+    }
+}
+
 impl Pretty for HeadingNode {
     fn pretty(&self, p: &mut Printer) {
         for _ in 0 .. self.level() {
@@ -253,12 +267,9 @@ impl Pretty for ArrayExpr {
 impl Pretty for DictExpr {
     fn pretty(&self, p: &mut Printer) {
         p.push('(');
-
-        let mut items = self.items().peekable();
-        if items.peek().is_none() {
+        let len = p.join(self.items(), ", ", |named, p| named.pretty(p));
+        if len == 0 {
             p.push(':');
-        } else {
-            p.join(items, ", ", |named, p| named.pretty(p));
         }
         p.push(')');
     }
@@ -291,13 +302,11 @@ impl Pretty for GroupExpr {
 impl Pretty for BlockExpr {
     fn pretty(&self, p: &mut Printer) {
         p.push('{');
-
-        let exprs: Vec<_> = self.exprs().collect();
-        if exprs.len() > 1 {
+        if self.exprs().count() > 1 {
             p.push(' ');
         }
-        p.join(&exprs, "; ", |expr, p| expr.pretty(p));
-        if exprs.len() > 1 {
+        let len = p.join(self.exprs(), "; ", |expr, p| expr.pretty(p));
+        if len > 1 {
             p.push(' ');
         }
         p.push('}');
@@ -348,17 +357,17 @@ impl Pretty for CallExpr {
         };
 
         let args: Vec<_> = self.args().items().collect();
-
-        if let Some(Expr::Template(template)) = args
-            .last()
-            .and_then(|x| if let CallArg::Pos(arg) = x { Some(arg) } else { None })
-        {
-            if args.len() > 1 {
-                write_args(&args[0 .. args.len() - 1]);
+        match args.as_slice() {
+            // This can be moved behind the arguments.
+            //
+            // Example: Transforms "#v(a, [b])" => "#v(a)[b]".
+            [head @ .., CallArg::Pos(Expr::Template(template))] => {
+                if !head.is_empty() {
+                    write_args(head);
+                }
+                template.pretty(p);
             }
-            template.pretty(p);
-        } else {
-            write_args(&args);
+            items => write_args(items),
         }
     }
 }
@@ -423,12 +432,12 @@ impl Pretty for LetExpr {
     fn pretty(&self, p: &mut Printer) {
         p.push_str("let ");
         self.binding().pretty(p);
-        if let Some(Expr::Closure(closure)) = &self.init() {
+        if let Some(Expr::Closure(closure)) = self.init() {
             p.push('(');
             p.join(closure.params(), ", ", |item, p| item.pretty(p));
             p.push_str(") = ");
             closure.body().pretty(p);
-        } else if let Some(init) = &self.init() {
+        } else if let Some(init) = self.init() {
             p.push_str(" = ");
             init.pretty(p);
         }
@@ -441,7 +450,7 @@ impl Pretty for IfExpr {
         self.condition().pretty(p);
         p.push(' ');
         self.if_body().pretty(p);
-        if let Some(expr) = &self.else_body() {
+        if let Some(expr) = self.else_body() {
             p.push_str(" else ");
             expr.pretty(p);
         }
@@ -525,7 +534,7 @@ mod tests {
     #[track_caller]
     fn test_parse(src: &str, expected: &str) {
         let source = SourceFile::detached(src);
-        let ast: Markup = source.ast().unwrap();
+        let ast = source.ast().unwrap();
         let found = pretty(&ast);
         if found != expected {
             println!("tree:     {:#?}", ast);
@@ -563,6 +572,11 @@ mod tests {
         test_parse("``` 1```", "`1`");
         test_parse("``` 1 ```", "`1 `");
         test_parse("```` ` ````", "``` ` ```");
+
+        // Math node.
+        roundtrip("$$");
+        roundtrip("$a+b$");
+        roundtrip("$[ a^2 + b^2 = c^2 ]$");
     }
 
     #[test]