feature: add support for typecasting (#27)

* feature: add support for typecasting

* fix: changes to  int typecasting

* fix: add support for args and kwargs in typecasting)

* fix: add unit tests for int typecasting

* fix: update typecasting tests

* fix: remove unused dependency

* fix: update operator to force casting in the correct location

* fix: type hint

* fix: add cosmetic changes

* Update test/unit_tests/autoqasm/test_operators.py

* fix: add xfail test

* fix: clarify test name

---------

Co-authored-by: Ryan Shaffer <[email protected]>

src/autoqasm/converters/typecast.py

@@ -0,0 +1,67 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+
+"""Converters for integer casting nodes."""
+
+import ast
+
+from malt.core import ag_ctx, converter
+from malt.pyct import templates
+
+TYPECASTING_OPERATORS = {
+    "int": "ag__.int_",
+}
+
+
+class TypecastTransformer(converter.Base):
+    def visit_Call(self, node: ast.stmt) -> ast.stmt:
+        """Converts type casting operations to their AutoQASM counterpart.
+
+        Args:
+            node (ast.stmt): AST node to transform.
+
+        Returns:
+            ast.stmt: Transformed node.
+        """
+        node = self.generic_visit(node)
+
+        if (
+            hasattr(node, "func")
+            and hasattr(node.func, "id")
+            and node.func.id in TYPECASTING_OPERATORS
+        ):
+            template = f"{TYPECASTING_OPERATORS[node.func.id]}(argument)"
+            new_node = templates.replace(
+                template,
+                argument=node.args,
+                original=node,
+            )
+            new_node = new_node[0].value
+        else:
+            new_node = node
+        return new_node
+
+
+def transform(node: ast.stmt, ctx: ag_ctx.ControlStatusCtx) -> ast.stmt:
+    """Transform int cast nodes.
+
+    Args:
+        node (ast.stmt): AST node to transform.
+        ctx (ag_ctx.ControlStatusCtx): Transformer context.
+
+    Returns:
+        ast.stmt: Transformed node.
+    """
+
+    return TypecastTransformer(ctx).visit(node)

src/autoqasm/operators/__init__.py

@@ -40,3 +40,4 @@
 from .logical import or_  # noqa: F401
 from .return_statements import return_output_from_main  # noqa: F401
 from .slices import GetItemOpts, get_item, set_item  # noqa: F401
+from .typecast import int_  # noqa: F401

src/autoqasm/operators/typecast.py

@@ -0,0 +1,48 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+
+"""Operators for int cast statements."""
+from __future__ import annotations
+
+from typing import Any
+
+from autoqasm import program
+from autoqasm import types as aq_types
+
+
+def int_(argument: Any, *args, **kwargs) -> aq_types.IntVar | int:
+    """Functional form of "int".
+
+    Args:
+        argument (Any): object to be converted into an int.
+
+    Returns:
+        IntVar | int : IntVar object if argument is QASM type, else int.
+    """
+    if aq_types.is_qasm_type(argument):
+        return _oqpy_int(argument)
+    else:
+        return _py_int(argument, *args, **kwargs)
+
+
+def _oqpy_int(argument: Any) -> aq_types.IntVar:
+    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    result = aq_types.IntVar()
+    oqpy_program.declare(result)
+    oqpy_program.set(result, argument)
+    return result
+
+
+def _py_int(argument: Any, *args, **kwargs) -> int:
+    return int(argument, *args, **kwargs)

src/autoqasm/transpiler/transpiler.py

@@ -54,6 +54,7 @@
     break_statements,
     comparisons,
     return_statements,
+    typecast,
 )
 
 
@@ -134,6 +135,7 @@ def transform_ast(
         # canonicalization creates.
         node = continue_statements.transform(node, ctx)
         node = return_statements.transform(node, ctx)
+        node = typecast.transform(node, ctx)
         node = assignments.transform(node, ctx)
         node = lists.transform(node, ctx)
         node = slices.transform(node, ctx)

test/unit_tests/autoqasm/test_operators.py

@@ -932,6 +932,68 @@ def test_list_ops():
     assert test_list_ops.build().to_ir()
 
 
+def test_int_typecasting_on_measure():
+    @aq.main(num_qubits=2)
+    def main():
+        test = int(measure([0, 1]))  # noqa: F841
+
+    expected_ir = """OPENQASM 3.0;
+int[32] test;
+qubit[2] __qubits__;
+bit[2] __bit_0__ = "00";
+__bit_0__[0] = measure __qubits__[0];
+__bit_0__[1] = measure __qubits__[1];
+int[32] __int_1__;
+__int_1__ = __bit_0__;
+test = __int_1__;"""
+    assert main.build().to_ir() == expected_ir
+
+
+def test_int_typecasting_on_python_string_not_captured():
+    @aq.main(num_qubits=2)
+    def main():
+        test = int("101", 2)  # noqa: F841
+
+    expected_ir = """OPENQASM 3.0;
+qubit[2] __qubits__;"""
+    assert main.build().to_ir() == expected_ir
+
+
+@pytest.mark.xfail(reason="Bug: assignments do not work as expected when operators are nested")
+def test_nested_int_typecasting_without_return():
+    @aq.main(num_qubits=2)
+    def main():
+        test = 2 * int(measure([0, 1]))  # noqa: F841
+
+    expected_ir = """OPENQASM 3.0;
+qubit[2] __qubits__;
+bit[2] __bit_0__ = "00";
+__bit_0__[0] = measure __qubits__[0];
+__bit_0__[1] = measure __qubits__[1];
+int[32] __int_1__;
+__int_1__ = __bit_0__;
+test = 2 * __int_1__;"""
+    assert main.build().to_ir() == expected_ir
+
+
+def test_nested_int_typecasting_with_return():
+    @aq.main(num_qubits=2)
+    def main():
+        test = 2 * int(measure([0, 1]))  # noqa: F841
+        return test
+
+    expected_ir = """OPENQASM 3.0;
+output int[32] test;
+qubit[2] __qubits__;
+bit[2] __bit_0__ = "00";
+__bit_0__[0] = measure __qubits__[0];
+__bit_0__[1] = measure __qubits__[1];
+int[32] __int_1__;
+__int_1__ = __bit_0__;
+test = 2 * __int_1__;"""
+    assert main.build().to_ir() == expected_ir
+
+
 def test_integer_division_on_intvars():
     @aq.main(num_qubits=2)
     def main():