Improve type inference for variadic parameters

3 files changed, 27 insertions, 5 deletions

diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index adcbf7684..e3a1b2575 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -107,7 +107,7 @@ void check_init_variables(CheckerContext *ctx, Entity **lhs, isize lhs_count, Ar
 	// an extra allocation
 	auto operands = array_make<Operand>(ctx->allocator, 0, 2*lhs_count);
 	defer (array_free(&operands));
-	check_unpack_arguments(ctx, lhs, lhs_count, &operands, inits, true);
+	check_unpack_arguments(ctx, lhs, lhs_count, &operands, inits, true, false);
 
 	isize rhs_count = operands.count;
 	for_array(i, operands) {
diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index de409a54e..2193af114 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -4425,7 +4425,7 @@ bool check_assignment_arguments(CheckerContext *ctx, Array<Operand> const &lhs,
 
 
 
-bool check_unpack_arguments(CheckerContext *ctx, Entity **lhs, isize lhs_count, Array<Operand> *operands, Array<Ast *> const &rhs, bool allow_ok) {
+bool check_unpack_arguments(CheckerContext *ctx, Entity **lhs, isize lhs_count, Array<Operand> *operands, Array<Ast *> const &rhs, bool allow_ok, bool is_variadic) {
 	bool optional_ok = false;
 	isize tuple_index = 0;
 	for_array(i, rhs) {
@@ -4443,6 +4443,25 @@ bool check_unpack_arguments(CheckerContext *ctx, Entity **lhs, isize lhs_count,
 				DeclInfo *decl = decl_info_of_entity(e);
 				if (decl) c->decl = decl;
 				type_hint = e->type;
+				if (e->flags & EntityFlag_Ellipsis) {
+					GB_ASSERT(is_type_slice(e->type));
+					GB_ASSERT(e->type->kind == Type_Slice);
+					type_hint = e->type->Slice.elem;
+				}
+			}
+		}
+		if (lhs != nullptr && tuple_index >= lhs_count && is_variadic) {
+			// NOTE(bill): override DeclInfo for dependency
+			Entity *e = lhs[lhs_count-1];
+			if (e != nullptr) {
+				DeclInfo *decl = decl_info_of_entity(e);
+				if (decl) c->decl = decl;
+				type_hint = e->type;
+				if (e->flags & EntityFlag_Ellipsis) {
+					GB_ASSERT(is_type_slice(e->type));
+					GB_ASSERT(e->type->kind == Type_Slice);
+					type_hint = e->type->Slice.elem;
+				}
 			}
 		}
 
@@ -4939,8 +4958,11 @@ CallArgumentData check_call_arguments(CheckerContext *c, Operand *operand, Type
 		operands = array_make<Operand>(heap_allocator(), 0, 2*ce->args.count);
 		Entity **lhs = nullptr;
 		isize lhs_count = -1;
+		bool is_variadic = false;
 		if (proc_type != nullptr && is_type_proc(proc_type)) {
 			TypeProc *pt = &base_type(proc_type)->Proc;
+			is_variadic = pt->variadic;
+
 			if (!pt->is_polymorphic || pt->is_poly_specialized) {
 				if (pt->params != nullptr) {
 					lhs = pt->params->Tuple.variables.data;
@@ -4959,7 +4981,7 @@ CallArgumentData check_call_arguments(CheckerContext *c, Operand *operand, Type
 			}
 		}
 
-		check_unpack_arguments(c, lhs, lhs_count, &operands, ce->args, false);
+		check_unpack_arguments(c, lhs, lhs_count, &operands, ce->args, false, is_variadic);
 	}
 
 	if (operand->mode == Addressing_ProcGroup) {
@@ -5212,7 +5234,7 @@ CallArgumentError check_polymorphic_record_type(CheckerContext *c, Operand *oper
 
 	} else {
 		operands = array_make<Operand>(heap_allocator(), 0, 2*ce->args.count);
-		check_unpack_arguments(c, nullptr, -1, &operands, ce->args, false);
+		check_unpack_arguments(c, nullptr, -1, &operands, ce->args, false, false);
 	}
 
 	CallArgumentError err = CallArgumentError_None;
diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index 2ef0a6e56..9949e0bb7 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -1247,7 +1247,7 @@ void check_stmt_internal(CheckerContext *ctx, Ast *node, u32 flags) {
 		auto operands = array_make<Operand>(heap_allocator(), 0, 2*rs->results.count);
 		defer (array_free(&operands));
 
-		check_unpack_arguments(ctx, result_entities, result_count, &operands, rs->results, true);
+		check_unpack_arguments(ctx, result_entities, result_count, &operands, rs->results, true, false);
 
 		if (result_count == 0 && rs->results.count > 0) {
 			error(rs->results[0], "No return values expected");