From 38afcb275bc8b904e3ebf6eae65a44ccfc51095c Mon Sep 17 00:00:00 2001
From: Benjamin Bannier <benjamin.bannier@corelight.com>
Date: Fri, 10 Feb 2023 18:46:47 +0100
Subject: [PATCH] Expose a number simple statement as nodes

---
 corpus/test | 41 +++++++++++++++++++++++------------------
 grammar.js  | 24 ++++++++++++++++++------
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/corpus/test b/corpus/test
index 0a12a12..d742df1 100644
--- a/corpus/test
+++ b/corpus/test
@@ -220,26 +220,31 @@ function foo(xs: vector of int): int
             (type))))
       (type))
     (stmt_list
-      (id)
-      (expr
-        (id))
-      (id)
-      (expr
-        (id))
-      (expr
-        (constant
-          (boolean)))
-      (expr
+      (for
+        (id)
+        (expr
+          (id)))
+      (for
+        (id)
+        (expr
+          (id)))
+      (while
         (expr
           (constant
-            (integer))))
-      (case_list
-        (expr_list
+            (boolean))))
+      (switch
+        (expr
           (expr
             (constant
               (integer))))
-        (stmt_list)
-        (stmt_list))
-      (expr
-        (constant
-          (integer))))))
+        (case_list
+          (expr_list
+            (expr
+              (constant
+                (integer))))
+          (stmt_list)
+          (stmt_list)))
+      (return
+        (expr
+          (constant
+            (integer)))))))
diff --git a/grammar.js b/grammar.js
index df523ce..0dfd2e3 100644
--- a/grammar.js
+++ b/grammar.js
@@ -67,13 +67,12 @@ module.exports = grammar({
             seq('{', optional($.stmt_list), '}'),
             seq('print', $.expr_list, ';'),
             seq('event', $.event_hdr, ';'),
-            prec_r(seq('if', '(', $.expr, ')', $._stmt, optional(seq('else', $._stmt)))),
-            seq('switch', $.expr, '{', optional($.case_list), '}'),
-            seq('for', '(', $.id, optional(seq(',', $.id)), 'in', $.expr, ')', $._stmt),
-            seq('for', '(', '[', list1($.id, ','), ']', optional(seq(',', $.id)), 'in', $.expr, ')', $._stmt),
-            seq('while', '(', $.expr, ')', $._stmt),
+            $.switch,
+            $.if,
+            $.for,
+            $.while,
+            $.return,
             seq(choice('next', 'break', 'fallthrough'), ';'),
-            seq('return', optional($.expr), ';'),
             seq(choice('add', 'delete'), $.expr, ';'),
             // Precedence works around ambiguity with `var_decl` in `_decl` at `source_file` scope.
             prec(-1, $.var_decl),
@@ -92,6 +91,19 @@ module.exports = grammar({
             ';',
         ),
 
+        if: $ => prec_r(seq('if', '(', $.expr, ')', $._stmt, optional(seq('else', $._stmt)))),
+
+        for: $ => choice(
+            seq('for', '(', $.id, optional(seq(',', $.id)), 'in', $.expr, ')', $._stmt),
+            seq('for', '(', '[', list1($.id, ','), ']', optional(seq(',', $.id)), 'in', $.expr, ')', $._stmt),
+        ),
+
+        while: $ => seq('while', '(', $.expr, ')', $._stmt),
+
+        switch: $ => seq('switch', $.expr, '{', optional($.case_list), '}'),
+
+        return: $ => seq('return', optional($.expr), ';'),
+
         stmt_list: $ => repeat1($._stmt),
 
         case_list: $ => repeat1(