core: Add gen_if_callback and gen_if_expr_callback

nac3core/src/codegen/stmt.rs

@@ -762,6 +762,125 @@ pub fn gen_if<G: CodeGenerator>(
     Ok(())
 }
 
+/// Generates an `if` construct using lambdas, similar to the following C code:
+///
+/// ```c
+/// if (cond()) {
+///     then();
+/// } else {
+///     if (orelse) orelse();
+/// }
+/// ```
+///
+/// * `cond` - A lambda containing IR statements checking whether the `if` body should be executed.
+/// The result value must be an `i1` indicating if the loop should continue.
+/// * `then` - A lambda containing IR statements executed when `cond` evaluates to `true`.
+/// * `orelse` - A lambda containing IR statements executed when `cond` evaluates to `false`. If set
+/// to [`None`], implies that the `else` block does not exist.
+pub fn gen_if_callback<'ctx, 'a, G, CondFn, BodyFn>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    cond: CondFn,
+    then: BodyFn,
+    orelse: Option<BodyFn>,
+) -> Result<(), String>
+    where
+        G: CodeGenerator + ?Sized,
+        CondFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        BodyFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<(), String>,
+{
+    let current = ctx.builder.get_insert_block().and_then(BasicBlock::get_parent).unwrap();
+    let then_bb = ctx.ctx.append_basic_block(current, "if.then");
+    let else_bb = if orelse.is_some() {
+        Some(ctx.ctx.append_basic_block(current, "if.else"))
+    } else {
+        None
+    };
+    let cont_bb = ctx.ctx.append_basic_block(current, "if.end");
+
+    let cond_v = cond(generator, ctx)?;
+    ctx.builder
+        .build_conditional_branch(generator.bool_to_i1(ctx, cond_v), then_bb, else_bb.unwrap_or(cont_bb))
+        .unwrap();
+
+    ctx.builder.position_at_end(then_bb);
+    then(generator, ctx)?;
+    ctx.builder.build_unconditional_branch(cont_bb).unwrap();
+
+    if let (Some(orelse), Some(else_bb)) = (orelse, else_bb) {
+        ctx.builder.position_at_end(else_bb);
+        orelse(generator, ctx)?;
+        ctx.builder.build_unconditional_branch(cont_bb).unwrap();
+    }
+
+    ctx.builder.position_at_end(cont_bb);
+
+    Ok(())
+}
+
+/// Generates an `if` construct with a value using lambdas, similar to the following C code:
+///
+/// ```c
+/// llvm_phi_t res;
+/// if (cond()) {
+///     res = then();
+/// } else {
+///     res = orelse();
+/// }
+/// ```
+/// 
+/// The type of both branches must be the same, and must also match `llvm_phi_t`.
+///
+/// * `llvm_phi_t` - The type of the resultant variable.
+/// * `cond` - A lambda containing IR statements checking whether the `if` body should be executed.
+/// The result value must be an `i1` indicating if the loop should continue.
+/// * `then` - A lambda containing IR statements executed when `cond` evaluates to `true`.
+/// * `orelse` - A lambda containing IR statements executed when `cond` evaluates to `false`.
+pub fn gen_if_expr_callback<'ctx, 'a, G, R, CondFn, BodyFn>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, 'a>,
+    llvm_phi_t: BasicTypeEnum<'ctx>,
+    cond: CondFn,
+    then: BodyFn,
+    orelse: BodyFn,
+) -> Result<BasicValueEnum<'ctx>, String>
+    where
+        G: CodeGenerator + ?Sized,
+        R: BasicValue<'ctx>,
+        CondFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<IntValue<'ctx>, String>,
+        BodyFn: FnOnce(&mut G, &mut CodeGenContext<'ctx, 'a>) -> Result<R, String>,
+{
+
+    let current = ctx.builder.get_insert_block().and_then(BasicBlock::get_parent).unwrap();
+    let then_bb = ctx.ctx.append_basic_block(current, "if.then");
+    let else_bb = ctx.ctx.append_basic_block(current, "if.else");
+    let cont_bb = ctx.ctx.append_basic_block(current, "if.end");
+
+    let cond_v = cond(generator, ctx)?;
+    ctx.builder
+        .build_conditional_branch(generator.bool_to_i1(ctx, cond_v), then_bb, else_bb)
+        .unwrap();
+
+    ctx.builder.position_at_end(then_bb);
+    let then_val = then(generator, ctx)?;
+    ctx.builder.build_unconditional_branch(cont_bb).unwrap();
+
+    ctx.builder.position_at_end(else_bb);
+    let else_val = orelse(generator, ctx)?;
+    ctx.builder.build_unconditional_branch(cont_bb).unwrap();
+
+    ctx.builder.position_at_end(cont_bb);
+    assert_eq!(llvm_phi_t, then_val.as_basic_value_enum().get_type());
+    assert_eq!(llvm_phi_t, else_val.as_basic_value_enum().get_type());
+    let phi_val = ctx.builder.build_phi(llvm_phi_t, "").unwrap();
+    phi_val.add_incoming(&[
+        (&then_val, then_bb),
+        (&else_val, else_bb),
+    ]);
+
+    Ok(phi_val.as_basic_value())
+}
+
 pub fn final_proxy<'ctx>(
     ctx: &mut CodeGenContext<'ctx, '_>,
     target: BasicBlock<'ctx>,