Merge pull request #15 from braised-babbage/allow-annotations-pragmas-in-cal-blocks

Allow annotations & pragmas in cal blocks

Author: psivaraj

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/ast.py

@@ -12,7 +12,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Union
 
 # Re-export the existing AST classes from openqasm3
 # pylint: disable=unused-import
@@ -22,6 +22,7 @@
 # pylint: disable=unused-import
 from openqasm3.ast import ExternArgument
 
+
 # From Pulse grammar
 class WaveformType(ClassicalType):
     """
@@ -72,7 +73,7 @@ class CalibrationStatement(Statement):
         }
     """
 
-    body: List[Statement]
+    body: List[Union[Statement, Pragma]]
 
 
 # Override the class from openqasm3

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

@@ -1,17 +1,16 @@
 import dataclasses
 
 import pytest
-
-from openqasm3.visitor import QASMVisitor
-
-from openpulse.parser import parse
 from openpulse.ast import (
     AngleType,
+    Annotation,
     ArrayLiteral,
     ArrayType,
+    AssignmentOperator,
     CalibrationDefinition,
     CalibrationStatement,
     ClassicalArgument,
+    ClassicalAssignment,
     ClassicalDeclaration,
     ComplexType,
     DurationType,
@@ -21,21 +20,24 @@
     FloatLiteral,
     FloatType,
     ForInLoop,
+    FrameType,
     FunctionCall,
     Identifier,
     IntegerLiteral,
     IntType,
+    PortType,
+    Pragma,
     Program,
     QASMNode,
     QuantumBarrier,
     RangeDefinition,
     ReturnStatement,
     UnaryExpression,
     UnaryOperator,
-    FrameType,
-    PortType,
     WaveformType,
 )
+from openpulse.parser import parse
+from openqasm3.visitor import QASMVisitor
 
 
 class SpanGuard(QASMVisitor):
@@ -268,6 +270,65 @@ def test_array():
     SpanGuard().visit(program)
 
 
+def test_annotation_in_cal_block():
+    p = """
+    cal {
+      @word1
+      int x = 10;
+
+      @word1 command1
+      @word2 command2 32f%^&
+      x = 0;
+
+      @word1 @not_a_separate_annotation uint x;
+      z();
+    }
+    """
+
+    xdecl = ClassicalDeclaration(
+        type=IntType(),
+        identifier=Identifier("x"),
+        init_expression=IntegerLiteral(10),
+    )
+    xdecl.annotations = [Annotation("word1")]
+
+    xassign = ClassicalAssignment(Identifier("x"), AssignmentOperator["="], IntegerLiteral(0))
+    xassign.annotations = [Annotation("word1", "command1"), Annotation("word2", "command2 32f%^&")]
+
+    zcall = ExpressionStatement(FunctionCall(Identifier("z"), []))
+    zcall.annotations = [Annotation("word1", "@not_a_separate_annotation uint x;")]
+    program = parse(p)
+    expected = Program(statements=[CalibrationStatement(body=[xdecl, xassign, zcall])])
+
+    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",
     [