add default return declaration for function literal.v0.22.2

5 files changed, 71 insertions, 10 deletions

diff --git a/src/yuescript/yue_ast.cpp b/src/yuescript/yue_ast.cpp
index faa175d..fbf7a63 100644
--- a/src/yuescript/yue_ast.cpp
+++ b/src/yuescript/yue_ast.cpp
@@ -1279,6 +1279,16 @@ std::string FunLit_t::to_string(void* ud) const {
         if (argsDef) {
                 line = argsDef->to_string(ud);
         }
+        if (defaultReturn) {
+                if (defaultReturn.is<DefaultValue_t>()) {
+                        line += ':';
+                } else {
+                        line += ": "s + defaultReturn->to_string(ud);
+                }
+        }
+        if (!line.empty()) {
+                line += ' ';
+        }
         line += arrow->to_string(ud);
         if (body) {
                 if (body->content.is<Statement_t>()) {
@@ -1302,7 +1312,7 @@ std::string MacroLit_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
         std::string line;
         if (argsDef) {
-                line = '(' + argsDef->to_string(ud) + ')';
+                line = '(' + argsDef->to_string(ud) + ") "s;
         }
         line += "->"s;
         if (body->content.is<Statement_t>()) {
@@ -1323,7 +1333,7 @@ std::string Macro_t::to_string(void* ud) const {
 }
 std::string MacroInPlace_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
-        auto line = "$->"s;
+        auto line = "$ ->"s;
         if (body->content.is<Statement_t>()) {
                 line += ' ' + body->to_string(ud);
         } else {
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index 90dce7a..73ffc9a 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -756,9 +756,10 @@ AST_END(FnArrow, "fn_arrow"sv)
 
 AST_NODE(FunLit)
         ast_ptr<false, FnArgsDef_t> argsDef;
+        ast_sel<false, ExpListLow_t, DefaultValue_t> defaultReturn;
         ast_ptr<true, FnArrow_t> arrow;
         ast_ptr<false, Body_t> body;
-        AST_MEMBER(FunLit, &argsDef, &arrow, &body)
+        AST_MEMBER(FunLit, &argsDef, &defaultReturn, &arrow, &body)
 AST_END(FunLit, "fun_lit"sv)
 
 AST_NODE(MacroName)
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 0d38fa5..f23d5f0 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -75,7 +75,7 @@ static std::unordered_set<std::string> Metamethods = {
         "close"s // Lua 5.4
 };
 
-const std::string_view version = "0.22.1"sv;
+const std::string_view version = "0.22.2"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -3787,11 +3787,44 @@ private:
                 str_list temp;
                 if (auto argsDef = funLit->argsDef.get()) {
                         transformFnArgsDef(argsDef, temp);
+                }
+                if (funLit->defaultReturn) {
+                        auto newBlock = funLit->new_ptr<Block_t>();
+                        if (funLit->body) {
+                                auto last = lastStatementFrom(funLit->body);
+                                if (!last->appendix && last->content.is<Return_t>() && !funLit->defaultReturn.is<DefaultValue_t>()) {
+                                        throw CompileError("duplicated return statement", last->content);
+                                }
+                                auto content = funLit->body->content.get();
+                                switch (content->get_id()) {
+                                        case id<Block_t>(): {
+                                                auto block = static_cast<Block_t*>(content);
+                                                newBlock->statements.dup(block->statements);
+                                                break;
+                                        }
+                                        case id<Statement_t>(): {
+                                                newBlock->statements.push_back(content);
+                                                break;
+                                        }
+                                        default: YUEE("AST node mismatch", content); break;
+                                }
+                        }
+                        if (funLit->defaultReturn.is<ExpListLow_t>()) {
+                                auto returnNode = newBlock->new_ptr<Return_t>();
+                                returnNode->valueList.set(funLit->defaultReturn);
+                                auto stmt = newBlock->new_ptr<Statement_t>();
+                                stmt->content.set(returnNode);
+                                newBlock->statements.push_back(stmt);
+                        }
+                        transformBlock(newBlock, temp, ExpUsage::Common);
+                } else {
                         if (funLit->body) {
                                 transformBody(funLit->body, temp, ExpUsage::Return);
                         } else {
                                 temp.push_back(Empty);
                         }
+                }
+                if (auto argsDef = funLit->argsDef.get()) {
                         auto it = temp.begin();
                         auto& args = *it;
                         auto& initArgs = *(++it);
@@ -3811,11 +3844,6 @@ private:
                                 _buf << " end"sv;
                         }
                 } else {
-                        if (funLit->body) {
-                                transformBody(funLit->body, temp, ExpUsage::Return);
-                        } else {
-                                temp.push_back(Empty);
-                        }
                         auto& bodyCodes = temp.back();
                         _buf << "function("sv << (isFatArrow ? "self"s : Empty) << ')';
                         if (!bodyCodes.empty()) {
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index d138979..3ffaf18 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -416,6 +416,18 @@ YueParser::YueParser() {
                 return true;
         });
 
+        disable_fun_lit = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->fnArrowAvailable = false;
+                return true;
+        });
+
+        enable_fun_lit = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                st->fnArrowAvailable = true;
+                return true;
+        });
+
         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);
@@ -846,7 +858,14 @@ YueParser::YueParser() {
 
         FnArgsDef = '(' >> *space_break >> -FnArgDefList >> -(white >> OuterVarShadow) >> white >> ')';
         FnArrow = expr("->") | "=>";
-        FunLit = -FnArgsDef >> space >> FnArrow >> -(space >> Body);
+        FunLit = pl::user(true_(), [](const item_t& item) {
+                State* st = reinterpret_cast<State*>(item.user_data);
+                return st->fnArrowAvailable;
+        }) >> -(FnArgsDef >>
+                -(':' >> space >>
+                        disable_fun_lit >> ensure(ExpListLow | DefaultValue, enable_fun_lit)
+                )
+        ) >> space >> FnArrow >> -(space >> Body);
 
         MacroName = '$' >> UnicodeName;
         macro_args_def = '(' >> white >> -FnArgDefList >> white >> ')';
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index 39b7cce..05aa9e6 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -105,6 +105,7 @@ protected:
                 bool exportMacro = false;
                 bool exportMetatable = false;
                 bool exportMetamethod = false;
+                bool fnArrowAvailable = true;
                 int exportCount = 0;
                 int expLevel = 0;
                 size_t stringOpen = 0;
@@ -198,6 +199,8 @@ private:
         NONE_AST_RULE(enable_arg_table_block);
         NONE_AST_RULE(disable_for);
         NONE_AST_RULE(enable_for);
+        NONE_AST_RULE(enable_fun_lit);
+        NONE_AST_RULE(disable_fun_lit);
         NONE_AST_RULE(switch_else);
         NONE_AST_RULE(switch_block);
         NONE_AST_RULE(if_else_if);