8 files changed, 135 insertions, 60 deletions

diff --git a/src/yuescript/ast.cpp b/src/yuescript/ast.cpp
index 239aec6..6db2c21 100644
--- a/src/yuescript/ast.cpp
+++ b/src/yuescript/ast.cpp
@@ -138,4 +138,25 @@ ast_node* parse(input& i, rule& g, error_list& el, void* ud) {
         return st.front();
 }
 
+/** check if the start part of given input matches grammar.
+        The parse procedures of each rule parsed are executed
+        before this function returns, if parsing succeeds.
+        @param i input.
+        @param g root rule of grammar.
+        @param ud user data, passed to the parse procedures.
+        @return true on parsing success, false on failure.
+*/
+ast_node* start_with(input& i, rule& g, error_list& el, void* ud) {
+        ast_stack st;
+        if (!start_with(i, g, el, &st, ud)) {
+                for (auto node : st) {
+                        delete node;
+                }
+                st.clear();
+                return nullptr;
+        }
+        assert(st.size() == 1);
+        return st.front();
+}
+
 } // namespace parserlib
diff --git a/src/yuescript/ast.hpp b/src/yuescript/ast.hpp
index c6da312..f59b50f 100644
--- a/src/yuescript/ast.hpp
+++ b/src/yuescript/ast.hpp
@@ -597,4 +597,14 @@ private:
 */
 ast_node* parse(input& i, rule& g, error_list& el, void* ud);
 
+/** check if the start part of given input matches grammar.
+        The parse procedures of each rule parsed are executed
+        before this function returns, if parsing succeeds.
+        @param i input.
+        @param g root rule of grammar.
+        @param ud user data, passed to the parse procedures.
+        @return true on parsing success, false on failure.
+*/
+ast_node* start_with(input& i, rule& g, error_list& el, void* ud);
+
 } // namespace parserlib
diff --git a/src/yuescript/parser.cpp b/src/yuescript/parser.cpp
index 25b735e..04b7ffd 100644
--- a/src/yuescript/parser.cpp
+++ b/src/yuescript/parser.cpp
@@ -1567,10 +1567,11 @@ expr user(const expr& e, const user_handler& handler) {
         @param i input.
         @param g root rule of grammar.
         @param el list of errors.
-        @param d user data, passed to the parse procedures.
+        @param st ast object stack.
+        @param ud user data, passed to the parse procedures.
         @return true on parsing success, false on failure.
 */
-bool parse(input& i, rule& g, error_list& el, void* d, void* ud) {
+bool parse(input& i, rule& g, error_list& el, void* st, void* ud) {
         // prepare context
         _context con(i, ud);
 
@@ -1591,7 +1592,32 @@ bool parse(input& i, rule& g, error_list& el, void* d, void* ud) {
         }
 
         // success; execute the parse procedures
-        con.do_parse_procs(d);
+        con.do_parse_procs(st);
+        return true;
+}
+
+/** check the start part of given input.
+        The parse procedures of each rule parsed are executed
+        before this function returns, if parsing succeeds.
+        @param i input.
+        @param g root rule of grammar.
+        @param el list of errors.
+        @param st ast object stack.
+        @param ud user data, passed to the parse procedures.
+        @return true on parsing success, false on failure.
+*/
+bool start_with(input& i, rule& g, error_list& el, void* st, void* ud) {
+        // prepare context
+        _context con(i, ud);
+
+        // parse grammar
+        if (!con.parse_non_term(g)) {
+                el.push_back(_syntax_error(con));
+                return false;
+        }
+
+        // success; execute the parse procedures
+        con.do_parse_procs(st);
         return true;
 }
 
diff --git a/src/yuescript/parser.hpp b/src/yuescript/parser.hpp
index 83682ae..d2612cd 100644
--- a/src/yuescript/parser.hpp
+++ b/src/yuescript/parser.hpp
@@ -394,10 +394,23 @@ expr user(const expr& e, const user_handler& handler);
         @param i input.
         @param g root rule of grammar.
         @param el list of errors.
-        @param d user data, passed to the parse procedures.
+        @param st ast object stack.
+        @param ud user data, passed to the parse procedures.
         @return true on parsing success, false on failure.
 */
-bool parse(input& i, rule& g, error_list& el, void* d, void* ud);
+bool parse(input& i, rule& g, error_list& el, void* st, void* ud);
+
+/** check if the start part of given input matches grammar.
+        The parse procedures of each rule parsed are executed
+        before this function returns, if parsing succeeds.
+        @param i input.
+        @param g root rule of grammar.
+        @param el list of errors.
+        @param st ast object stack.
+        @param ud user data, passed to the parse procedures.
+        @return true on parsing success, false on failure.
+*/
+bool start_with(input& i, rule& g, error_list& el, void* st, void* ud);
 
 /** output the specific input range to the specific stream.
         @param stream stream.
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index acb7221..e95e35c 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -875,8 +875,7 @@ AST_NODE(Statement)
                 Backcall_t, LocalAttrib_t, PipeBody_t, ExpListAssign_t, ChainAssign_t
         > content;
         ast_ptr<false, StatementAppendix_t> appendix;
-        ast_ptr<false, StatementSep_t> needSep;
-        AST_MEMBER(Statement, &sep, &comments, &content, &appendix, &needSep)
+        AST_MEMBER(Statement, &sep, &comments, &content, &appendix)
 AST_END(Statement, "statement"sv)
 
 AST_NODE(Body)
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 98d19a7..8574e6c 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -72,7 +72,7 @@ static std::unordered_set<std::string> Metamethods = {
         "close"s // Lua 5.4
 };
 
-const std::string_view version = "0.17.0"sv;
+const std::string_view version = "0.17.1"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -1497,25 +1497,6 @@ private:
                         case id<ChainAssign_t>(): transformChainAssign(static_cast<ChainAssign_t*>(content), out); break;
                         default: YUEE("AST node mismatch", content); break;
                 }
-                if (statement->needSep && !out.empty() && !out.back().empty()) {
-                        auto index = std::string::npos;
-                        if (_config.reserveLineNumber) {
-                                index = out.back().rfind(" -- "sv);
-                        } else {
-                                index = out.back().find_last_not_of('\n');
-                                if (index != std::string::npos) index++;
-                        }
-                        if (index != std::string::npos) {
-                                auto ending = out.back().substr(0, index);
-                                auto ind = ending.find_last_of(" \t\n"sv);
-                                if (ind != std::string::npos) {
-                                        ending = ending.substr(ind + 1);
-                                }
-                                if (LuaKeywords.find(ending) == LuaKeywords.end()) {
-                                        out.back().insert(index, ";"sv);
-                                }
-                        }
-                }
         }
 
         str_list getAssignVars(ExpListAssign_t* assignment) {
@@ -3543,8 +3524,6 @@ private:
                                 }
                                 lst->appendix.set(stmt->appendix);
                                 stmt->appendix.set(nullptr);
-                                lst->needSep.set(stmt->needSep);
-                                stmt->needSep.set(nullptr);
                                 auto exp = lastExpFromStatement(lst);
                                 BREAK_IF(!exp);
                                 for (auto val : pipeBody->values.objects()) {
@@ -3753,22 +3732,6 @@ private:
                                 returnNode->valueList.set(expListLow);
                                 returnNode->allowBlockMacroReturn = true;
                                 last->content.set(returnNode);
-                                last->needSep.set(nullptr);
-                                auto bLast = ++nodes.rbegin();
-                                if (bLast != nodes.rend()) {
-                                        bool isMacro = false;
-                                        BLOCK_START
-                                        BREAK_IF(expListLow->exprs.size() != 1);
-                                        auto exp = static_cast<Exp_t*>(expListLow->exprs.back());
-                                        BREAK_IF(!exp->opValues.empty());
-                                        auto chainValue = exp->getByPath<UnaryExp_t, Value_t, ChainValue_t>();
-                                        BREAK_IF(!chainValue);
-                                        isMacro = isMacroChain(chainValue);
-                                        BLOCK_END
-                                        if (!isMacro) {
-                                                ast_to<Statement_t>(*bLast)->needSep.set(nullptr);
-                                        }
-                                }
                                 BLOCK_END
                                 break;
                         }
@@ -3802,11 +3765,6 @@ private:
                                         }
                                         newAssignment->action.set(assign);
                                         last->content.set(newAssignment);
-                                        last->needSep.set(nullptr);
-                                        auto bLast = ++nodes.rbegin();
-                                        if (bLast != nodes.rend()) {
-                                                static_cast<Statement_t*>(*bLast)->needSep.set(nullptr);
-                                        }
                                 } else if (!last->content.is<BreakLoop_t>()) {
                                         throw CompileError("expecting assignable statement or break loop"sv, last);
                                 }
@@ -3818,6 +3776,28 @@ private:
                         str_list temp;
                         for (auto node : nodes) {
                                 transformStatement(static_cast<Statement_t*>(node), temp);
+                                if (_parser.startWith<StatementSep_t>(temp.back())) {
+                                        auto rit = ++temp.rbegin();
+                                        if (rit != temp.rend() && !rit->empty()) {
+                                                auto index = std::string::npos;
+                                                if (_config.reserveLineNumber) {
+                                                        index = rit->rfind(" -- "sv);
+                                                } else {
+                                                        index = rit->find_last_not_of('\n');
+                                                        if (index != std::string::npos) index++;
+                                                }
+                                                if (index != std::string::npos) {
+                                                        auto ending = rit->substr(0, index);
+                                                        auto ind = ending.find_last_of(" \t\n"sv);
+                                                        if (ind != std::string::npos) {
+                                                                ending = ending.substr(ind + 1);
+                                                        }
+                                                        if (LuaKeywords.find(ending) == LuaKeywords.end()) {
+                                                                rit->insert(index, ";"sv);
+                                                        }
+                                                }
+                                        }
+                                }
                         }
                         out.push_back(join(temp));
                 } else {
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index f383275..f72407d 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -853,13 +853,6 @@ YueParser::YueParser() {
         ChainAssign = Seperator >> Exp >> +(space >> '=' >> space >> Exp >> space >> and_('=')) >> space >> Assign;
 
         StatementAppendix = (IfLine | WhileLine | CompInner) >> space;
-        StatementSep = and_(
-                *space_break >> check_indent_match >> space >> (
-                        set("($'\"") |
-                        "[[" |
-                        "[="
-                )
-        );
         Statement =
                 Seperator >>
                 -(
@@ -876,8 +869,9 @@ YueParser::YueParser() {
                         StatementAppendix >> empty_block_error
                 ) >>
                 space >>
-                -StatementAppendix >>
-                -StatementSep;
+                -StatementAppendix;
+
+        StatementSep = white >> (set("('\"") | "[[" | "[=");
 
         Body = in_block | Statement;
 
@@ -905,6 +899,32 @@ YueParser::YueParser() {
 }
 // clang-format on
 
+bool YueParser::startWith(std::string_view codes, rule& r) {
+        std::unique_ptr<input> converted;
+        if (codes.substr(0, 3) == "\xEF\xBB\xBF"sv) {
+                codes = codes.substr(3);
+        }
+        try {
+                if (!codes.empty()) {
+                        converted = std::make_unique<input>(_converter.from_bytes(&codes.front(), &codes.back() + 1));
+                } else {
+                        converted = std::make_unique<input>();
+                }
+        } catch (const std::range_error&) {
+                return false;
+        }
+        error_list errors;
+        try {
+                State state;
+                return ::yue::start_with(*converted, r, errors, &state);
+        } catch (const ParserError&) {
+                return false;
+        } catch (const std::logic_error&) {
+                return false;
+        }
+        return true;
+}
+
 ParseInfo YueParser::parse(std::string_view codes, rule& r) {
         ParseInfo res;
         if (codes.substr(0, 3) == "\xEF\xBB\xBF"sv) {
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index d4f66eb..d7fbf90 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -84,11 +84,17 @@ public:
                 return parse(codes, rEnd).node;
         }
 
+        template <class AST>
+        bool startWith(std::string_view codes) {
+                return startWith(codes, getRule<AST>());
+        }
+
         std::string toString(ast_node* node);
         std::string toString(input::iterator begin, input::iterator end);
 
 protected:
         ParseInfo parse(std::string_view codes, rule& r);
+        bool startWith(std::string_view codes, rule& r);
 
         struct State {
                 State() {
@@ -387,8 +393,8 @@ private:
         AST_RULE(WhileLine)
         AST_RULE(BreakLoop)
         AST_RULE(StatementAppendix)
-        AST_RULE(StatementSep)
         AST_RULE(Statement)
+        AST_RULE(StatementSep);
         AST_RULE(YueLineComment)
         AST_RULE(MultilineCommentInner)
         AST_RULE(YueMultilineComment)