Simple OpenQASM 3.0 expressions and casts

pennylane/io/qasm_interpreter.py

@@ -7,25 +7,37 @@
 from functools import partial
 from typing import Any, Iterable
 
+import numpy as np
+from numpy import uint
 from openqasm3.ast import (
     AliasStatement,
     ArrayLiteral,
+    ArrayType,
+    BinaryExpression,
     BitstringLiteral,
     BooleanLiteral,
+    BoolType,
+    Cast,
     ClassicalAssignment,
     ClassicalDeclaration,
+    ComplexType,
     ConstantDeclaration,
     DurationLiteral,
     EndStatement,
     Expression,
+    ExpressionStatement,
     FloatLiteral,
+    FloatType,
     Identifier,
     ImaginaryLiteral,
     IndexExpression,
     IntegerLiteral,
+    IntType,
     QuantumGate,
     QubitDeclaration,
     RangeDefinition,
+    UintType,
+    UnaryExpression,
 )
 from openqasm3.visitor import QASMNode
 
@@ -66,6 +78,64 @@
     "CRZ": ops.CRZ,
 }
 
+EQUALS = "="
+ARROW = "->"
+PLUS = "+"
+DOUBLE_PLUS = "++"
+MINUS = "-"
+ASTERISK = "*"
+DOUBLE_ASTERISK = "**"
+SLASH = "/"
+PERCENT = "%"
+PIPE = "|"
+DOUBLE_PIPE = "||"
+AMPERSAND = "&"
+DOUBLE_AMPERSAND = "&&"
+CARET = "^"
+AT = "@"
+TILDE = "~"
+EXCLAMATION_POINT = "!"
+EQUALITY_OPERATORS = ["==", "!=", "~="]
+COMPOUND_ASSIGNMENT_OPERATORS = [
+    "+=",
+    "-=",
+    "*=",
+    "/=",
+    "&=",
+    "|=",
+    "^=",
+    "<<=",
+    ">>=",
+    "%=",
+    "**=",
+]
+COMPARISON_OPERATORS = [">", "<", ">=", "<="]
+BIT_SHIFT_OPERATORS = [">>", "<<"]
+
+NON_ASSIGNMENT_CLASSICAL_OPERATORS = (
+    EQUALITY_OPERATORS
+    + COMPARISON_OPERATORS
+    + BIT_SHIFT_OPERATORS
+    + [
+        PLUS,
+        MINUS,
+        ASTERISK,
+        DOUBLE_ASTERISK,
+        SLASH,
+        PERCENT,
+        PIPE,
+        DOUBLE_PIPE,
+        AMPERSAND,
+        DOUBLE_AMPERSAND,
+        CARET,
+        AT,
+        TILDE,
+        EXCLAMATION_POINT,
+    ]
+)
+
+ASSIGNMENT_CLASSICAL_OPERATORS = [EQUALS, DOUBLE_PLUS] + COMPOUND_ASSIGNMENT_OPERATORS
+
 
 @dataclass
 class Variable:
@@ -130,7 +200,7 @@ def update_var(
         value: any,
         name: str,
         node: QASMNode,
-    ):
+    ):  # pylint: disable=too-many-branches
         """
         Updates a variable, or raises if it is constant.
         Args:
@@ -144,8 +214,40 @@ def update_var(
                     f"Attempt to mutate a constant {name} on line {node.span.start_line} that was "
                     f"defined on line {self.vars[name].line}"
                 )
-            self.vars[name].val = value
-            self.vars[name].line = node.span.start_line
+            if node.op.name in ASSIGNMENT_CLASSICAL_OPERATORS:
+                if node.op.name == "=":
+                    self.vars[name].val = value
+                if node.op.name == "++":
+                    self.vars[name].val = self.vars[name].val + 1
+                if node.op.name == "+=":
+                    self.vars[name].val += value
+                if node.op.name == "-=":
+                    self.vars[name].val -= value
+                if node.op.name == "*=":
+                    self.vars[name].val = self.vars[name].val * value
+                if node.op.name == "/=":
+                    self.vars[name].val = self.vars[name].val / value
+                if node.op.name == "&=":
+                    self.vars[name].val = self.vars[name].val & value
+                if node.op.name == "|=":
+                    self.vars[name].val = self.vars[name].val | value
+                if node.op.name == "^=":
+                    self.vars[name].val = self.vars[name].val ^ value
+                if node.op.name == "<<=":
+                    self.vars[name].val = self.vars[name].val << value
+                if node.op.name == ">>=":
+                    self.vars[name].val = self.vars[name].val >> value
+                if node.op.name == "%=":
+                    self.vars[name].val = self.vars[name].val % value
+                if node.op.name == "**=":
+                    self.vars[name].val = self.vars[name].val ** value
+                self.vars[name].line = node.span.start_line
+            else:
+                # we shouldn't ever get thi error if the parser did its job right
+                raise SyntaxError(  # pragma: no cover
+                    f"Invalid operator {node.op.name} encountered in assignment expression "
+                    f"on line {node.span.start_line}."
+                )  # pragma: no cover
         else:
             raise TypeError(f"Attempt to use undeclared variable {name} in {self.name}")
 
@@ -205,6 +307,24 @@ def _resolve_name(node: QASMNode):
     return node.name if isinstance(node.name, str) else node.name.name
 
 
+def preprocess_operands(operand):
+    """
+    Interprets a string operand as an appropriate type.
+
+    Args:
+        operand (str): the string operand to interpret.
+
+    Returns:
+        The interpreted operand as an appropriate type.
+    """
+    if isinstance(operand, str):
+        if operand.isdigit():
+            operand = int(operand)
+        elif operand.isnumeric():
+            operand = float(operand)
+    return operand
+
+
 class EndProgram(Exception):
     """Exception raised when it encounters an end statement in the QASM circuit."""
 
@@ -335,7 +455,6 @@ def visit_classical_assignment(self, node: QASMNode, context: Context):
         # references to an unresolved value see a func for now
         name = _resolve_name(node.lvalue)
         res = self.visit(node.rvalue, context)
-        # TODO: different types of assignments
         context.update_var(res, name, node)
 
     @visit.register(AliasStatement)
@@ -513,6 +632,145 @@ def apply_modifier(self, mod: QuantumGate, previous: Operator, context: Context,
 
         return next, wires
 
+    @visit.register(ExpressionStatement)
+    def visit_expression_statement(self, node: ExpressionStatement, context: Context):
+        """
+        Registers an expression statement.
+        Args:
+            node (ExpressionStatement): The expression statement.
+            context (Context): The current context.
+        """
+        return self.visit(node.expression, context)
+
+    @visit.register(Cast)
+    def visit_cast(self, node: Cast, context: Context):
+        """
+        Registers a Cast expression.
+
+        Args:
+            node (Cast): The Cast expression.
+            context (Context): The current context.
+
+        Returns:
+            Any: The argument cast to the appropriate type.
+
+        Raises:
+            TypeError: If the cast cannot be made.
+        """
+        arg = self.visit(node.argument, context)
+        ret = arg
+        try:
+            if isinstance(node.type, IntType):
+                ret = int(arg)
+            if isinstance(node.type, UintType):
+                ret = uint(arg)
+            if isinstance(node.type, FloatType):
+                ret = float(arg)
+            if isinstance(node.type, ComplexType):
+                ret = complex(arg)
+            if isinstance(node.type, BoolType):
+                ret = bool(arg)
+            if isinstance(node.type, ArrayType):
+                ret = arg.val
+            # TODO: durations, angles, etc.
+        except TypeError as e:
+            raise TypeError(
+                f"Unable to cast {arg.__class__.__name__} to {node.type.__class__.__name__}: {str(e)}"
+            ) from e
+        return ret
+
+    @visit.register(BinaryExpression)
+    def visit_binary_expression(
+        self, node: BinaryExpression, context: Context
+    ):  # pylint: disable=too-many-branches
+        """
+        Registers a binary expression.
+
+        Args:
+            node (BinaryExpression): The binary expression.
+            context (Context): The current context.
+
+        Returns:
+            The result of the evaluated expression.
+        """
+        lhs = preprocess_operands(self.visit(node.lhs, context))
+        rhs = preprocess_operands(self.visit(node.rhs, context))
+        ret = None
+        if node.op.name in NON_ASSIGNMENT_CLASSICAL_OPERATORS:
+            match node.op.name:
+                case "==":
+                    ret = lhs == rhs
+                case "!=":
+                    ret = lhs != rhs
+                case "~=":
+                    ret = np.isclose(lhs, rhs)
+                case ">":
+                    ret = lhs > rhs
+                case "<":
+                    ret = lhs < rhs
+                case ">=":
+                    ret = lhs >= rhs
+                case "<=":
+                    ret = lhs <= rhs
+                case ">>":
+                    ret = lhs >> rhs
+                case "<<":
+                    ret = lhs << rhs
+                case "+":
+                    ret = lhs + rhs
+                case "-":
+                    ret = lhs - rhs
+                case "*":
+                    ret = lhs * rhs
+                case "**":
+                    ret = lhs**rhs
+                case "/":
+                    ret = lhs / rhs
+                case "%":
+                    ret = lhs % rhs
+                case "|":
+                    ret = lhs | rhs
+                case "||":
+                    ret = lhs or rhs
+                case "&":
+                    ret = lhs & rhs
+                case "&&":
+                    ret = lhs and rhs
+                case "^":
+                    ret = lhs ^ rhs
+            return ret
+        # we shouldn't ever get thi error if the parser did its job right
+        raise SyntaxError(  # pragma: no cover
+            f"Invalid operator {node.op.name} encountered in binary expression "
+            f"on line {node.span.start_line}."
+        )  # pragma: no cover
+
+    @visit.register(UnaryExpression)
+    def visit_unary_expression(self, node: UnaryExpression, context: Context):
+        """
+        Registers a unary expression.
+
+        Args:
+            node (UnaryExpression): The unary expression.
+            context (Context): The current context.
+
+        Returns:
+            The result of the evaluated expression.
+        """
+        if node.op.name in NON_ASSIGNMENT_CLASSICAL_OPERATORS:
+            operand = preprocess_operands(self.visit(node.expression, context))
+            if node.op.name == "!":
+                return not operand
+            if node.op.name == "-":
+                return -operand
+            if node.op.name == "~":
+                return ~operand  # pylint: disable=invalid-unary-operand-type
+        # we shouldn't ever get thi error if the parser did its job right
+        raise SyntaxError(  # pragma: no cover
+            f"Invalid operator {node.op.name} encountered in unary expression "
+            f"on line {node.span.start_line}."
+        )  # pragma: no cover
+
     @visit.register(IndexExpression)
     def visit_index_expression(
         self, node: IndexExpression, context: Context, aliasing: bool = False

tests/io/qasm_interpreter/assignment.qasm

@@ -0,0 +1,25 @@
+int a = 0;
+int b = 1;
+int c = 2;
+int d = 3;
+int e = 4;
+bool f = true;
+bool g = true;
+bit i = "10";
+bit j = "10";
+bit k = "11";
+int l = 5;
+int m = 6;
+int n = 2;
+a = 1;
+b += 2;
+c -= 3;
+d *= 4;
+e /= 5;
+f &= true;
+g |= false;
+i ^= 9;
+j <<= 10;
+k >>= 1;
+l %= 2;
+m **= 13;

tests/io/qasm_interpreter/binary_expressions.qasm

@@ -0,0 +1,32 @@
+const int lhs = 0;
+const int rhs = 1;
+bool a = lhs == rhs;
+bool b = lhs != rhs;
+bool c = lhs > rhs;
+bool d = lhs < rhs;
+bool e = lhs >= rhs;
+bool f = lhs <= rhs;
+
+bit lh = "10";
+bit rh = "01";
+bit g = lh >> rh;
+bit h = lh << rh;
+bit i = lh | rh;
+bit j = lh ^ rh;
+bit k = lh & rh;
+
+int l = 3;
+int r = 2;
+int m = l + r;
+int n = l / r;
+int o = l - r;
+int p = l * r;
+int q = l ** r;
+int s = l % r;
+
+bool left = true;
+bool right = false;
+int t = left || right;
+int u = left && right;
+
+bool v = lhs ~= rhs;