allow annotations & statements in cal blocks

Author: braised-babbage

source/grammar/openpulseParser.g4

@@ -15,29 +15,32 @@ calibrationBlock: openpulseStatement*;
 
 openpulseStatement:
     // Statements can be used in the cal or defcal body
-    aliasDeclarationStatement
-    | assignmentStatement
-    | barrierStatement
-    | boxStatement
-    | breakStatement
-    | classicalDeclarationStatement
-    | constDeclarationStatement
-    | continueStatement
-    | defStatement
-    | delayStatement
-    | endStatement
-    | expressionStatement
-    | externStatement
-    | forStatement
-    | gateCallStatement
-    | ifStatement
-    | includeStatement
-    | ioDeclarationStatement
-    | quantumDeclarationStatement
-    | resetStatement
-    | returnStatement
-    | whileStatement
-    ;
+    pragma
+    | annotation* (
+        aliasDeclarationStatement
+        | assignmentStatement
+        | barrierStatement
+        | boxStatement
+        | breakStatement
+        | classicalDeclarationStatement
+        | constDeclarationStatement
+        | continueStatement
+        | defStatement
+        | delayStatement
+        | endStatement
+        | expressionStatement
+        | externStatement
+        | forStatement
+        | gateCallStatement
+        | ifStatement
+        | includeStatement
+        | ioDeclarationStatement
+        | quantumDeclarationStatement
+        | resetStatement
+        | returnStatement
+        | whileStatement
+    )
+;
 
 /** In the following we extend existing OpenQASM nodes. Need to refresh whenever OpenQASM is updated. **/
 // We extend the scalarType with WAVEFORM, PORT and FRAME

source/openpulse/openpulse/parser.py

@@ -216,6 +216,14 @@ def visitReturnStatement(self, ctx: qasm3Parser.ReturnStatementContext):
             expression = None
         return ast.ReturnStatement(expression=expression)
 
+    @span
+    def visitOpenpulseStatement(self, ctx: openpulseParser.OpenpulseStatementContext):
+        if ctx.pragma():
+            return self.visit(ctx.pragma())
+        out = self.visit(ctx.getChild(-1))
+        out.annotations = [self.visit(annotation) for annotation in ctx.annotation()]
+        return out
+
 
 # Reuse some QASMNodeVisitor methods in OpenPulseNodeVisitor
 # The following methods are overridden in OpenPulseNodeVisitor and thus not imported:

source/openpulse/tests/test_openpulse_parser.py

@@ -7,6 +7,7 @@
 from openpulse.parser import parse
 from openpulse.ast import (
     AngleType,
+    Annotation,
     ArrayLiteral,
     ArrayType,
     CalibrationDefinition,
@@ -25,6 +26,7 @@
     Identifier,
     IntegerLiteral,
     IntType,
+    Pragma,
     Program,
     QASMNode,
     QuantumBarrier,
@@ -268,6 +270,59 @@ def test_array():
     SpanGuard().visit(program)
 
 
+def test_annotation_in_cal_block():
+    p = """
+    cal {
+      @foo
+      int x = 10;
+    }
+    """
+
+    program = parse(p)
+    expected = Program(
+        statements=[
+            CalibrationStatement(
+                body=[
+                    ClassicalDeclaration(
+                        type=IntType(),
+                        identifier=Identifier("x"),
+                        init_expression=IntegerLiteral(10),
+                    )
+                ]
+            )
+        ]
+    )
+    # annotate the declaration
+    expected.statements[0].body[0].annotations = [Annotation("foo")]
+    assert _remove_spans(program) == expected
+
+
+def test_pragma_in_cal_block():
+    p = """
+    cal {
+      #pragma foo bar
+      int x = 10;
+    }
+    """
+
+    program = parse(p)
+    expected = Program(
+        statements=[
+            CalibrationStatement(
+                body=[
+                    Pragma(command="foo bar"),
+                    ClassicalDeclaration(
+                        type=IntType(),
+                        identifier=Identifier("x"),
+                        init_expression=IntegerLiteral(10),
+                    ),
+                ]
+            )
+        ]
+    )
+    assert _remove_spans(program) == expected
+
+
 @pytest.mark.parametrize(
     "p",
     [