5 files changed, 205 insertions, 29 deletions

diff --git a/spec/inputs/with.yue b/spec/inputs/with.yue
index 19b7be1..3fee48e 100644
--- a/spec/inputs/with.yue
+++ b/spec/inputs/with.yue
@@ -152,4 +152,13 @@ do
                 return with {}
                         return [123]
 
+do
+        f with item
+                if .id > 0
+                        break .content
+
+        a = with tb
+                if .v
+                        break .a
+
 nil
diff --git a/spec/outputs/codes_from_doc.lua b/spec/outputs/codes_from_doc.lua
index 177c33d..3ee20bb 100644
--- a/spec/outputs/codes_from_doc.lua
+++ b/spec/outputs/codes_from_doc.lua
@@ -694,6 +694,7 @@ local some_string = "Here is a string\n  that has a line break in it."
 print("I am " .. tostring(math.random() * 100) .. "% sure.")
 local integer = 1000000
 local hex = 0xEFBBBF
+local binary = 19
 local my_function
 my_function = function() end
 my_function()
@@ -2228,6 +2229,38 @@ local inventory = {
                 }
         }
 }
+local map
+map = function(arr, action)
+        local _accum_0 = { }
+        local _len_0 = 1
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                _accum_0[_len_0] = action(item)
+                _len_0 = _len_0 + 1
+        end
+        return _accum_0
+end
+local filter
+filter = function(arr, cond)
+        local _accum_0 = { }
+        local _len_0 = 1
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                if cond(item) then
+                        _accum_0[_len_0] = item
+                        _len_0 = _len_0 + 1
+                end
+        end
+        return _accum_0
+end
+local reduce
+reduce = function(arr, init, action)
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                init = action(init, item)
+        end
+        return init
+end
 print(reduce(filter(map({
         1,
         2,
diff --git a/spec/outputs/codes_from_doc_zh.lua b/spec/outputs/codes_from_doc_zh.lua
index 6349010..52204b7 100644
--- a/spec/outputs/codes_from_doc_zh.lua
+++ b/spec/outputs/codes_from_doc_zh.lua
@@ -694,6 +694,7 @@ local some_string = "这是一个字符串\n  并包括一个换行。"
 print("我有" .. tostring(math.random() * 100) .. "%的把握。")
 local integer = 1000000
 local hex = 0xEFBBBF
+local binary = 19
 local my_function
 my_function = function() end
 my_function()
@@ -2222,6 +2223,38 @@ local inventory = {
                 }
         }
 }
+local map
+map = function(arr, action)
+        local _accum_0 = { }
+        local _len_0 = 1
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                _accum_0[_len_0] = action(item)
+                _len_0 = _len_0 + 1
+        end
+        return _accum_0
+end
+local filter
+filter = function(arr, cond)
+        local _accum_0 = { }
+        local _len_0 = 1
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                if cond(item) then
+                        _accum_0[_len_0] = item
+                        _len_0 = _len_0 + 1
+                end
+        end
+        return _accum_0
+end
+local reduce
+reduce = function(arr, init, action)
+        for _index_0 = 1, #arr do
+                local item = arr[_index_0]
+                init = action(init, item)
+        end
+        return init
+end
 print(reduce(filter(map({
         1,
         2,
diff --git a/spec/outputs/with.lua b/spec/outputs/with.lua
index 1a795c1..20c5d44 100644
--- a/spec/outputs/with.lua
+++ b/spec/outputs/with.lua
@@ -187,4 +187,31 @@ do
                 return _with_0[123]
         end
 end
+do
+        f((function()
+                local _with_0 = item
+                do
+                        local _accum_0
+                        while true do
+                                if _with_0.id > 0 then
+                                        _accum_0 = _with_0.content
+                                        break
+                                end
+                        end
+                        _with_0 = _accum_0
+                end
+                return _with_0
+        end)())
+        local a
+        local _with_0 = tb
+        local _accum_0
+        while true do
+                if _with_0.v then
+                        _accum_0 = _with_0.a
+                        break
+                end
+        end
+        _with_0 = _accum_0
+        a = _with_0
+end
 return nil
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 60ed7ac..7abc929 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.2"sv;
+const std::string_view version = "0.28.3"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -2064,6 +2064,9 @@ private:
         }
 
         bool transformAssignment(ExpListAssign_t* assignment, str_list& out, bool optionalDestruct = false) {
+                if (assignment->action.is<SubBackcall_t>()) {
+                        YUEE("AST node mismatch", assignment->action);
+                }
                 checkAssignable(assignment->expList);
                 BLOCK_START
                 auto assign = ast_cast<Assign_t>(assignment->action);
@@ -8195,18 +8198,18 @@ private:
                 transformForHead(forNode->varName, forNode->startValue, forNode->stopValue, forNode->stepValue, out);
         }
 
-        void transform_plain_body(ast_node* body, str_list& out, ExpUsage usage, ExpList_t* assignList = nullptr) {
-                switch (body->get_id()) {
+        void transform_plain_body(ast_node* bodyOrStmt, str_list& out, ExpUsage usage, ExpList_t* assignList = nullptr) {
+                switch (bodyOrStmt->get_id()) {
                         case id<Block_t>():
-                                transformBlock(static_cast<Block_t*>(body), out, usage, assignList);
+                                transformBlock(static_cast<Block_t*>(bodyOrStmt), out, usage, assignList);
                                 break;
                         case id<Statement_t>(): {
-                                auto newBlock = body->new_ptr<Block_t>();
-                                newBlock->statements.push_back(body);
+                                auto newBlock = bodyOrStmt->new_ptr<Block_t>();
+                                newBlock->statements.push_back(bodyOrStmt);
                                 transformBlock(newBlock, out, usage, assignList);
                                 break;
                         }
-                        default: YUEE("AST node mismatch", body); break;
+                        default: YUEE("AST node mismatch", bodyOrStmt); break;
                 }
         }
 
@@ -8477,12 +8480,17 @@ private:
                 _buf << indent() << "local "sv << len << " = 1"sv << nll(forNode);
                 auto& lenAssign = out.emplace_back(clearBuf());
                 transformForHead(forNode, out);
-                auto expList = toAst<ExpList_t>(accum + '[' + len + ']', x);
-                auto followStmt = toAst<Statement_t>(len + "+=1"s, forNode->body);
-                expList->followStmt = followStmt.get();
-                transformLoopBody(forNode->body, out, breakLoopType, ExpUsage::Assignment, expList);
-                if (!expList->followStmtProcessed) {
+                if (hasBreakWithValue(breakLoopType)) {
                         lenAssign.clear();
+                        transformLoopBody(forNode->body, out, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto expList = toAst<ExpList_t>(accum + '[' + len + ']', x);
+                        auto followStmt = toAst<Statement_t>(len + "+=1"s, forNode->body);
+                        expList->followStmt = followStmt.get();
+                        transformLoopBody(forNode->body, out, breakLoopType, ExpUsage::Assignment, expList);
+                        if (!expList->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
                 }
                 popScope();
                 out.push_back(indent() + "end"s + nlr(forNode));
@@ -8584,12 +8592,17 @@ private:
                 _buf << indent() << "local "sv << len << " = 1"sv << nll(forEach);
                 auto& lenAssign = out.emplace_back(clearBuf());
                 transformForEachHead(forEach->nameList, forEach->loopValue, out, true);
-                auto expList = toAst<ExpList_t>(accum + '[' + len + ']', x);
-                auto followStmt = toAst<Statement_t>(len + "+=1"s, forEach->body);
-                expList->followStmt = followStmt.get();
-                transformLoopBody(forEach->body, out, breakLoopType, ExpUsage::Assignment, expList);
-                if (!expList->followStmtProcessed) {
+                if (hasBreakWithValue(breakLoopType)) {
                         lenAssign.clear();
+                        transformLoopBody(forEach->body, out, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto expList = toAst<ExpList_t>(accum + '[' + len + ']', x);
+                        auto followStmt = toAst<Statement_t>(len + "+=1"s, forEach->body);
+                        expList->followStmt = followStmt.get();
+                        transformLoopBody(forEach->body, out, breakLoopType, ExpUsage::Assignment, expList);
+                        if (!expList->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
                 }
                 popScope();
                 out.push_back(indent() + "end"s + nlr(forEach));
@@ -9476,15 +9489,66 @@ private:
                         }
                 }
                 _withVars.push(withVar);
+                std::string breakWithVar;
+                if (assignList || returnValue) {
+                        auto breakLoopType = getBreakLoopType(with->body, withVar);
+                        if (hasBreakWithValue(breakLoopType)) {
+                                breakWithVar = withVar;
+                        }
+                }
                 if (with->eop) {
                         auto ifNode = x->new_ptr<If_t>();
                         ifNode->type.set(toAst<IfType_t>("if"sv, x));
                         ifNode->nodes.push_back(toAst<IfCond_t>(withVar + "~=nil"s, x));
                         ifNode->nodes.push_back(with->body);
-                        transformIf(ifNode, temp, ExpUsage::Common);
+                        if (breakWithVar.empty()) {
+                                transformIf(ifNode, temp, ExpUsage::Common);
+                        } else {
+                                auto simpleValue = x->new_ptr<SimpleValue_t>();
+                                simpleValue->value.set(ifNode);
+                                auto exp = newExp(simpleValue, x);
+                                auto expList = x->new_ptr<ExpList_t>();
+                                expList->exprs.push_back(exp);
+                                auto expListAssign = x->new_ptr<ExpListAssign_t>();
+                                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 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);
+                                auto sVal = x->new_ptr<SimpleValue_t>();
+                                sVal->value.set(whileNode);
+                                auto asmt = assignmentFrom(toAst<Exp_t>(breakWithVar, x), newExp(sVal, x), x);
+                                transformAssignment(asmt, temp);
+                        }
                 } else {
                         bool transformed = false;
-                        if (!extraScope && assignList) {
+                        if (!breakWithVar.empty()) {
+                                auto whileNode = toAst<While_t>("while true do break"s, x);
+                                auto block = x->new_ptr<Block_t>();
+                                if (auto blk = with->body.as<Block_t>()) {
+                                        block->statements.dup(blk->statements);
+                                } else {
+                                        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);
+                                auto sVal = x->new_ptr<SimpleValue_t>();
+                                sVal->value.set(whileNode);
+                                auto asmt = assignmentFrom(toAst<Exp_t>(breakWithVar, x), newExp(sVal, x), x);
+                                transformAssignment(asmt, temp);
+                                transformed = true;
+                        } else if (!extraScope && assignList) {
                                 if (auto block = with->body.as<Block_t>()) {
                                         if (!block->statements.empty()) {
                                                 Statement_t* stmt = static_cast<Statement_t*>(block->statements.back());
@@ -10543,12 +10607,17 @@ private:
                 auto condStr = transformCondExp(whileNode->condition, isUntil);
                 temp.push_back(indent() + "while "s + condStr + " do"s + nll(whileNode));
                 pushScope();
-                auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
-                auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, whileNode);
-                assignLeft->followStmt = followStmt.get();
-                transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
-                if (!assignLeft->followStmtProcessed) {
+                if (hasBreakWithValue(breakLoopType)) {
                         lenAssign.clear();
+                        transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
+                        auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, whileNode);
+                        assignLeft->followStmt = followStmt.get();
+                        transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
+                        if (!assignLeft->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
                 }
                 popScope();
                 temp.push_back(indent() + "end"s + nlr(whileNode));
@@ -10593,12 +10662,17 @@ private:
                 auto condStr = transformCondExp(whileNode->condition, isUntil);
                 temp.push_back(indent() + "while "s + condStr + " do"s + nll(whileNode));
                 pushScope();
-                auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
-                auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, whileNode);
-                assignLeft->followStmt = followStmt.get();
-                transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
-                if (!assignLeft->followStmtProcessed) {
+                if (hasBreakWithValue(breakLoopType)) {
                         lenAssign.clear();
+                        transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Common);
+                } else {
+                        auto assignLeft = toAst<ExpList_t>(accumVar + '[' + lenVar + ']', x);
+                        auto followStmt = toAst<Statement_t>(lenVar + "+=1"s, whileNode);
+                        assignLeft->followStmt = followStmt.get();
+                        transformLoopBody(whileNode->body, temp, breakLoopType, ExpUsage::Assignment, assignLeft);
+                        if (!assignLeft->followStmtProcessed) {
+                                lenAssign.clear();
+                        }
                 }
                 popScope();
                 temp.push_back(indent() + "end"s + nlr(whileNode));