diff --git a/tree-sitter-p4/grammar.js b/tree-sitter-p4/grammar.js
index 59b8440..5f33005 100644
--- a/tree-sitter-p4/grammar.js
+++ b/tree-sitter-p4/grammar.js
@@ -84,7 +84,7 @@ export default grammar({
 
         // Control declarations
         control_declaration: $ => seq($.control, $.identifier, $.parameters, '{', repeat(choice($.table_declaration, $.action_declaration)), $.apply_statement, '}'),
-        action_declaration: $ => seq($.action, $.identifier, $.parameters, $.statement),
+        action_declaration: $ => seq($.action, $.identifier, $.parameters, $.blockStatement),
         table_declaration: $ => seq($.table, $.identifier, '{', optional($.table_property_list), '}'),
 
         // Table property list
@@ -107,7 +107,7 @@ export default grammar({
         assignmentStatement: $=> seq($.expression, $.assignment, $.expression, $._semicolon),
         return_statement: $=> seq($._return, $.expression, $._semicolon),
         exit_statement: $=> seq($._exit, $._semicolon),
-        apply_statement: $=> seq($.apply, $.statement),
+        apply_statement: $=> seq($.apply, $.blockStatement),
 
         // Parser statements
         parserStatements: $ => repeat1($.parserStatement),
diff --git a/tree-sitter-p4/test/corpus/control.txt b/tree-sitter-p4/test/corpus/control.txt
index 30dc98b..04c5075 100644
--- a/tree-sitter-p4/test/corpus/control.txt
+++ b/tree-sitter-p4/test/corpus/control.txt
@@ -24,9 +24,7 @@ control simple() {
     (action)
     (identifier)
     (parameters)
-    (statement
-     (blockStatement)
-    )
+    (blockStatement)
    )
    (table_declaration
     (table)
@@ -51,9 +49,7 @@ control simple() {
    )
    (apply_statement
     (apply)
-    (statement
-     (blockStatement)
-    )
+    (blockStatement)
    )
   )
  )
diff --git a/tree-sitter-p4/test/corpus/statements.txt b/tree-sitter-p4/test/corpus/statements.txt
index 70f0855..a1ce8a2 100644
--- a/tree-sitter-p4/test/corpus/statements.txt
+++ b/tree-sitter-p4/test/corpus/statements.txt
@@ -15,12 +15,10 @@ control T() {
    (parameters)
    (apply_statement
     (apply)
-    (statement
-     (blockStatement
-      (statements
-       (statement
-        (exit_statement)
-       )
+    (blockStatement
+     (statements
+      (statement
+       (exit_statement)
       )
      )
     )