fix some ambiguous syntax.

5 files changed, 104 insertions, 62 deletions

diff --git a/spec/inputs/goto.mp b/spec/inputs/goto.mp
index 3410ca3..1197251 100644
--- a/spec/inputs/goto.mp
+++ b/spec/inputs/goto.mp
@@ -7,7 +7,7 @@ do
         ::done::
 
 do
-        for z = 1, 10 for y = 1, 10 for x = 1, 10
+        for z = 1, 10 do for y = 1, 10 do for x = 1, 10
                 if x^2 + y^2 == z^2
                         print 'found a Pythagorean triple:', x, y, z
                         goto done
@@ -15,7 +15,7 @@ do
 
 do
         for z = 1, 10
-                for y = 1, 10 for x = 1, 10
+                for y = 1, 10 do for x = 1, 10
                                 if x^2 + y^2 == z^2
                                         print 'found a Pythagorean triple:', x, y, z
                                         print 'now trying next z...'
@@ -24,7 +24,7 @@ do
 
 do
         ::redo::
-        for x = 1, 10 for y = 1, 10
+        for x = 1, 10 do for y = 1, 10
                 if not f x, y then goto continue
                 if not g x, y then goto skip
                 if not h x, y then goto redo
diff --git a/src/MoonP/moon_ast.h b/src/MoonP/moon_ast.h
index e2f54de..6b738d4 100644
--- a/src/MoonP/moon_ast.h
+++ b/src/MoonP/moon_ast.h
@@ -223,19 +223,20 @@ AST_END(Return)
 
 class existential_op_t;
 class Assign_t;
-class Body_t;
+class Block_t;
+class Statement_t;
 
 AST_NODE(With)
         ast_ptr<false, existential_op_t> eop;
         ast_ptr<true, ExpList_t> valueList;
         ast_ptr<false, Assign_t> assigns;
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         AST_MEMBER(With, &eop, &valueList, &assigns, &body)
 AST_END(With)
 
 AST_NODE(SwitchCase)
         ast_ptr<true, ExpList_t> valueList;
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         AST_MEMBER(SwitchCase, &valueList, &body)
 AST_END(SwitchCase)
 
@@ -243,7 +244,7 @@ AST_NODE(Switch)
         ast_ptr<true, Exp_t> target;
         ast_ptr<true, Seperator_t> sep;
         ast_list<true, SwitchCase_t> branches;
-        ast_ptr<false, Body_t> lastBranch;
+        ast_sel<false, Block_t, Statement_t> lastBranch;
         AST_MEMBER(Switch, &target, &sep, &branches, &lastBranch)
 AST_END(Switch)
 
@@ -255,22 +256,24 @@ AST_END(IfCond)
 
 AST_NODE(If)
         ast_ptr<true, Seperator_t> sep;
-        ast_sel_list<true, IfCond_t, Body_t> nodes;
+        ast_sel_list<true, IfCond_t, Block_t, Statement_t> nodes;
         AST_MEMBER(If, &sep, &nodes)
 AST_END(If)
 
 AST_NODE(Unless)
         ast_ptr<true, Seperator_t> sep;
-        ast_sel_list<true, IfCond_t, Body_t> nodes;
+        ast_sel_list<true, IfCond_t, Block_t, Statement_t> nodes;
         AST_MEMBER(Unless, &sep, &nodes)
 AST_END(Unless)
 
 AST_NODE(While)
         ast_ptr<true, Exp_t> condition;
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         AST_MEMBER(While, &condition, &body)
 AST_END(While)
 
+class Body_t;
+
 AST_NODE(Repeat)
         ast_ptr<true, Body_t> body;
         ast_ptr<true, Exp_t> condition;
@@ -287,7 +290,7 @@ AST_NODE(For)
         ast_ptr<true, Exp_t> startValue;
         ast_ptr<true, Exp_t> stopValue;
         ast_ptr<false, for_step_value_t> stepValue;
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         AST_MEMBER(For, &varName, &startValue, &stopValue, &stepValue, &body)
 AST_END(For)
 
@@ -297,7 +300,7 @@ class star_exp_t;
 AST_NODE(ForEach)
         ast_ptr<true, AssignableNameList_t> nameList;
         ast_sel<true, star_exp_t, ExpList_t> loopValue;
-        ast_ptr<true, Body_t> body;
+        ast_sel<true, Block_t, Statement_t> body;
         AST_MEMBER(ForEach, &nameList, &loopValue, &body)
 AST_END(ForEach)
 
diff --git a/src/MoonP/moon_compiler.cpp b/src/MoonP/moon_compiler.cpp
index 6b721f1..24e4a97 100644
--- a/src/MoonP/moon_compiler.cpp
+++ b/src/MoonP/moon_compiler.cpp
@@ -53,7 +53,7 @@ inline std::string s(std::string_view sv) {
         return std::string(sv);
 }
 
-const std::string_view version = "0.6.2"sv;
+const std::string_view version = "0.6.3"sv;
 const std::string_view extension = "mp"sv;
 
 class MoonCompilerImpl {
@@ -516,6 +516,18 @@ private:
                 return nullptr;
         }
 
+        Statement_t* lastStatementFrom(ast_node* body) const {
+                switch (body->getId()) {
+                        case id<Block_t>():
+                                return lastStatementFrom(static_cast<Block_t*>(body));
+                        case id<Statement_t>(): {
+                                return static_cast<Statement_t*>(body);
+                        }
+                        default: assert(false); break;
+                }
+                return nullptr;
+        }
+
         Statement_t* lastStatementFrom(const node_container& stmts) const {
                 if (!stmts.empty()) {
                         auto it = stmts.end(); --it;
@@ -803,9 +815,7 @@ private:
 
                                         auto stmt = x->new_ptr<Statement_t>();
                                         stmt->content.set(statement->content);
-                                        auto body = x->new_ptr<Body_t>();
-                                        body->content.set(stmt);
-                                        ifNode->nodes.push_back(body);
+                                        ifNode->nodes.push_back(stmt);
 
                                         statement->appendix.set(nullptr);
                                         auto simpleValue = x->new_ptr<SimpleValue_t>();
@@ -830,9 +840,7 @@ private:
 
                                         auto stmt = x->new_ptr<Statement_t>();
                                         stmt->content.set(statement->content);
-                                        auto body = x->new_ptr<Body_t>();
-                                        body->content.set(stmt);
-                                        unless->nodes.push_back(body);
+                                        unless->nodes.push_back(stmt);
 
                                         statement->appendix.set(nullptr);
                                         auto simpleValue = x->new_ptr<SimpleValue_t>();
@@ -1593,9 +1601,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);
-                                        ns.push_back(body.get());
+                                        ns.push_back(stmt.get());
                                 }
                         }
                 }
@@ -1614,15 +1620,16 @@ private:
                         pushScope();
                         _enableReturn.push(true);
                 }
-                std::list<std::pair<IfCond_t*, Body_t*>> ifCondPairs;
+                std::list<std::pair<IfCond_t*, ast_node*>> ifCondPairs;
                 ifCondPairs.emplace_back();
                 for (auto node : nodes) {
                         switch (node->getId()) {
                                 case id<IfCond_t>():
                                         ifCondPairs.back().first = static_cast<IfCond_t*>(node);
                                         break;
-                                case id<Body_t>():
-                                        ifCondPairs.back().second = static_cast<Body_t*>(node);
+                                case id<Block_t>():
+                                case id<Statement_t>():
+                                        ifCondPairs.back().second = node;
                                         ifCondPairs.emplace_back();
                                         break;
                                 default: assert(false); break;
@@ -1718,7 +1725,7 @@ private:
                                 if (pair == ifCondPairs.front() && extraAssignment) {
                                         transformAssignment(extraAssignment, temp);
                                 }
-                                transformBody(pair.second, temp, usage, assignList);
+                                transform_plain_body(pair.second, temp, usage, assignList);
                                 popScope();
                         }
                         if (!pair.first) {
@@ -3945,7 +3952,22 @@ private:
                 out.push_back(clearBuf());
         }
 
-        void transformLoopBody(Body_t* body, str_list& out, const std::string& appendContent, ExpUsage usage, ExpList_t* assignList = nullptr) {
+        void transform_plain_body(ast_node* body, str_list& out, ExpUsage usage, ExpList_t* assignList = nullptr) {
+                switch (body->getId()) {
+                        case id<Block_t>():
+                                transformBlock(static_cast<Block_t*>(body), out, usage, assignList);
+                                break;
+                        case id<Statement_t>(): {
+                                auto newBlock = body->new_ptr<Block_t>();
+                                newBlock->statements.push_back(body);
+                                transformBlock(newBlock, out, usage, assignList);
+                                break;
+                        }
+                        default: assert(false); break;
+                }
+        }
+
+        void transformLoopBody(ast_node* body, str_list& out, const std::string& appendContent, ExpUsage usage, ExpList_t* assignList = nullptr) {
                 str_list temp;
                 bool withContinue = traversal::Stop == body->traverse([&](ast_node* node) {
                         if (auto stmt = ast_cast<Statement_t>(node)) {
@@ -3968,7 +3990,7 @@ private:
                         _continueVars.push(continueVar);
                         pushScope();
                 }
-                transformBody(body, temp, usage, assignList);
+                transform_plain_body(body, temp, usage, assignList);
                 if (withContinue) {
                         if (!appendContent.empty()) {
                                 _buf << indent() << appendContent;
@@ -4739,7 +4761,7 @@ private:
                         ifNode->nodes.push_back(with->body);
                         transformIf(ifNode, temp, ExpUsage::Common);
                 } else {
-                        transformBody(with->body, temp, ExpUsage::Common);
+                        transform_plain_body(with->body, temp, ExpUsage::Common);
                 }
                 _withVars.pop();
                 if (assignList) {
@@ -5401,7 +5423,7 @@ private:
         void transformRepeat(Repeat_t* repeat, str_list& out) {
                 str_list temp;
                 pushScope();
-                transformLoopBody(repeat->body, temp, Empty, ExpUsage::Common);
+                transformLoopBody(repeat->body->content, temp, Empty, ExpUsage::Common);
                 transformExp(repeat->condition, temp, ExpUsage::Closure);
                 popScope();
                 _buf << indent() << "repeat"sv << nll(repeat);
@@ -5442,13 +5464,13 @@ private:
                         }
                         temp.back().append(s(" then"sv) + nll(branch));
                         pushScope();
-                        transformBody(branch->body, temp, usage, assignList);
+                        transform_plain_body(branch->body, temp, usage, assignList);
                         popScope();
                 }
                 if (switchNode->lastBranch) {
                         temp.push_back(indent() + s("else"sv) + nll(switchNode->lastBranch));
                         pushScope();
-                        transformBody(switchNode->lastBranch, temp, usage, assignList);
+                        transform_plain_body(switchNode->lastBranch, temp, usage, assignList);
                         popScope();
                 }
                 temp.push_back(indent() + s("end"sv) + nlr(switchNode));
diff --git a/src/MoonP/moon_parser.cpp b/src/MoonP/moon_parser.cpp
index aafe730..a581e15 100644
--- a/src/MoonP/moon_parser.cpp
+++ b/src/MoonP/moon_parser.cpp
@@ -70,8 +70,13 @@ MoonParser::MoonParser() {
 
         #define sym(str) (Space >> str)
         #define symx(str) expr(str)
-        #define ensure(patt, finally) (((patt) >> (finally)) | ((finally) >> (Cut)))
+        #define ensure(patt, finally) ((patt) >> (finally) | (finally) >> Cut)
         #define key(str) (Space >> str >> not_(AlphaNum))
+        #define disable_do(patt) (DisableDo >> ((patt) >> EnableDo | EnableDo >> Cut))
+        #define disable_chain(patt) (DisableChain >> ((patt) >> EnableChain | EnableChain >> Cut))
+        #define disable_do_chain(patt) (DisableDoChain >> ((patt) >> EnableDoChain | EnableDoChain >> Cut))
+        #define plain_body_with(str) (-key(str) >> InBlock | key(str) >> Statement)
+        #define plain_body (InBlock | Statement)
 
         Variable = pl::user(Name, [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
@@ -169,7 +174,7 @@ MoonParser::MoonParser() {
                 return true;
         });
 
-        InBlock = Advance >> ensure(Block, PopIndent);
+        InBlock = +SpaceBreak >> Advance >> ensure(Block, PopIndent);
 
         local_flag = expr('*') | expr('^');
         local_values = NameList >> -(sym('=') >> (TableBlock | ExpListLow));
@@ -219,11 +224,11 @@ MoonParser::MoonParser() {
 
         Return = key("return") >> -ExpListLow;
 
-        WithExp = DisableChainBlock >> ensure(ExpList >> -Assign, PopChainBlock);
+        WithExp = ExpList >> -Assign;
 
-        With = key("with") >> -existential_op >> DisableDo >> ensure(WithExp, PopDo) >> -key("do") >> Body;
-        SwitchCase = key("when") >> ExpList >> -key("then") >> Body;
-        SwitchElse = key("else") >> Body;
+        With = key("with") >> -existential_op >> disable_do_chain(WithExp) >> plain_body_with("do");
+        SwitchCase = key("when") >> disable_chain(ExpList) >> plain_body_with("then");
+        SwitchElse = key("else") >> plain_body;
 
         SwitchBlock = *EmptyLine >>
                 Advance >> Seperator >>
@@ -232,31 +237,27 @@ MoonParser::MoonParser() {
                 -(+SpaceBreak >> SwitchElse) >>
                 PopIndent;
 
-        Switch = key("switch") >>
-                DisableDo >> ensure(Exp, PopDo) >>
-                -key("do") >> -Space >> Break >> SwitchBlock;
+        Switch = key("switch") >> disable_do(Exp) >> -key("do")
+                >> -Space >> Break >> SwitchBlock;
 
-        IfCond = Exp >> -Assign;
-        IfElseIf = -(Break >> *EmptyLine >> CheckIndent) >> key("elseif") >> IfCond >> -key("then") >> Body;
-        IfElse = -(Break >> *EmptyLine >> CheckIndent) >> key("else") >> Body;
-        If = key("if") >> Seperator >> IfCond >> -key("then") >> Body >> *IfElseIf >> -IfElse;
-        Unless = key("unless") >> Seperator >> IfCond >> -key("then") >> Body >> *IfElseIf >> -IfElse;
-
-        While = key("while") >> DisableDo >> ensure(Exp, PopDo) >> -key("do") >> Body;
+        IfCond = disable_chain(Exp >> -Assign);
+        IfElseIf = -(Break >> *EmptyLine >> CheckIndent) >> key("elseif") >> IfCond >> plain_body_with("then");
+        IfElse = -(Break >> *EmptyLine >> CheckIndent) >> key("else") >> plain_body;
+        If = key("if") >> Seperator >> IfCond >> plain_body_with("then") >> *IfElseIf >> -IfElse;
+        Unless = key("unless") >> Seperator >> IfCond >> plain_body_with("then") >> *IfElseIf >> -IfElse;
+        
+        While = key("while") >> disable_do_chain(Exp) >> plain_body_with("do");
         Repeat = key("repeat") >> Body >> Break >> *EmptyLine >> CheckIndent >> key("until") >> Exp;
 
         for_step_value = sym(',') >> Exp;
         for_args = Space >> Variable >> sym('=') >> Exp >> sym(',') >> Exp >> -for_step_value;
 
-        For = key("for") >> DisableDo >>
-                ensure(for_args, PopDo) >>
-                -key("do") >> Body;
+        For = key("for") >> disable_do_chain(for_args) >> plain_body_with("do");
 
         for_in = star_exp | ExpList;
 
         ForEach = key("for") >> AssignableNameList >> key("in") >>
-                DisableDo >> ensure(for_in, PopDo) >>
-                -key("do") >> Body;
+                disable_do_chain(for_in) >> plain_body_with("do");
 
         Do = pl::user(key("do"), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
@@ -269,12 +270,26 @@ MoonParser::MoonParser() {
                 return true;
         });
 
-        PopDo = pl::user(true_(), [](const item_t& item) {
+        EnableDo = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
                 st->doStack.pop();
                 return true;
         });
 
+        DisableDoChain = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->doStack.push(false);
+                st->chainBlockStack.push(false);
+                return true;
+        });
+
+        EnableDoChain = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->doStack.pop();
+                st->chainBlockStack.pop();
+                return true;
+        });
+
         Comprehension = sym('[') >> Exp >> CompInner >> sym(']');
         comp_value = sym(',') >> Exp;
         TblComprehension = sym('{') >> Exp >> -comp_value >> CompInner >> sym('}');
@@ -338,13 +353,13 @@ MoonParser::MoonParser() {
         exp_op_value = Space >> BinaryOperator >> *SpaceBreak >> backcall_exp;
         Exp = Seperator >> backcall_exp >> *exp_op_value;
 
-        DisableChainBlock = pl::user(true_(), [](const item_t& item) {
+        DisableChain = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
                 st->chainBlockStack.push(false);
                 return true;
         });
 
-        PopChainBlock = pl::user(true_(), [](const item_t& item) {
+        EnableChain = pl::user(true_(), [](const item_t& item) {
                 State* st = reinterpret_cast<State*>(item.user_data);
                 st->chainBlockStack.pop();
                 return true;
@@ -356,7 +371,7 @@ MoonParser::MoonParser() {
                 return st->chainBlockStack.empty() || st->chainBlockStack.top();
         }) >> +SpaceBreak >> Advance >> ensure(
                 chain_line >> *(+SpaceBreak >> chain_line), PopIndent);
-        ChainValue = Seperator >> (Chain | Callable) >> -existential_op >> (chain_block | -InvokeArgs);
+        ChainValue = Seperator >> (Chain | Callable) >> -existential_op >> -(InvokeArgs | chain_block);
 
         simple_table = Seperator >> KeyValue >> *(sym(',') >> KeyValue);
         Value = SimpleValue | simple_table | ChainValue | String;
@@ -586,10 +601,10 @@ MoonParser::MoonParser() {
         ) >> Space >>
         -statement_appendix >> -statement_sep;
 
-        Body = Space >> Break >> *EmptyLine >> InBlock | Statement;
+        Body = InBlock | Statement;
 
         empty_line_stop = Space >> and_(Stop);
-        Line = CheckIndent >> Statement | and_(Space >> BackcallOperator) >> Advance >> ensure(Statement, PopIndent) | empty_line_stop;
+        Line = and_(check_indent >> Space >> not_(BackcallOperator)) >> Statement | Advance >> ensure(and_(Space >> BackcallOperator) >> Statement, PopIndent) | empty_line_stop;
         Block = Seperator >> Line >> *(+Break >> Line);
 
         Shebang = expr("#!") >> *(not_(Stop) >> Any);
diff --git a/src/MoonP/moon_parser.h b/src/MoonP/moon_parser.h
index 66f2cd9..b165070 100644
--- a/src/MoonP/moon_parser.h
+++ b/src/MoonP/moon_parser.h
@@ -131,7 +131,11 @@ private:
         rule import_tab_lines;
         rule WithExp;
         rule DisableDo;
-        rule PopDo;
+        rule EnableDo;
+        rule DisableChain;
+        rule EnableChain;
+        rule DisableDoChain;
+        rule EnableDoChain;
         rule SwitchElse;
         rule SwitchBlock;
         rule IfElseIf;
@@ -158,8 +162,6 @@ private:
         rule ChainItem;
         rule chain_line;
         rule chain_block;
-        rule DisableChainBlock;
-        rule PopChainBlock;
         rule Index;
         rule invoke_chain;
         rule TableValue;