From 60774555568dbc64866d088cb51e45f9434c0f89 Mon Sep 17 00:00:00 2001
From: Li Jin <dragon-fly@qq.com>
Date: Tue, 8 Aug 2023 18:12:06 +0800
Subject: make vararg assignment work with line decorator.

---
 spec/inputs/vararg.yue         | 22 ++++++++++++
 spec/outputs/vararg.lua        | 45 +++++++++++++++++++++++
 src/yuescript/yue_compiler.cpp | 82 +++++++++++++++++++++++++++++++++++++-----
 3 files changed, 141 insertions(+), 8 deletions(-)

diff --git a/spec/inputs/vararg.yue b/spec/inputs/vararg.yue
index c7b8ebf..466a4a6 100644
--- a/spec/inputs/vararg.yue
+++ b/spec/inputs/vararg.yue
@@ -52,5 +52,27 @@ join = (...) ->
 	f_colon f!\func
 	f_colon f(...)\func
 
+	_ = ->
+		list = {1, 2, 3, 4, 5}
+		fn = (ok) ->
+		  ok, table.unpack list
+		ok, ... = fn true
+		print ok, ...
+
+		fn_many_args = ->
+		  10, nil, 20, nil, 30
+
+		... = fn_many_args!
+		print select "#", ...
+		print ...
+
+	do
+		... = 1, 2 if true
+		print ...
+
+	do
+		a, ... = 1, 2 unless true
+		print a, select '#', ...
+
 	nil
 
diff --git a/spec/outputs/vararg.lua b/spec/outputs/vararg.lua
index 56e2011..6cfaccb 100644
--- a/spec/outputs/vararg.lua
+++ b/spec/outputs/vararg.lua
@@ -220,5 +220,50 @@ join = function(...)
 			return _fn_0(_base_0, ...)
 		end
 	end)(...))
+	local _
+	_ = function()
+		local list = {
+			1,
+			2,
+			3,
+			4,
+			5
+		}
+		local fn
+		fn = function(ok)
+			return ok, table.unpack(list)
+		end
+		return (function(_arg_0, ...)
+			local ok = _arg_0
+			print(ok, ...)
+			local fn_many_args
+			fn_many_args = function()
+				return 10, nil, 20, nil, 30
+			end
+			return (function(...)
+				print(select("#", ...))
+				return print(...)
+			end)(fn_many_args())
+		end)(fn(true))
+	end
+	do
+		(function(...)
+			return print(...)
+		end)((function()
+			if true then
+				return 1, 2
+			end
+		end)())
+	end
+	do
+		(function(_arg_0, ...)
+			local a = _arg_0
+			return print(a, select('#', ...))
+		end)((function()
+			if not true then
+				return 1, 2
+			end
+		end)())
+	end
 	return nil
 end
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 9384214..6644ca5 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -3619,14 +3619,82 @@ private:
 				auto varArg = simpleValue->value.as<VarArg_t>();
 				BREAK_IF(!varArg);
 				auto assignAction = static_cast<Assign_t*>(expListAssign->action.get());
-				bool isExps = true;
-				for (auto item : assignAction->values.objects()) {
-					if (!ast_is<Exp_t, TableBlock_t>(item)) {
-						isExps = false;
-						break;
+				auto newInvoke = assignAction->new_ptr<InvokeArgs_t>();
+				if (auto appendix = stmt->appendix.get()) {
+					switch (appendix->item->get_id()) {
+						case id<IfLine_t>(): {
+							auto if_line = static_cast<IfLine_t*>(appendix->item.get());
+							auto ifNode = appendix->new_ptr<If_t>();
+							ifNode->type.set(if_line->type);
+							ifNode->nodes.push_back(if_line->condition);
+							auto expList = appendix->new_ptr<ExpList_t>();
+							for (auto val : assignAction->values.objects()) {
+								switch (val->get_id()) {
+									case id<If_t>():
+									case id<Switch_t>():
+									case id<With_t>(): {
+										auto simpleValue = val->new_ptr<SimpleValue_t>();
+										simpleValue->value.set(val);
+										auto exp = newExp(simpleValue, val);
+										expList->exprs.push_back(exp);
+										break;
+									}
+									case id<TableBlock_t>(): {
+										auto tableBlock = static_cast<TableBlock_t*>(val);
+										auto tabLit = val->new_ptr<TableLit_t>();
+										tabLit->values.dup(tableBlock->values);
+										auto simpleValue = val->new_ptr<SimpleValue_t>();
+										simpleValue->value.set(tabLit);
+										auto exp = newExp(simpleValue, val);
+										expList->exprs.push_back(exp);
+										break;
+									}
+									case id<Exp_t>(): {
+										expList->exprs.push_back(val);
+										break;
+									}
+									default: YUEE("AST node mismatch", val); break;
+								}
+							}
+							auto newExpListAssign = assignAction->new_ptr<ExpListAssign_t>();
+							newExpListAssign->expList.set(expList);
+							auto newStmt = assignAction->new_ptr<Statement_t>();
+							newStmt->content.set(newExpListAssign);
+							ifNode->nodes.push_back(newStmt);
+							auto newSVal = ifNode->new_ptr<SimpleValue_t>();
+							newSVal->value.set(ifNode);
+							newInvoke->args.push_back(newExp(newSVal, newSVal));
+							break;
+						}
+						case id<WhileLine_t>(): {
+							throw CompileError("while-loop line decorator is not supported here"sv, appendix->item.get());
+							break;
+						}
+						case id<CompInner_t>(): {
+							throw CompileError("for-loop line decorator is not supported here"sv, appendix->item.get());
+							break;
+						}
+						default: YUEE("AST node mismatch", appendix->item.get()); break;
+					}
+				}
+				if (newInvoke->args.empty()) {
+					for (auto item : assignAction->values.objects()) {
+						switch (item->get_id()) {
+							case id<With_t>():
+							case id<If_t>():
+							case id<Switch_t>(): {
+								auto sVal = item->new_ptr<SimpleValue_t>();
+								sVal->value.set(item);
+								newInvoke->args.push_back(newExp(sVal, item));
+								break;
+							}
+							case id<TableBlock_t>():
+							case id<Exp_t>():
+								newInvoke->args.push_back(item);
+								break;
+						}
 					}
 				}
-				BREAK_IF(!isExps);
 				auto x = *nodes.begin();
 				auto newBlock = x->new_ptr<Block_t>();
 				if (it != nodes.begin()) {
@@ -3685,8 +3753,6 @@ private:
 				newCallable->item.set(newParens);
 				auto newChainValue = x->new_ptr<ChainValue_t>();
 				newChainValue->items.push_back(newCallable);
-				auto newInvoke = x->new_ptr<InvokeArgs_t>();
-				newInvoke->args.dup(assignAction->values);
 				newChainValue->items.push_back(newInvoke);
 				auto newItem = newExp(newChainValue, x);
 				auto newItemList = x->new_ptr<ExpList_t>();
-- 
cgit v1.2.3-55-g6feb