fix multi assign (#348)

I just realize I misinterpreted what it means when there are multiple
targets in `ast.Assign` it actually means the following

```py
x = y.z = 1
```

this also fix:
- statement with multiple return values. Allow one create such
statements without using a tuple.
- named expr, e.g `x := 1`
- `ast.AugAssign`, e.g `x += 1`

Author: Roger-luo

src/kirin/dialects/py/assign.py

@@ -133,24 +133,42 @@ def lower_Assign(self, state: lowering.State, node: ast.Assign) -> lowering.Resu
                 stmt.result.name = lhs_name
                 current_frame.defs[lhs_name] = current_frame.push(stmt).result
             case _:
-                for target, value in zip(node.targets, result.data):
-                    self.assign_item(state, target, value)
+                for target in node.targets:
+                    self.assign_item(state, target, result)
 
     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)
+        self.assign_item_value(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())
+        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())
+        self.assign_item_value(
+            state,
+            node.target,
+            state.lower(ast.BinOp(rhs, node.op, node.value)).expect_one(),
+        )
 
-    @staticmethod
-    def assign_item(state: lowering.State, target, value: ir.SSAValue):
+    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
+
+    @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()):
@@ -168,9 +186,15 @@ def assign_item(state: lowering.State, target, value: ir.SSAValue):
             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
+    @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())

test/lowering/test_assign.py

@@ -0,0 +1,84 @@
+import pytest
+
+from kirin import ir, lowering
+from kirin.decl import info, statement
+from kirin.prelude import basic_no_opt
+from kirin.dialects import cf, py
+
+dialect = ir.Dialect("test")
+
+
+@statement(dialect=dialect)
+class MultiResult(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    result_a: ir.ResultValue = info.result()
+    result_b: ir.ResultValue = info.result()
+
+
+dummy_dialect = basic_no_opt.add(dialect)
+
+
+def test_multi_result():
+    @dummy_dialect
+    def multi_assign():
+        (x, y) = MultiResult()  # type: ignore
+        return x, y
+
+    stmt = multi_assign.callable_region.blocks[0].stmts.at(0)
+    assert isinstance(stmt, MultiResult)
+    assert stmt.result_a.name == "x"
+    assert stmt.result_b.name == "y"
+
+    with pytest.raises(lowering.BuildError):
+
+        @dummy_dialect
+        def multi_assign_error():
+            (x, y, z) = MultiResult()  # type: ignore
+            return x, y, z
+
+
+def test_chain_assign_setattr():
+
+    @dummy_dialect
+    def chain_assign(y):
+        x = y.z = 1
+        return x, y
+
+    stmt = chain_assign.callable_region.blocks[0].stmts.at(1)
+    assert isinstance(stmt, py.assign.SetAttribute)
+    assert stmt.obj.name == "y"
+    assert stmt.attr == "z"
+    assert stmt.value.name == "x"
+
+
+def test_aug_assign():
+    @dummy_dialect
+    def aug_assign(y):
+        y += 1
+        return y
+
+    y = aug_assign.callable_region.blocks[0].args[1]
+    const = aug_assign.callable_region.blocks[0].stmts.at(0)
+    assert isinstance(const, py.Constant)
+    assert const.value.unwrap() == 1
+    add = aug_assign.callable_region.blocks[0].stmts.at(1)
+    assert isinstance(add, py.binop.Add)
+    assert add.lhs is y
+    assert add.rhs is const.result
+
+
+def test_named_expr():
+
+    @dummy_dialect
+    def named_expr(y):
+        if y := y + 1:
+            return y
+        return y
+
+    stmt = named_expr.callable_region.blocks[0].stmts.at(1)
+    y = named_expr.callable_region.blocks[0].args[1]
+    assert isinstance(stmt, py.binop.Add)
+    assert stmt.lhs is y
+    br = named_expr.callable_region.blocks[0].stmts.at(2)
+    assert isinstance(br, cf.ConditionalBranch)
+    assert stmt.result is br.cond