Skip to content

Commit

Permalink
multi-assign tuple unrolling
Browse files Browse the repository at this point in the history
  • Loading branch information
dakk committed Oct 28, 2023
1 parent d55a0de commit 3084751
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 0 deletions.
22 changes: 22 additions & 0 deletions qlasskit/ast2ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,28 @@ def visit_FunctionDef(self, node):
return super().generic_visit(node)

def visit_Assign(self, node):
# Transform multi-target assign to single target assigns
if len(node.targets) == 1 and hasattr(node.targets[0], "elts"):
_temptup = self.visit(
ast.Assign(targets=[ast.Name(id="_temptup")], value=node.value)
)

single_assigns = [
self.visit(
ast.Assign(
targets=[ast.Name(id=node.targets[0].elts[i].id)],
value=ast.Subscript(
value=ast.Name(id="_temptup"),
slice=ast.Index(
value=ast.Constant(value=i), ctx=ast.Load()
),
),
)
)
for i in range(len(node.targets[0].elts))
]
return [_temptup] + single_assigns

was_known = node.targets[0].id in self.env

if isinstance(node.value, ast.Name) and node.value.id in self.env:
Expand Down
15 changes: 15 additions & 0 deletions test/test_qlassf_tuple.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,21 @@ def test_tuple_assign3(self):
)
compute_and_compare_results(self, qf)

def test_multi_assign(self):
f = "def test(a: bool) -> bool:\n\tc, d = a, not a\n\treturn c and d"
qf = qlassf(f, to_compile=COMPILATION_ENABLED)
compute_and_compare_results(self, qf)

def test_multi_assign2(self):
f = "def test() -> Qint4:\n\tc, d = 1, 2\n\treturn c+d"
qf = qlassf(f, to_compile=COMPILATION_ENABLED)
compute_and_compare_results(self, qf)

def test_multi_assign3(self):
f = "def test() -> Qint4:\n\tc, d, e = 1, 2, 0xa\n\treturn c+d+e"
qf = qlassf(f, to_compile=COMPILATION_ENABLED)
compute_and_compare_results(self, qf)

def test_tuple_result(self):
f = "def test(a: bool, b: bool) -> Tuple[bool,bool]:\n\treturn a,b"
qf = qlassf(f, to_compile=COMPILATION_ENABLED)
Expand Down

0 comments on commit 3084751

Please sign in to comment.