core/typecheck: fix missing ExprKind::Asterisk in fix_assignment_target_context

- distinguish between setitem and getitem
- allow starred assignment targets, but the assigned value would be a tuple
- allow both [...] and (...) to be target lists

nac3core/src/codegen/expr.rs

@@ -995,8 +995,10 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
     ctx.builder.position_at_end(init_bb);
 
     let Comprehension { target, iter, ifs, .. } = &generators[0];
+
+    let iter_ty = iter.custom.unwrap();
     let iter_val = if let Some(v) = generator.gen_expr(ctx, iter)? {
-        v.to_basic_value_enum(ctx, generator, iter.custom.unwrap())?
+        v.to_basic_value_enum(ctx, generator, iter_ty)?
     } else {
         for bb in [test_bb, body_bb, cont_bb] {
             ctx.builder.position_at_end(bb);
@@ -1014,96 +1016,120 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
     ctx.builder.build_store(index, zero_size_t).unwrap();
 
     let elem_ty = ctx.get_llvm_type(generator, elt.custom.unwrap());
-    let is_range = ctx.unifier.unioned(iter.custom.unwrap(), ctx.primitives.range);
     let list;
 
-    if is_range {
+    match &*ctx.unifier.get_ty(iter_ty) {
-        let iter_val = RangeValue::from_ptr_val(iter_val.into_pointer_value(), Some("range"));
+        TypeEnum::TObj { obj_id, .. }
-        let (start, stop, step) = destructure_range(ctx, iter_val);
+            if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
-        let diff = ctx.builder.build_int_sub(stop, start, "diff").unwrap();
+        {
-        // add 1 to the length as the value is rounded to zero
+            let iter_val = RangeValue::from_ptr_val(iter_val.into_pointer_value(), Some("range"));
-        // the length may be 1 more than the actual length if the division is exact, but the
+            let (start, stop, step) = destructure_range(ctx, iter_val);
-        // length is a upper bound only anyway so it does not matter.
+            let diff = ctx.builder.build_int_sub(stop, start, "diff").unwrap();
-        let length = ctx.builder.build_int_signed_div(diff, step, "div").unwrap();
+            // add 1 to the length as the value is rounded to zero
-        let length = ctx.builder.build_int_add(length, int32.const_int(1, false), "add1").unwrap();
+            // the length may be 1 more than the actual length if the division is exact, but the
-        // in case length is non-positive
+            // length is a upper bound only anyway so it does not matter.
-        let is_valid =
+            let length = ctx.builder.build_int_signed_div(diff, step, "div").unwrap();
-            ctx.builder.build_int_compare(IntPredicate::SGT, length, zero_32, "check").unwrap();
+            let length =
+                ctx.builder.build_int_add(length, int32.const_int(1, false), "add1").unwrap();
+            // in case length is non-positive
+            let is_valid =
+                ctx.builder.build_int_compare(IntPredicate::SGT, length, zero_32, "check").unwrap();
 
-        let list_alloc_size = ctx
+            let list_alloc_size = ctx
-            .builder
+                .builder
-            .build_select(
+                .build_select(
-                is_valid,
+                    is_valid,
-                ctx.builder.build_int_z_extend_or_bit_cast(length, size_t, "z_ext_len").unwrap(),
+                    ctx.builder
-                zero_size_t,
+                        .build_int_z_extend_or_bit_cast(length, size_t, "z_ext_len")
-                "listcomp.alloc_size",
+                        .unwrap(),
-            )
+                    zero_size_t,
-            .unwrap();
+                    "listcomp.alloc_size",
-        list = allocate_list(
+                )
-            generator,
+                .unwrap();
-            ctx,
+            list = allocate_list(
-            Some(elem_ty),
+                generator,
-            list_alloc_size.into_int_value(),
+                ctx,
-            Some("listcomp.addr"),
+                Some(elem_ty),
-        );
+                list_alloc_size.into_int_value(),
+                Some("listcomp.addr"),
+            );
 
-        let i = generator.gen_store_target(ctx, target, Some("i.addr"))?.unwrap();
+            let i = generator.gen_store_target(ctx, target, Some("i.addr"))?.unwrap();
-        ctx.builder
+            ctx.builder
-            .build_store(i, ctx.builder.build_int_sub(start, step, "start_init").unwrap())
+                .build_store(i, ctx.builder.build_int_sub(start, step, "start_init").unwrap())
-            .unwrap();
+                .unwrap();
 
-        ctx.builder
+            ctx.builder
-            .build_conditional_branch(gen_in_range_check(ctx, start, stop, step), test_bb, cont_bb)
+                .build_conditional_branch(
-            .unwrap();
+                    gen_in_range_check(ctx, start, stop, step),
+                    test_bb,
+                    cont_bb,
+                )
+                .unwrap();
 
-        ctx.builder.position_at_end(test_bb);
+            ctx.builder.position_at_end(test_bb);
-        // add and test
+            // add and test
-        let tmp = ctx
+            let tmp = ctx
-            .builder
+                .builder
-            .build_int_add(
+                .build_int_add(
-                ctx.builder.build_load(i, "i").map(BasicValueEnum::into_int_value).unwrap(),
+                    ctx.builder.build_load(i, "i").map(BasicValueEnum::into_int_value).unwrap(),
-                step,
+                    step,
-                "start_loop",
+                    "start_loop",
-            )
+                )
-            .unwrap();
+                .unwrap();
-        ctx.builder.build_store(i, tmp).unwrap();
+            ctx.builder.build_store(i, tmp).unwrap();
-        ctx.builder
+            ctx.builder
-            .build_conditional_branch(gen_in_range_check(ctx, tmp, stop, step), body_bb, cont_bb)
+                .build_conditional_branch(
-            .unwrap();
+                    gen_in_range_check(ctx, tmp, stop, step),
+                    body_bb,
+                    cont_bb,
+                )
+                .unwrap();
 
-        ctx.builder.position_at_end(body_bb);
+            ctx.builder.position_at_end(body_bb);
-    } else {
+        }
-        let length = ctx
+        TypeEnum::TObj { obj_id, .. }
-            .build_gep_and_load(
+            if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
-                iter_val.into_pointer_value(),
+        {
-                &[zero_size_t, int32.const_int(1, false)],
+            let length = ctx
-                Some("length"),
+                .build_gep_and_load(
-            )
+                    iter_val.into_pointer_value(),
-            .into_int_value();
+                    &[zero_size_t, int32.const_int(1, false)],
-        list = allocate_list(generator, ctx, Some(elem_ty), length, Some("listcomp"));
+                    Some("length"),
+                )
+                .into_int_value();
+            list = allocate_list(generator, ctx, Some(elem_ty), length, Some("listcomp"));
 
-        let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?;
+            let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?;
-        // counter = -1
+            // counter = -1
-        ctx.builder.build_store(counter, size_t.const_all_ones()).unwrap();
+            ctx.builder.build_store(counter, size_t.const_all_ones()).unwrap();
-        ctx.builder.build_unconditional_branch(test_bb).unwrap();
+            ctx.builder.build_unconditional_branch(test_bb).unwrap();
 
-        ctx.builder.position_at_end(test_bb);
+            ctx.builder.position_at_end(test_bb);
-        let tmp = ctx.builder.build_load(counter, "i").map(BasicValueEnum::into_int_value).unwrap();
+            let tmp =
-        let tmp = ctx.builder.build_int_add(tmp, size_t.const_int(1, false), "inc").unwrap();
+                ctx.builder.build_load(counter, "i").map(BasicValueEnum::into_int_value).unwrap();
-        ctx.builder.build_store(counter, tmp).unwrap();
+            let tmp = ctx.builder.build_int_add(tmp, size_t.const_int(1, false), "inc").unwrap();
-        let cmp = ctx.builder.build_int_compare(IntPredicate::SLT, tmp, length, "cmp").unwrap();
+            ctx.builder.build_store(counter, tmp).unwrap();
-        ctx.builder.build_conditional_branch(cmp, body_bb, cont_bb).unwrap();
+            let cmp = ctx.builder.build_int_compare(IntPredicate::SLT, tmp, length, "cmp").unwrap();
+            ctx.builder.build_conditional_branch(cmp, body_bb, cont_bb).unwrap();
 
-        ctx.builder.position_at_end(body_bb);
+            ctx.builder.position_at_end(body_bb);
-        let arr_ptr = ctx
+            let arr_ptr = ctx
-            .build_gep_and_load(
+                .build_gep_and_load(
-                iter_val.into_pointer_value(),
+                    iter_val.into_pointer_value(),
-                &[zero_size_t, zero_32],
+                    &[zero_size_t, zero_32],
-                Some("arr.addr"),
+                    Some("arr.addr"),
-            )
+                )
-            .into_pointer_value();
+                .into_pointer_value();
-        let val = ctx.build_gep_and_load(arr_ptr, &[tmp], Some("val"));
+            let val = ctx.build_gep_and_load(arr_ptr, &[tmp], Some("val"));
-        generator.gen_assign(ctx, target, val.into())?;
+            generator.gen_assign(ctx, target, val.into(), elt.custom.unwrap())?;
+        }
+        _ => {
+            panic!(
+                "unsupported list comprehension iterator type: {}",
+                ctx.unifier.stringify(iter_ty)
+            );
+        }
     }
 
     // Emits the content of `cont_bb`

nac3core/src/codegen/generator.rs

@@ -123,11 +123,45 @@ pub trait CodeGenerator {
         ctx: &mut CodeGenContext<'ctx, '_>,
         target: &Expr<Option<Type>>,
         value: ValueEnum<'ctx>,
+        value_ty: Type,
     ) -> Result<(), String>
     where
         Self: Sized,
     {
-        gen_assign(self, ctx, target, value)
+        gen_assign(self, ctx, target, value, value_ty)
+    }
+
+    /// Generate code for an assignment expression where LHS is a `"target_list"`.
+    ///
+    /// See <https://docs.python.org/3/reference/simple_stmts.html#assignment-statements>.
+    fn gen_assign_target_list<'ctx>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        targets: &Vec<Expr<Option<Type>>>,
+        value: ValueEnum<'ctx>,
+        value_ty: Type,
+    ) -> Result<(), String>
+    where
+        Self: Sized,
+    {
+        gen_assign_target_list(self, ctx, targets, value, value_ty)
+    }
+
+    /// Generate code for an item assignment.
+    ///
+    /// i.e., `target[key] = value`
+    fn gen_setitem<'ctx>(
+        &mut self,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        target: &Expr<Option<Type>>,
+        key: &Expr<Option<Type>>,
+        value: ValueEnum<'ctx>,
+        value_ty: Type,
+    ) -> Result<(), String>
+    where
+        Self: Sized,
+    {
+        gen_setitem(self, ctx, target, key, value, value_ty)
     }
 
     /// Generate code for a while expression.

nac3core/src/codegen/stmt.rs

@@ -10,10 +10,10 @@ use crate::{
         expr::gen_binop_expr,
         gen_in_range_check,
     },
-    toplevel::{helper::PrimDef, numpy::unpack_ndarray_var_tys, DefinitionId, TopLevelDef},
+    toplevel::{DefinitionId, TopLevelDef},
     typecheck::{
         magic_methods::Binop,
-        typedef::{FunSignature, Type, TypeEnum},
+        typedef::{iter_type_vars, FunSignature, Type, TypeEnum},
     },
 };
 use inkwell::{
@@ -23,10 +23,10 @@ use inkwell::{
     values::{BasicValue, BasicValueEnum, FunctionValue, IntValue, PointerValue},
     IntPredicate,
 };
+use itertools::{izip, Itertools};
 use nac3parser::ast::{
     Constant, ExcepthandlerKind, Expr, ExprKind, Location, Stmt, StmtKind, StrRef,
 };
-use std::convert::TryFrom;
 
 /// See [`CodeGenerator::gen_var_alloc`].
 pub fn gen_var<'ctx>(
@@ -97,8 +97,6 @@ pub fn gen_store_target<'ctx, G: CodeGenerator>(
     pattern: &Expr<Option<Type>>,
     name: Option<&str>,
 ) -> Result<Option<PointerValue<'ctx>>, String> {
-    let llvm_usize = generator.get_size_type(ctx.ctx);
-
     // very similar to gen_expr, but we don't do an extra load at the end
     // and we flatten nested tuples
     Ok(Some(match &pattern.node {
@@ -137,65 +135,6 @@ pub fn gen_store_target<'ctx, G: CodeGenerator>(
             }
             .unwrap()
         }
-        ExprKind::Subscript { value, slice, .. } => {
-            match ctx.unifier.get_ty_immutable(value.custom.unwrap()).as_ref() {
-                TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::List.id() => {
-                    let v = generator
-                        .gen_expr(ctx, value)?
-                        .unwrap()
-                        .to_basic_value_enum(ctx, generator, value.custom.unwrap())?
-                        .into_pointer_value();
-                    let v = ListValue::from_ptr_val(v, llvm_usize, None);
-                    let len = v.load_size(ctx, Some("len"));
-                    let raw_index = generator
-                        .gen_expr(ctx, slice)?
-                        .unwrap()
-                        .to_basic_value_enum(ctx, generator, slice.custom.unwrap())?
-                        .into_int_value();
-                    let raw_index = ctx
-                        .builder
-                        .build_int_s_extend(raw_index, generator.get_size_type(ctx.ctx), "sext")
-                        .unwrap();
-                    // handle negative index
-                    let is_negative = ctx
-                        .builder
-                        .build_int_compare(
-                            IntPredicate::SLT,
-                            raw_index,
-                            generator.get_size_type(ctx.ctx).const_zero(),
-                            "is_neg",
-                        )
-                        .unwrap();
-                    let adjusted = ctx.builder.build_int_add(raw_index, len, "adjusted").unwrap();
-                    let index = ctx
-                        .builder
-                        .build_select(is_negative, adjusted, raw_index, "index")
-                        .map(BasicValueEnum::into_int_value)
-                        .unwrap();
-                    // unsigned less than is enough, because negative index after adjustment is
-                    // bigger than the length (for unsigned cmp)
-                    let bound_check = ctx
-                        .builder
-                        .build_int_compare(IntPredicate::ULT, index, len, "inbound")
-                        .unwrap();
-                    ctx.make_assert(
-                        generator,
-                        bound_check,
-                        "0:IndexError",
-                        "index {0} out of bounds 0:{1}",
-                        [Some(raw_index), Some(len), None],
-                        slice.location,
-                    );
-                    v.data().ptr_offset(ctx, generator, &index, name)
-                }
-
-                TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
-                    todo!()
-                }
-
-                _ => unreachable!(),
-            }
-        }
         _ => unreachable!(),
     }))
 }
@@ -206,70 +145,20 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
     ctx: &mut CodeGenContext<'ctx, '_>,
     target: &Expr<Option<Type>>,
     value: ValueEnum<'ctx>,
+    value_ty: Type,
 ) -> Result<(), String> {
-    let llvm_usize = generator.get_size_type(ctx.ctx);
+    // See https://docs.python.org/3/reference/simple_stmts.html#assignment-statements.
-
     match &target.node {
-        ExprKind::Tuple { elts, .. } => {
+        ExprKind::Subscript { value: target, slice: key, .. } => {
-            let BasicValueEnum::StructValue(v) =
+            // Handle "slicing" or "subscription"
-                value.to_basic_value_enum(ctx, generator, target.custom.unwrap())?
+            generator.gen_setitem(ctx, target, key, value, value_ty)?;
-            else {
-                unreachable!()
-            };
-
-            for (i, elt) in elts.iter().enumerate() {
-                let v = ctx
-                    .builder
-                    .build_extract_value(v, u32::try_from(i).unwrap(), "struct_elem")
-                    .unwrap();
-                generator.gen_assign(ctx, elt, v.into())?;
-            }
         }
-        ExprKind::Subscript { value: ls, slice, .. }
+        ExprKind::Tuple { elts, .. } | ExprKind::List { elts, .. } => {
-            if matches!(&slice.node, ExprKind::Slice { .. }) =>
+            // Fold on `"[" [target_list] "]"` and `"(" [target_list] ")"`
-        {
+            generator.gen_assign_target_list(ctx, elts, value, value_ty)?;
-            let ExprKind::Slice { lower, upper, step } = &slice.node else { unreachable!() };
-
-            let ls = generator
-                .gen_expr(ctx, ls)?
-                .unwrap()
-                .to_basic_value_enum(ctx, generator, ls.custom.unwrap())?
-                .into_pointer_value();
-            let ls = ListValue::from_ptr_val(ls, llvm_usize, None);
-            let Some((start, end, step)) =
-                handle_slice_indices(lower, upper, step, ctx, generator, ls.load_size(ctx, None))?
-            else {
-                return Ok(());
-            };
-            let value = value
-                .to_basic_value_enum(ctx, generator, target.custom.unwrap())?
-                .into_pointer_value();
-            let value = ListValue::from_ptr_val(value, llvm_usize, None);
-            let ty = match &*ctx.unifier.get_ty_immutable(target.custom.unwrap()) {
-                TypeEnum::TObj { obj_id, params, .. } if *obj_id == PrimDef::List.id() => {
-                    *params.iter().next().unwrap().1
-                }
-                TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
-                    unpack_ndarray_var_tys(&mut ctx.unifier, target.custom.unwrap()).0
-                }
-                _ => unreachable!(),
-            };
-
-            let ty = ctx.get_llvm_type(generator, ty);
-            let Some(src_ind) = handle_slice_indices(
-                &None,
-                &None,
-                &None,
-                ctx,
-                generator,
-                value.load_size(ctx, None),
-            )?
-            else {
-                return Ok(());
-            };
-            list_slice_assignment(generator, ctx, ty, ls, (start, end, step), value, src_ind);
         }
         _ => {
+            // Handle attribute and direct variable assignments.
             let name = if let ExprKind::Name { id, .. } = &target.node {
                 format!("{id}.addr")
             } else {
@@ -293,6 +182,233 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
     Ok(())
 }
 
+/// See [`CodeGenerator::gen_assign_target_list`].
+pub fn gen_assign_target_list<'ctx, G: CodeGenerator>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    targets: &Vec<Expr<Option<Type>>>,
+    value: ValueEnum<'ctx>,
+    value_ty: Type,
+) -> Result<(), String> {
+    // Deconstruct the tuple `value`
+    let BasicValueEnum::StructValue(tuple) = value.to_basic_value_enum(ctx, generator, value_ty)?
+    else {
+        unreachable!()
+    };
+
+    // NOTE: Currently, RHS's type is forced to be a Tuple by the type inferencer.
+    let TypeEnum::TTuple { ty: tuple_tys } = &*ctx.unifier.get_ty(value_ty) else {
+        unreachable!();
+    };
+
+    assert_eq!(tuple.get_type().count_fields() as usize, tuple_tys.len());
+
+    let tuple = (0..tuple.get_type().count_fields())
+        .map(|i| ctx.builder.build_extract_value(tuple, i, "item").unwrap())
+        .collect_vec();
+
+    // Find the starred target if it exists.
+    let mut starred_target_index: Option<usize> = None; // Index of the "starred" target. If it exists, there may only be one.
+    for (i, target) in targets.iter().enumerate() {
+        if matches!(target.node, ExprKind::Starred { .. }) {
+            assert!(starred_target_index.is_none()); // The typechecker ensures this
+            starred_target_index = Some(i);
+        }
+    }
+
+    if let Some(starred_target_index) = starred_target_index {
+        assert!(tuple_tys.len() >= targets.len() - 1); // The typechecker ensures this
+
+        let a = starred_target_index; // Number of RHS values before the starred target
+        let b = tuple_tys.len() - (targets.len() - 1 - starred_target_index); // Number of RHS values after the starred target
+                                                                              // Thus `tuple[a..b]` is assigned to the starred target.
+
+        // Handle assignment before the starred target
+        for (target, val, val_ty) in
+            izip!(&targets[..starred_target_index], &tuple[..a], &tuple_tys[..a])
+        {
+            generator.gen_assign(ctx, target, ValueEnum::Dynamic(*val), *val_ty)?;
+        }
+
+        // Handle assignment to the starred target
+        if let ExprKind::Starred { value: target, .. } = &targets[starred_target_index].node {
+            let vals = &tuple[a..b];
+            let val_tys = &tuple_tys[a..b];
+
+            // Create a sub-tuple from `value` for the starred target.
+            let sub_tuple_ty = ctx
+                .ctx
+                .struct_type(&vals.iter().map(BasicValueEnum::get_type).collect_vec(), false);
+            let psub_tuple_val =
+                ctx.builder.build_alloca(sub_tuple_ty, "starred_target_value_ptr").unwrap();
+            for (i, val) in vals.iter().enumerate() {
+                let pitem = ctx
+                    .builder
+                    .build_struct_gep(psub_tuple_val, i as u32, "starred_target_value_item")
+                    .unwrap();
+                ctx.builder.build_store(pitem, *val).unwrap();
+            }
+            let sub_tuple_val =
+                ctx.builder.build_load(psub_tuple_val, "starred_target_value").unwrap();
+
+            // Create the typechecker type of the sub-tuple
+            let sub_tuple_ty = ctx.unifier.add_ty(TypeEnum::TTuple { ty: val_tys.to_vec() });
+
+            // Now assign with that sub-tuple to the starred target.
+            generator.gen_assign(ctx, target, ValueEnum::Dynamic(sub_tuple_val), sub_tuple_ty)?;
+        } else {
+            unreachable!() // The typechecker ensures this
+        }
+
+        // Handle assignment after the starred target
+        for (target, val, val_ty) in
+            izip!(&targets[starred_target_index + 1..], &tuple[b..], &tuple_tys[b..])
+        {
+            generator.gen_assign(ctx, target, ValueEnum::Dynamic(*val), *val_ty)?;
+        }
+    } else {
+        assert_eq!(tuple_tys.len(), targets.len()); // The typechecker ensures this
+
+        for (target, val, val_ty) in izip!(targets, tuple, tuple_tys) {
+            generator.gen_assign(ctx, target, ValueEnum::Dynamic(val), *val_ty)?;
+        }
+    }
+    Ok(())
+}
+
+/// See [`CodeGenerator::gen_setitem`].
+pub fn gen_setitem<'ctx, G: CodeGenerator>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    target: &Expr<Option<Type>>,
+    key: &Expr<Option<Type>>,
+    value: ValueEnum<'ctx>,
+    value_ty: Type,
+) -> Result<(), String> {
+    let target_ty = target.custom.unwrap();
+    let key_ty = key.custom.unwrap();
+
+    match &*ctx.unifier.get_ty(target_ty) {
+        TypeEnum::TObj { obj_id, params: list_params, .. }
+            if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+        {
+            // Handle list item assignment
+            let llvm_usize = generator.get_size_type(ctx.ctx);
+            let target_item_ty = iter_type_vars(list_params).next().unwrap().ty;
+
+            let target = generator
+                .gen_expr(ctx, target)?
+                .unwrap()
+                .to_basic_value_enum(ctx, generator, target_ty)?
+                .into_pointer_value();
+            let target = ListValue::from_ptr_val(target, llvm_usize, None);
+
+            if let ExprKind::Slice { .. } = &key.node {
+                // Handle assigning to a slice
+                let ExprKind::Slice { lower, upper, step } = &key.node else { unreachable!() };
+                let Some((start, end, step)) = handle_slice_indices(
+                    lower,
+                    upper,
+                    step,
+                    ctx,
+                    generator,
+                    target.load_size(ctx, None),
+                )?
+                else {
+                    return Ok(());
+                };
+
+                let value =
+                    value.to_basic_value_enum(ctx, generator, value_ty)?.into_pointer_value();
+                let value = ListValue::from_ptr_val(value, llvm_usize, None);
+
+                let target_item_ty = ctx.get_llvm_type(generator, target_item_ty);
+                let Some(src_ind) = handle_slice_indices(
+                    &None,
+                    &None,
+                    &None,
+                    ctx,
+                    generator,
+                    value.load_size(ctx, None),
+                )?
+                else {
+                    return Ok(());
+                };
+                list_slice_assignment(
+                    generator,
+                    ctx,
+                    target_item_ty,
+                    target,
+                    (start, end, step),
+                    value,
+                    src_ind,
+                );
+            } else {
+                // Handle assigning to an index
+                let len = target.load_size(ctx, Some("len"));
+
+                let index = generator
+                    .gen_expr(ctx, key)?
+                    .unwrap()
+                    .to_basic_value_enum(ctx, generator, key_ty)?
+                    .into_int_value();
+                let index = ctx
+                    .builder
+                    .build_int_s_extend(index, generator.get_size_type(ctx.ctx), "sext")
+                    .unwrap();
+
+                // handle negative index
+                let is_negative = ctx
+                    .builder
+                    .build_int_compare(
+                        IntPredicate::SLT,
+                        index,
+                        generator.get_size_type(ctx.ctx).const_zero(),
+                        "is_neg",
+                    )
+                    .unwrap();
+                let adjusted = ctx.builder.build_int_add(index, len, "adjusted").unwrap();
+                let index = ctx
+                    .builder
+                    .build_select(is_negative, adjusted, index, "index")
+                    .map(BasicValueEnum::into_int_value)
+                    .unwrap();
+
+                // unsigned less than is enough, because negative index after adjustment is
+                // bigger than the length (for unsigned cmp)
+                let bound_check = ctx
+                    .builder
+                    .build_int_compare(IntPredicate::ULT, index, len, "inbound")
+                    .unwrap();
+                ctx.make_assert(
+                    generator,
+                    bound_check,
+                    "0:IndexError",
+                    "index {0} out of bounds 0:{1}",
+                    [Some(index), Some(len), None],
+                    key.location,
+                );
+
+                // Write value to index on list
+                let item_ptr =
+                    target.data().ptr_offset(ctx, generator, &index, Some("list_item_ptr"));
+                let value = value.to_basic_value_enum(ctx, generator, value_ty)?;
+                ctx.builder.build_store(item_ptr, value).unwrap();
+            }
+        }
+        TypeEnum::TObj { obj_id, .. }
+            if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() =>
+        {
+            // Handle NDArray item assignment
+            todo!("ndarray subscript assignment is not yet implemented");
+        }
+        _ => {
+            panic!("encountered unknown target type: {}", ctx.unifier.stringify(target_ty));
+        }
+    }
+    Ok(())
+}
+
 /// See [`CodeGenerator::gen_for`].
 pub fn gen_for<G: CodeGenerator>(
     generator: &mut G,
@@ -315,9 +431,6 @@ pub fn gen_for<G: CodeGenerator>(
     let orelse_bb =
         if orelse.is_empty() { cont_bb } else { ctx.ctx.append_basic_block(current, "for.orelse") };
 
-    // Whether the iterable is a range() expression
-    let is_iterable_range_expr = ctx.unifier.unioned(iter.custom.unwrap(), ctx.primitives.range);
-
     // The BB containing the increment expression
     let incr_bb = ctx.ctx.append_basic_block(current, "for.incr");
     // The BB containing the loop condition check
@@ -326,113 +439,132 @@ pub fn gen_for<G: CodeGenerator>(
     // store loop bb information and restore it later
     let loop_bb = ctx.loop_target.replace((incr_bb, cont_bb));
 
+    let iter_ty = iter.custom.unwrap();
     let iter_val = if let Some(v) = generator.gen_expr(ctx, iter)? {
-        v.to_basic_value_enum(ctx, generator, iter.custom.unwrap())?
+        v.to_basic_value_enum(ctx, generator, iter_ty)?
     } else {
         return Ok(());
     };
-    if is_iterable_range_expr {
-        let iter_val = RangeValue::from_ptr_val(iter_val.into_pointer_value(), Some("range"));
-        // Internal variable for loop; Cannot be assigned
-        let i = generator.gen_var_alloc(ctx, int32.into(), Some("for.i.addr"))?;
-        // Variable declared in "target" expression of the loop; Can be reassigned *or* shadowed
-        let Some(target_i) = generator.gen_store_target(ctx, target, Some("for.target.addr"))?
-        else {
-            unreachable!()
-        };
-        let (start, stop, step) = destructure_range(ctx, iter_val);
-
-        ctx.builder.build_store(i, start).unwrap();
-
-        // Check "If step is zero, ValueError is raised."
-        let rangenez =
-            ctx.builder.build_int_compare(IntPredicate::NE, step, int32.const_zero(), "").unwrap();
-        ctx.make_assert(
-            generator,
-            rangenez,
-            "ValueError",
-            "range() arg 3 must not be zero",
-            [None, None, None],
-            ctx.current_loc,
-        );
-        ctx.builder.build_unconditional_branch(cond_bb).unwrap();
 
+    match &*ctx.unifier.get_ty(iter_ty) {
+        TypeEnum::TObj { obj_id, .. }
+            if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
         {
-            ctx.builder.position_at_end(cond_bb);
+            let iter_val = RangeValue::from_ptr_val(iter_val.into_pointer_value(), Some("range"));
-            ctx.builder
+            // Internal variable for loop; Cannot be assigned
-                .build_conditional_branch(
+            let i = generator.gen_var_alloc(ctx, int32.into(), Some("for.i.addr"))?;
-                    gen_in_range_check(
+            // Variable declared in "target" expression of the loop; Can be reassigned *or* shadowed
-                        ctx,
+            let Some(target_i) =
-                        ctx.builder.build_load(i, "").map(BasicValueEnum::into_int_value).unwrap(),
+                generator.gen_store_target(ctx, target, Some("for.target.addr"))?
-                        stop,
+            else {
-                        step,
+                unreachable!()
-                    ),
+            };
-                    body_bb,
+            let (start, stop, step) = destructure_range(ctx, iter_val);
-                    orelse_bb,
+
+            ctx.builder.build_store(i, start).unwrap();
+
+            // Check "If step is zero, ValueError is raised."
+            let rangenez = ctx
+                .builder
+                .build_int_compare(IntPredicate::NE, step, int32.const_zero(), "")
+                .unwrap();
+            ctx.make_assert(
+                generator,
+                rangenez,
+                "ValueError",
+                "range() arg 3 must not be zero",
+                [None, None, None],
+                ctx.current_loc,
+            );
+            ctx.builder.build_unconditional_branch(cond_bb).unwrap();
+
+            {
+                ctx.builder.position_at_end(cond_bb);
+                ctx.builder
+                    .build_conditional_branch(
+                        gen_in_range_check(
+                            ctx,
+                            ctx.builder
+                                .build_load(i, "")
+                                .map(BasicValueEnum::into_int_value)
+                                .unwrap(),
+                            stop,
+                            step,
+                        ),
+                        body_bb,
+                        orelse_bb,
+                    )
+                    .unwrap();
+            }
+
+            ctx.builder.position_at_end(incr_bb);
+            let next_i = ctx
+                .builder
+                .build_int_add(
+                    ctx.builder.build_load(i, "").map(BasicValueEnum::into_int_value).unwrap(),
+                    step,
+                    "inc",
                 )
                 .unwrap();
+            ctx.builder.build_store(i, next_i).unwrap();
+            ctx.builder.build_unconditional_branch(cond_bb).unwrap();
+
+            ctx.builder.position_at_end(body_bb);
+            ctx.builder
+                .build_store(
+                    target_i,
+                    ctx.builder.build_load(i, "").map(BasicValueEnum::into_int_value).unwrap(),
+                )
+                .unwrap();
+            generator.gen_block(ctx, body.iter())?;
         }
+        TypeEnum::TObj { obj_id, params: list_params, .. }
+            if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+        {
+            let index_addr = generator.gen_var_alloc(ctx, size_t.into(), Some("for.index.addr"))?;
+            ctx.builder.build_store(index_addr, size_t.const_zero()).unwrap();
+            let len = ctx
+                .build_gep_and_load(
+                    iter_val.into_pointer_value(),
+                    &[zero, int32.const_int(1, false)],
+                    Some("len"),
+                )
+                .into_int_value();
+            ctx.builder.build_unconditional_branch(cond_bb).unwrap();
 
-        ctx.builder.position_at_end(incr_bb);
+            ctx.builder.position_at_end(cond_bb);
-        let next_i = ctx
+            let index = ctx
-            .builder
+                .builder
-            .build_int_add(
+                .build_load(index_addr, "for.index")
-                ctx.builder.build_load(i, "").map(BasicValueEnum::into_int_value).unwrap(),
+                .map(BasicValueEnum::into_int_value)
-                step,
+                .unwrap();
-                "inc",
+            let cmp = ctx.builder.build_int_compare(IntPredicate::SLT, index, len, "cond").unwrap();
-            )
+            ctx.builder.build_conditional_branch(cmp, body_bb, orelse_bb).unwrap();
-            .unwrap();
-        ctx.builder.build_store(i, next_i).unwrap();
-        ctx.builder.build_unconditional_branch(cond_bb).unwrap();
 
-        ctx.builder.position_at_end(body_bb);
+            ctx.builder.position_at_end(incr_bb);
-        ctx.builder
+            let index =
-            .build_store(
+                ctx.builder.build_load(index_addr, "").map(BasicValueEnum::into_int_value).unwrap();
-                target_i,
+            let inc = ctx.builder.build_int_add(index, size_t.const_int(1, true), "inc").unwrap();
-                ctx.builder.build_load(i, "").map(BasicValueEnum::into_int_value).unwrap(),
+            ctx.builder.build_store(index_addr, inc).unwrap();
-            )
+            ctx.builder.build_unconditional_branch(cond_bb).unwrap();
-            .unwrap();
-        generator.gen_block(ctx, body.iter())?;
-    } else {
-        let index_addr = generator.gen_var_alloc(ctx, size_t.into(), Some("for.index.addr"))?;
-        ctx.builder.build_store(index_addr, size_t.const_zero()).unwrap();
-        let len = ctx
-            .build_gep_and_load(
-                iter_val.into_pointer_value(),
-                &[zero, int32.const_int(1, false)],
-                Some("len"),
-            )
-            .into_int_value();
-        ctx.builder.build_unconditional_branch(cond_bb).unwrap();
 
-        ctx.builder.position_at_end(cond_bb);
+            ctx.builder.position_at_end(body_bb);
-        let index = ctx
+            let arr_ptr = ctx
-            .builder
+                .build_gep_and_load(iter_val.into_pointer_value(), &[zero, zero], Some("arr.addr"))
-            .build_load(index_addr, "for.index")
+                .into_pointer_value();
-            .map(BasicValueEnum::into_int_value)
+            let index = ctx
-            .unwrap();
+                .builder
-        let cmp = ctx.builder.build_int_compare(IntPredicate::SLT, index, len, "cond").unwrap();
+                .build_load(index_addr, "for.index")
-        ctx.builder.build_conditional_branch(cmp, body_bb, orelse_bb).unwrap();
+                .map(BasicValueEnum::into_int_value)
-
+                .unwrap();
-        ctx.builder.position_at_end(incr_bb);
+            let val = ctx.build_gep_and_load(arr_ptr, &[index], Some("val"));
-        let index =
+            let val_ty = iter_type_vars(list_params).next().unwrap().ty;
-            ctx.builder.build_load(index_addr, "").map(BasicValueEnum::into_int_value).unwrap();
+            generator.gen_assign(ctx, target, val.into(), val_ty)?;
-        let inc = ctx.builder.build_int_add(index, size_t.const_int(1, true), "inc").unwrap();
+            generator.gen_block(ctx, body.iter())?;
-        ctx.builder.build_store(index_addr, inc).unwrap();
+        }
-        ctx.builder.build_unconditional_branch(cond_bb).unwrap();
+        _ => {
-
+            panic!("unsupported for loop iterator type: {}", ctx.unifier.stringify(iter_ty));
-        ctx.builder.position_at_end(body_bb);
+        }
-        let arr_ptr = ctx
-            .build_gep_and_load(iter_val.into_pointer_value(), &[zero, zero], Some("arr.addr"))
-            .into_pointer_value();
-        let index = ctx
-            .builder
-            .build_load(index_addr, "for.index")
-            .map(BasicValueEnum::into_int_value)
-            .unwrap();
-        let val = ctx.build_gep_and_load(arr_ptr, &[index], Some("val"));
-        generator.gen_assign(ctx, target, val.into())?;
-        generator.gen_block(ctx, body.iter())?;
     }
 
     for (k, (_, _, counter)) in &var_assignment {
@@ -1588,14 +1720,14 @@ pub fn gen_stmt<G: CodeGenerator>(
         }
         StmtKind::AnnAssign { target, value, .. } => {
             if let Some(value) = value {
-                let Some(value) = generator.gen_expr(ctx, value)? else { return Ok(()) };
+                let Some(value_enum) = generator.gen_expr(ctx, value)? else { return Ok(()) };
-                generator.gen_assign(ctx, target, value)?;
+                generator.gen_assign(ctx, target, value_enum, value.custom.unwrap())?;
             }
         }
         StmtKind::Assign { targets, value, .. } => {
-            let Some(value) = generator.gen_expr(ctx, value)? else { return Ok(()) };
+            let Some(value_enum) = generator.gen_expr(ctx, value)? else { return Ok(()) };
             for target in targets {
-                generator.gen_assign(ctx, target, value.clone())?;
+                generator.gen_assign(ctx, target, value_enum.clone(), value.custom.unwrap())?;
             }
         }
         StmtKind::Continue { .. } => {
@@ -1609,15 +1741,16 @@ pub fn gen_stmt<G: CodeGenerator>(
         StmtKind::For { .. } => generator.gen_for(ctx, stmt)?,
         StmtKind::With { .. } => generator.gen_with(ctx, stmt)?,
         StmtKind::AugAssign { target, op, value, .. } => {
-            let value = gen_binop_expr(
+            let value_enum = gen_binop_expr(
                 generator,
                 ctx,
                 target,
                 Binop::aug_assign(*op),
                 value,
                 stmt.location,
-            )?;
+            )?
-            generator.gen_assign(ctx, target, value.unwrap())?;
+            .unwrap();
+            generator.gen_assign(ctx, target, value_enum, value.custom.unwrap())?;
         }
         StmtKind::Try { .. } => gen_try(generator, ctx, stmt)?,
         StmtKind::Raise { exc, .. } => {

nac3core/src/typecheck/function_check.rs

@@ -34,13 +34,18 @@ impl<'a> Inferencer<'a> {
                 self.should_have_value(pattern)?;
                 Ok(())
             }
-            ExprKind::Tuple { elts, .. } => {
+            ExprKind::List { elts, .. } | ExprKind::Tuple { elts, .. } => {
                 for elt in elts {
                     self.check_pattern(elt, defined_identifiers)?;
                     self.should_have_value(elt)?;
                 }
                 Ok(())
             }
+            ExprKind::Starred { value, .. } => {
+                self.check_pattern(value, defined_identifiers)?;
+                self.should_have_value(value)?;
+                Ok(())
+            }
             ExprKind::Subscript { value, slice, .. } => {
                 self.check_expr(value, defined_identifiers)?;
                 self.should_have_value(value)?;

nac3standalone/demo/src/assignment.py

@@ -0,0 +1,66 @@
+@extern
+def output_int32(x: int32):
+    ...
+
+@extern
+def output_bool(x: bool):
+    ...
+
+def example1():
+    x, *ys, z = (1, 2, 3, 4, 5)
+    output_int32(x)
+    output_int32(ys[0])
+    output_int32(ys[1])
+    output_int32(ys[2])
+    output_int32(z)
+
+def example2():
+    x, y, *zs = (1, 2, 3, 4, 5)
+    output_int32(x)
+    output_int32(y)
+    output_int32(zs[0])
+    output_int32(zs[1])
+    output_int32(zs[2])
+
+def example3():
+    *xs, y, z = (1, 2, 3, 4, 5)
+    output_int32(xs[0])
+    output_int32(xs[1])
+    output_int32(xs[2])
+    output_int32(y)
+    output_int32(z)
+
+def example4():
+    # Example from: https://docs.python.org/3/reference/simple_stmts.html#assignment-statements
+    x = [0, 1]
+    i = 0
+    i, x[i] = 1, 2 # i is updated, then x[i] is updated
+    output_int32(i)
+    output_int32(x[0])
+    output_int32(x[1])
+
+class A:
+    value: int32
+    def __init__(self):
+        self.value = 1000
+
+def example5():
+    ws = [88, 7, 8]
+    a = A()
+    x, [y, *ys, a.value], ws[0], (ws[0],) = 1, (2, False, 4, 5), 99, (6,)
+    output_int32(x)
+    output_int32(y)
+    output_bool(ys[0])
+    output_int32(ys[1])
+    output_int32(a.value)
+    output_int32(ws[0])
+    output_int32(ws[1])
+    output_int32(ws[2])
+
+def run() -> int32:
+    example1()
+    example2()
+    example3()
+    example4()
+    example5()
+    return 0