Modify ExternStatement AST node (#410)

* Change ExternalDeclaration's return type node * Export ExternArgument * Add tests for extern * Update externStatement parsing tests
openqasm · Dec 12, 2022 · bcd5348 · bcd5348
1 parent 25f8730
commit bcd5348
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 1 deletion.
diff --git a/source/openqasm/openqasm3/ast.py b/source/openqasm/openqasm3/ast.py
@@ -52,6 +52,7 @@
     "EndStatement",
     "Expression",
     "ExpressionStatement",
+    "ExternArgument",
     "ExternDeclaration",
     "FloatLiteral",
     "FloatType",
@@ -231,7 +232,7 @@ class ExternDeclaration(Statement):
 
     name: Identifier
     arguments: List[ExternArgument]
-    return_type: Optional[ExternArgument] = None
+    return_type: Optional[ClassicalType] = None
 
 
 class Expression(QASMNode):

diff --git a/source/openqasm/tests/test_qasm_parser.py b/source/openqasm/tests/test_qasm_parser.py
@@ -35,6 +35,8 @@
     DurationType,
     EndStatement,
     ExpressionStatement,
+    ExternArgument,
+    ExternDeclaration,
     FloatLiteral,
     FloatType,
     ForInLoop,
@@ -406,6 +408,63 @@ def test_array_declaration():
     )
 
 
+def test_extern_declarations():
+    p = """
+    extern f();
+    extern f() -> bool;
+    extern f(bool);
+    extern f(int[32], uint[32]);
+    extern f(mutable array[complex[float[64]], N_ELEMENTS]) -> int[2 * INT_SIZE];
+    """.strip()
+    program = parse(p)
+    assert _remove_spans(program) == Program(
+        statements=[
+            ExternDeclaration(
+                name=Identifier(name="f"),
+                arguments=[],
+            ),
+            ExternDeclaration(
+                name=Identifier(name="f"),
+                arguments=[],
+                return_type=BoolType(),
+            ),
+            ExternDeclaration(
+                name=Identifier(name="f"),
+                arguments=[
+                    ExternArgument(type=BoolType()),
+                ],
+            ),
+            ExternDeclaration(
+                name=Identifier(name="f"),
+                arguments=[
+                    ExternArgument(type=IntType(IntegerLiteral(32))),
+                    ExternArgument(type=UintType(IntegerLiteral(32))),
+                ],
+            ),
+            ExternDeclaration(
+                name=Identifier(name="f"),
+                arguments=[
+                    ExternArgument(
+                        type=ArrayReferenceType(
+                            base_type=ComplexType(FloatType(IntegerLiteral(64))),
+                            dimensions=[Identifier(name="N_ELEMENTS")],
+                        ),
+                        access=AccessControl["mutable"],
+                    ),
+                ],
+                return_type=IntType(
+                    size=BinaryExpression(
+                        op=BinaryOperator["*"],
+                        lhs=IntegerLiteral(2),
+                        rhs=Identifier(name="INT_SIZE"),
+                    )
+                ),
+            ),
+        ]
+    )
+    SpanGuard().visit(program)
+
+
 def test_single_gatecall():
     p = """
     h q;