5 files changed, 183 insertions, 54 deletions

diff --git a/src/yuescript/yue_ast.cpp b/src/yuescript/yue_ast.cpp
index 83e5985..9d68161 100644
--- a/src/yuescript/yue_ast.cpp
+++ b/src/yuescript/yue_ast.cpp
@@ -527,10 +527,10 @@ std::string While_t::to_string(void* ud) const {
 }
 std::string Repeat_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
-        str_list temp;
-        if (body->content.is<Statement_t>()) {
-                temp.emplace_back("repeat "s + body->to_string(ud));
+        if (body.is<Statement_t>()) {
+                return "repeat "s + body->to_string(ud) + " until "s + condition->to_string(ud);
         } else {
+                str_list temp;
                 temp.emplace_back("repeat"s);
                 info->pushScope();
                 temp.emplace_back(body->to_string(ud));
@@ -538,9 +538,9 @@ std::string Repeat_t::to_string(void* ud) const {
                         temp.back() = info->ind() + "--"s;
                 }
                 info->popScope();
+                temp.emplace_back(info->ind() + "until "s + condition->to_string(ud));
+                return join(temp, "\n"sv);
         }
-        temp.emplace_back(info->ind() + "until "s + condition->to_string(ud));
-        return join(temp, "\n"sv);
 }
 std::string ForStepValue_t::to_string(void* ud) const {
         return value->to_string(ud);
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index 946a587..d3e6368 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -346,7 +346,7 @@ AST_NODE(While)
 AST_END(While)
 
 AST_NODE(Repeat)
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         ast_ptr<true, Exp_t> condition;
         AST_MEMBER(Repeat, &body, &condition)
 AST_END(Repeat)
@@ -550,8 +550,8 @@ AST_NODE(SimpleValue)
         ast_sel<true,
                 TableLit_t, ConstValue_t,
                 If_t, Switch_t, With_t, ClassDecl_t,
-                ForEach_t, For_t, While_t, Do_t, Try_t,
-                UnaryValue_t,
+                ForEach_t, For_t, While_t, Repeat_t,
+                Do_t, Try_t, UnaryValue_t,
                 TblComprehension_t, Comprehension_t,
                 FunLit_t, Num_t, VarArg_t> value;
         AST_MEMBER(SimpleValue, &value)
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 9d2037f..2312025 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.4"sv;
+const std::string_view version = "0.28.5"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -1846,6 +1846,7 @@ private:
                                                                 case id<ForEach_t>(): transformForEach(static_cast<ForEach_t*>(value), out); break;
                                                                 case id<For_t>(): transformFor(static_cast<For_t*>(value), out); break;
                                                                 case id<While_t>(): transformWhile(static_cast<While_t*>(value), out); break;
+                                                                case id<Repeat_t>(): transformRepeat(static_cast<Repeat_t*>(value), out); break;
                                                                 case id<Do_t>(): transformDo(static_cast<Do_t*>(value), out, ExpUsage::Common); break;
                                                                 case id<Try_t>(): transformTry(static_cast<Try_t*>(value), out, ExpUsage::Common); break;
                                                                 case id<Comprehension_t>(): {
@@ -2400,6 +2401,13 @@ private:
                                 out.back().insert(0, preDefine);
                                 return false;
                         }
+                        case id<Repeat_t>(): {
+                                auto expList = assignment->expList.get();
+                                std::string preDefine = getPreDefineLine(assignment);
+                                transformRepeatInPlace(static_cast<Repeat_t*>(value), out, expList);
+                                out.back().insert(0, preDefine);
+                                return false;
+                        }
                         case id<TableLit_t>(): {
                                 auto tableLit = static_cast<TableLit_t*>(value);
                                 if (hasSpreadExp(tableLit->values.objects())) {
@@ -4471,6 +4479,7 @@ private:
                         case id<ForEach_t>(): transformForEachClosure(static_cast<ForEach_t*>(value), out); break;
                         case id<For_t>(): transformForClosure(static_cast<For_t*>(value), out); break;
                         case id<While_t>(): transformWhileClosure(static_cast<While_t*>(value), out); break;
+                        case id<Repeat_t>(): transformRepeatClosure(static_cast<Repeat_t*>(value), out); break;
                         case id<Do_t>(): transformDo(static_cast<Do_t*>(value), out, ExpUsage::Closure); break;
                         case id<Try_t>(): transformTry(static_cast<Try_t*>(value), out, ExpUsage::Closure); break;
                         case id<UnaryValue_t>(): transformUnaryValue(static_cast<UnaryValue_t*>(value), out); break;
@@ -5445,6 +5454,9 @@ private:
                                                 case id<While_t>():
                                                         transformWhileInPlace(static_cast<While_t*>(value), out);
                                                         return;
+                                                case id<Repeat_t>():
+                                                        transformRepeatInPlace(static_cast<Repeat_t*>(value), out);
+                                                        return;
                                                 case id<For_t>():
                                                         transformForInPlace(static_cast<For_t*>(value), out);
                                                         return;
@@ -8386,7 +8398,7 @@ private:
         std::string transformRepeatBody(Repeat_t* repeatNode, str_list& out) {
                 str_list temp;
                 bool extraDo = false;
-                auto body = repeatNode->body->content.get();
+                auto body = repeatNode->body.get();
                 auto breakLoopType = getBreakLoopType(body, Empty);
                 bool withContinue = hasContinue(breakLoopType);
                 std::string conditionVar;
@@ -9513,22 +9525,19 @@ private:
                                 expListAssign->expList.set(expList);
                                 auto stmt = x->new_ptr<Statement_t>();
                                 stmt->content.set(expListAssign);
-                                auto whileNode = toAst<While_t>("while true do break"s, x);
+                                auto repeatNode = toAst<Repeat_t>("repeat\n\t--\nuntil true"s, x);
                                 auto block = x->new_ptr<Block_t>();
                                 block->statements.push_back(stmt);
-                                block->statements.push_back(whileNode->body);
-                                auto body = x->new_ptr<Body_t>();
-                                body->content.set(block);
-                                whileNode->body.set(block);
+                                repeatNode->body.set(block);
                                 auto sVal = x->new_ptr<SimpleValue_t>();
-                                sVal->value.set(whileNode);
+                                sVal->value.set(repeatNode);
                                 auto asmt = assignmentFrom(toAst<Exp_t>(breakWithVar, x), newExp(sVal, x), x);
                                 transformAssignment(asmt, temp);
                         }
                 } else {
                         bool transformed = false;
                         if (!breakWithVar.empty()) {
-                                auto whileNode = toAst<While_t>("while true do break"s, x);
+                                auto repeatNode = toAst<Repeat_t>("repeat\n\t--\nuntil true"s, x);
                                 auto block = x->new_ptr<Block_t>();
                                 if (auto blk = with->body.as<Block_t>()) {
                                         block->statements.dup(blk->statements);
@@ -9536,15 +9545,9 @@ private:
                                         auto stmt = with->body.to<Statement_t>();
                                         block->statements.push_back(stmt);
                                 }
-                                auto breakLoop = whileNode->body.to<Statement_t>()->content.as<BreakLoop_t>();
-                                if (!(breakLoop && breakLoop->type.is<Break_t>())) {
-                                        block->statements.push_back(whileNode->body);
-                                }
-                                auto body = x->new_ptr<Body_t>();
-                                body->content.set(block);
-                                whileNode->body.set(block);
+                                repeatNode->body.set(block);
                                 auto sVal = x->new_ptr<SimpleValue_t>();
-                                sVal->value.set(whileNode);
+                                sVal->value.set(repeatNode);
                                 auto asmt = assignmentFrom(toAst<Exp_t>(breakWithVar, x), newExp(sVal, x), x);
                                 transformAssignment(asmt, temp);
                                 transformed = true;
@@ -10708,9 +10711,7 @@ private:
                         expListAssign->expList.set(expList);
                         auto stmt = x->new_ptr<Statement_t>();
                         stmt->content.set(expListAssign);
-                        auto body = x->new_ptr<Body_t>();
-                        body->content.set(stmt);
-                        repeat->body.set(body);
+                        repeat->body.set(stmt);
                         transformRepeat(repeat, out);
                         return;
                 }
@@ -10727,6 +10728,106 @@ private:
                 out.push_back(clearBuf());
         }
 
+        void transformRepeatInPlace(Repeat_t* repeatNode, str_list& out, ExpList_t* expList = nullptr) {
+                auto x = repeatNode;
+                str_list temp;
+                bool extraScope = false;
+                if (expList) {
+                        if (!currentScope().lastStatement) {
+                                extraScope = true;
+                                temp.push_back(indent() + "do"s + nll(repeatNode));
+                                pushScope();
+                        }
+                }
+                auto accumVar = getUnusedName("_accum_"sv);
+                addToScope(accumVar);
+                auto lenVar = getUnusedName("_len_"sv);
+                addToScope(lenVar);
+                auto breakLoopType = getBreakLoopType(repeatNode->body, accumVar);
+                _buf << indent() << "local "sv << accumVar << (hasBreakWithValue(breakLoopType) ? ""sv : " = { }"sv) << nll(repeatNode);
+                temp.emplace_back(clearBuf());
+                _buf << indent() << "local "s << lenVar << " = 1"s << nll(repeatNode);
+                auto& lenAssign = temp.emplace_back(clearBuf());
+                auto condStr = transformCondExp(repeatNode->condition, false);
+                temp.push_back(indent() + "repeat"s + nll(repeatNode));
+                pushScope();
+                if (hasBreakWithValue(breakLoopType)) {
+                        lenAssign.clear();
+                        transformLoopBody(repeatNode->body, temp, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
+                        auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, repeatNode);
+                        assignLeft->followStmt = followStmt.get();
+                        transformLoopBody(repeatNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
+                        if (!assignLeft->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
+                }
+                popScope();
+                temp.push_back(indent() + "until "s + condStr + nlr(repeatNode));
+                if (expList) {
+                        auto assign = x->new_ptr<Assign_t>();
+                        assign->values.push_back(toAst<Exp_t>(accumVar, x));
+                        auto assignment = x->new_ptr<ExpListAssign_t>();
+                        assignment->expList.set(expList);
+                        assignment->action.set(assign);
+                        transformAssignment(assignment, temp);
+                        if (extraScope) popScope();
+                } else {
+                        temp.push_back(indent() + "return "s + accumVar + nlr(repeatNode));
+                }
+                if (expList && extraScope) {
+                        temp.push_back(indent() + "end"s + nlr(repeatNode));
+                }
+                out.push_back(join(temp));
+        }
+
+        void transformRepeatClosure(Repeat_t* repeatNode, str_list& out) {
+                auto x = repeatNode;
+                auto simpleValue = x->new_ptr<SimpleValue_t>();
+                simpleValue->value.set(repeatNode);
+                if (transformAsUpValueFunc(newExp(simpleValue, x), out)) {
+                        return;
+                }
+                str_list temp;
+                pushAnonFunctionScope();
+                pushAnonVarArg();
+                std::string& funcStart = temp.emplace_back();
+                pushScope();
+                auto accumVar = getUnusedName("_accum_"sv);
+                addToScope(accumVar);
+                auto lenVar = getUnusedName("_len_"sv);
+                addToScope(lenVar);
+                auto breakLoopType = getBreakLoopType(repeatNode->body, accumVar);
+                _buf << indent() << "local "sv << accumVar << (hasBreakWithValue(breakLoopType) ? ""sv : " = { }"sv) << nll(repeatNode);
+                temp.emplace_back(clearBuf());
+                auto& lenAssign = temp.emplace_back(indent() + "local "s + lenVar + " = 1"s + nll(repeatNode));
+                auto condStr = transformCondExp(repeatNode->condition, false);
+                temp.push_back(indent() + "repeat"s + nll(repeatNode));
+                pushScope();
+                if (hasBreakWithValue(breakLoopType)) {
+                        lenAssign.clear();
+                        transformLoopBody(repeatNode->body, temp, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
+                        auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, repeatNode);
+                        assignLeft->followStmt = followStmt.get();
+                        transformLoopBody(repeatNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
+                        if (!assignLeft->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
+                }
+                popScope();
+                temp.push_back(indent() + "until "s + condStr + nlr(repeatNode));
+                temp.push_back(indent() + "return "s + accumVar + nlr(repeatNode));
+                popScope();
+                funcStart = anonFuncStart() + nll(repeatNode);
+                temp.push_back(indent() + anonFuncEnd());
+                popAnonVarArg();
+                popFunctionScope();
+                out.push_back(join(temp));
+        }
+
         void transformRepeat(Repeat_t* repeat, str_list& out) {
                 str_list temp;
                 pushScope();
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index 0271e07..e5bdc26 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -162,6 +162,13 @@ YueParser::YueParser() {
                 ) \
         )
 
+        #define disable_until_rule(patt) ( \
+                disable_until >> ( \
+                        (patt) >> enable_until | \
+                        enable_until >> cut \
+                ) \
+        )
+
         #define body_with(str) ( \
                 key(str) >> space >> (in_block | Statement) | \
                 in_block | \
@@ -389,13 +396,19 @@ YueParser::YueParser() {
         IfType = (expr("if") | "unless") >> not_alpha_num;
         If = IfType >> space >> IfCond >> space >> opt_body_with("then") >> *if_else_if >> -if_else;
 
-        WhileType = (expr("while") | "until") >> not_alpha_num;
-        While = WhileType >> space >> disable_do_chain_arg_table_block_rule(Exp >> -(space >> Assignment)) >> space >> opt_body_with("do");
-        Repeat = key("repeat") >> space >> Body >> line_break >> *space_break >> check_indent_match >> space >> key("until") >> space >> Exp;
+        WhileType = (expr("while") | pl::user("until", [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                return st->noUntilStack.empty() || !st->noUntilStack.back();
+        })) >> not_alpha_num;
+        While = key(WhileType) >> space >> disable_do_chain_arg_table_block_rule(Exp >> -(space >> Assignment)) >> space >> opt_body_with("do");
+        Repeat = key("repeat") >> space >> (
+                in_block >> line_break >> *space_break >> check_indent_match |
+                disable_until_rule(Statement)
+        ) >> space >> key("until") >> space >> Exp;
 
         for_key = pl::user(key("for"), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                return st->noForStack.empty() || !st->noForStack.top();
+                return st->noForStack.empty() || !st->noForStack.back();
         });
         ForStepValue = ',' >> space >> Exp;
         for_args = Variable >> space >> '=' >> space >> Exp >> space >> ',' >> space >> Exp >> space >> -ForStepValue;
@@ -409,18 +422,18 @@ YueParser::YueParser() {
 
         Do = pl::user(key("do"), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                return st->noDoStack.empty() || !st->noDoStack.top();
+                return st->noDoStack.empty() || !st->noDoStack.back();
         }) >> space >> Body;
 
         disable_do = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noDoStack.push(true);
+                st->noDoStack.push_back(true);
                 return true;
         });
 
         enable_do = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noDoStack.pop();
+                st->noDoStack.pop_back();
                 return true;
         });
 
@@ -438,41 +451,53 @@ YueParser::YueParser() {
 
         disable_do_chain_arg_table_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noDoStack.push(true);
-                st->noChainBlockStack.push(true);
-                st->noTableBlockStack.push(true);
+                st->noDoStack.push_back(true);
+                st->noChainBlockStack.push_back(true);
+                st->noTableBlockStack.push_back(true);
                 return true;
         });
 
         enable_do_chain_arg_table_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noDoStack.pop();
-                st->noChainBlockStack.pop();
-                st->noTableBlockStack.pop();
+                st->noDoStack.pop_back();
+                st->noChainBlockStack.pop_back();
+                st->noTableBlockStack.pop_back();
                 return true;
         });
 
         disable_arg_table_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noTableBlockStack.push(true);
+                st->noTableBlockStack.push_back(true);
                 return true;
         });
 
         enable_arg_table_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noTableBlockStack.pop();
+                st->noTableBlockStack.pop_back();
                 return true;
         });
 
         disable_for = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noForStack.push(true);
+                st->noForStack.push_back(true);
                 return true;
         });
 
         enable_for = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noForStack.pop();
+                st->noForStack.pop_back();
+                return true;
+        });
+
+        disable_until = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->noUntilStack.push_back(true);
+                return true;
+        });
+
+        enable_until = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->noUntilStack.pop_back();
                 return true;
         });
 
@@ -564,20 +589,20 @@ YueParser::YueParser() {
 
         disable_chain = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noChainBlockStack.push(true);
+                st->noChainBlockStack.push_back(true);
                 return true;
         });
 
         enable_chain = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                st->noChainBlockStack.pop();
+                st->noChainBlockStack.pop_back();
                 return true;
         });
 
         chain_line = check_indent_match >> space >> (chain_dot_chain | colon_chain) >> -InvokeArgs;
         chain_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                return st->noChainBlockStack.empty() || !st->noChainBlockStack.top();
+                return st->noChainBlockStack.empty() || !st->noChainBlockStack.back();
         }) >> +space_break >> advance_match >> ensure(
                 chain_line >> *(+space_break >> chain_line), pop_indent);
         ChainValue =
@@ -904,7 +929,7 @@ YueParser::YueParser() {
 
         arg_table_block = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
-                return st->noTableBlockStack.empty() || !st->noTableBlockStack.top();
+                return st->noTableBlockStack.empty() || !st->noTableBlockStack.back();
         }) >> TableBlock;
 
         invoke_args_with_table =
@@ -944,7 +969,7 @@ YueParser::YueParser() {
 
         SimpleValue =
                 TableLit | ConstValue | If | Switch | Try | With |
-                ClassDecl | ForEach | For | While | Do |
+                ClassDecl | ForEach | For | While | Repeat | Do |
                 UnaryValue | TblComprehension | Comprehension |
                 FunLit | Num | VarArg;
 
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index 6263857..773bdbe 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -119,10 +119,11 @@ protected:
                 size_t stringOpen = 0;
                 std::string buffer;
                 std::stack<int> indents;
-                std::stack<bool> noDoStack;
-                std::stack<bool> noChainBlockStack;
-                std::stack<bool> noTableBlockStack;
-                std::stack<bool> noForStack;
+                std::vector<bool> noDoStack;
+                std::vector<bool> noChainBlockStack;
+                std::vector<bool> noTableBlockStack;
+                std::vector<bool> noForStack;
+                std::vector<bool> noUntilStack;
                 std::unordered_set<std::string> usedNames;
         };
 
@@ -217,6 +218,8 @@ private:
         NONE_AST_RULE(enable_for);
         NONE_AST_RULE(enable_fun_lit);
         NONE_AST_RULE(disable_fun_lit);
+        NONE_AST_RULE(disable_until);
+        NONE_AST_RULE(enable_until);
         NONE_AST_RULE(switch_else);
         NONE_AST_RULE(switch_block);
         NONE_AST_RULE(if_else_if);