From 967b6154657b40767e6241903f0776322c17bd2f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20=C4=8Euri=C5=A1?=
 <tomas.duris@gendigital.com>
Date: Tue, 23 Jan 2024 10:48:19 +0100
Subject: [PATCH] add better variable support

---
 example.yar                            |  4 ++--
 src/parser/grammar/expressions.rs      | 23 ++++++++++++++++-------
 src/parser/grammar/expressions/atom.rs | 13 ++++++++++---
 src/parser/syntaxkind.rs               |  1 +
 4 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/example.yar b/example.yar
index 12f2d0b..9f68a9b 100644
--- a/example.yar
+++ b/example.yar
@@ -10,6 +10,6 @@ rule test
 		$a = "foo"
 		$b = "bar"
 	condition:
-		$a and
-		$b
+		$a or
+		$b and true
 }
diff --git a/src/parser/grammar/expressions.rs b/src/parser/grammar/expressions.rs
index 83f038c..da80b41 100644
--- a/src/parser/grammar/expressions.rs
+++ b/src/parser/grammar/expressions.rs
@@ -45,25 +45,34 @@ fn condition(p: &mut Parser) {
     m.complete(p, CONDITION);
 }
 
+const VARIABLE_RECOVERY_SET: TokenSet = TokenSet::new(&[VARIABLE]);
+
 pub(super) fn strings_body(p: &mut Parser) {
     // add support for meta also
     while !p.at(EOF) && !p.at(STRINGS) && !p.at(CONDITION) && !p.at(RBRACE) {
-        assert!(p.at(VARIABLE));
         let m = p.start();
-        p.bump(VARIABLE);
+        if p.at(VARIABLE) {
+            let m = p.start();
+            p.bump(VARIABLE);
+            m.complete(p, VARIABLE);
+        } else {
+            p.err_recover("expected a variable", VARIABLE_RECOVERY_SET);
+        }
         p.expect(ASSIGN);
         // so far only strings are supported, later add match for hex strings and regex
         string(p);
-        m.complete(p, VARIABLE);
+        m.complete(p, VARIABLE_STMT);
     }
 }
 
-// do the same for hex and regex strings
+// add support for hex and regex strings later on
 fn string(p: &mut Parser) {
-    assert!(p.at(STRING));
     let m = p.start();
-    p.bump(STRING);
-    // add plain string modifiers
+    match p.current() {
+        STRING => p.bump(STRING),
+        _ => p.err_and_bump("expected a string"),
+    }
+    // add string modifiers
     m.complete(p, STRING);
 }
 
diff --git a/src/parser/grammar/expressions/atom.rs b/src/parser/grammar/expressions/atom.rs
index ae0dc39..4a5fab0 100644
--- a/src/parser/grammar/expressions/atom.rs
+++ b/src/parser/grammar/expressions/atom.rs
@@ -13,11 +13,18 @@ pub(crate) fn literal(p: &mut Parser) -> Option<CompletedMarker> {
     Some(m.complete(p, LITERAL))
 }
 
+const EXPR_RECOVERY_SET: TokenSet = TokenSet::new(&[VARIABLE, TRUE, FALSE]);
+
 // add support for while/for loops, if/else statements, etc.
 pub(super) fn atom_expr(p: &mut Parser) -> Option<CompletedMarker> {
     if let Some(m) = literal(p) {
-        Some(m)
-    } else {
-        todo!("add support for other atoms")
+        return Some(m);
     }
+
+    let _done = match p.current() {
+        _ => {
+            p.err_recover("expected expression", EXPR_RECOVERY_SET);
+            return None;
+        }
+    };
 }
diff --git a/src/parser/syntaxkind.rs b/src/parser/syntaxkind.rs
index fea3024..fc47c2c 100644
--- a/src/parser/syntaxkind.rs
+++ b/src/parser/syntaxkind.rs
@@ -33,6 +33,7 @@ pub enum SyntaxKind {
     LITERAL,
     EXPRESSION,
     EXPRESSION_STMT,
+    VARIABLE_STMT,
     __LAST,
 }