WIP

nac3artiq/src/codegen.rs

@@ -15,7 +15,7 @@ use pyo3::{
 use super::{symbol_resolver::InnerResolver, timeline::TimeFns, SpecialPythonId};
 use nac3core::{
     codegen::{
-        expr::{create_fn_and_call, destructure_range, gen_call, infer_and_call_function},
+        expr::{create_fn_and_call, gen_call, infer_and_call_function},
         llvm_intrinsics::{call_int_smax, call_memcpy, call_stackrestore, call_stacksave},
         stmt::{gen_block, gen_for_callback_incrementing, gen_if_callback, gen_with},
         type_aligned_alloca,
@@ -1414,7 +1414,7 @@ fn polymorphic_print<'ctx>(
 
                 let val = RangeType::new(ctx).map_pointer_value(value.into_pointer_value(), None);
 
-                let (start, stop, step) = destructure_range(ctx, val);
+                let (start, stop, step) = val.load_values(ctx);
 
                 polymorphic_print(
                     ctx,

nac3core/src/codegen/builtin_fns.rs

@@ -6,7 +6,6 @@ use inkwell::{
 use itertools::Itertools;
 
 use super::{
-    expr::destructure_range,
     extern_fns, irrt,
     irrt::calculate_len_for_slice_range,
     llvm_intrinsics,
@@ -48,7 +47,7 @@ pub fn call_len<'ctx, G: CodeGenerator + ?Sized>(
 
     Ok(if ctx.unifier.unioned(arg_ty, range_ty) {
         let arg = RangeType::new(ctx).map_pointer_value(arg.into_pointer_value(), Some("range"));
-        let (start, end, step) = destructure_range(ctx, arg);
+        let (start, end, step) = arg.load_values(ctx);
         calculate_len_for_slice_range(generator, ctx, start, end, step)
     } else {
         match &*ctx.unifier.get_ty_immutable(arg_ty) {

nac3core/src/codegen/expr.rs

@@ -29,16 +29,15 @@ use super::{
     macros::codegen_unreachable,
     need_sret,
     stmt::{
-        gen_for_callback_incrementing, gen_if_callback, gen_if_else_expr_callback, gen_raise,
-        gen_var,
+        gen_for_callback, gen_for_callback_incrementing, gen_if_callback,
+        gen_if_else_expr_callback, gen_raise, gen_var,
     },
     types::{
         ndarray::NDArrayType, ExceptionType, ListType, OptionType, RangeType, StringType, TupleType,
     },
     values::{
         ndarray::{NDArrayOut, RustNDIndex, ScalarOrNDArray},
-        ArrayLikeIndexer, ArrayLikeValue, ListValue, ProxyValue, RangeValue,
-        UntypedArrayLikeAccessor,
+        ArrayLikeIndexer, ArrayLikeValue, ListValue, ProxyValue, UntypedArrayLikeAccessor,
     },
     CodeGenContext, CodeGenTask, CodeGenerator,
 };
@@ -1043,18 +1042,6 @@ pub fn gen_call<'ctx, G: CodeGenerator>(
     Ok(ctx.build_call_or_invoke(fun_val, &param_vals, "call"))
 }
 
-/// Generates three LLVM variables representing the start, stop, and step values of a [range] class
-/// respectively.
-pub fn destructure_range<'ctx>(
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    range: RangeValue<'ctx>,
-) -> (IntValue<'ctx>, IntValue<'ctx>, IntValue<'ctx>) {
-    let start = range.load_start(ctx, None);
-    let end = range.load_end(ctx, None);
-    let step = range.load_step(ctx, None);
-    (start, end, step)
-}
-
 /// Generates LLVM IR for a [list comprehension expression][expr].
 pub fn gen_comprehension<'ctx, G: CodeGenerator>(
     generator: &mut G,
@@ -1063,110 +1050,56 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
 ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
     let ExprKind::ListComp { elt, generators } = &expr.node else { codegen_unreachable!(ctx) };
 
-    let current = ctx.builder.get_insert_block().unwrap().get_parent().unwrap();
-
-    let init_bb = ctx.ctx.append_basic_block(current, "listcomp.init");
-    let test_bb = ctx.ctx.append_basic_block(current, "listcomp.test");
-    let body_bb = ctx.ctx.append_basic_block(current, "listcomp.body");
-    let cont_bb = ctx.ctx.append_basic_block(current, "listcomp.cont");
-
-    ctx.builder.build_unconditional_branch(init_bb).unwrap();
-
-    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_ty)?
     } else {
-        for bb in [test_bb, body_bb, cont_bb] {
-            ctx.builder.position_at_end(bb);
-            ctx.builder.build_unreachable().unwrap();
-        }
-
         return Ok(None);
     };
-    let int32 = ctx.ctx.i32_type();
-    let size_t = ctx.get_size_type();
-    let zero_size_t = size_t.const_zero();
-    let zero_32 = int32.const_zero();
 
-    let index = generator.gen_var_alloc(ctx, size_t.into(), Some("index.addr"))?;
-    ctx.builder.build_store(index, zero_size_t).unwrap();
+    let llvm_i32 = ctx.ctx.i32_type();
+    let llvm_usize = ctx.get_size_type();
 
     let elem_ty = ctx.get_llvm_type(generator, elt.custom.unwrap());
-    let list;
-
-    match &*ctx.unifier.get_ty(iter_ty) {
+    let list = match &*ctx.unifier.get_ty(iter_ty) {
         TypeEnum::TObj { obj_id, .. }
             if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
         {
             let iter_val =
                 RangeType::new(ctx).map_pointer_value(iter_val.into_pointer_value(), Some("range"));
-            let (start, stop, step) = destructure_range(ctx, iter_val);
+            let (start, stop, step) = iter_val.load_values(ctx);
             let diff = ctx.builder.build_int_sub(stop, start, "diff").unwrap();
             // add 1 to the length as the value is rounded to zero
             // the length may be 1 more than the actual length if the division is exact, but the
             // length is a upper bound only anyway so it does not matter.
             let length = ctx.builder.build_int_signed_div(diff, step, "div").unwrap();
             let length =
-                ctx.builder.build_int_add(length, int32.const_int(1, false), "add1").unwrap();
+                ctx.builder.build_int_add(length, llvm_i32.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 is_valid = ctx
+                .builder
+                .build_int_compare(IntPredicate::SGT, length, llvm_i32.const_zero(), "check")
+                .unwrap();
 
             let list_alloc_size = ctx
                 .builder
                 .build_select(
                     is_valid,
                     ctx.builder
-                        .build_int_z_extend_or_bit_cast(length, size_t, "z_ext_len")
+                        .build_int_z_extend_or_bit_cast(length, llvm_usize, "z_ext_len")
                         .unwrap(),
-                    zero_size_t,
+                    llvm_usize.const_zero(),
                     "listcomp.alloc_size",
                 )
                 .unwrap();
-            list = ListType::new(ctx, &elem_ty).construct(
+            ListType::new(ctx, &elem_ty).construct(
                 generator,
                 ctx,
                 list_alloc_size.into_int_value(),
                 Some("listcomp"),
-            );
-
-            let i = generator.gen_store_target(ctx, target, Some("i.addr"))?.unwrap();
-            ctx.builder
-                .build_store(i, ctx.builder.build_int_sub(start, step, "start_init").unwrap())
-                .unwrap();
-
-            ctx.builder
-                .build_conditional_branch(
-                    gen_in_range_check(ctx, start, stop, step),
-                    test_bb,
-                    cont_bb,
-                )
-                .unwrap();
-
-            ctx.builder.position_at_end(test_bb);
-            // add and test
-            let tmp = ctx
-                .builder
-                .build_int_add(
-                    ctx.builder.build_load(i, "i").map(BasicValueEnum::into_int_value).unwrap(),
-                    step,
-                    "start_loop",
-                )
-                .unwrap();
-            ctx.builder.build_store(i, tmp).unwrap();
-            ctx.builder
-                .build_conditional_branch(
-                    gen_in_range_check(ctx, tmp, stop, step),
-                    body_bb,
-                    cont_bb,
-                )
-                .unwrap();
-
-            ctx.builder.position_at_end(body_bb);
+            )
         }
         TypeEnum::TObj { obj_id, .. }
             if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
@@ -1174,35 +1107,11 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
             let length = ctx
                 .build_gep_and_load(
                     iter_val.into_pointer_value(),
-                    &[zero_size_t, int32.const_int(1, false)],
+                    &[llvm_usize.const_zero(), llvm_i32.const_int(1, false)],
                     Some("length"),
                 )
                 .into_int_value();
-            list = ListType::new(ctx, &elem_ty).construct(generator, ctx, length, Some("listcomp"));
-
-            let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?;
-            // counter = -1
-            ctx.builder.build_store(counter, size_t.const_all_ones()).unwrap();
-            ctx.builder.build_unconditional_branch(test_bb).unwrap();
-
-            ctx.builder.position_at_end(test_bb);
-            let tmp =
-                ctx.builder.build_load(counter, "i").map(BasicValueEnum::into_int_value).unwrap();
-            let tmp = ctx.builder.build_int_add(tmp, size_t.const_int(1, false), "inc").unwrap();
-            ctx.builder.build_store(counter, tmp).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);
-            let arr_ptr = ctx
-                .build_gep_and_load(
-                    iter_val.into_pointer_value(),
-                    &[zero_size_t, zero_32],
-                    Some("arr.addr"),
-                )
-                .into_pointer_value();
-            let val = ctx.build_gep_and_load(arr_ptr, &[tmp], Some("val"));
-            generator.gen_assign(ctx, target, val.into(), elt.custom.unwrap())?;
+            ListType::new(ctx, &elem_ty).construct(generator, ctx, length, Some("listcomp"))
         }
         _ => {
             panic!(
@@ -1210,54 +1119,211 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
                 ctx.unifier.stringify(iter_ty)
             );
         }
-    }
-
-    // Emits the content of `cont_bb`
-    let emit_cont_bb = |ctx: &CodeGenContext<'ctx, '_>, list: ListValue<'ctx>| {
-        ctx.builder.position_at_end(cont_bb);
-        list.store_size(
-            ctx,
-            ctx.builder.build_load(index, "index").map(BasicValueEnum::into_int_value).unwrap(),
-        );
     };
 
-    for cond in ifs {
-        let result = if let Some(v) = generator.gen_expr(ctx, cond)? {
-            v.to_basic_value_enum(ctx, generator, cond.custom.unwrap())?.into_int_value()
-        } else {
-            // Bail if the predicate is an ellipsis - Emit cont_bb contents in case the
-            // no element matches the predicate
-            emit_cont_bb(ctx, list);
+    gen_for_callback(
+        generator,
+        ctx,
+        Some("listcomp"),
+        |generator, ctx| {
+            // index with respect to the new list
+            let list_idx = generator.gen_var_alloc(ctx, llvm_usize.into(), Some("index.addr"))?;
+            ctx.builder.build_store(list_idx, llvm_usize.const_zero()).unwrap();
 
-            return Ok(None);
-        };
-        let result = generator.bool_to_i1(ctx, result);
-        let succ = ctx.ctx.append_basic_block(current, "then");
-        ctx.builder.build_conditional_branch(result, succ, test_bb).unwrap();
+            // index with respect to the iterable value
+            let iter_idx = match &*ctx.unifier.get_ty(iter_ty) {
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let iter_val = RangeType::new(ctx)
+                        .map_pointer_value(iter_val.into_pointer_value(), Some("range"));
+                    let start = iter_val.load_start(ctx, Some("range"));
 
-        ctx.builder.position_at_end(succ);
-    }
+                    let i = generator.gen_store_target(ctx, target, Some("i.addr"))?.unwrap();
+                    ctx.builder.build_store(i, start).unwrap();
 
-    let Some(elem) = generator.gen_expr(ctx, elt)? else {
-        // Similarly, bail if the generator expression is an ellipsis, but keep cont_bb contents
-        emit_cont_bb(ctx, list);
+                    i
+                }
 
-        return Ok(None);
-    };
-    let i = ctx.builder.build_load(index, "i").map(BasicValueEnum::into_int_value).unwrap();
-    let elem_ptr =
-        unsafe { list.data().ptr_offset_unchecked(ctx, generator, &i, Some("elem_ptr")) };
-    let val = elem.to_basic_value_enum(ctx, generator, elt.custom.unwrap())?;
-    ctx.builder.build_store(elem_ptr, val).unwrap();
-    ctx.builder
-        .build_store(
-            index,
-            ctx.builder.build_int_add(i, size_t.const_int(1, false), "inc").unwrap(),
-        )
-        .unwrap();
-    ctx.builder.build_unconditional_branch(test_bb).unwrap();
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let i = generator.gen_var_alloc(ctx, llvm_usize.into(), Some("counter"))?;
+                    ctx.builder.build_store(i, llvm_usize.const_zero()).unwrap();
 
-    emit_cont_bb(ctx, list);
+                    i
+                }
+
+                _ => unreachable!(),
+            };
+
+            Ok((list_idx, iter_idx))
+        },
+        |generator, ctx, (_, piter_idx)| {
+            Ok(match &*ctx.unifier.get_ty(iter_ty) {
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let iter_val = RangeType::new(ctx)
+                        .map_pointer_value(iter_val.into_pointer_value(), Some("range"));
+
+                    let iter_idx = ctx
+                        .builder
+                        .build_load(piter_idx, "start_loop")
+                        .map(BasicValueEnum::into_int_value)
+                        .unwrap();
+                    gen_in_range_check(
+                        ctx,
+                        iter_idx,
+                        iter_val.load_end(ctx, None),
+                        iter_val.load_step(ctx, None),
+                    )
+                }
+
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let list = ListType::from_unifier_type(generator, ctx, iter_ty)
+                        .map_pointer_value(iter_val.into_pointer_value(), None);
+                    let length = list.load_size(ctx, Some("length"));
+
+                    let iter_idx = ctx
+                        .builder
+                        .build_load(piter_idx, "i")
+                        .map(BasicValueEnum::into_int_value)
+                        .unwrap();
+                    ctx.builder
+                        .build_int_compare(IntPredicate::SLT, iter_idx, length, "cmp")
+                        .unwrap()
+                }
+
+                _ => unreachable!(),
+            })
+        },
+        |generator, ctx, hooks, (plist_idx, piter_idx)| {
+            // If the iterable is a list, store its value into the `target` variable first
+            match &*ctx.unifier.get_ty(iter_ty) {
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let iter_idx = ctx
+                        .builder
+                        .build_load(piter_idx, "")
+                        .map(BasicValueEnum::into_int_value)
+                        .unwrap();
+
+                    let arr_ptr = ctx
+                        .build_gep_and_load(
+                            iter_val.into_pointer_value(),
+                            &[llvm_usize.const_zero(), llvm_i32.const_zero()],
+                            Some("arr.addr"),
+                        )
+                        .into_pointer_value();
+                    let val = ctx.build_gep_and_load(arr_ptr, &[iter_idx], Some("val"));
+                    generator.gen_assign(ctx, target, val.into(), elt.custom.unwrap())?;
+                }
+
+                _ => {}
+            }
+
+            for cond in ifs {
+                let result = if let Some(v) = generator.gen_expr(ctx, cond)? {
+                    v.to_basic_value_enum(ctx, generator, cond.custom.unwrap())?.into_int_value()
+                } else {
+                    // Bail if the predicate is an ellipsis
+                    hooks.build_continue_branch(&ctx.builder);
+
+                    return Ok(());
+                };
+
+                gen_if_callback(
+                    generator,
+                    ctx,
+                    |generator, ctx| Ok(generator.bool_to_i1(ctx, result)),
+                    |_, _| Ok(()),
+                    |_, ctx| {
+                        hooks.build_continue_branch(&ctx.builder);
+                        Ok(())
+                    },
+                )?;
+            }
+
+            let Some(elem) = generator.gen_expr(ctx, elt)? else {
+                // Similarly, bail if the generator expression is an ellipsis
+                hooks.build_continue_branch(&ctx.builder);
+
+                return Ok(());
+            };
+            let list_idx =
+                ctx.builder.build_load(plist_idx, "i").map(BasicValueEnum::into_int_value).unwrap();
+            let elem_ptr = unsafe {
+                list.data().ptr_offset_unchecked(ctx, generator, &list_idx, Some("elem_ptr"))
+            };
+            let val = elem.to_basic_value_enum(ctx, generator, elt.custom.unwrap())?;
+            ctx.builder.build_store(elem_ptr, val).unwrap();
+            ctx.builder
+                .build_store(
+                    plist_idx,
+                    ctx.builder
+                        .build_int_add(list_idx, llvm_usize.const_int(1, false), "inc")
+                        .unwrap(),
+                )
+                .unwrap();
+
+            Ok(())
+        },
+        |_, ctx, (plist_idx, piter_idx)| {
+            list.store_size(
+                ctx,
+                ctx.builder
+                    .build_load(plist_idx, "index")
+                    .map(BasicValueEnum::into_int_value)
+                    .unwrap(),
+            );
+
+            match &*ctx.unifier.get_ty(iter_ty) {
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.range.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let iter_val = RangeType::new(ctx)
+                        .map_pointer_value(iter_val.into_pointer_value(), Some("range"));
+                    let step = iter_val.load_step(ctx, Some("range"));
+
+                    let iter_idx = ctx
+                        .builder
+                        .build_int_add(
+                            ctx.builder
+                                .build_load(piter_idx, "i")
+                                .map(BasicValueEnum::into_int_value)
+                                .unwrap(),
+                            step,
+                            "start_loop",
+                        )
+                        .unwrap();
+                    ctx.builder.build_store(piter_idx, iter_idx).unwrap();
+                }
+
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
+                {
+                    let iter_idx = ctx
+                        .builder
+                        .build_load(piter_idx, "i")
+                        .map(BasicValueEnum::into_int_value)
+                        .unwrap();
+                    let tmp = ctx
+                        .builder
+                        .build_int_add(iter_idx, llvm_usize.const_int(1, false), "inc")
+                        .unwrap();
+                    ctx.builder.build_store(piter_idx, tmp).unwrap();
+                }
+
+                _ => unreachable!(),
+            };
+
+            Ok(())
+        },
+    )?;
 
     Ok(Some(list.as_abi_value(ctx).into()))
 }
@@ -2515,75 +2581,52 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
         }
         ExprKind::BoolOp { op, values } => {
             // requires conditional branches for short-circuiting...
+
+            let llvm_i1 = ctx.ctx.bool_type();
+
             let left = if let Some(v) = generator.gen_expr(ctx, &values[0])? {
                 v.to_basic_value_enum(ctx, generator, values[0].custom.unwrap())?.into_int_value()
             } else {
                 return Ok(None);
             };
-            let left = generator.bool_to_i1(ctx, left);
-            let current = ctx.builder.get_insert_block().unwrap().get_parent().unwrap();
-            let a_begin_bb = ctx.ctx.append_basic_block(current, "a_begin");
-            let a_end_bb = ctx.ctx.append_basic_block(current, "a_end");
-            let b_begin_bb = ctx.ctx.append_basic_block(current, "b_begin");
-            let b_end_bb = ctx.ctx.append_basic_block(current, "b_end");
-            let cont_bb = ctx.ctx.append_basic_block(current, "cont");
-            ctx.builder.build_conditional_branch(left, a_begin_bb, b_begin_bb).unwrap();
 
-            ctx.builder.position_at_end(a_end_bb);
-            ctx.builder.build_unconditional_branch(cont_bb).unwrap();
-            ctx.builder.position_at_end(b_end_bb);
-            ctx.builder.build_unconditional_branch(cont_bb).unwrap();
-            let (a, b) = match op {
-                Boolop::Or => {
-                    ctx.builder.position_at_end(a_begin_bb);
-                    let a = ctx.ctx.i8_type().const_int(1, false);
-                    ctx.builder.build_unconditional_branch(a_end_bb).unwrap();
-
-                    ctx.builder.position_at_end(b_begin_bb);
-                    let b = if let Some(v) = generator.gen_expr(ctx, &values[1])? {
-                        let b = v
-                            .to_basic_value_enum(ctx, generator, values[1].custom.unwrap())?
-                            .into_int_value();
-                        let b = generator.bool_to_i8(ctx, b);
-                        Some(b)
-                    } else {
-                        None
-                    };
-                    ctx.builder.build_unconditional_branch(b_end_bb).unwrap();
-
-                    (Some(a), b)
-                }
-                Boolop::And => {
-                    ctx.builder.position_at_end(a_begin_bb);
-                    let a = if let Some(v) = generator.gen_expr(ctx, &values[1])? {
-                        let a = v
-                            .to_basic_value_enum(ctx, generator, values[1].custom.unwrap())?
-                            .into_int_value();
-                        let a = generator.bool_to_i8(ctx, a);
-                        Some(a)
-                    } else {
-                        None
-                    };
-                    ctx.builder.build_unconditional_branch(a_end_bb).unwrap();
-
-                    ctx.builder.position_at_end(b_begin_bb);
-                    let b = ctx.ctx.i8_type().const_zero();
-                    ctx.builder.build_unconditional_branch(b_end_bb).unwrap();
-
-                    (a, Some(b))
-                }
+            let gen_right_expr = |generator: &mut G, ctx: &mut _| -> Result<_, String> {
+                Ok(if let Some(v) = generator.gen_expr(ctx, &values[1])? {
+                    Some(
+                        v.to_basic_value_enum(ctx, generator, values[1].custom.unwrap())?
+                            .into_int_value(),
+                    )
+                } else {
+                    None
+                })
             };
 
-            ctx.builder.position_at_end(cont_bb);
-            match (a, b) {
-                (Some(a), Some(b)) => {
-                    let phi = ctx.builder.build_phi(ctx.ctx.i8_type(), "").unwrap();
-                    phi.add_incoming(&[(&a, a_end_bb), (&b, b_end_bb)]);
-                    phi.as_basic_value().into()
-                }
-                (Some(a), None) => a.into(),
-                (None, Some(b)) => b.into(),
-                (None, None) => codegen_unreachable!(ctx),
+            let result = gen_if_else_expr_callback(
+                generator,
+                ctx,
+                |generator, ctx| Ok(generator.bool_to_i1(ctx, left)),
+                |generator, ctx| {
+                    Ok(match op {
+                        Boolop::And => gen_right_expr(generator, ctx)?
+                            .map(|right| generator.bool_to_i1(ctx, right)),
+
+                        Boolop::Or => Some(llvm_i1.const_all_ones()),
+                    })
+                },
+                |generator, ctx| {
+                    Ok(match op {
+                        Boolop::Or => gen_right_expr(generator, ctx)?
+                            .map(|right| generator.bool_to_i1(ctx, right)),
+
+                        Boolop::And => Some(llvm_i1.const_zero()),
+                    })
+                },
+            )?;
+
+            if let Some(result) = result {
+                generator.bool_to_i8(ctx, result.into_int_value()).into()
+            } else {
+                return Ok(None);
             }
         }
         ExprKind::BinOp { op, left, right } => {
@@ -2600,50 +2643,27 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
                 }
                 None => return Ok(None),
             };
-            let test = generator.bool_to_i1(ctx, test);
-            let body_ty = body.custom.unwrap();
-            let is_none = ctx.unifier.get_representative(body_ty) == ctx.primitives.none;
-            let result = if is_none {
-                None
-            } else {
-                let llvm_ty = ctx.get_llvm_type(generator, body_ty);
-                Some(ctx.builder.build_alloca(llvm_ty, "if_exp_result").unwrap())
-            };
-            let current = ctx.builder.get_insert_block().unwrap().get_parent().unwrap();
-            let then_bb = ctx.ctx.append_basic_block(current, "then");
-            let else_bb = ctx.ctx.append_basic_block(current, "else");
-            let cont_bb = ctx.ctx.append_basic_block(current, "cont");
-            ctx.builder.build_conditional_branch(test, then_bb, else_bb).unwrap();
 
-            ctx.builder.position_at_end(then_bb);
-            let a = generator.gen_expr(ctx, body)?;
-            if let Some(a) = a {
-                match result {
-                    None => None,
-                    Some(v) => {
-                        let a = a.to_basic_value_enum(ctx, generator, body.custom.unwrap())?;
-                        Some(ctx.builder.build_store(v, a))
-                    }
-                };
-                ctx.builder.build_unconditional_branch(cont_bb).unwrap();
-            }
+            let result = gen_if_else_expr_callback(
+                generator,
+                ctx,
+                |generator, ctx| Ok(generator.bool_to_i1(ctx, test)),
+                |generator, ctx| {
+                    generator
+                        .gen_expr(ctx, body)?
+                        .map(|a| a.to_basic_value_enum(ctx, generator, body.custom.unwrap()))
+                        .transpose()
+                },
+                |generator, ctx| {
+                    generator
+                        .gen_expr(ctx, orelse)?
+                        .map(|b| b.to_basic_value_enum(ctx, generator, body.custom.unwrap()))
+                        .transpose()
+                },
+            )?;
 
-            ctx.builder.position_at_end(else_bb);
-            let b = generator.gen_expr(ctx, orelse)?;
-            if let Some(b) = b {
-                match result {
-                    None => None,
-                    Some(v) => {
-                        let b = b.to_basic_value_enum(ctx, generator, orelse.custom.unwrap())?;
-                        Some(ctx.builder.build_store(v, b))
-                    }
-                };
-                ctx.builder.build_unconditional_branch(cont_bb).unwrap();
-            }
-
-            ctx.builder.position_at_end(cont_bb);
             if let Some(v) = result {
-                ctx.builder.build_load(v, "if_exp_val_load").map(Into::into).unwrap()
+                v.into()
             } else {
                 return Ok(None);
             }

nac3core/src/codegen/snapshots/nac3core__codegen__test__test_primitives.snap

@@ -8,12 +8,12 @@ target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(none)
-define i32 @testing(i32 %0, i32 %1) local_unnamed_addr #0 !dbg !4 {
+define i32 @main(i32 %0, i32 %1) local_unnamed_addr #0 !dbg !4 {
 init:
   %add = add i32 %1, %0, !dbg !9
   %cmp = icmp eq i32 %add, 1, !dbg !10
-  %. = select i1 %cmp, i32 %0, i32 0, !dbg !11
-  ret i32 %., !dbg !12
+  %2 = select i1 %cmp, i32 %0, i32 0, !dbg !10
+  ret i32 %2, !dbg !11
 }
 
 attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) }
@@ -25,12 +25,11 @@ attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memor
 !1 = !{i32 2, !"Dwarf Version", i32 4}
 !2 = distinct !DICompileUnit(language: DW_LANG_Python, file: !3, producer: "NAC3", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug)
 !3 = !DIFile(filename: "unknown", directory: "")
-!4 = distinct !DISubprogram(name: "testing", linkageName: "testing", scope: null, file: !3, line: 1, type: !5, scopeLine: 1, flags: DIFlagPublic, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !8)
+!4 = distinct !DISubprogram(name: "main", linkageName: "main", scope: null, file: !3, line: 1, type: !5, scopeLine: 1, flags: DIFlagPublic, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !2, retainedNodes: !8)
 !5 = !DISubroutineType(flags: DIFlagPublic, types: !6)
 !6 = !{!7}
 !7 = !DIBasicType(name: "_", flags: DIFlagPublic)
 !8 = !{}
 !9 = !DILocation(line: 1, column: 9, scope: !4)
 !10 = !DILocation(line: 2, column: 15, scope: !4)
-!11 = !DILocation(line: 2, scope: !4)
-!12 = !DILocation(line: 3, column: 8, scope: !4)
+!11 = !DILocation(line: 3, column: 8, scope: !4)

nac3core/src/codegen/stmt.rs

@@ -13,7 +13,7 @@ use nac3parser::ast::{
 };
 
 use super::{
-    expr::{destructure_range, gen_binop_expr},
+    expr::gen_binop_expr,
     gen_in_range_check,
     irrt::{handle_slice_indices, list_slice_assignment},
     macros::codegen_unreachable,
@@ -467,6 +467,83 @@ pub fn gen_setitem<'ctx, G: CodeGenerator>(
     Ok(())
 }
 
+pub fn gen_for_pythonic<'ctx, 'a, G, I, InitFn, CondFn, BodyFn, UpdateFn, OrElseFn>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    label: Option<&str>,
+    init: InitFn,
+    cond: CondFn,
+    body: BodyFn,
+    update: UpdateFn,
+    orelse: Option<OrElseFn>,
+) -> Result<(), String>
+where
+    G: CodeGenerator + ?Sized,
+    I: Clone,
+    InitFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<I, String>,
+    CondFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>, I) -> Result<IntValue<'ctx>, String>,
+    BodyFn: FnOnce(
+        &mut G,
+        &mut CodeGenContext<'ctx, 'a>,
+        LoopHooks<'ctx>,
+        I,
+    ) -> Result<(), String>,
+    UpdateFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>, I) -> Result<(), String>,
+    OrElseFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+{
+    // var_assignment static values may be changed in another branch
+    // if so, remove the static value as it may not be correct in this branch
+    let var_assignment = ctx.var_assignment.clone();
+
+    // if there is no orelse, we just go to cont_bb
+    let current = ctx.builder.get_insert_block().and_then(BasicBlock::get_parent).unwrap();
+    let orelse_bb = orelse.as_ref().map(|_| ctx.ctx.append_basic_block(current, "for.orelse"));
+
+    gen_for_callback(
+        generator,
+        ctx,
+        label,
+        init,
+        |generator, ctx, i| todo!(),
+        |generator, ctx, hooks, i| {
+            body(generator, ctx, hooks, i)?;
+
+            for (k, (_, _, counter)) in &var_assignment {
+                let (_, static_val, counter2) = ctx.var_assignment.get_mut(k).unwrap();
+                if counter != counter2 {
+                    *static_val = None;
+                }
+            }
+
+            Ok(())
+        },
+        update,
+    )?;
+
+    if let (Some(orelse), Some(orelse_bb)) = (orelse, orelse_bb) {
+        let cont_bb = ctx.builder.get_insert_block().unwrap();
+
+        ctx.builder.position_at_end(orelse_bb);
+        orelse(generator, ctx)?;
+        if !ctx.is_terminated() {
+            ctx.builder.build_unconditional_branch(cont_bb).unwrap();
+        }
+
+        orelse_bb.move_before(cont_bb).unwrap();
+
+        for (k, (_, _, counter)) in &var_assignment {
+            let (_, static_val, counter2) = ctx.var_assignment.get_mut(k).unwrap();
+            if counter != counter2 {
+                *static_val = None;
+            }
+        }
+
+        ctx.builder.position_at_end(cont_bb);
+    }
+
+    Ok(())
+}
+
 /// See [`CodeGenerator::gen_for`].
 pub fn gen_for<G: CodeGenerator>(
     generator: &mut G,
@@ -520,7 +597,7 @@ pub fn gen_for<G: CodeGenerator>(
             else {
                 codegen_unreachable!(ctx)
             };
-            let (start, stop, step) = destructure_range(ctx, iter_val);
+            let (start, stop, step) = iter_val.load_values(ctx);
 
             ctx.builder.build_store(i, start).unwrap();
 
@@ -661,7 +738,7 @@ pub fn gen_for<G: CodeGenerator>(
 }
 
 #[derive(PartialEq, Eq, Debug, Clone, Copy, Hash)]
-pub struct BreakContinueHooks<'ctx> {
+pub struct LoopHooks<'ctx> {
     /// The [exit block][`BasicBlock`] to branch to when `break`-ing out of a loop.
     exit_bb: BasicBlock<'ctx>,
 
@@ -670,7 +747,7 @@ pub struct BreakContinueHooks<'ctx> {
     latch_bb: BasicBlock<'ctx>,
 }
 
-impl<'ctx> BreakContinueHooks<'ctx> {
+impl<'ctx> LoopHooks<'ctx> {
     /// Creates a [`br` instruction][Builder::build_unconditional_branch] to the exit
     /// [`BasicBlock`], as if by calling `break`.
     pub fn build_break_branch(&self, builder: &Builder<'ctx>) {
@@ -715,7 +792,7 @@ where
     BodyFn: FnOnce(
         &mut G,
         &mut CodeGenContext<'ctx, 'a>,
-        BreakContinueHooks<'ctx>,
+        LoopHooks<'ctx>,
         I,
     ) -> Result<(), String>,
     UpdateFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>, I) -> Result<(), String>,
@@ -750,7 +827,7 @@ where
     }
 
     ctx.builder.position_at_end(body_bb);
-    let hooks = BreakContinueHooks { exit_bb: cont_bb, latch_bb: update_bb };
+    let hooks = LoopHooks { exit_bb: cont_bb, latch_bb: update_bb };
     body(generator, ctx, hooks, loop_var.clone())?;
     if !ctx.is_terminated() {
         ctx.builder.build_unconditional_branch(update_bb).unwrap();
@@ -797,7 +874,7 @@ where
     BodyFn: FnOnce(
         &mut G,
         &mut CodeGenContext<'ctx, 'a>,
-        BreakContinueHooks<'ctx>,
+        LoopHooks<'ctx>,
         IntValue<'ctx>,
     ) -> Result<(), String>,
 {
@@ -877,7 +954,7 @@ where
     BodyFn: FnOnce(
         &mut G,
         &mut CodeGenContext<'ctx, 'a>,
-        BreakContinueHooks<'ctx>,
+        LoopHooks<'ctx>,
         IntValue<'ctx>,
     ) -> Result<(), String>,
 {

nac3core/src/codegen/test.rs

@@ -89,15 +89,18 @@ impl SymbolResolver for Resolver {
     }
 }
 
-#[test]
-#[named]
-fn test_primitives() {
-    let source = indoc! { "
-        c = a + b
-        d = a if c == 1 else 0
-        return d
-        "};
-    let statements = parse_program(source, FileName::default()).unwrap();
+/// Tests compilation of a function to its IR representation.
+///
+/// - `main`: A tuple containing the [`FunSignature`] and body of the function.
+/// - `codegen_opts_override`: Optional lambda to override compilation defaults, which is
+///   `-O2 -march=native`.
+/// - `test_name`: The name of the test case; usually `function_name!()`.
+fn test_compile_to_ir(
+    main: (fn(&PrimitiveStore) -> FunSignature, &str),
+    codegen_opts_override: Option<fn(CodeGenLLVMOptions) -> CodeGenLLVMOptions>,
+    test_name: &'static str,
+) {
+    let statements = parse_program(main.1, FileName::default()).unwrap();
 
     let context = inkwell::context::Context::create();
     let composer = TopLevelComposer::new(Vec::new(), Vec::new(), ComposerConfig::default(), 64).0;
@@ -111,7 +114,97 @@ fn test_primitives() {
             as Arc<dyn SymbolResolver + Send + Sync>;
 
     let threads = vec![DefaultCodeGenerator::new("test".into(), context.i64_type()).into()];
-    let signature = FunSignature {
+    let signature = (main.0)(&primitives);
+
+    let mut store = ConcreteTypeStore::new();
+    let mut cache = HashMap::new();
+    let signature_ty = store.from_signature(&mut unifier, &primitives, &signature, &mut cache);
+    let signature_ty = store.add_cty(signature_ty);
+
+    let mut function_data = FunctionData {
+        resolver: resolver.clone(),
+        bound_variables: Vec::new(),
+        return_type: Some(signature.ret),
+    };
+    let mut virtual_checks = Vec::new();
+    let mut calls = HashMap::new();
+    let mut identifiers: HashMap<_, _> = signature
+        .args
+        .iter()
+        .map(|arg| arg.name)
+        .map(|id| (id, IdentifierInfo::default()))
+        .collect();
+    let mut inferencer = Inferencer {
+        top_level: &top_level,
+        function_data: &mut function_data,
+        unifier: &mut unifier,
+        variable_mapping: HashMap::default(),
+        primitives: &primitives,
+        virtual_checks: &mut virtual_checks,
+        calls: &mut calls,
+        defined_identifiers: identifiers.clone(),
+        in_handler: false,
+    };
+    for arg in signature.args {
+        inferencer.variable_mapping.insert(arg.name, arg.ty);
+    }
+
+    let statements = statements
+        .into_iter()
+        .map(|v| inferencer.fold_stmt(v))
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+
+    inferencer.check_block(&statements, &mut identifiers).unwrap();
+    let top_level = Arc::new(TopLevelContext {
+        definitions: Arc::new(RwLock::new(std::mem::take(&mut *top_level.definitions.write()))),
+        unifiers: Arc::new(RwLock::new(vec![(unifier.get_shared_unifier(), primitives)])),
+        personality_symbol: None,
+    });
+
+    let task = CodeGenTask {
+        subst: Vec::default(),
+        symbol_name: "main".into(),
+        body: Arc::new(statements),
+        unifier_index: 0,
+        calls: Arc::new(calls),
+        resolver,
+        store,
+        signature: signature_ty,
+        id: 0,
+    };
+    let f = Arc::new(WithCall::new(Box::new(move |module| {
+        insta::assert_snapshot!(
+            test_name,
+            module.print_to_string().to_str().map(str::trim).unwrap()
+        );
+    })));
+
+    Target::initialize_all(&InitializationConfig::default());
+
+    let llvm_options = CodeGenLLVMOptions {
+        opt_level: OptimizationLevel::Default,
+        target: CodeGenTargetMachineOptions::from_host_triple(),
+    };
+    let llvm_options = if let Some(opt_override) = codegen_opts_override {
+        opt_override(llvm_options)
+    } else {
+        llvm_options
+    };
+    let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
+    registry.add_task(task);
+    registry.wait_tasks_complete(handles);
+}
+
+#[test]
+#[named]
+fn test_primitives() {
+    let source = indoc! { "
+        c = a + b
+        d = a if c == 1 else 0
+        return d
+        "};
+    let signature = |primitives: &PrimitiveStore| FunSignature {
         args: vec![
             FuncArg {
                 name: "a".into(),
@@ -130,74 +223,7 @@ fn test_primitives() {
         vars: VarMap::new(),
     };
 
-    let mut store = ConcreteTypeStore::new();
-    let mut cache = HashMap::new();
-    let signature = store.from_signature(&mut unifier, &primitives, &signature, &mut cache);
-    let signature = store.add_cty(signature);
-
-    let mut function_data = FunctionData {
-        resolver: resolver.clone(),
-        bound_variables: Vec::new(),
-        return_type: Some(primitives.int32),
-    };
-    let mut virtual_checks = Vec::new();
-    let mut calls = HashMap::new();
-    let mut identifiers: HashMap<_, _> =
-        ["a".into(), "b".into()].map(|id| (id, IdentifierInfo::default())).into();
-    let mut inferencer = Inferencer {
-        top_level: &top_level,
-        function_data: &mut function_data,
-        unifier: &mut unifier,
-        variable_mapping: HashMap::default(),
-        primitives: &primitives,
-        virtual_checks: &mut virtual_checks,
-        calls: &mut calls,
-        defined_identifiers: identifiers.clone(),
-        in_handler: false,
-    };
-    inferencer.variable_mapping.insert("a".into(), inferencer.primitives.int32);
-    inferencer.variable_mapping.insert("b".into(), inferencer.primitives.int32);
-
-    let statements = statements
-        .into_iter()
-        .map(|v| inferencer.fold_stmt(v))
-        .collect::<Result<Vec<_>, _>>()
-        .unwrap();
-
-    inferencer.check_block(&statements, &mut identifiers).unwrap();
-    let top_level = Arc::new(TopLevelContext {
-        definitions: Arc::new(RwLock::new(std::mem::take(&mut *top_level.definitions.write()))),
-        unifiers: Arc::new(RwLock::new(vec![(unifier.get_shared_unifier(), primitives)])),
-        personality_symbol: None,
-    });
-
-    let task = CodeGenTask {
-        subst: Vec::default(),
-        symbol_name: "testing".into(),
-        body: Arc::new(statements),
-        unifier_index: 0,
-        calls: Arc::new(calls),
-        resolver,
-        store,
-        signature,
-        id: 0,
-    };
-    let f = Arc::new(WithCall::new(Box::new(|module| {
-        insta::assert_snapshot!(
-            function_name!(),
-            module.print_to_string().to_str().map(str::trim).unwrap()
-        );
-    })));
-
-    Target::initialize_all(&InitializationConfig::default());
-
-    let llvm_options = CodeGenLLVMOptions {
-        opt_level: OptimizationLevel::Default,
-        target: CodeGenTargetMachineOptions::from_host_triple(),
-    };
-    let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
-    registry.add_task(task);
-    registry.wait_tasks_complete(handles);
+    test_compile_to_ir((signature, source), None, function_name!());
 }
 
 #[test]

nac3core/src/codegen/values/ndarray/fold.rs

@@ -2,7 +2,7 @@ use inkwell::values::{BasicValue, BasicValueEnum};
 
 use super::{NDArrayValue, NDIterValue, ScalarOrNDArray};
 use crate::codegen::{
-    stmt::{gen_for_callback, BreakContinueHooks},
+    stmt::{gen_for_callback, LoopHooks},
     types::ndarray::NDIterType,
     CodeGenContext, CodeGenerator,
 };
@@ -11,7 +11,7 @@ impl<'ctx> NDArrayValue<'ctx> {
     /// Folds the elements of this ndarray into an accumulator value by applying `f`, returning the
     /// final value.
     ///
-    /// `f` has access to [`BreakContinueHooks`] to short-circuit the `fold` operation, an instance
+    /// `f` has access to [`LoopHooks`] to short-circuit the `fold` operation, an instance
     /// of `V` representing the current accumulated value, and an [`NDIterValue`] to get the
     /// properties of the current iterated element.
     pub fn fold<'a, G, V, F>(
@@ -28,7 +28,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         F: FnOnce(
             &mut G,
             &mut CodeGenContext<'ctx, 'a>,
-            BreakContinueHooks<'ctx>,
+            LoopHooks<'ctx>,
             V,
             NDIterValue<'ctx>,
         ) -> Result<V, String>,
@@ -83,7 +83,7 @@ impl<'ctx> ScalarOrNDArray<'ctx> {
         F: FnOnce(
             &mut G,
             &mut CodeGenContext<'ctx, 'a>,
-            Option<&BreakContinueHooks<'ctx>>,
+            Option<&LoopHooks<'ctx>>,
             V,
             BasicValueEnum<'ctx>,
         ) -> Result<V, String>,

nac3core/src/codegen/values/ndarray/nditer.rs

@@ -7,7 +7,7 @@ use inkwell::{
 use super::NDArrayValue;
 use crate::codegen::{
     irrt,
-    stmt::{gen_for_callback, BreakContinueHooks},
+    stmt::{gen_for_callback, LoopHooks},
     types::{
         ndarray::NDIterType,
         structure::{StructField, StructProxyType},
@@ -155,7 +155,7 @@ impl<'ctx> From<NDIterValue<'ctx>> for PointerValue<'ctx> {
 impl<'ctx> NDArrayValue<'ctx> {
     /// Iterate through every element in the ndarray.
     ///
-    /// `body` has access to [`BreakContinueHooks`] to short-circuit and [`NDIterValue`] to
+    /// `body` has access to [`LoopHooks`] to short-circuit and [`NDIterValue`] to
     /// get properties of the current iteration (e.g., the current element, indices, etc.)
     pub fn foreach<'a, G, F>(
         &self,
@@ -168,7 +168,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         F: FnOnce(
             &mut G,
             &mut CodeGenContext<'ctx, 'a>,
-            BreakContinueHooks<'ctx>,
+            LoopHooks<'ctx>,
             NDIterValue<'ctx>,
         ) -> Result<(), String>,
     {

nac3core/src/codegen/values/range.rs

@@ -154,6 +154,18 @@ impl<'ctx> RangeValue<'ctx> {
             .map(BasicValueEnum::into_int_value)
             .unwrap()
     }
+
+    /// Return a tuple of LLVM values representing the start, stop, and step values of this `range`
+    /// respectively.
+    pub fn load_values(
+        &self,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) -> (IntValue<'ctx>, IntValue<'ctx>, IntValue<'ctx>) {
+        let start = self.load_start(ctx, None);
+        let end = self.load_end(ctx, None);
+        let step = self.load_step(ctx, None);
+        (start, end, step)
+    }
 }
 
 impl<'ctx> ProxyValue<'ctx> for RangeValue<'ctx> {