Basic copy elision support for multiple return values

4 files changed, 160 insertions, 76 deletions

diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index 4b5b24b71..59881735d 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -254,6 +254,10 @@ struct lbTargetList {
 };
 
 
+struct lbTupleFix {
+	Slice<lbValue> values;
+};
+
 enum lbProcedureFlag : u32 {
 	lbProcedureFlag_WithoutMemcpyPass = 1<<0,
 	lbProcedureFlag_DebugAllocaCopy = 1<<1,
@@ -306,6 +310,7 @@ struct lbProcedure {
 
 	PtrMap<Ast *, lbValue> selector_values;
 	PtrMap<Ast *, lbAddr>  selector_addr;
+	PtrMap<LLVMValueRef, lbTupleFix> tuple_fix_map;
 };
 
 
@@ -360,6 +365,7 @@ lbValue lb_emit_epi(lbModule *m, lbValue const &value, isize index);
 lbValue lb_emit_array_epi(lbModule *m, lbValue s, isize index);
 lbValue lb_emit_struct_ep(lbProcedure *p, lbValue s, i32 index);
 lbValue lb_emit_struct_ev(lbProcedure *p, lbValue s, i32 index);
+lbValue lb_emit_tuple_ev(lbProcedure *p, lbValue value, i32 index);
 lbValue lb_emit_array_epi(lbProcedure *p, lbValue value, isize index);
 lbValue lb_emit_array_ep(lbProcedure *p, lbValue s, lbValue index);
 lbValue lb_emit_deep_field_gep(lbProcedure *p, lbValue e, Selection sel);
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index d330c90a5..a518e501f 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -123,6 +123,7 @@ lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool ignore_body)
 	p->scope_stack.allocator   = a;
 	map_init(&p->selector_values,  a, 0);
 	map_init(&p->selector_addr,    a, 0);
+	map_init(&p->tuple_fix_map,    a, 0);
 
 	if (p->is_foreign) {
 		lb_add_foreign_library_path(p->module, entity->Procedure.foreign_library);
@@ -711,15 +712,8 @@ Array<lbValue> lb_value_to_array(lbProcedure *p, lbValue value) {
 	if (t == nullptr) {
 		// Do nothing
 	} else if (is_type_tuple(t)) {
-		GB_ASSERT(t->kind == Type_Tuple);
-		auto *rt = &t->Tuple;
-		if (rt->variables.count > 0) {
-			array = array_make<lbValue>(permanent_allocator(), rt->variables.count);
-			for_array(i, rt->variables) {
-				lbValue elem = lb_emit_struct_ev(p, value, cast(i32)i);
-				array[i] = elem;
-			}
-		}
+		array = array_make<lbValue>(permanent_allocator(), 0, t->Tuple.variables.count);
+		lb_append_tuple_values(p, &array, value);
 	} else {
 		array = array_make<lbValue>(permanent_allocator(), 1);
 		array[0] = value;
@@ -1034,18 +1028,43 @@ lbValue lb_emit_call(lbProcedure *p, lbValue value, Array<lbValue> const &args,
 
 		if (original_rt != rt) {
 			GB_ASSERT(split_returns);
-			GB_ASSERT(original_rt->kind == Type_Tuple);
+			GB_ASSERT(is_type_tuple(original_rt));
+
+			// IMPORTANT NOTE(bill, 2022-11-24)
+			// result_ptr is a dummy value which is only used to reference a tuple
+			// value for the "tuple-fix"
+			//
+			// The reason for the fake stack allocation is to have a unique pointer
+			// for the value to be used as a key within the procedure itself
+
 			lbValue result_ptr = lb_add_local_generated(p, original_rt, false).addr;
 			isize ret_count = original_rt->Tuple.variables.count;
+
+			auto tuple_fix_values = slice_make<lbValue>(permanent_allocator(), ret_count);
+			auto tuple_geps = slice_make<lbValue>(permanent_allocator(), ret_count);
+
 			isize offset = ft->original_arg_count;
 			for (isize j = 0; j < ret_count-1; j++) {
 				lbValue ret_arg_ptr = processed_args[offset + j];
 				lbValue ret_arg = lb_emit_load(p, ret_arg_ptr);
-				lb_emit_store(p, lb_emit_struct_ep(p, result_ptr, cast(i32)j), ret_arg);
+				tuple_fix_values[j] = ret_arg;
 			}
-			lb_emit_store(p, lb_emit_struct_ep(p, result_ptr, cast(i32)(ret_count-1)), result);
+			tuple_fix_values[ret_count-1] = result;
+
+		#if 0
+			for (isize j = 0; j < ret_count; j++) {
+				tuple_geps[j] = lb_emit_struct_ep(p, result_ptr, cast(i32)j);
+			}
+			for (isize j = 0; j < ret_count; j++) {
+				lb_emit_store(p, tuple_geps[j], tuple_fix_values[j]);
+			}
+		#endif
 
 			result = lb_emit_load(p, result_ptr);
+
+			lbTupleFix tf = {tuple_fix_values};
+			map_set(&p->tuple_fix_map, result_ptr.value, tf);
+			map_set(&p->tuple_fix_map, result.value, tf);
 		}
 
 	}
@@ -2338,7 +2357,7 @@ lbValue lb_build_builtin_proc(lbProcedure *p, Ast *expr, TypeAndValue const &tv,
 		);
 		LLVMSetWeak(value, weak);
 
-		if (tv.type->kind == Type_Tuple) {
+		if (is_type_tuple(tv.type)) {
 			Type *fix_typed = alloc_type_tuple();
 			slice_init(&fix_typed->Tuple.variables, permanent_allocator(), 2);
 			fix_typed->Tuple.variables[0] = tv.type->Tuple.variables[0];
@@ -3082,7 +3101,7 @@ lbValue lb_build_call_expr_internal(lbProcedure *p, Ast *expr) {
 		GB_ASSERT_MSG(tav.mode != Addressing_Invalid, "%s %s %d", expr_to_string(arg), expr_to_string(expr), tav.mode);
 		GB_ASSERT_MSG(tav.mode != Addressing_ProcGroup, "%s", expr_to_string(arg));
 		Type *at = tav.type;
-		if (at->kind == Type_Tuple) {
+		if (is_type_tuple(at)) {
 			arg_count += at->Tuple.variables.count;
 		} else {
 			arg_count++;
@@ -3122,9 +3141,16 @@ lbValue lb_build_call_expr_internal(lbProcedure *p, Ast *expr) {
 			lbValue a = lb_build_expr(p, arg);
 			Type *at = a.type;
 			if (at->kind == Type_Tuple) {
-				for_array(i, at->Tuple.variables) {
-					lbValue v = lb_emit_struct_ev(p, a, cast(i32)i);
-					args[arg_index++] = v;
+				lbTupleFix *tf = map_get(&p->tuple_fix_map, a.value);
+				if (tf) {
+					for_array(j, tf->values) {
+						args[arg_index++] = tf->values[j];
+					}
+				} else {
+					for_array(j, at->Tuple.variables) {
+						lbValue v = lb_emit_struct_ev(p, a, cast(i32)j);
+						args[arg_index++] = v;
+					}
 				}
 			} else {
 				args[arg_index++] = a;
diff --git a/src/llvm_backend_stmt.cpp b/src/llvm_backend_stmt.cpp
index 949ceed7d..d026d2ecd 100644
--- a/src/llvm_backend_stmt.cpp
+++ b/src/llvm_backend_stmt.cpp
@@ -698,13 +698,13 @@ void lb_build_range_tuple(lbProcedure *p, Ast *expr, Type *val0_type, Type *val1
 	i32 tuple_count = cast(i32)tuple->Tuple.variables.count;
 	i32 cond_index = tuple_count-1;
 
-	lbValue cond = lb_emit_struct_ev(p, tuple_value, cond_index);
+	lbValue cond = lb_emit_tuple_ev(p, tuple_value, cond_index);
 	lb_emit_if(p, cond, body, done);
 	lb_start_block(p, body);
 
 
-	if (val0_) *val0_ = lb_emit_struct_ev(p, tuple_value, 0);
-	if (val1_) *val1_ = lb_emit_struct_ev(p, tuple_value, 1);
+	if (val0_) *val0_ = lb_emit_tuple_ev(p, tuple_value, 0);
+	if (val1_) *val1_ = lb_emit_tuple_ev(p, tuple_value, 1);
 	if (loop_) *loop_ = loop;
 	if (done_) *done_ = done;
 }
@@ -1543,6 +1543,24 @@ void lb_build_static_variables(lbProcedure *p, AstValueDecl *vd) {
 		lb_add_member(p->module, mangled_name, global_val);
 	}
 }
+void lb_append_tuple_values(lbProcedure *p, Array<lbValue> *dst_values, lbValue src_value) {
+	Type *t = src_value.type;
+	if (t->kind == Type_Tuple) {
+		lbTupleFix *tf = map_get(&p->tuple_fix_map, src_value.value);
+		if (tf) {
+			for_array(j, tf->values) {
+				array_add(dst_values, tf->values[j]);
+			}
+		} else {
+			for_array(i, t->Tuple.variables) {
+				lbValue v = lb_emit_tuple_ev(p, src_value, cast(i32)i);
+				array_add(dst_values, v);
+			}
+		}
+	} else {
+		array_add(dst_values, src_value);
+	}
+}
 
 
 void lb_build_assignment(lbProcedure *p, Array<lbAddr> &lvals, Slice<Ast *> const &values) {
@@ -1554,18 +1572,8 @@ void lb_build_assignment(lbProcedure *p, Array<lbAddr> &lvals, Slice<Ast *> cons
 
 	for_array(i, values) {
 		Ast *rhs = values[i];
-		if (is_type_tuple(type_of_expr(rhs))) {
-			lbValue init = lb_build_expr(p, rhs);
-			Type *t = init.type;
-			GB_ASSERT(t->kind == Type_Tuple);
-			for_array(i, t->Tuple.variables) {
-				lbValue v = lb_emit_struct_ev(p, init, cast(i32)i);
-				array_add(&inits, v);
-			}
-		} else {
-			lbValue init = lb_build_expr(p, rhs);
-			array_add(&inits, init);
-		}
+		lbValue init = lb_build_expr(p, rhs);
+		lb_append_tuple_values(p, &inits, init);
 	}
 
 	GB_ASSERT(lvals.count == inits.count);
@@ -1655,15 +1663,7 @@ void lb_build_return_stmt(lbProcedure *p, Slice<Ast *> const &return_results) {
 		if (res_count != 0) {
 			for (isize res_index = 0; res_index < res_count; res_index++) {
 				lbValue res = lb_build_expr(p, return_results[res_index]);
-				Type *t = res.type;
-				if (t->kind == Type_Tuple) {
-					for_array(i, t->Tuple.variables) {
-						lbValue v = lb_emit_struct_ev(p, res, cast(i32)i);
-						array_add(&results, v);
-					}
-				} else {
-					array_add(&results, res);
-				}
+				lb_append_tuple_values(p, &results, res);
 			}
 		} else {
 			for (isize res_index = 0; res_index < return_count; res_index++) {
diff --git a/src/llvm_backend_utility.cpp b/src/llvm_backend_utility.cpp
index fa8433a36..3ae9aba8f 100644
--- a/src/llvm_backend_utility.cpp
+++ b/src/llvm_backend_utility.cpp
@@ -339,16 +339,16 @@ void lb_emit_try_lhs_rhs(lbProcedure *p, Ast *arg, TypeAndValue const &tv, lbVal
 	if (is_type_tuple(value.type)) {
 		i32 n = cast(i32)(value.type->Tuple.variables.count-1);
 		if (value.type->Tuple.variables.count == 2) {
-			lhs = lb_emit_struct_ev(p, value, 0);
+			lhs = lb_emit_tuple_ev(p, value, 0);
 		} else {
 			lbAddr lhs_addr = lb_add_local_generated(p, tv.type, false);
 			lbValue lhs_ptr = lb_addr_get_ptr(p, lhs_addr);
 			for (i32 i = 0; i < n; i++) {
-				lb_emit_store(p, lb_emit_struct_ep(p, lhs_ptr, i), lb_emit_struct_ev(p, value, i));
+				lb_emit_store(p, lb_emit_struct_ep(p, lhs_ptr, i), lb_emit_tuple_ev(p, value, i));
 			}
 			lhs = lb_addr_load(p, lhs_addr);
 		}
-		rhs = lb_emit_struct_ev(p, value, n);
+		rhs = lb_emit_tuple_ev(p, value, n);
 	} else {
 		rhs = value;
 	}
@@ -943,6 +943,54 @@ char const *llvm_type_kinds[] = {
 	"LLVMBFloatTypeKind",
 };
 
+gb_internal lbValue lb_emit_struct_ep_internal(lbProcedure *p, lbValue s, i32 index, Type *result_type) {
+	Type *t = base_type(type_deref(s.type));
+
+	i32 original_index = index;
+	index = lb_convert_struct_index(p->module, t, index);
+
+	if (lb_is_const(s)) {
+		// NOTE(bill): this cannot be replaced with lb_emit_epi
+		lbModule *m = p->module;
+		lbValue res = {};
+		LLVMValueRef indices[2] = {llvm_zero(m), LLVMConstInt(lb_type(m, t_i32), index, false)};
+		res.value = LLVMConstGEP2(lb_type(m, type_deref(s.type)), s.value, indices, gb_count_of(indices));
+		res.type = alloc_type_pointer(result_type);
+		return res;
+	} else {
+		lbValue res = {};
+		LLVMTypeRef st = lb_type(p->module, type_deref(s.type));
+		// gb_printf_err("%s\n", type_to_string(s.type));
+		// gb_printf_err("%s\n", LLVMPrintTypeToString(LLVMTypeOf(s.value)));
+		// gb_printf_err("%d\n", index);
+		GB_ASSERT_MSG(LLVMGetTypeKind(st) == LLVMStructTypeKind, "%s", llvm_type_kinds[LLVMGetTypeKind(st)]);
+		unsigned count = LLVMCountStructElementTypes(st);
+		GB_ASSERT_MSG(count >= cast(unsigned)index, "%u %d %d", count, index, original_index);
+
+		res.value = LLVMBuildStructGEP2(p->builder, st, s.value, cast(unsigned)index, "");
+		res.type = alloc_type_pointer(result_type);
+		return res;
+	}
+}
+
+lbValue lb_emit_tuple_ep(lbProcedure *p, lbValue ptr, i32 index) {
+	Type *t = type_deref(ptr.type);
+	GB_ASSERT(is_type_tuple(t));
+	Type *result_type = t->Tuple.variables[index]->type;
+
+	lbValue res = {};
+	lbTupleFix *tf = map_get(&p->tuple_fix_map, ptr.value);
+	if (tf) {
+		res = tf->values[index];
+		GB_ASSERT(are_types_identical(res.type, result_type));
+		res = lb_address_from_load_or_generate_local(p, res);
+	} else {
+		res = lb_emit_struct_ep_internal(p, ptr, index, result_type);
+	}
+	return res;
+}
+
+
 lbValue lb_emit_struct_ep(lbProcedure *p, lbValue s, i32 index) {
 	GB_ASSERT(is_type_pointer(s.type));
 	Type *t = base_type(type_deref(s.type));
@@ -958,8 +1006,7 @@ lbValue lb_emit_struct_ep(lbProcedure *p, lbValue s, i32 index) {
 		GB_ASSERT(index == -1);
 		return lb_emit_union_tag_ptr(p, s);
 	} else if (is_type_tuple(t)) {
-		GB_ASSERT(t->Tuple.variables.count > 0);
-		result_type = t->Tuple.variables[index]->type;
+		return lb_emit_tuple_ep(p, s, index);
 	} else if (is_type_complex(t)) {
 		Type *ft = base_complex_elem_type(t);
 		switch (index) {
@@ -1024,34 +1071,45 @@ lbValue lb_emit_struct_ep(lbProcedure *p, lbValue s, i32 index) {
 
 	GB_ASSERT_MSG(result_type != nullptr, "%s %d", type_to_string(t), index);
 	
-	i32 original_index = index;
-	index = lb_convert_struct_index(p->module, t, index);
-	
-	if (lb_is_const(s)) {
-		// NOTE(bill): this cannot be replaced with lb_emit_epi
-		lbModule *m = p->module;
-		lbValue res = {};
-		LLVMValueRef indices[2] = {llvm_zero(m), LLVMConstInt(lb_type(m, t_i32), index, false)};
-		res.value = LLVMConstGEP2(lb_type(m, type_deref(s.type)), s.value, indices, gb_count_of(indices));
-		res.type = alloc_type_pointer(result_type);
-		return res;
+	return lb_emit_struct_ep_internal(p, s, index, result_type);
+}
+
+lbValue lb_emit_tuple_ev(lbProcedure *p, lbValue value, i32 index) {
+	Type *t = value.type;
+	GB_ASSERT(is_type_tuple(t));
+	Type *result_type = t->Tuple.variables[index]->type;
+
+	lbValue res = {};
+	lbTupleFix *tf = map_get(&p->tuple_fix_map, value.value);
+	if (tf) {
+		res = tf->values[index];
+		GB_ASSERT(are_types_identical(res.type, result_type));
 	} else {
-		lbValue res = {};
-		LLVMTypeRef st = lb_type(p->module, type_deref(s.type));
-		// gb_printf_err("%s\n", type_to_string(s.type));
-		// gb_printf_err("%s\n", LLVMPrintTypeToString(LLVMTypeOf(s.value)));
-		// gb_printf_err("%d\n", index);
-		GB_ASSERT_MSG(LLVMGetTypeKind(st) == LLVMStructTypeKind, "%s", llvm_type_kinds[LLVMGetTypeKind(st)]);
-		unsigned count = LLVMCountStructElementTypes(st);
-		GB_ASSERT_MSG(count >= cast(unsigned)index, "%u %d %d", count, index, original_index);
-		
-		res.value = LLVMBuildStructGEP2(p->builder, st, s.value, cast(unsigned)index, "");
-		res.type = alloc_type_pointer(result_type);
-		return res;
+		if (t->Tuple.variables.count == 1) {
+			GB_ASSERT(index == 0);
+			// value.type = result_type;
+			return value;
+		}
+		if (LLVMIsALoadInst(value.value)) {
+			lbValue res = {};
+			res.value = LLVMGetOperand(value.value, 0);
+			res.type = alloc_type_pointer(value.type);
+			lbValue ptr = lb_emit_struct_ep(p, res, index);
+			return lb_emit_load(p, ptr);
+		}
+
+		res.value = LLVMBuildExtractValue(p->builder, value.value, cast(unsigned)index, "");
+		res.type = result_type;
 	}
+	return res;
 }
 
 lbValue lb_emit_struct_ev(lbProcedure *p, lbValue s, i32 index) {
+	Type *t = base_type(s.type);
+	if (is_type_tuple(t)) {
+		return lb_emit_tuple_ev(p, s, index);
+	}
+
 	if (LLVMIsALoadInst(s.value)) {
 		lbValue res = {};
 		res.value = LLVMGetOperand(s.value, 0);
@@ -1060,7 +1118,6 @@ lbValue lb_emit_struct_ev(lbProcedure *p, lbValue s, i32 index) {
 		return lb_emit_load(p, ptr);
 	}
 
-	Type *t = base_type(s.type);
 	Type *result_type = nullptr;
 
 	switch (t->kind) {
@@ -1113,12 +1170,7 @@ lbValue lb_emit_struct_ev(lbProcedure *p, lbValue s, i32 index) {
 		GB_PANIC("lb_emit_union_tag_value");
 
 	case Type_Tuple:
-		GB_ASSERT(t->Tuple.variables.count > 0);
-		result_type = t->Tuple.variables[index]->type;
-		if (t->Tuple.variables.count == 1) {
-			return s;
-		}
-		break;
+		return lb_emit_tuple_ev(p, s, index);
 	case Type_Slice:
 		switch (index) {
 		case 0: result_type = alloc_type_pointer(t->Slice.elem); break;