Skip to content

Commit

Permalink
IR: Ensure Conditional.else_bodies does not return empty tuples
Browse files Browse the repository at this point in the history
  • Loading branch information
mlange05 committed Oct 9, 2024
1 parent 2ee867e commit ddc67d1
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
2 changes: 1 addition & 1 deletion loki/ir/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -712,7 +712,7 @@ def else_bodies(self):
"""
if self.has_elseif:
return (self.else_body[0].body,) + self.else_body[0].else_bodies
return (self.else_body,)
return (self.else_body,) if self.else_body else ()


@dataclass_strict(frozen=True)
Expand Down
19 changes: 19 additions & 0 deletions loki/ir/tests/test_ir_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,25 @@ def test_multi_conditional(scope, one, i, n, a_i):
assert isinstance(else_bodies[2][0], ir.Assignment)
assert else_bodies[2][0].lhs == 'a(i)' and else_bodies[2][0].rhs == '42.0'

# Not try without the final else
multicond = ir.Conditional(
condition=sym.Comparison(i, '==', sym.IntLiteral(1)),
body=ir.Assignment(lhs=a_i, rhs=sym.Literal(1.0)),
)
for idx in range(2, 4):
multicond = ir.Conditional(
condition=sym.Comparison(i, '==', sym.IntLiteral(idx)),
body=ir.Assignment(lhs=a_i, rhs=sym.Literal(float(idx))),
else_body=multicond, has_elseif=True
)
else_bodies = multicond.else_bodies
assert len(else_bodies) == 2
assert all(isinstance(b, tuple) for b in else_bodies)
assert isinstance(else_bodies[0][0], ir.Assignment)
assert else_bodies[0][0].lhs == 'a(i)' and else_bodies[0][0].rhs == '2.0'
assert isinstance(else_bodies[1][0], ir.Assignment)
assert else_bodies[1][0].lhs == 'a(i)' and else_bodies[1][0].rhs == '1.0'


def test_section(scope, one, i, n, a_n, a_i):
"""
Expand Down

0 comments on commit ddc67d1

Please sign in to comment.