7 files changed, 176 insertions, 10 deletions

diff --git a/spec/inputs/funcs.yue b/spec/inputs/funcs.yue
index 0e45ff1..e647edc 100644
--- a/spec/inputs/funcs.yue
+++ b/spec/inputs/funcs.yue
@@ -158,4 +158,39 @@ args = (f,g,m
         return 1 if n == 0
         n * @(n-1)
 
+do
+        items.every (item) ->
+                if item.field
+                        value = item.field.get "abc"
+                        if value
+                                switch value\get!
+                                        when 123
+                                                return false
+                                        when 456
+                                                handle item
+                true
+
+        items.every (item): true ->
+                if item.field
+                        value = item.field.get "abc"
+                        if value
+                                switch value\get!
+                                        when 123
+                                                return false
+                                        when 456
+                                                -- prevent implicit return for next line
+                                                handle item
+
+        HttpServer\post "/login", (req): success: false ->
+                switch req when {:name, :pwd}
+                        if name ~= ""
+                                if user := DB\queryUser name, pwd
+                                        if user.status == "available"
+                                                return success: true
+
+        check = (num) -> return num
+        -- func without implicit return
+        func = (): -> check 123
+        print func! -- get nil
+
 nil
diff --git a/spec/outputs/funcs.lua b/spec/outputs/funcs.lua
index 03281e4..c07989e 100644
--- a/spec/outputs/funcs.lua
+++ b/spec/outputs/funcs.lua
@@ -219,4 +219,74 @@ self = function(n)
         end
         return n * self(n - 1)
 end
+do
+        items.every(function(item)
+                if item.field then
+                        local value = item.field.get("abc")
+                        if value then
+                                do
+                                        local _exp_0 = value:get()
+                                        if 123 == _exp_0 then
+                                                return false
+                                        elseif 456 == _exp_0 then
+                                                handle(item)
+                                        end
+                                end
+                        end
+                end
+                return true
+        end)
+        items.every(function(item)
+                if item.field then
+                        local value = item.field.get("abc")
+                        if value then
+                                do
+                                        local _exp_0 = value:get()
+                                        if 123 == _exp_0 then
+                                                return false
+                                        elseif 456 == _exp_0 then
+                                                handle(item)
+                                        end
+                                end
+                        end
+                end
+                return true
+        end)
+        HttpServer:post("/login", function(req)
+                do
+                        local _type_0 = type(req)
+                        local _tab_0 = "table" == _type_0 or "userdata" == _type_0
+                        if _tab_0 then
+                                local name = req.name
+                                local pwd = req.pwd
+                                if name ~= nil and pwd ~= nil then
+                                        if name ~= "" then
+                                                do
+                                                        local user = DB:queryUser(name, pwd)
+                                                        if user then
+                                                                if user.status == "available" then
+                                                                        return {
+                                                                                success = true
+                                                                        }
+                                                                end
+                                                        end
+                                                end
+                                        end
+                                end
+                        end
+                end
+                return {
+                        success = false
+                }
+        end)
+        local check
+        check = function(num)
+                return num
+        end
+        local func
+        func = function()
+                check(123)
+        end
+        print(func())
+end
 return nil
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);