5 files changed, 46 insertions, 25 deletions

diff --git a/src/yuescript/yue_ast.cpp b/src/yuescript/yue_ast.cpp
index fdb1d20..e69dd37 100644
--- a/src/yuescript/yue_ast.cpp
+++ b/src/yuescript/yue_ast.cpp
@@ -372,8 +372,8 @@ std::string With_t::to_string(void* ud) const {
         str_list temp{
                 eop ? "with?"s : "with"s,
                 valueList->to_string(ud)};
-        if (assigns) {
-                temp.push_back(assigns->to_string(ud));
+        if (assign) {
+                temp.push_back(':' + assign->to_string(ud));
         }
         if (body.is<Statement_t>()) {
                 return join(temp, " "sv) + " do "s + body->to_string(ud);
@@ -419,6 +419,9 @@ std::string SwitchCase_t::to_string(void* ud) const {
 std::string Switch_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
         str_list temp{"switch "s + target->to_string(ud)};
+        if (assignment) {
+                temp.back().append(" :"s + assignment->to_string(ud));
+        }
         info->pushScope();
         for (auto branch : branches.objects()) {
                 temp.emplace_back(info->ind() + branch->to_string(ud));
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index 0c15fac..946a587 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -287,9 +287,9 @@ AST_END(Return)
 AST_NODE(With)
         ast_ptr<false, ExistentialOp_t> eop;
         ast_ptr<true, ExpList_t> valueList;
-        ast_ptr<false, Assign_t> assigns;
+        ast_ptr<false, Assign_t> assign;
         ast_sel<true, Block_t, Statement_t> body;
-        AST_MEMBER(With, &eop, &valueList, &assigns, &body)
+        AST_MEMBER(With, &eop, &valueList, &assign, &body)
 AST_END(With)
 
 AST_NODE(SwitchList)
@@ -304,20 +304,21 @@ AST_NODE(SwitchCase)
         AST_MEMBER(SwitchCase, &condition, &body)
 AST_END(SwitchCase)
 
+AST_NODE(Assignment)
+        ast_ptr<false, ExpList_t> expList;
+        ast_ptr<true, Assign_t> assign;
+        AST_MEMBER(Assignment, &expList, &assign)
+AST_END(Assignment)
+
 AST_NODE(Switch)
         ast_ptr<true, Exp_t> target;
+        ast_ptr<false, Assignment_t> assignment;
         ast_ptr<true, Seperator_t> sep;
         ast_list<true, SwitchCase_t> branches;
         ast_sel<false, Block_t, Statement_t> lastBranch;
-        AST_MEMBER(Switch, &target, &sep, &branches, &lastBranch)
+        AST_MEMBER(Switch, &target, &assignment, &sep, &branches, &lastBranch)
 AST_END(Switch)
 
-AST_NODE(Assignment)
-        ast_ptr<false, ExpList_t> expList;
-        ast_ptr<true, Assign_t> assign;
-        AST_MEMBER(Assignment, &expList, &assign)
-AST_END(Assignment)
-
 AST_NODE(IfCond)
         ast_ptr<true, Exp_t> condition;
         ast_ptr<false, Assignment_t> assignment;
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 7abc929..9d2037f 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -78,7 +78,7 @@ static std::unordered_set<std::string> Metamethods = {
         "close"s // Lua 5.4
 };
 
-const std::string_view version = "0.28.3"sv;
+const std::string_view version = "0.28.4"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -9370,11 +9370,11 @@ private:
                 std::string withVar;
                 bool needScope = !currentScope().lastStatement && !returnValue;
                 bool extraScope = false;
-                if (with->assigns) {
+                if (with->assign) {
                         auto vars = getAssignVars(with);
                         if (vars.front().empty() || isDeclaredAsGlobal(vars.front())) {
-                                if (with->assigns->values.objects().size() == 1) {
-                                        auto var = singleVariableFrom(with->assigns->values.objects().front(), AccessType::Read);
+                                if (with->assign->values.objects().size() == 1) {
+                                        auto var = singleVariableFrom(with->assign->values.objects().front(), AccessType::Read);
                                         if (!var.empty() && isLocal(var)) {
                                                 withVar = var;
                                         }
@@ -9384,7 +9384,7 @@ private:
                                         auto assignment = x->new_ptr<ExpListAssign_t>();
                                         assignment->expList.set(toAst<ExpList_t>(withVar, x));
                                         auto assign = x->new_ptr<Assign_t>();
-                                        assign->values.push_back(with->assigns->values.objects().front());
+                                        assign->values.push_back(with->assign->values.objects().front());
                                         assignment->action.set(assign);
                                         if (needScope) {
                                                 extraScope = true;
@@ -9398,7 +9398,7 @@ private:
                                 auto assign = x->new_ptr<Assign_t>();
                                 assign->values.push_back(toAst<Exp_t>(withVar, x));
                                 bool skipFirst = true;
-                                for (auto value : with->assigns->values.objects()) {
+                                for (auto value : with->assign->values.objects()) {
                                         if (skipFirst) {
                                                 skipFirst = false;
                                                 continue;
@@ -9411,7 +9411,7 @@ private:
                                 withVar = vars.front();
                                 auto assignment = x->new_ptr<ExpListAssign_t>();
                                 assignment->expList.set(with->valueList);
-                                assignment->action.set(with->assigns);
+                                assignment->action.set(with->assign);
                                 if (needScope) {
                                         extraScope = true;
                                         temp.push_back(indent() + "do"s + nll(with));
@@ -10760,10 +10760,26 @@ private:
                         pushScope();
                 }
                 bool extraScope = false;
+                if (switchNode->assignment) {
+                        if (needScope) {
+                                extraScope = true;
+                                temp.push_back(indent() + "do"s + nll(x));
+                                pushScope();
+                        }
+                        auto asmt = x->new_ptr<ExpListAssign_t>();
+                        auto expList = x->new_ptr<ExpList_t>();
+                        expList->exprs.push_back(switchNode->target);
+                        if (switchNode->assignment->expList) {
+                                expList->exprs.dup(switchNode->assignment->expList->exprs);
+                        }
+                        asmt->expList.set(expList);
+                        asmt->action.set(switchNode->assignment->assign);
+                        transformAssignment(asmt, temp);
+                }
                 auto objVar = singleVariableFrom(switchNode->target, AccessType::Read);
                 if (objVar.empty() || !isLocal(objVar)) {
                         if (usage == ExpUsage::Common || usage == ExpUsage::Assignment) {
-                                if (needScope) {
+                                if (needScope && !extraScope) {
                                         extraScope = true;
                                         temp.push_back(indent() + "do"s + nll(x));
                                         pushScope();
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index 078509c..0271e07 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -118,8 +118,8 @@ YueParser::YueParser() {
                 return false;
         });
 
-        if_assignment_syntax_error = pl::user(true_(), [](const item_t& item) {
-                throw ParserError("use := for if-assignment expression"sv, item.begin);
+        assignment_expression_syntax_error = pl::user(true_(), [](const item_t& item) {
+                throw ParserError("use := for assignment expression"sv, item.begin);
                 return false;
         });
 
@@ -356,7 +356,7 @@ YueParser::YueParser() {
 
         Return = key("return") >> -(space >> (TableBlock | ExpListLow));
 
-        with_exp = ExpList >> -(space >> Assign);
+        with_exp = ExpList >> -(space >> (':' >> Assign | and_('=') >> assignment_expression_syntax_error));
 
         With = key("with") >> -ExistentialOp >> space >> disable_do_chain_arg_table_block_rule(with_exp) >> space >> body_with("do");
         SwitchCase = key("when") >> space >> disable_chain_rule(disable_arg_table_block_rule(SwitchList)) >> space >> body_with("then");
@@ -372,7 +372,8 @@ YueParser::YueParser() {
                 and_(SimpleTable | TableLit) >> Exp |
                 exp_not_tab >> *(space >> ',' >> space >> exp_not_tab)
         );
-        Switch = key("switch") >> space >> Exp >>
+        Switch = key("switch") >> space >>
+                Exp >> -(space >> Assignment) >>
                 space >> Seperator >> (
                         SwitchCase >> space >> (
                                 switch_block |
@@ -381,7 +382,7 @@ YueParser::YueParser() {
                         +space_break >> advance_match >> space >> SwitchCase >> switch_block >> pop_indent
                 );
 
-        Assignment = -(',' >> space >> ExpList >> space) >> (':' >> Assign | and_('=') >> if_assignment_syntax_error);
+        Assignment = -(',' >> space >> ExpList >> space) >> (':' >> Assign | and_('=') >> assignment_expression_syntax_error);
         IfCond = disable_chain_rule(disable_arg_table_block_rule(Exp >> -(space >> Assignment)));
         if_else_if = -(line_break >> *space_break >> check_indent_match) >> space >> key("elseif") >> space >> IfCond >> space >> body_with("then");
         if_else = -(line_break >> *space_break >> check_indent_match) >> space >> key("else") >> space >> body;
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index 4488685..6263857 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -156,7 +156,7 @@ private:
         NONE_AST_RULE(invalid_interpolation_error);
         NONE_AST_RULE(confusing_unary_not_error);
         NONE_AST_RULE(table_key_pair_error);
-        NONE_AST_RULE(if_assignment_syntax_error);
+        NONE_AST_RULE(assignment_expression_syntax_error);
 
         NONE_AST_RULE(inc_exp_level);
         NONE_AST_RULE(dec_exp_level);