IR: implement runtime enum init and switch on enum with variable

5 files changed, 90 insertions, 19 deletions

diff --git a/src/analyze.cpp b/src/analyze.cpp
index 221868debc..738c0b36d2 100644
--- a/src/analyze.cpp
+++ b/src/analyze.cpp
@@ -749,6 +749,7 @@ TypeTableEntry *get_fn_type(CodeGen *g, FnTypeId *fn_type_id) {
     // next, loop over the parameters again and compute debug information
     // and codegen information
     if (!skip_debug_info) {
+        ensure_complete_type(g, fn_type_id->return_type);
         bool first_arg_return = !fn_type_id->is_extern && handle_is_ptr(fn_type_id->return_type);
         // +1 for maybe making the first argument the return value
         LLVMTypeRef *gen_param_types = allocate<LLVMTypeRef>(1 + fn_type_id->param_count);
@@ -2534,6 +2535,7 @@ bool handle_is_ptr(TypeTableEntry *type_entry) {
         case TypeTableEntryIdErrorUnion:
              return type_has_bits(type_entry->data.error.child_type);
         case TypeTableEntryIdEnum:
+             assert(type_entry->data.enumeration.complete);
              return type_entry->data.enumeration.gen_field_count != 0;
         case TypeTableEntryIdMaybe:
              return type_entry->data.maybe.child_type->id != TypeTableEntryIdPointer &&
diff --git a/src/codegen.cpp b/src/codegen.cpp
index e60e641b2a..49a6f4739b 100644
--- a/src/codegen.cpp
+++ b/src/codegen.cpp
@@ -645,7 +645,7 @@ static LLVMValueRef gen_struct_memcpy(CodeGen *g, LLVMValueRef src, LLVMValueRef
     return LLVMBuildCall(g->builder, g->memcpy_fn_val, params, 5, "");
 }
 
-static LLVMValueRef gen_assign_raw(CodeGen *g, AstNode *source_node,
+static LLVMValueRef gen_assign_raw(CodeGen *g,
         LLVMValueRef target_ref, LLVMValueRef value,
         TypeTableEntry *op1_type, TypeTableEntry *op2_type)
 {
@@ -698,8 +698,7 @@ static LLVMValueRef ir_render_return(CodeGen *g, IrExecutable *executable, IrIns
             LLVMBuildRet(g->builder, by_val_value);
         } else {
             assert(g->cur_ret_ptr);
-            gen_assign_raw(g, return_instruction->base.source_node, g->cur_ret_ptr, value,
-                    return_type, return_instruction->value->type_entry);
+            gen_assign_raw(g, g->cur_ret_ptr, value, return_type, return_instruction->value->type_entry);
             LLVMBuildRetVoid(g->builder);
         }
     } else {
@@ -1232,8 +1231,7 @@ static LLVMValueRef ir_render_decl_var(CodeGen *g, IrExecutable *executable,
         want_zeroes = true;
 
     if (have_init_expr) {
-        gen_assign_raw(g, init_value->source_node, var->value_ref,
-                ir_llvm_value(g, init_value), var->type, init_value->type_entry);
+        gen_assign_raw(g, var->value_ref, ir_llvm_value(g, init_value), var->type, init_value->type_entry);
     } else {
         bool ignore_uninit = false;
         // handle runtime stack allocation
@@ -2078,8 +2076,7 @@ static LLVMValueRef ir_render_maybe_wrap(CodeGen *g, IrExecutable *executable, I
     assert(instruction->tmp_ptr);
 
     LLVMValueRef val_ptr = LLVMBuildStructGEP(g->builder, instruction->tmp_ptr, maybe_child_index, "");
-    gen_assign_raw(g, instruction->base.source_node, val_ptr, payload_val, child_type, instruction->value->type_entry);
-
+    gen_assign_raw(g, val_ptr, payload_val, child_type, instruction->value->type_entry);
     LLVMValueRef maybe_ptr = LLVMBuildStructGEP(g->builder, instruction->tmp_ptr, maybe_null_index, "");
     LLVMBuildStore(g->builder, LLVMConstAllOnes(LLVMInt1Type()), maybe_ptr);
 
@@ -2125,7 +2122,7 @@ static LLVMValueRef ir_render_err_wrap_payload(CodeGen *g, IrExecutable *executa
     LLVMBuildStore(g->builder, ok_err_val, err_tag_ptr);
 
     LLVMValueRef payload_ptr = LLVMBuildStructGEP(g->builder, instruction->tmp_ptr, err_union_payload_index, "");
-    gen_assign_raw(g, instruction->base.source_node, payload_ptr, payload_val, child_type, instruction->value->type_entry);
+    gen_assign_raw(g, payload_ptr, payload_val, child_type, instruction->value->type_entry);
 
     return instruction->tmp_ptr;
 }
@@ -2145,7 +2142,30 @@ static LLVMValueRef ir_render_enum_tag(CodeGen *g, IrExecutable *executable, IrI
 }
 
 static LLVMValueRef ir_render_init_enum(CodeGen *g, IrExecutable *executable, IrInstructionInitEnum *instruction) {
-    zig_panic("TODO ir_render_init_enum");
+    TypeTableEntry *enum_type = instruction->enum_type;
+    uint32_t value = instruction->field->value;
+    LLVMTypeRef tag_type_ref = enum_type->data.enumeration.tag_type->type_ref;
+    LLVMValueRef tag_value = LLVMConstInt(tag_type_ref, value, false);
+
+    if (enum_type->data.enumeration.gen_field_count == 0)
+        return tag_value;
+
+    LLVMValueRef tmp_struct_ptr = instruction->tmp_ptr;
+
+    LLVMValueRef tag_field_ptr = LLVMBuildStructGEP(g->builder, tmp_struct_ptr, enum_gen_tag_index, "");
+    LLVMBuildStore(g->builder, tag_value, tag_field_ptr);
+
+    TypeTableEntry *union_val_type = instruction->field->type_entry;
+    if (type_has_bits(union_val_type)) {
+        LLVMValueRef new_union_val = ir_llvm_value(g, instruction->init_value);
+        LLVMValueRef union_field_ptr = LLVMBuildStructGEP(g->builder, tmp_struct_ptr, enum_gen_union_index, "");
+        LLVMValueRef bitcasted_union_field_ptr = LLVMBuildBitCast(g->builder, union_field_ptr,
+                LLVMPointerType(union_val_type->type_ref, 0), "");
+
+        gen_assign_raw(g, bitcasted_union_field_ptr, new_union_val, union_val_type, union_val_type);
+    }
+
+    return tmp_struct_ptr;
 }
 
 static void set_debug_location(CodeGen *g, IrInstruction *instruction) {
@@ -2750,7 +2770,9 @@ static void do_code_gen(CodeGen *g) {
 
         if (!type_has_bits(fn_type->data.fn.fn_type_id.return_type)) {
             // nothing to do
-        } else if (fn_type->data.fn.fn_type_id.return_type->id == TypeTableEntryIdPointer) {
+        } else if (fn_type->data.fn.fn_type_id.return_type->id == TypeTableEntryIdPointer ||
+                   fn_type->data.fn.fn_type_id.return_type->id == TypeTableEntryIdFn)
+        {
             ZigLLVMAddNonNullAttr(fn_val, 0);
         } else if (handle_is_ptr(fn_type->data.fn.fn_type_id.return_type) &&
                 !fn_type->data.fn.fn_type_id.is_extern)
diff --git a/src/ir.cpp b/src/ir.cpp
index ad05c8e23e..ae629e0910 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -4324,6 +4324,13 @@ static ErrorMsg *ir_add_error(IrAnalyze *ira, IrInstruction *source_instruction,
     return ir_add_error_node(ira, source_instruction->source_node, msg);
 }
 
+static void ir_add_typedef_err_note(IrAnalyze *ira, ErrorMsg *msg, TypeTableEntry *type_entry) {
+    if (type_entry->id == TypeTableEntryIdTypeDecl) {
+        // requires tracking source_node in the typedecl type
+        zig_panic("TODO add error note about typedecls");
+    }
+}
+
 static IrInstruction *ir_exec_const_result(IrExecutable *exec) {
     if (exec->basic_block_list.length != 1)
         return nullptr;
@@ -5451,8 +5458,7 @@ static TypeTableEntry *ir_analyze_ref(IrAnalyze *ira, IrInstruction *source_inst
     if (value->type_entry->id == TypeTableEntryIdInvalid)
         return ira->codegen->builtin_types.entry_invalid;
 
-    bool is_inline = ir_should_inline(&ira->new_irb);
-    if (is_inline || instr_is_comptime(value)) {
+    if (instr_is_comptime(value)) {
         ConstExprValue *val = ir_resolve_const(ira, value, UndefBad);
         if (!val)
             return ira->codegen->builtin_types.entry_invalid;
@@ -8223,7 +8229,7 @@ static TypeTableEntry *ir_analyze_instruction_switch_target(IrAnalyze *ira,
         case TypeTableEntryIdUnion:
         case TypeTableEntryIdBlock:
         case TypeTableEntryIdBoundFn:
-            ir_add_error_node(ira, switch_target_instruction->base.source_node,
+            ir_add_error(ira, &switch_target_instruction->base,
                 buf_sprintf("invalid switch target type '%s'", buf_ptr(&target_type->name)));
             // TODO if this is a typedecl, add error note showing the declaration of the type decl
             return ira->codegen->builtin_types.entry_invalid;
@@ -8231,10 +8237,36 @@ static TypeTableEntry *ir_analyze_instruction_switch_target(IrAnalyze *ira,
     zig_unreachable();
 }
 
-static TypeTableEntry *ir_analyze_instruction_switch_var(IrAnalyze *ira,
-        IrInstructionSwitchVar *switch_var_instruction)
-{
-    zig_panic("TODO switch var analyze");
+static TypeTableEntry *ir_analyze_instruction_switch_var(IrAnalyze *ira, IrInstructionSwitchVar *instruction) {
+    IrInstruction *target_value_ptr = instruction->target_value_ptr->other;
+    if (target_value_ptr->type_entry->id == TypeTableEntryIdInvalid)
+        return ira->codegen->builtin_types.entry_invalid;
+
+    IrInstruction *prong_value = instruction->prong_value->other;
+    if (prong_value->type_entry->id == TypeTableEntryIdInvalid)
+        return ira->codegen->builtin_types.entry_invalid;
+
+    assert(target_value_ptr->type_entry->id == TypeTableEntryIdPointer);
+    TypeTableEntry *target_type = target_value_ptr->type_entry->data.pointer.child_type;
+    if (target_type->id == TypeTableEntryIdEnum) {
+        ConstExprValue *prong_val = ir_resolve_const(ira, prong_value, UndefBad);
+        if (!prong_val)
+            return ira->codegen->builtin_types.entry_invalid;
+
+        TypeEnumField *field = &target_type->data.enumeration.fields[prong_val->data.x_bignum.data.x_uint];
+        if (instr_is_comptime(target_value_ptr)) {
+            zig_panic("TODO comptime switch var");
+        }
+
+        ir_build_enum_field_ptr_from(&ira->new_irb, &instruction->base, target_value_ptr, field);
+        return get_pointer_to_type(ira->codegen, field->type_entry,
+                target_value_ptr->type_entry->data.pointer.is_const);
+    } else {
+        ErrorMsg *msg = ir_add_error(ira, &instruction->base,
+            buf_sprintf("switch on type '%s' provides no expression parameter", buf_ptr(&target_type->name)));
+        ir_add_typedef_err_note(ira, msg, target_type);
+        return ira->codegen->builtin_types.entry_invalid;
+    }
 }
 
 static TypeTableEntry *ir_analyze_instruction_enum_tag(IrAnalyze *ira, IrInstructionEnumTag *enum_tag_instruction) {
diff --git a/src/ir_print.cpp b/src/ir_print.cpp
index c723da4726..7acbbb0207 100644
--- a/src/ir_print.cpp
+++ b/src/ir_print.cpp
@@ -919,9 +919,9 @@ static void ir_print_test_comptime(IrPrint *irp, IrInstructionTestComptime *inst
 }
 
 static void ir_print_init_enum(IrPrint *irp, IrInstructionInitEnum *instruction) {
-    fprintf(irp->f, "%s.%s { ", buf_ptr(&instruction->enum_type->name), buf_ptr(instruction->field->name));
+    fprintf(irp->f, "%s.%s {", buf_ptr(&instruction->enum_type->name), buf_ptr(instruction->field->name));
     ir_print_other_instruction(irp, instruction->init_value);
-    fprintf(irp->f, "{");
+    fprintf(irp->f, "}");
 }
 
 static void ir_print_instruction(IrPrint *irp, IrInstruction *instruction) {
diff --git a/test/cases3/enum.zig b/test/cases3/enum.zig
index e588a19091..9c89414958 100644
--- a/test/cases3/enum.zig
+++ b/test/cases3/enum.zig
@@ -12,6 +12,16 @@ fn enumType() {
     assert(@sizeOf(Foo) == expected_foo_size);
     assert(@sizeOf(Bar) == 1);
 }
+
+fn enumAsReturnValue () {
+    @setFnTest(this);
+
+    switch (returnAnInt(13)) {
+        Foo.One => |value| assert(value == 13),
+        else => @unreachable(),
+    }
+}
+
 const Point = struct {
     x: u64,
     y: u64,
@@ -28,6 +38,11 @@ const Bar = enum {
     D,
 };
 
+fn returnAnInt(x: i32) -> Foo {
+    Foo.One { x }
+}
+
+
 fn assert(ok: bool) {
     if (!ok)
         @unreachable();