Add shorthand for root operation (#929)

4 files changed, 43 insertions, 0 deletions

diff --git a/src/syntax/ast.rs b/src/syntax/ast.rs
index 53e92949..6e55d106 100644
--- a/src/syntax/ast.rs
+++ b/src/syntax/ast.rs
@@ -126,6 +126,8 @@ pub enum Expr {
     MathAttach(MathAttach),
     /// A fraction in math: `x/2`.
     MathFrac(MathFrac),
+    /// A root in math: `√x`, `∛x` or `∜x`.
+    MathRoot(MathRoot),
     /// An identifier: `left`.
     Ident(Ident),
     /// The `none` literal.
@@ -223,6 +225,7 @@ impl AstNode for Expr {
             SyntaxKind::MathDelimited => node.cast().map(Self::MathDelimited),
             SyntaxKind::MathAttach => node.cast().map(Self::MathAttach),
             SyntaxKind::MathFrac => node.cast().map(Self::MathFrac),
+            SyntaxKind::MathRoot => node.cast().map(Self::MathRoot),
             SyntaxKind::Ident => node.cast().map(Self::Ident),
             SyntaxKind::None => node.cast().map(Self::None),
             SyntaxKind::Auto => node.cast().map(Self::Auto),
@@ -283,6 +286,7 @@ impl AstNode for Expr {
             Self::MathDelimited(v) => v.as_untyped(),
             Self::MathAttach(v) => v.as_untyped(),
             Self::MathFrac(v) => v.as_untyped(),
+            Self::MathRoot(v) => v.as_untyped(),
             Self::Ident(v) => v.as_untyped(),
             Self::None(v) => v.as_untyped(),
             Self::Auto(v) => v.as_untyped(),
@@ -857,6 +861,28 @@ impl MathFrac {
 }
 
 node! {
+    /// A root in math: `√x`, `∛x` or `∜x`.
+    MathRoot
+}
+
+impl MathRoot {
+    /// The index of the root.
+    pub fn index(&self) -> Option<usize> {
+        match self.0.children().next().map(|node| node.text().as_str()) {
+            Some("∜") => Some(4),
+            Some("∛") => Some(3),
+            Some("√") => Option::None,
+            _ => Option::None,
+        }
+    }
+
+    /// The radicand.
+    pub fn radicand(&self) -> Expr {
+        self.0.cast_first_match().unwrap_or_default()
+    }
+}
+
+node! {
     /// An identifier: `it`.
     Ident
 }
diff --git a/src/syntax/kind.rs b/src/syntax/kind.rs
index 0717e16c..111fb260 100644
--- a/src/syntax/kind.rs
+++ b/src/syntax/kind.rs
@@ -67,6 +67,8 @@ pub enum SyntaxKind {
     MathAttach,
     /// A fraction in math: `x/2`.
     MathFrac,
+    /// A root in math: `√x`, `∛x` or `∜x`.
+    MathRoot,
 
     /// A hashtag that switches into code mode: `#`.
     Hashtag,
@@ -134,6 +136,8 @@ pub enum SyntaxKind {
     Dots,
     /// An arrow between a closure's parameters and body: `=>`.
     Arrow,
+    /// A root: `√`, `∛` or `∜`.
+    Root,
 
     /// The `not` operator.
     Not,
@@ -347,6 +351,7 @@ impl SyntaxKind {
             Self::MathDelimited => "delimited math",
             Self::MathAttach => "math attachments",
             Self::MathFrac => "math fraction",
+            Self::MathRoot => "math root",
             Self::Hashtag => "hashtag",
             Self::LeftBrace => "opening brace",
             Self::RightBrace => "closing brace",
@@ -378,6 +383,7 @@ impl SyntaxKind {
             Self::SlashEq => "divide-assign operator",
             Self::Dots => "dots",
             Self::Arrow => "arrow",
+            Self::Root => "root",
             Self::Not => "operator `not`",
             Self::And => "operator `and`",
             Self::Or => "operator `or`",
diff --git a/src/syntax/lexer.rs b/src/syntax/lexer.rs
index dc75a902..eb19d8d9 100644
--- a/src/syntax/lexer.rs
+++ b/src/syntax/lexer.rs
@@ -437,6 +437,7 @@ impl Lexer<'_> {
             '/' => SyntaxKind::Slash,
             '^' => SyntaxKind::Hat,
             '&' => SyntaxKind::MathAlignPoint,
+            '√' | '∛' | '∜' => SyntaxKind::Root,
 
             // Identifiers.
             c if is_math_id_start(c) && self.s.at(is_math_id_continue) => {
diff --git a/src/syntax/parser.rs b/src/syntax/parser.rs
index 1198774b..13c7bd47 100644
--- a/src/syntax/parser.rs
+++ b/src/syntax/parser.rs
@@ -282,6 +282,16 @@ fn math_expr_prec(p: &mut Parser, min_prec: usize, stop: SyntaxKind) {
             p.eat();
         }
 
+        SyntaxKind::Root => {
+            if min_prec < 3 {
+                p.eat();
+                let m2 = p.marker();
+                math_expr_prec(p, 2, stop);
+                math_unparen(p, m2);
+                p.wrap(m, SyntaxKind::MathRoot);
+            }
+        }
+
         _ => p.expected("expression"),
     }