core/expr: Implement negative indices for ndarray

For generating if-constructs in IR.

nac3core/src/codegen/expr.rs

@@ -44,6 +44,8 @@ use itertools::{chain, izip, Itertools, Either};
 use nac3parser::ast::{
     self, Boolop, Comprehension, Constant, Expr, ExprKind, Location, Operator, StrRef,
 };
+use crate::codegen::classes::TypedArrayLikeAccessor;
+use crate::codegen::stmt::gen_if_else_expr_callback;
 
 use super::{CodeGenerator, llvm_intrinsics::call_memcpy_generic, need_sret};
 
@@ -1692,21 +1694,55 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
         slice.location,
     );
 
+    if let ExprKind::Slice { .. } = &slice.node {
+        return Err(String::from("subscript operator for ndarray not implemented"))
+    }
+
+    let index = if let Some(index) = generator.gen_expr(ctx, slice)? {
+        let index = index.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?.into_int_value();
+
+        gen_if_else_expr_callback(
+            generator,
+            ctx,
+            |_, ctx| {
+                Ok(ctx.builder.build_int_compare(
+                    IntPredicate::SGE,
+                    index,
+                    index.get_type().const_zero(),
+                    "",
+                ).unwrap())
+            },
+            |_, _| Ok(Some(index)),
+            |generator, ctx| {
+                let llvm_i32 = ctx.ctx.i32_type();
+
+                let len = unsafe {
+                    v.dim_sizes().get_typed_unchecked(
+                        ctx,
+                        generator,
+                        llvm_usize.const_zero(),
+                        None,
+                    )
+                };
+
+                let index = ctx.builder.build_int_add(
+                    len,
+                    ctx.builder.build_int_s_extend(index, llvm_usize, "").unwrap(),
+                    "",
+                ).unwrap();
+
+                Ok(Some(ctx.builder.build_int_truncate(index, llvm_i32, "").unwrap()))
+            },
+        )?.map(BasicValueEnum::into_int_value).unwrap()
+    } else {
+        return Ok(None)
+    };
+    let index_addr = generator.gen_var_alloc(ctx, index.get_type().into(), None)?;
+    ctx.builder.build_store(index_addr, index).unwrap();
+
     if ndims.len() == 1 && ndims[0] == 1 {
         // Accessing an element from a 1-dimensional `ndarray`
 
-        if let ExprKind::Slice { .. } = &slice.node {
-            return Err(String::from("subscript operator for ndarray not implemented"))
-        }
-
-        let index = if let Some(v) = generator.gen_expr(ctx, slice)? {
-            v.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?.into_int_value()
-        } else {
-            return Ok(None)
-        };
-        let index_addr = generator.gen_var_alloc(ctx, index.get_type().into(), None)?;
-        ctx.builder.build_store(index_addr, index).unwrap();
-
         Ok(Some(v.data()
             .get(
                 ctx,
@@ -1718,18 +1754,6 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
     } else {
         // Accessing an element from a multi-dimensional `ndarray`
 
-        if let ExprKind::Slice { .. } = &slice.node {
-            return Err(String::from("subscript operator for ndarray not implemented"))
-        }
-
-        let index = if let Some(v) = generator.gen_expr(ctx, slice)? {
-            v.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?.into_int_value()
-        } else {
-            return Ok(None)
-        };
-        let index_addr = generator.gen_var_alloc(ctx, index.get_type().into(), None)?;
-        ctx.builder.build_store(index_addr, index).unwrap();
-
         // Create a new array, remove the top dimension from the dimension-size-list, and copy the
         // elements over
         let subscripted_ndarray = generator.gen_var_alloc(

nac3core/src/codegen/stmt.rs

@@ -29,6 +29,7 @@ use nac3parser::ast::{
     Constant, ExcepthandlerKind, Expr, ExprKind, Location, Stmt, StmtKind, StrRef,
 };
 use std::convert::TryFrom;
+use itertools::Itertools;
 
 /// See [`CodeGenerator::gen_var_alloc`].
 pub fn gen_var<'ctx>(
@@ -509,28 +510,32 @@ pub fn gen_for_callback<'ctx, 'a, G, I, InitFn, CondFn, BodyFn, UpdateFn>(
 
     ctx.builder.build_unconditional_branch(init_bb).unwrap();
 
-    let loop_var = {
-        ctx.builder.position_at_end(init_bb);
-        let result = init(generator, ctx)?;
+    ctx.builder.position_at_end(init_bb);
+    let loop_var = init(generator, ctx)?;
+    if !ctx.is_terminated() {
         ctx.builder.build_unconditional_branch(cond_bb).unwrap();
-
-        result
-    };
+    }
 
     ctx.builder.position_at_end(cond_bb);
     let cond = cond(generator, ctx, loop_var.clone())?;
     assert_eq!(cond.get_type().get_bit_width(), ctx.ctx.bool_type().get_bit_width());
-    ctx.builder
-        .build_conditional_branch(cond, body_bb, cont_bb)
-        .unwrap();
+    if !ctx.is_terminated() {
+        ctx.builder
+            .build_conditional_branch(cond, body_bb, cont_bb)
+            .unwrap();
+    }
 
     ctx.builder.position_at_end(body_bb);
     body(generator, ctx, loop_var.clone())?;
-    ctx.builder.build_unconditional_branch(update_bb).unwrap();
+    if !ctx.is_terminated() {
+        ctx.builder.build_unconditional_branch(update_bb).unwrap();
+    }
 
     ctx.builder.position_at_end(update_bb);
     update(generator, ctx, loop_var)?;
-    ctx.builder.build_unconditional_branch(cond_bb).unwrap();
+    if !ctx.is_terminated() {
+        ctx.builder.build_unconditional_branch(cond_bb).unwrap();
+    }
 
     ctx.builder.position_at_end(cont_bb);
     ctx.loop_target = loop_bb;
@@ -690,6 +695,243 @@ pub fn gen_while<G: CodeGenerator>(
     Ok(())
 }
 
+/// Generates a C-style chained-`if` construct using lambdas, similar to the following C code:
+/// 
+/// ```c
+/// T val;
+/// if (ifts[0].cond()) {
+///   val = ifts[0].then();
+/// } else if (ifts[1].cond()) {
+///   val = ifts[1].then();
+/// } else if /* ... */
+/// else {
+///   if (else_fn) {
+///     val = else_fn();
+///   } else {
+///     __builtin_unreachable();
+///   }
+/// }
+/// ```
+/// 
+/// - `ifts` - A slice of tuples containing the condition and body of a branch respectively. The 
+/// branches will be generated in the order as appears in the slice.
+/// - `else_fn` - The body to generate if no other branches evaluates to `true`. If [`None`], a call
+/// to `__builtin_unreachable` will be generated instead.
+pub fn gen_chained_if_expr_callback<'ctx, 'a, G, CondFn, ThenFn, ElseFn, R>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    ifts: &[(CondFn, ThenFn)],
+    else_fn: Option<ElseFn>,
+) -> Result<Option<BasicValueEnum<'ctx>>, String>
+    where
+        G: CodeGenerator + ?Sized,
+        CondFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        ThenFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<Option<R>, String>,
+        ElseFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<Option<R>, String>,
+        R: BasicValue<'ctx>,
+{
+    assert!(!ifts.is_empty());
+
+    let current_bb = ctx.builder.get_insert_block().unwrap();
+    let current_fn = current_bb.get_parent().unwrap();
+
+    let end_bb = ctx.ctx.append_basic_block(current_fn, "if.end");
+
+    let vals = {
+        let mut vals = ifts.iter()
+            .map(|(cond, then)| -> Result<_, String> {
+                let then_bb = ctx.ctx.insert_basic_block_after(current_bb, "if.then");
+                let else_bb = ctx.ctx.insert_basic_block_after(current_bb, "if.else");
+
+                let cond = cond(generator, ctx)?;
+                assert_eq!(cond.get_type().get_bit_width(), ctx.ctx.bool_type().get_bit_width());
+                ctx.builder.build_conditional_branch(cond, then_bb, else_bb).unwrap();
+
+                ctx.builder.position_at_end(then_bb);
+                let val = then(generator, ctx)?;
+
+                if !ctx.is_terminated() {
+                    ctx.builder.build_unconditional_branch(end_bb).unwrap();
+                }
+
+                ctx.builder.position_at_end(else_bb);
+
+                Ok((val, then_bb))
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        if let Some(else_fn) = else_fn {
+            let else_bb = ctx.builder.get_insert_block().unwrap();
+            let else_val = else_fn(generator, ctx)?;
+            vals.push((else_val, else_bb));
+
+            if !ctx.is_terminated() {
+                ctx.builder.build_unconditional_branch(end_bb).unwrap();
+            }
+        } else {
+            ctx.builder.build_unreachable().unwrap();
+        }
+
+        vals
+    };
+
+    ctx.builder.position_at_end(end_bb);
+    let phi = if vals.is_empty() {
+        None
+    } else {
+        let llvm_val_ty = vals.iter()
+            .filter_map(|(val, _)| val.as_ref().map(|v| v.as_basic_value_enum().get_type()))
+            .reduce(|acc, ty| {
+                assert_eq!(acc, ty);
+                acc
+            })
+            .unwrap();
+
+        let phi = ctx.builder.build_phi(llvm_val_ty, "").unwrap();
+        vals.into_iter()
+            .filter_map(|(val, bb)| val.map(|v| (v, bb)))
+            .for_each(|(val, bb)| phi.add_incoming(&[(&val.as_basic_value_enum(), bb)]));
+
+        Some(phi.as_basic_value())
+    };
+
+    Ok(phi)
+}
+
+/// Generates a C-style chained-`if` construct using lambdas, similar to the following C code:
+///
+/// ```c
+/// if (ifts[0].cond()) {
+///   ifts[0].then();
+/// } else if (ifts[1].cond()) {
+///   ifts[1].then();
+/// } else if /* ... */
+/// else {
+///   if (else_fn) {
+///     else_fn();
+///   } else {
+///     __builtin_unreachable();
+///   }
+/// }
+/// ```
+/// 
+/// This function mainly serves as an abstraction over [`gen_chained_if_expr_callback`] when a value
+/// does not need to be returned from the `if` construct.
+///
+/// - `ifts` - A slice of tuples containing the condition and body of a branch respectively. The 
+/// branches will be generated in the order as appears in the slice.
+/// - `else_fn` - The body to generate if no other branches evaluates to `true`. If [`None`], a call
+/// to `__builtin_unreachable` will be generated instead.
+pub fn gen_chained_if_callback<'ctx, 'a, G, CondFn, ThenFn, ElseFn>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    ifts: &[(CondFn, ThenFn)],
+    else_fn: &Option<ElseFn>,
+) -> Result<(), String>
+    where
+        G: CodeGenerator + ?Sized,
+        CondFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        ThenFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+        ElseFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+{
+    let res = gen_chained_if_expr_callback(
+        generator,
+        ctx,
+        ifts.iter()
+            .map(|(cond, then)| {
+                (
+                    cond,
+                    |generator: &mut G, ctx: &mut CodeGenContext<'ctx, 'a>| {
+                        then(generator, ctx)?;
+                        Ok(None::<BasicValueEnum<'ctx>>)
+                    },
+                )
+            })
+            .collect_vec()
+            .as_slice(),
+        else_fn
+            .as_ref()
+            .map(|else_fn| |generator: &mut G, ctx: &mut CodeGenContext<'ctx, 'a>| {
+                else_fn(generator, ctx)?;
+                Ok(None)
+            }),
+    )?;
+
+    assert!(res.is_none());
+
+    Ok(())
+}
+
+/// Generates a C-style chained-`if` construct using lambdas, similar to the following C code:
+///
+/// ```c
+/// T val;
+/// if (cond_fn()) {
+///   val = then_fn();
+/// } else {
+///   val = else_fn();
+/// }
+/// ```
+///
+/// This function mainly serves as an abstraction over [`gen_chained_if_expr_callback`] for a basic
+/// `if`-`else` construct that returns a value.
+pub fn gen_if_else_expr_callback<'ctx, 'a, G, CondFn, ThenFn, ElseFn, R>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    cond_fn: CondFn,
+    then_fn: ThenFn,
+    else_fn: ElseFn,
+) -> Result<Option<BasicValueEnum<'ctx>>, String>
+    where
+        G: CodeGenerator + ?Sized,
+        CondFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        ThenFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<Option<R>, String>,
+        ElseFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<Option<R>, String>,
+        R: BasicValue<'ctx>,
+{
+    gen_chained_if_expr_callback(
+        generator,
+        ctx,
+        &[(cond_fn, then_fn)],
+        Some(else_fn),
+    )
+}
+
+/// Generates a C-style chained-`if` construct using lambdas, similar to the following C code:
+///
+/// ```c
+/// if (cond_fn()) {
+///   then_fn();
+/// } else {
+///   if (else_fn) {
+///     else_fn();
+///   }
+/// }
+/// ```
+///
+/// This function mainly serves as an abstraction over [`gen_chained_if_expr_callback`] for a basic
+/// `if`-`else` construct that does not return a value.
+pub fn gen_if_callback<'ctx, 'a, G, CondFn, ThenFn, ElseFn>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    cond_fn: CondFn,
+    then_fn: ThenFn,
+    else_fn: &Option<ElseFn>,
+) -> Result<(), String>
+    where
+        G: CodeGenerator + ?Sized,
+        CondFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        ThenFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+        ElseFn: Fn(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+{
+    gen_chained_if_callback(
+        generator,
+        ctx,
+        &[(cond_fn, then_fn)],
+        else_fn,
+    )
+}
+
 /// See [`CodeGenerator::gen_if`].
 pub fn gen_if<G: CodeGenerator>(
     generator: &mut G,

nac3standalone/demo/src/ndarray.py

@@ -79,6 +79,13 @@ def test_ndarray_copy():
     output_ndarray_float_2(x)
     output_ndarray_float_2(y)
 
+def test_ndarray_neg_idx():
+    x = np_identity(2)
+
+    for i in range(-1, -3, -1):
+        for j in range(-1, -3, -1):
+            output_float64(x[i][j])
+
 def test_ndarray_add():
     x = np_identity(2)
     y = x + np_ones([2, 2])
@@ -639,6 +646,7 @@ def run() -> int32:
     test_ndarray_identity()
     test_ndarray_fill()
     test_ndarray_copy()
+    test_ndarray_neg_idx()
     test_ndarray_add()
     test_ndarray_add_broadcast()
     test_ndarray_add_broadcast_lhs_scalar()