support annotated assign (#344)

closes #324

Author: Roger-luo

src/kirin/dialects/ilist/interp.py

@@ -17,12 +17,7 @@ def _range(self, interp, frame: Frame, stmt: Range):
 
     @impl(New)
     def new(self, interp, frame: Frame, stmt: New):
-        elem_type = types.Any
-        if stmt.values:
-            elem_type = stmt.values[0].type
-            for each in stmt.values[1:]:
-                elem_type = elem_type.join(each.type)
-        return (IList(list(frame.get_values(stmt.values)), elem=elem_type),)
+        return (IList(list(frame.get_values(stmt.values)), elem=stmt.elem_type),)
 
     @impl(Len, types.PyClass(IList))
     def len(self, interp, frame: Frame, stmt: Len):

src/kirin/dialects/ilist/lowering.py

@@ -19,4 +19,32 @@ def lower_List(self, state: lowering.State, node: ast.List) -> lowering.Result:
         else:
             typ = types.Any
 
-        return state.current_frame.push(ilist.New(values=tuple(elts)))
+        return state.current_frame.push(ilist.New(values=tuple(elts), elem_type=typ))
+
+    @lowering.akin(ilist.IList)
+    def lower_Call_IList(
+        self, state: lowering.State, node: ast.Call
+    ) -> lowering.Result:
+        if len(node.args) != 1:
+            raise lowering.BuildError("IList constructor only takes one argument")
+        value = node.args[0]
+        if not isinstance(value, ast.List):
+            raise lowering.BuildError("IList constructor only takes a list")
+
+        if len(node.keywords) > 1:
+            raise lowering.BuildError(
+                "IList constructor only takes one keyword argument"
+            )
+
+        if len(node.keywords) == 1:
+            kw = node.keywords[0]
+            if kw.arg != "elem":
+                raise lowering.BuildError(
+                    "IList constructor only takes elem keyword argument"
+                )
+            elem = self.get_hint(state, kw.value)
+            elts = tuple(state.lower(each).expect_one() for each in value.elts)
+            stmt = ilist.New(values=tuple(elts), elem_type=elem)
+            return state.current_frame.push(stmt)
+        else:
+            return self.lower_List(state, value)

src/kirin/dialects/ilist/stmts.py

@@ -25,25 +25,22 @@ class Range(ir.Statement):
 class New(ir.Statement):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     values: tuple[ir.SSAValue, ...] = info.argument(ElemT)
+    elem_type: types.TypeAttribute = info.attribute()
     result: ir.ResultValue = info.result(IListType[ElemT])
 
     def __init__(
         self,
         values: Sequence[ir.SSAValue],
+        elem_type: types.TypeAttribute | None = None,
     ) -> None:
-        # get elem type
-        if not values:
+        if not elem_type:
             elem_type = types.Any
-        else:
-            elem_type = values[0].type
-            for v in values:
-                elem_type = elem_type.join(v.type)
-
         result_type = IListType[elem_type, types.Literal(len(values))]
         super().__init__(
             args=values,
             result_types=(result_type,),
             args_slice={"values": slice(0, len(values))},
+            attributes={"elem_type": elem_type},
         )
 
 

src/kirin/dialects/lowering/func.py

@@ -115,16 +115,3 @@ def assert_simple_arguments(self, node: ast.arguments) -> None:
 
         if node.posonlyargs:
             raise lowering.BuildError("positional-only arguments are not supported")
-
-    @staticmethod
-    def get_hint(state: lowering.State, node: ast.expr | None):
-        if node is None:
-            return types.Any
-
-        try:
-            t = state.get_global(node).data
-            return types.hint2type(t)
-        except Exception as e:  # noqa: E722
-            raise lowering.BuildError(
-                f"expect a type hint, got {ast.unparse(node)}"
-            ) from e

src/kirin/dialects/lowering/range.py

@@ -15,6 +15,7 @@
 @py.register
 class PyLowering(lowering.FromPythonAST):
 
+    @lowering.akin(range)
     def lower_Call_range(
         self, state: lowering.State, node: ast.Call
     ) -> lowering.Result:
@@ -24,6 +25,7 @@ def lower_Call_range(
 @ilist.register
 class IListLowering(lowering.FromPythonAST):
 
+    @lowering.akin(range)
     def lower_Call_range(
         self, state: lowering.State, node: ast.Call
     ) -> lowering.Result:

src/kirin/dialects/py/assign.py

@@ -49,6 +49,31 @@ class SetItem(ir.Statement):
     index: ir.SSAValue = info.argument(print=False)
 
 
+@statement(dialect=dialect)
+class SetAttribute(ir.Statement):
+    name = "setattr"
+    traits = frozenset({lowering.FromPythonCall()})
+    obj: ir.SSAValue = info.argument(print=False)
+    attr: str = info.attribute()
+    value: ir.SSAValue = info.argument(print=False)
+
+
+@statement(dialect=dialect)
+class TypeAssert(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    got: ir.SSAValue = info.argument(print=False)
+    expected: types.TypeAttribute = info.attribute()
+    result: ir.ResultValue = info.result()
+
+    def __init__(self, got: ir.SSAValue, *, expected: types.TypeAttribute):
+        super().__init__(
+            args=(got,),
+            attributes={"expected": expected},
+            result_types=(expected,),
+            args_slice={"got": 0},
+        )
+
+
 @dialect.register
 class Concrete(interp.MethodTable):
 
@@ -59,7 +84,37 @@ def alias(self, interp, frame: interp.Frame, stmt: Alias):
     @interp.impl(SetItem)
     def setindex(self, interp, frame: interp.Frame, stmt: SetItem):
         frame.get(stmt.obj)[frame.get(stmt.index)] = frame.get(stmt.value)
-        return (None,)
+
+    @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)
+
+    # NOTE: we don't do much runtime type checking here, object with generic
+    # types will unlikely work here.
+    # TODO: consider runtime type checking by boxing the 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 interp.WrapException(
+                TypeError(f"Expected {stmt.expected}, got {got_type}")
+            )
+        return (frame.get(stmt.got),)
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+    @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,)
 
 
 @dialect.register
@@ -79,16 +134,43 @@ def lower_Assign(self, state: lowering.State, node: ast.Assign) -> lowering.Resu
                 current_frame.defs[lhs_name] = current_frame.push(stmt).result
             case _:
                 for target, value in zip(node.targets, result.data):
-                    match target:
-                        # NOTE: if the name exists new ssa value will be
-                        # used in the future to shadow the old one
-                        case ast.Name(name, ast.Store()):
-                            value.name = name
-                            current_frame.defs[name] = value
-                        case ast.Subscript(obj, slice):
-                            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}")
+                    self.assign_item(state, target, value)
+
+    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(state, node.target, stmt.result)
+
+    def lower_AugAssign(
+        self, state: lowering.State, node: ast.AugAssign
+    ) -> lowering.Result:
+        self.assign_item(state, node.target, state.lower(node.value).expect_one())
+
+    @staticmethod
+    def assign_item(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}")
+
+    @staticmethod
+    def assert_assign_value_type(value: ir.SSAValue, type_hint: types.TypeAttribute):
+        value_type = value.type.meet(type_hint)
+        if value_type is value_type.bottom():
+            raise lowering.BuildError(f"Cannot assign {value.type} to {type_hint}")
+        return value_type

src/kirin/dialects/py/attr.py

@@ -40,15 +40,14 @@ class Lowering(lowering.FromPythonAST):
     def lower_Attribute(
         self, state: lowering.State, node: ast.Attribute
     ) -> lowering.Result:
-        from kirin.dialects.py import Constant
 
         if not isinstance(node.ctx, ast.Load):
             raise lowering.BuildError(f"unsupported attribute context {node.ctx}")
 
         # NOTE: eagerly load global variables
         value = state.get_global(node, no_raise=True)
         if value is not None:
-            return state.current_frame.push(Constant(value.data))
+            return state.lower(ast.Constant(value.data)).expect_one()
 
         value = state.lower(node.value).expect_one()
         return state.current_frame.push(GetAttr(obj=value, attrname=node.attr))

src/kirin/dialects/py/builtin.py

@@ -39,9 +39,11 @@ class Sum(ir.Statement):
 @dialect.register
 class Lowering(lowering.FromPythonAST):
 
+    @lowering.akin(abs)
     def lower_Call_abs(self, state: lowering.State, node: Call) -> lowering.Result:
         return state.current_frame.push(Abs(state.lower(node.args[0]).expect_one()))
 
+    @lowering.akin(sum)
     def lower_Call_sum(self, state: lowering.State, node: Call) -> lowering.Result:
         return state.current_frame.push(Sum(state.lower(node.args[0]).expect_one()))
 

src/kirin/dialects/py/iterable.py

@@ -66,13 +66,15 @@ def next_(self, interp, frame: interp.Frame, stmt: Next):
 @dialect.register
 class Lowering(lowering.FromPythonAST):
 
+    @lowering.akin(iter)
     def lower_Call_iter(self, state: lowering.State, node: Call) -> lowering.Result:
         if len(node.args) != 1:
             raise lowering.BuildError("iter() takes exactly 1 argument")
         return state.current_frame.push(
             Iter(state.lower(node.args[0]).expect_one()),
         )
 
+    @lowering.akin(next)
     def lower_Call_next(self, state: lowering.State, node: Call) -> lowering.Result:
         if len(node.args) == 2:
             raise lowering.BuildError(

src/kirin/dialects/py/len.py

@@ -49,5 +49,6 @@ def len(self, interp, frame: interp.Frame, stmt: Len):
 @dialect.register
 class Lowering(lowering.FromPythonAST):
 
+    @lowering.akin(len)
     def lower_Call_len(self, state: lowering.State, node: ast.Call) -> lowering.Result:
         return state.current_frame.push(Len(state.lower(node.args[0]).expect_one()))