Core Python Dialects

def lower_Expr(self, state: lowering.State, node: ast.Expr) -> lowering.Result:
    return state.parent.visit(state, node.value)

def lower_Name(self, state: lowering.State, node: ast.Name) -> lowering.Result:
    name = node.id
    if isinstance(node.ctx, ast.Load):
        value = state.current_frame.get(name)
        if value is None:
            raise lowering.BuildError(f"{name} is not defined")
        return value
    elif isinstance(node.ctx, ast.Store):
        raise lowering.BuildError("unhandled store operation")
    else:  # Del
        raise lowering.BuildError("unhandled del operation")

@interp.impl(Constant)
def constant(self, interp, frame: interp.Frame, stmt: Constant):
    return (stmt.value.unwrap(),)

def check_type(self) -> None:
    if not isinstance(self.result.type, types.TypeAttribute):
        raise TypeError(
            f"Expected result type to be PyType, got {self.result.type}"
        )

def print_impl(self, printer: Printer) -> None:
    printer.print_name(self)
    printer.plain_print(" ")
    if isinstance(self.value, ir.PyAttr):
        printer.plain_print(repr(self.value.data))
    else:  # other attributes
        printer.plain_print(repr(self.value))
    with printer.rich(style="comment"):
        printer.plain_print(" : ")
        printer.print(self.result.type)

@interp.impl(Constant)
def emit_Constant(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Constant):
    return (emit.emit_attribute(stmt.value),)

def lower_Constant(
    self, state: lowering.State, node: ast.Constant
) -> lowering.Result:
    return state.current_frame.push(
        Constant(node.value),
    )

def print_impl(self, printer: Printer) -> None:
    with printer.rich(style="keyword"):
        printer.print_name(self)

    printer.plain_print(" ")
    printer.print(self.condition)

    if self.message:
        printer.plain_print(", ")
        printer.print(self.message)

@interp.impl(Assert)
def assert_stmt(
    self, interp_: interp.Interpreter, frame: interp.Frame, stmt: Assert
):
    if frame.get(stmt.condition) is True:
        return ()

    if stmt.message:
        raise AssertionError(frame.get(stmt.message))
    else:
        raise AssertionError("Assertion failed")

@interp.impl(Assert)
def emit_assert(self, interp: julia.EmitJulia, frame: EmitStrFrame, stmt: Assert):
    interp.writeln(
        frame, f"@assert {frame.get(stmt.condition)} {frame.get(stmt.message)}"
    )
    return ()

def lower_Assert(self, state: lowering.State, node: ast.Assert) -> lowering.Result:
    from kirin.dialects.py.constant import Constant

    cond = state.lower(node.test).expect_one()
    if node.msg:
        message = state.lower(node.msg).expect_one()
        state.current_frame.push(Assert(condition=cond, message=message))
    else:
        message_stmt = state.current_frame.push(Constant(""))
        state.current_frame.push(
            Assert(condition=cond, message=message_stmt.result)
        )

@interp.impl(Assert)
def assert_stmt(self, interp, frame, stmt: Assert):
    return (types.Bottom,)

def print_impl(self, printer: Printer) -> None:
    printer.print_name(self)
    printer.plain_print(" ")
    with printer.rich(style="symbol"):
        printer.plain_print(self.target.data)

    with printer.rich(style="keyword"):
        printer.plain_print(" = ")

    printer.print(self.value)

@interp.impl(Alias)
def alias(self, interp, frame: interp.Frame, stmt: Alias):
    return (frame.get(stmt.value),)

@interp.impl(SetAttribute)
def set_attribute(self, interp, frame: interp.Frame, stmt: SetAttribute):
    obj = frame.get(stmt.obj)
    value = frame.get(stmt.value)
    setattr(obj, stmt.attr, value)

@interp.impl(SetItem)
def setindex(self, interp, frame: interp.Frame, stmt: SetItem):
    frame.get(stmt.obj)[frame.get(stmt.index)] = frame.get(stmt.value)

@interp.impl(TypeAssert)
def type_assert(self, interp_, frame: interp.Frame, stmt: TypeAssert):
    got = frame.get(stmt.got)
    got_type = types.PyClass(type(got))
    if not got_type.is_subseteq(stmt.expected):
        raise TypeError(f"Expected {stmt.expected}, got {got_type}")
    return (frame.get(stmt.got),)

@classmethod
def assign_item(cls, state: lowering.State, target, result: lowering.State.Result):
    match target:
        case ast.Tuple(elts, ast.Store()):
            if len(elts) != len(result.data):
                raise lowering.BuildError(
                    f"tuple assignment length mismatch: {len(elts)} != {len(result.data)}"
                )
            for target, value in zip(elts, result.data):
                cls.assign_item_value(state, target, value)
        case _:
            cls.assign_item_value(state, target, result.expect_one())

@classmethod
def assign_item_value(cls, state: lowering.State, target, value: ir.SSAValue):
    current_frame = state.current_frame
    match target:
        case ast.Name(name, ast.Store()):
            value.name = name
            current_frame.defs[name] = value
        case ast.Attribute(obj, attr, ast.Store()):
            obj = state.lower(obj).expect_one()
            stmt = SetAttribute(obj, value, attr=attr)
            current_frame.push(stmt)
        case ast.Subscript(obj, slice, ast.Store()):
            obj = state.lower(obj).expect_one()
            slice = state.lower(slice).expect_one()
            stmt = SetItem(obj=obj, index=slice, value=value)
            current_frame.push(stmt)
        case _:
            raise lowering.BuildError(f"unsupported target {target}")

def lower_AnnAssign(
    self, state: lowering.State, node: ast.AnnAssign
) -> lowering.Result:
    type_hint = self.get_hint(state, node.annotation)
    value = state.lower(node.value).expect_one()
    stmt = state.current_frame.push(TypeAssert(got=value, expected=type_hint))
    self.assign_item_value(state, node.target, stmt.result)

def lower_Assign(self, state: lowering.State, node: ast.Assign) -> lowering.Result:
    result = state.lower(node.value)
    current_frame = state.current_frame
    match node:
        case ast.Assign(
            targets=[ast.Name(lhs_name, ast.Store())], value=ast.Name(_, ast.Load())
        ):
            stmt = Alias(
                value=result.data[0], target=ir.PyAttr(lhs_name)
            )  # NOTE: this is guaranteed to be one result
            stmt.result.name = lhs_name
            current_frame.defs[lhs_name] = current_frame.push(stmt).result
        case _:
            for target in node.targets:
                self.assign_item(state, target, result)

def lower_AugAssign(
    self, state: lowering.State, node: ast.AugAssign
) -> lowering.Result:
    match node.target:
        case ast.Name(name, ast.Store()):
            rhs = ast.Name(name, ast.Load())
        case ast.Attribute(obj, attr, ast.Store()):
            rhs = ast.Attribute(obj, attr, ast.Load())
        case ast.Subscript(obj, slice, ast.Store()):
            rhs = ast.Subscript(obj, slice, ast.Load())
        case _:
            raise lowering.BuildError(f"unsupported target {node.target}")
    self.assign_item_value(
        state,
        node.target,
        state.lower(ast.BinOp(rhs, node.op, node.value)).expect_one(),
    )

def lower_NamedExpr(
    self, state: lowering.State, node: ast.NamedExpr
) -> lowering.Result:
    value = state.lower(node.value).expect_one()
    self.assign_item_value(state, node.target, value)
    return value

@interp.impl(TypeAssert)
def type_assert(
    self, interp_, frame: interp.Frame[types.TypeAttribute], stmt: TypeAssert
):
    got = frame.get(stmt.got)
    if got.is_subseteq(stmt.expected):
        return (got.meet(stmt.expected),)
    return (types.Bottom,)

@interp.impl(Unpack)
def unpack(self, interp: interp.Interpreter, frame: interp.Frame, stmt: Unpack):
    return tuple(frame.get(stmt.value))

@interp.impl(Unpack)
def unpack(self, interp, frame: interp.Frame[types.TypeAttribute], stmt: Unpack):
    value = frame.get(stmt.value)
    if isinstance(value, types.Generic) and value.is_subseteq(types.Tuple):
        if value.vararg:
            rest = tuple(value.vararg.typ for _ in stmt.names[len(value.vars) :])
            return tuple(value.vars) + rest
        else:
            return value.vars
    # TODO: support unpacking other types
    return tuple(types.Any for _ in stmt.names)

def print_impl(self, printer: Printer) -> None:
    printer.print_name(self)
    printer.plain_print(" ")
    printer.print(self.value)

@interp.impl(And)
def and_(self, interp, frame: interp.Frame, stmt: And):
    return (frame.get(stmt.lhs) and frame.get(stmt.rhs),)

@interp.impl(Or)
def or_(self, interp, frame: interp.Frame, stmt: Or):
    return (frame.get(stmt.lhs) or frame.get(stmt.rhs),)

@interp.impl(And)
def emit_And(self, emit: EmitJulia, frame: EmitStrFrame, stmt: And):
    return emit.emit_binaryop(frame, "&&", stmt.lhs, stmt.rhs, stmt.result)

@interp.impl(Or)
def emit_Or(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Or):
    return emit.emit_binaryop(frame, "||", stmt.lhs, stmt.rhs, stmt.result)

def lower_BoolOp(self, state: lowering.State, node: ast.BoolOp) -> lowering.Result:
    lhs = state.lower(node.values[0]).expect_one()
    match node.op:
        case ast.And():
            boolop = And
        case ast.Or():
            boolop = Or
        case _:
            raise lowering.BuildError(f"unsupported boolop {node.op}")

    for value in node.values[1:]:
        lhs = state.current_frame.push(
            boolop(lhs=lhs, rhs=state.lower(value).expect_one())
        ).result
    return lhs