From c37319808422551b707005965fcf298009565d6c Mon Sep 17 00:00:00 2001
From: Ross Andrews <ross.andrews@gmail.com>
Date: Sat, 8 Jul 2023 23:33:30 -0500
Subject: [PATCH] Prefix operators

---
 forge_core/src/ast.rs          |  9 ++++++++-
 forge_core/src/forge.pest      |  5 ++++-
 forge_core/src/forge_parser.rs | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/forge_core/src/ast.rs b/forge_core/src/ast.rs
index 1ce8eaa..5e5e86a 100644
--- a/forge_core/src/ast.rs
+++ b/forge_core/src/ast.rs
@@ -38,7 +38,7 @@ pub struct Member {
 #[derive(Eq, PartialEq, Clone, Debug)]
 pub struct Const {
     pub name: String,
-    pub value: Option<i32>,
+    pub value: Option<i32>, // todo: this should take a Node
     pub string: Option<String>,
 }
 
@@ -150,6 +150,12 @@ pub enum Operator {
     Rshift,
 }
 
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum Prefix {
+    Neg,
+    Not,
+}
+
 #[derive(PartialEq, Clone, Debug)]
 pub enum Node {
     Number(i32),
@@ -158,6 +164,7 @@ pub enum Node {
     Name(String),
     Address(String),
     Expr(BoxNode, Operator, BoxNode),
+    Prefix(Prefix, BoxNode),
 }
 
 impl From<i32> for Node {
diff --git a/forge_core/src/forge.pest b/forge_core/src/forge.pest
index 5da8603..5e74b11 100644
--- a/forge_core/src/forge.pest
+++ b/forge_core/src/forge.pest
@@ -5,6 +5,8 @@ name_char = { ASCII_ALPHA_LOWER | ASCII_ALPHA_UPPER | "_" | "$" }
 name = @{ name_char ~ (name_char | ASCII_DIGIT)* }
 
 dec_number = @{ ASCII_NONZERO_DIGIT ~ ASCII_DIGIT* }
+// Literal numbers use this rule, but in exprs it's the prefix operator. The difference
+// is that some things (const initialization) only accept literal numbers.
 neg_number = ${ "-" ~ dec_number }
 hex_number = ${ "0x" ~ ASCII_HEX_DIGIT+ }
 bin_number = ${ "0b" ~ ASCII_BIN_DIGIT+ }
@@ -38,6 +40,7 @@ eq = { "==" }
 ne = { "!=" }
 lshift = { "<<" }
 rshift = { ">>" }
+prefix = { "-" | "!" }
 
 operator = _{
  add | sub |
@@ -47,7 +50,7 @@ operator = _{
  lshift | rshift |
  lt | le | gt | ge |
  eq | ne }
-expr = { val ~ (operator ~ val)* }
+expr = { prefix? ~ val ~ (operator ~ prefix? ~ val)* }
 val = { // These are what get evaluated and left on the stack:
   number | // Literal numbers
   ("(" ~ expr ~ ")") | // Nested parenthesized exprs
diff --git a/forge_core/src/forge_parser.rs b/forge_core/src/forge_parser.rs
index 2061a88..c9ccee9 100644
--- a/forge_core/src/forge_parser.rs
+++ b/forge_core/src/forge_parser.rs
@@ -22,6 +22,7 @@ lazy_static::lazy_static! {
             .op(Op::infix(lshift, Left) | Op::infix(rshift, Left))
             .op(Op::infix(add, Left) | Op::infix(sub, Left))
             .op(Op::infix(mul, Left) | Op::infix(div, Left) | Op::infix(modulus, Left))
+            .op(Op::prefix(Rule::prefix))
     };
 }
 
@@ -467,6 +468,7 @@ impl AstNode for Node {
                     rule => unreachable!("Expr::parse expected atom, found {:?}", rule),
                 }
             })
+            .map_prefix(|op, rhs| Node::Prefix(Prefix::from_pair(op), rhs.into()))
             .map_infix(|lhs, op, rhs| Node::Expr(lhs.into(), Operator::from_pair(op), rhs.into()))
             .parse(pair.into_inner())
     }
@@ -501,6 +503,18 @@ impl AstNode for Operator {
     }
 }
 
+impl AstNode for Prefix {
+    const RULE: Rule = Rule::prefix;
+
+    fn from_pair(pair: Pair) -> Self {
+        match pair.as_str() {
+            "!" => Self::Not,
+            "-" => Self::Neg,
+            _ => unreachable!(),
+        }
+    }
+}
+
 ///////////////////////////////////////////////////////////////////////////////////////////
 
 impl AstNode for ArrayRef {
@@ -765,6 +779,7 @@ mod test {
             Ok(Expr(Expr(2.into(), Mul, 3.into()).into(), Add, 4.into()))
         );
 
+        // Various operators
         assert_eq!(
             Node::from_str("1 || 2 && 3"),
             Ok(Expr(1.into(), Or, Expr(2.into(), And, 3.into()).into()))
@@ -799,6 +814,23 @@ mod test {
             Ok(Expr(1.into(), Lshift, 6.into()))
         );
 
+        assert_eq!(
+            Node::from_str("!a - -3"),
+            Ok(Expr(
+                Prefix(crate::ast::Prefix::Not, Name("a".into()).into()).into(),
+                Sub,
+                Prefix(crate::ast::Prefix::Neg, 3.into()).into()
+            ))
+        );
+
+        assert_eq!(
+            Node::from_str("-(4 * 5)"),
+            Ok(Prefix(
+                crate::ast::Prefix::Neg,
+                Expr(4.into(), Mul, 5.into()).into()
+            ))
+        );
+
         // Parens
         assert_eq!(
             Node::from_str("(1 + 2) * 3"),