Fix bug that ignored defs in if when missing from else

[david-pl]

@Roger-luo this fixes the bug in #416 when there's a definition in the if path, but not in the else path.

However, the issue with the elif persists since the definition from the nested if isn't propagated upwards properly.

In code:

This fails on main, but works on this branch

@kernel
def main(n: int):
    x = 0

    if x == n:
        x = 1
    else:
        y = 2  # noqa: F841

    return x

main.print()

This fails both on main and this branch:

@kernel
def main_elif(n: int):
    x = 0

    if x == n:
        x = 3
    elif x == n + 1:
        x = 4

    return x

main_elif.print()

(or, equivalently)

@kernel
def main_elif(n: int):
    x = 0

    if x == n:
        x = 3
    else:
        if x == n + 1:
            x = 4

    return x

We somehow need to lower the bodies first in order to obtain the nested definitions in them, but I'm not sure how exactly.

[github-actions]

PR Preview Action v1.6.1
🚀 View preview at https://QuEraComputing.github.io/kirin/pr-preview/pr-417/
Built to branch gh-pages at 2025-06-13 07:07 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[david-pl]

pyright failure seems to be unrelated, I think.

[weinbe58]

@Roger-luo do we support elif as a nested ifelse statement in the else body?

[david-pl]

@Roger-luo do we support elif as a nested ifelse statement in the else body?

From what I saw looking into this: yes, that's how an elif is represented in the IR.

[Roger-luo]

actually why do we remove node.orelse here?

[david-pl]

actually why do we remove node.orelse here?

@Roger-luo you mean the if node.orelse? Because we need to check if the definitions are in either branch regardless of whether there's an orelse. Before, it would only check the else_frame if there was an orelse, which is wrong.

But also, this still doesn't work for all cases! See my question in the description.

[david-pl]

There we go: the issue that kept the elif from working was that the if nested in the else didn't check definitions present in its parent frame.

I added

if frame.parent is not None:
            frame.defs.update(frame.parent.defs)

and now everything seems to work.

The example from before

@kernel
def main_elif(n: int):
    x = 0

    if x == n:
        x = 3
    elif x == n + 1:
        x = 4

    return x

main_elif.print()

now prints

func.func @main_elif(n : !py.int) -> !Any {
  ^0(%main_elif_self, %n):
  │   %x = py.constant.constant 0 : !py.int
  │   %0 = py.cmp.eq(lhs=%x, rhs=%n : !py.int) : !py.bool
  │ %x_1 = scf.if %0 {
  │        │ ^1(%1):
  │        │ │ %x_2 = py.constant.constant 3 : !py.int
  │        │ │        scf.yield %x_2
  │        } else {
  │        │ ^2(%5):
  │        │ │   %2 = py.constant.constant 1 : !py.int
  │        │ │   %3 = py.binop.add(%n : !py.int, %2) : ~T
  │        │ │   %4 = py.cmp.eq(lhs=%x, rhs=%3) : !py.bool
  │        │ │ %x_3 = scf.if %4 {
  │        │ │        │ ^3(%6):
  │        │ │        │ │ %x_4 = py.constant.constant 4 : !py.int
  │        │ │        │ │        scf.yield %x_4
  │        │ │        } else {
  │        │ │        │ ^4(%7):
  │        │ │        │ │ scf.yield %x
  │        │ │        } -> purity=False
  │        │ │        scf.yield %x_3
  │        } -> purity=False
  │        func.return %x_1
} // func.func main_elif

That looks correct to me!

@Roger-luo ready for another round of review!

[Roger-luo]

I think this breaks some of the test case?

[david-pl]

Ah, damn, I need to take another look then.

[david-pl]

Looks like copying over the definitions from the parent frame was too much. Now, I just check for them and add them to the yields if needed.

[codecov]

Codecov Report

Attention: Patch coverage is 89.47368% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/kirin/dialects/scf/lowering.py	89.47%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
9686	8588	89%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
src/kirin/dialects/scf/lowering.py	95%	🟢
TOTAL	95%	🟢

updated for commit: 739169c by action🐍

[david-pl]

@Roger-luo I found another test case that didn't work when not copying over the defs with frame.defs.update(frame.parent.defs):

@kernel
def main_nested_if(n: int):
    x = 0

    if n > 0:
        if n > 1:
            if n == 3:
                x = 4

    return x

With multiple levels of nesting, the definitions are no longer carried over from the respective parent frame and I only checked up to one level above (just one frame.parent).

This would still work if we just copied over the definitions with frame.defs.update, but I couldn't figure out why some other tests mixing for loops and ifs broke then.

So, now I wrote a little helper function that walks through all parent frames to see if there's a definition there. That's probably not the cleanest way to do this, but at least all tests seem to work now.