From 4e9927c60191cc0cafae1bdcb3dafffd3166e911 Mon Sep 17 00:00:00 2001
From: Will Hawkins <hawkinsw@obs.cr>
Date: Tue, 10 Mar 2026 05:41:28 -0400
Subject: [PATCH] Add Tree Sitter Tests for Annotations (And General Test
 Updates)

Update Tests

Signed-off-by: Will Hawkins <hawkinsw@obs.cr>
---
 Tests/p4rseTests/ParserTests.swift         |   2 +-
 tree-sitter-p4/grammar.js                  |   4 +-
 tree-sitter-p4/test/corpus/annotations.txt | 101 +++++++++++++++++++++
 3 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 tree-sitter-p4/test/corpus/annotations.txt

diff --git a/Tests/p4rseTests/ParserTests.swift b/Tests/p4rseTests/ParserTests.swift
index f42ebef..9c27460 100644
--- a/Tests/p4rseTests/ParserTests.swift
+++ b/Tests/p4rseTests/ParserTests.swift
@@ -114,5 +114,5 @@ import P4Lang
     #RequireErrorResult<(EvaluatableStatement, LexicalScopes)>(
       Error(withMessage: "{2, 154}: Did not find assignment statement"),
       ParserAssignmentStatement.Compile(
-        node: result.rootNode!, inTree: result, withScopes: LexicalScopes())))
+        node: result.rootNode!, withTypesInScope: LexicalScopes())))
 }
diff --git a/tree-sitter-p4/grammar.js b/tree-sitter-p4/grammar.js
index 44b09a4..944fdf9 100644
--- a/tree-sitter-p4/grammar.js
+++ b/tree-sitter-p4/grammar.js
@@ -56,7 +56,7 @@ export default grammar({
         annotations: $ => repeat1($.annotation),
 
         //annotation: $ => choice(seq('@', "[a-z]+"), seq('@', "[a-z]+", '(', $.annotationBody, ')'), seq('@', "[a-z]+", '[', $.structuredAnnotationBody, ']')),
-        annotation: $ => choice(seq('@', "[a-z]+")),// seq('@', "[a-z]+", '(', /* empty for now*/ ')'), seq('@', "[a-z]+", '[', /* empty for now */ ']')),
+        annotation: $ => choice($.annotation_literal), // seq('@', "[a-z]+", '(', /* empty for now*/ ')'), seq('@', "[a-z]+", '[', /* empty for now */ ']')),
 
 
         // Instantiation
@@ -146,7 +146,7 @@ export default grammar({
         type_identifier: $ => /[a-z]+/,
         string_literal: $ => /".*"/,
         integer: $ => /[0-9]+/,
-
+        annotation_literal: $ => /@[a-z_]+/,
     },
 }
 );
diff --git a/tree-sitter-p4/test/corpus/annotations.txt b/tree-sitter-p4/test/corpus/annotations.txt
new file mode 100644
index 0000000..900e76d
--- /dev/null
+++ b/tree-sitter-p4/test/corpus/annotations.txt
@@ -0,0 +1,101 @@
+=========================
+Simple Annotated Declaration
+=========================
+@testing parser simple() {
+  state start {
+    bool l;
+    transition accept;
+  }
+};
+
+---
+(p4program
+ (declaration
+  (parserDeclaration
+   (parserType
+    (annotations
+     (annotation
+      (annotation_literal)
+     )
+    )
+    (parser)
+    (identifier)
+   )
+   (parserStates
+    (parserState
+     (state)
+     (identifier)
+     (parserStatements
+      (parserStatement
+       (variableDeclaration
+        (typeRef
+         (baseType
+          (bool)
+         )
+        )
+        (identifier)
+       )
+      )
+     )
+     (parserTransitionStatement
+      (transition)
+      (transitionSelectionExpression
+       (identifier)
+      )
+     )
+    )
+   )
+  )
+ )
+)
+
+=========================
+Simple Annotated State
+=========================
+parser simple() {
+  @testing state start {
+    bool l;
+    transition accept;
+  }
+};
+
+---
+(p4program
+ (declaration
+  (parserDeclaration
+   (parserType
+    (parser)
+    (identifier)
+   )
+   (parserStates
+    (parserState
+     (annotations
+      (annotation
+       (annotation_literal)
+      )
+     )
+     (state)
+     (identifier)
+     (parserStatements
+      (parserStatement
+       (variableDeclaration
+        (typeRef
+         (baseType
+          (bool)
+         )
+        )
+        (identifier)
+       )
+      )
+     )
+     (parserTransitionStatement
+      (transition)
+      (transitionSelectionExpression
+       (identifier)
+      )
+     )
+    )
+   )
+  )
+ )
+)