fix regression with cache

Author: jcjaskula-aws

oqpy/base.py

@@ -363,7 +363,7 @@ class CachedExpressionConvertible(Protocol):
 
     _oqpy_cache_key: Hashable
 
-    def _to_cached_oqpy_expression(self) -> HasToAst: ...  # pragma: no cover
+    def _to_cached_oqpy_expression(self) -> AstConvertible: ...  # pragma: no cover
 
 
 class OQPyUnaryExpression(OQPyExpression):
@@ -490,8 +490,10 @@ def to_ast(program: Program, item: AstConvertible) -> ast.Expression:
         if item._oqpy_cache_key is None:
             item._oqpy_cache_key = uuid.uuid1()
         if item._oqpy_cache_key not in program.expr_cache:
-            program.expr_cache[item._oqpy_cache_key] = item._to_cached_oqpy_expression()
-        item = program.expr_cache[item._oqpy_cache_key]
+            program.expr_cache[item._oqpy_cache_key] = to_ast(
+                program, item._to_cached_oqpy_expression()
+            )
+        return program.expr_cache[item._oqpy_cache_key]
     if isinstance(item, (complex, np.complexfloating)):
         if item.imag == 0:
             return to_ast(program, item.real)
@@ -507,7 +509,9 @@ def to_ast(program: Program, item: AstConvertible) -> ast.Expression:
                 ast.ImaginaryLiteral(-item.imag),
             )
         return ast.BinaryExpression(
-            ast.BinaryOperator["+"], ast.FloatLiteral(item.real), ast.ImaginaryLiteral(item.imag)
+            ast.BinaryOperator["+"],
+            ast.FloatLiteral(item.real),
+            ast.ImaginaryLiteral(item.imag),
         )
     if isinstance(item, (bool, np.bool_)):
         return ast.BooleanLiteral(item)

tests/test_directives.py

@@ -1665,6 +1665,7 @@ def _to_cached_oqpy_expression(self):
     assert dur.count == 2
     # This gets computed just once
     assert frame.count == 1
+    assert all(isinstance(v, ast.QASMNode) for v in prog.expr_cache.values())
 
 
 def test_waveform_extern_arg_passing():