Math delimiter grouping

3 files changed, 52 insertions, 19 deletions

diff --git a/src/syntax/ast.rs b/src/syntax/ast.rs
index d70c4ae4..ceda2d57 100644
--- a/src/syntax/ast.rs
+++ b/src/syntax/ast.rs
@@ -117,6 +117,9 @@ pub enum Expr {
     Math(Math),
     /// An atom in a math formula: `x`, `+`, `12`.
     Atom(Atom),
+    /// A subsection in a math formula that is surrounded by matched delimiters:
+    /// `[x + y]`.
+    Delimited(Delimited),
     /// A base with optional sub- and superscripts in a math formula: `a_1^2`.
     Script(Script),
     /// A fraction in a math formula: `x/2`.
@@ -216,6 +219,7 @@ impl AstNode for Expr {
             SyntaxKind::TermItem => node.cast().map(Self::Term),
             SyntaxKind::Math => node.cast().map(Self::Math),
             SyntaxKind::Atom => node.cast().map(Self::Atom),
+            SyntaxKind::Delimited => node.cast().map(Self::Delimited),
             SyntaxKind::Script => node.cast().map(Self::Script),
             SyntaxKind::Frac => node.cast().map(Self::Frac),
             SyntaxKind::AlignPoint => node.cast().map(Self::AlignPoint),
@@ -275,6 +279,7 @@ impl AstNode for Expr {
             Self::Term(v) => v.as_untyped(),
             Self::Math(v) => v.as_untyped(),
             Self::Atom(v) => v.as_untyped(),
+            Self::Delimited(v) => v.as_untyped(),
             Self::Script(v) => v.as_untyped(),
             Self::Frac(v) => v.as_untyped(),
             Self::AlignPoint(v) => v.as_untyped(),
@@ -658,6 +663,19 @@ impl Atom {
 }
 
 node! {
+    /// A subsection in a math formula that is surrounded by matched delimiters:
+    /// `[x + y]`.
+    Delimited
+}
+
+impl Delimited {
+    /// The contents, including the delimiters.
+    pub fn exprs(&self) -> impl DoubleEndedIterator<Item = Expr> + '_ {
+        self.0.children().filter_map(Expr::cast_with_space)
+    }
+}
+
+node! {
     /// A base with an optional sub- and superscript in a formula: `a_1^2`.
     Script
 }
@@ -673,8 +691,7 @@ impl Script {
         self.0
             .children()
             .skip_while(|node| !matches!(node.kind(), SyntaxKind::Underscore))
-            .nth(1)
-            .map(|node| node.cast().expect("script node has invalid subscript"))
+            .find_map(SyntaxNode::cast)
     }
 
     /// The superscript.
@@ -682,8 +699,7 @@ impl Script {
         self.0
             .children()
             .skip_while(|node| !matches!(node.kind(), SyntaxKind::Hat))
-            .nth(1)
-            .map(|node| node.cast().expect("script node has invalid superscript"))
+            .find_map(SyntaxNode::cast)
     }
 }
 
diff --git a/src/syntax/kind.rs b/src/syntax/kind.rs
index 5928fa0a..206df911 100644
--- a/src/syntax/kind.rs
+++ b/src/syntax/kind.rs
@@ -61,6 +61,9 @@ pub enum SyntaxKind {
     Math,
     /// An atom in math: `x`, `+`, `12`.
     Atom,
+    /// A subsection in a math formula that is surrounded by matched delimiters:
+    /// `[x + y]`.
+    Delimited,
     /// A base with optional sub- and superscripts in math: `a_1^2`.
     Script,
     /// A fraction in math: `x/2`.
@@ -336,6 +339,7 @@ impl SyntaxKind {
             Self::TermItem => "term list item",
             Self::TermMarker => "term marker",
             Self::Math => "math formula",
+            Self::Delimited => "delimited math",
             Self::Atom => "math atom",
             Self::Script => "script",
             Self::Frac => "fraction",
diff --git a/src/syntax/parser.rs b/src/syntax/parser.rs
index a379500c..15839e18 100644
--- a/src/syntax/parser.rs
+++ b/src/syntax/parser.rs
@@ -1,6 +1,8 @@
 use std::collections::HashSet;
 use std::ops::Range;
 
+use unicode_math_class::MathClass;
+
 use super::{ast, is_newline, ErrorPos, LexMode, Lexer, SyntaxKind, SyntaxNode};
 use crate::util::{format_eco, EcoString};
 
@@ -233,12 +235,13 @@ fn math_expr_prec(p: &mut Parser, min_prec: usize, stop: SyntaxKind) {
             }
         }
 
-        SyntaxKind::Atom => match p.current_text() {
-            "(" => math_delimited(p, ")"),
-            "{" => math_delimited(p, "}"),
-            "[" => math_delimited(p, "]"),
-            _ => p.eat(),
-        },
+        SyntaxKind::Atom if math_class(p.current_text()) == Some(MathClass::Fence) => {
+            math_delimited(p, MathClass::Fence)
+        }
+
+        SyntaxKind::Atom if math_class(p.current_text()) == Some(MathClass::Opening) => {
+            math_delimited(p, MathClass::Closing)
+        }
 
         SyntaxKind::Let
         | SyntaxKind::Set
@@ -254,7 +257,8 @@ fn math_expr_prec(p: &mut Parser, min_prec: usize, stop: SyntaxKind) {
         | SyntaxKind::LeftBrace
         | SyntaxKind::LeftBracket => embedded_code_expr(p),
 
-        SyntaxKind::Linebreak
+        SyntaxKind::Atom
+        | SyntaxKind::Linebreak
         | SyntaxKind::Escape
         | SyntaxKind::Shorthand
         | SyntaxKind::Symbol
@@ -288,21 +292,30 @@ fn math_expr_prec(p: &mut Parser, min_prec: usize, stop: SyntaxKind) {
     }
 }
 
-fn math_delimited(p: &mut Parser, closing: &str) {
+fn math_delimited(p: &mut Parser, stop: MathClass) {
     let m = p.marker();
     p.expect(SyntaxKind::Atom);
-    while !p.eof()
-        && !p.at(SyntaxKind::Dollar)
-        && (!p.at(SyntaxKind::Atom) || p.current_text() != closing)
-    {
+    while !p.eof() && !p.at(SyntaxKind::Dollar) {
+        if math_class(p.current_text()) == Some(stop) {
+            p.eat();
+            p.wrap(m, SyntaxKind::Delimited);
+            return;
+        }
+
         let prev = p.prev_end();
         math_expr(p);
         if !p.progress(prev) {
             p.unexpected();
         }
     }
-    p.expect(SyntaxKind::Atom);
-    p.wrap(m, SyntaxKind::Math);
+}
+
+fn math_class(text: &str) -> Option<MathClass> {
+    let mut chars = text.chars();
+    chars
+        .next()
+        .filter(|_| chars.next().is_none())
+        .and_then(unicode_math_class::class)
 }
 
 fn math_op(kind: SyntaxKind) -> Option<(SyntaxKind, SyntaxKind, ast::Assoc, usize)> {
@@ -324,7 +337,7 @@ fn math_args(p: &mut Parser) {
     p.expect(SyntaxKind::Atom);
     let m = p.marker();
     let mut m2 = p.marker();
-    while !p.eof() {
+    while !p.eof() && !p.at(SyntaxKind::Dollar) {
         match p.current_text() {
             ")" => break,
             "," => {