Fixing issue #178.

9 files changed, 94 insertions, 36 deletions

diff --git a/doc/docs/doc/README.md b/doc/docs/doc/README.md
index 56d3635..e1ba56e 100755
--- a/doc/docs/doc/README.md
+++ b/doc/docs/doc/README.md
@@ -278,12 +278,11 @@ macro yueFunc = (var) -> "local #{var} = ->"
 $yueFunc funcA
 funcA = -> "fail to assign to the Yue defined variable"
 
--- take care and let YueScript know the
+-- YueScript knows the
 -- local variables you declared in Lua code
 macro luaFunc = (var) -> {
   code: "local function #{var}() end"
   type: "lua"
-  locals: {var}
 }
 $luaFunc funcB
 funcB = -> "assign to the Lua defined variable"
@@ -307,12 +306,11 @@ macro yueFunc = (var) -> "local #{var} = ->"
 $yueFunc funcA
 funcA = -> "fail to assign to the Yue defined variable"
 
--- take care and let YueScript know the
+-- YueScript knows the
 -- local variables you declared in Lua codes
 macro luaFunc = (var) -> {
   code: "local function #{var}() end"
   type: "lua"
-  locals: {var}
 }
 $luaFunc funcB
 funcB = -> "assign to the Lua defined variable"
diff --git a/doc/docs/zh/doc/README.md b/doc/docs/zh/doc/README.md
index 0e5b532..5e28122 100755
--- a/doc/docs/zh/doc/README.md
+++ b/doc/docs/zh/doc/README.md
@@ -277,11 +277,10 @@ macro yueFunc = (var) -> "local #{var} = ->"
 $yueFunc funcA
 funcA = -> "无法访问宏生成月之脚本里定义的变量"
 
--- 让月之脚本知道你在Lua代码中声明的局部变量
+-- 月之脚本会知道你在 Lua 代码中声明的局部变量
 macro luaFunc = (var) -> {
   code: "local function #{var}() end"
   type: "lua"
-  locals: {var}
 }
 $luaFunc funcB
 funcB = -> "访问宏生成Lua代码里定义的变量"
@@ -305,11 +304,10 @@ macro yueFunc = (var) -> "local #{var} = ->"
 $yueFunc funcA
 funcA = -> "无法访问宏生成月之脚本里定义的变量"
 
--- 让月之脚本知道你在Lua代码中声明的局部变量
+-- 月之脚本会知道你在 Lua 代码中声明的局部变量
 macro luaFunc = (var) -> {
   code: "local function #{var}() end"
   type: "lua"
-  locals: {var}
 }
 $luaFunc funcB
 funcB = -> "访问宏生成Lua代码里定义的变量"
diff --git a/spec/inputs/macro.yue b/spec/inputs/macro.yue
index 5dafd84..5d5f1a9 100644
--- a/spec/inputs/macro.yue
+++ b/spec/inputs/macro.yue
@@ -340,6 +340,15 @@ _1 = ->
         print 2
         print 3
 
+do
+        macro foo = -> code: "tb:func(123)", type: "lua"
+        f = ->
+                x = $foo\bar 456
+                $foo!
+        f1 = ->
+                $foo!
+                return
+
 macro implicitReturnMacroIsAllowed = -> "print 'abc'\n123"
 
 $implicitReturnMacroIsAllowed
diff --git a/spec/outputs/macro.lua b/spec/outputs/macro.lua
index 83c6a90..7163077 100644
--- a/spec/outputs/macro.lua
+++ b/spec/outputs/macro.lua
@@ -330,6 +330,17 @@ _1 = function()
         return _accum_0
 end
 do
+        f = function()
+                x = (tb:func(123)):bar(456)
+                return (tb:func(123))
+        end
+        local f1
+        f1 = function()
+tb:func(123)
+                return
+        end
+end
+do
         print('abc')
         return 123
 end
diff --git a/src/yuescript/yue_ast.cpp b/src/yuescript/yue_ast.cpp
index 3dac699..fe6e726 100644
--- a/src/yuescript/yue_ast.cpp
+++ b/src/yuescript/yue_ast.cpp
@@ -167,6 +167,9 @@ std::string ExistentialOp_t::to_string(void*) const {
 std::string TableAppendingOp_t::to_string(void*) const {
         return "[]"s;
 }
+std::string PlainItem_t::to_string(void *) const {
+        return {};
+}
 std::string GlobalOp_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
         return info->convert(this);
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index 139f4b3..e670126 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -635,9 +635,13 @@ AST_END(ExistentialOp)
 AST_LEAF(TableAppendingOp)
 AST_END(TableAppendingOp)
 
+AST_LEAF(PlainItem)
+AST_END(PlainItem)
+
 AST_NODE(ChainValue)
         ast_ptr<true, Seperator_t> sep;
-        ast_sel_list<true, Callable_t, Invoke_t, DotChainItem_t, ColonChainItem_t, Slice_t, Exp_t, String_t, InvokeArgs_t, ExistentialOp_t, TableAppendingOp_t> items;
+        ast_sel_list<true, Callable_t, Invoke_t, DotChainItem_t, ColonChainItem_t, Slice_t, Exp_t, String_t, InvokeArgs_t, ExistentialOp_t, TableAppendingOp_t,
+                /*non-syntax-rule*/ PlainItem_t> items;
         AST_MEMBER(ChainValue, &sep, &items)
 AST_END(ChainValue)
 
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 369f7e8..446414d 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -24,6 +24,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI
 
 extern "C" {
 #include "lauxlib.h"
+#include "lparser.h"
 #include "lua.h"
 #include "lualib.h"
 } // extern "C"
@@ -6272,6 +6273,9 @@ private:
                                 case id<TableAppendingOp_t>():
                                         transform_table_appending_op(static_cast<TableAppendingOp_t*>(item), temp);
                                         break;
+                                case id<PlainItem_t>():
+                                        temp.push_back(_parser.toString(item));
+                                        break;
                                 default: YUEE("AST node mismatch", item); break;
                         }
                 }
@@ -6519,24 +6523,26 @@ private:
                                 throw CompileError("macro table must contain field \"type\" of value \"lua\" or \"text\""sv, x);
                         }
                         lua_pop(L, 1); // cur tab
-                        lua_getfield(L, -1, "locals"); // cur tab locals
-                        if (lua_istable(L, -1) != 0) {
-                                for (int i = 0; i < static_cast<int>(lua_objlen(L, -1)); i++) {
-                                        lua_rawgeti(L, -1, i + 1); // cur tab locals item
-                                        size_t len = 0;
-                                        if (lua_isstring(L, -1) == 0) {
-                                                throw CompileError("macro table field \"locals\" must be a table of strings"sv, x);
-                                        }
-                                        auto name = lua_tolstring(L, -1, &len);
-                                        if (auto varNode = toAst<Variable_t>({name, len}, x)) {
-                                                localVars.push_back(variableToString(varNode));
-                                        } else {
-                                                throw CompileError("macro table field \"locals\" must contain names for local variables, got \""s + std::string(name, len) + '"', x);
+                        if (type == "text"sv) {
+                                lua_getfield(L, -1, "locals"); // cur tab locals
+                                if (lua_istable(L, -1) != 0) {
+                                        for (int i = 0; i < static_cast<int>(lua_objlen(L, -1)); i++) {
+                                                lua_rawgeti(L, -1, i + 1); // cur tab locals item
+                                                size_t len = 0;
+                                                if (lua_isstring(L, -1) == 0) {
+                                                        throw CompileError("macro table field \"locals\" must be a table of strings"sv, x);
+                                                }
+                                                auto name = lua_tolstring(L, -1, &len);
+                                                if (auto varNode = toAst<Variable_t>({name, len}, x)) {
+                                                        localVars.push_back(variableToString(varNode));
+                                                } else {
+                                                        throw CompileError("macro table field \"locals\" must contain names for local variables, got \""s + std::string(name, len) + '"', x);
+                                                }
+                                                lua_pop(L, 1); // cur tab locals
                                         }
-                                        lua_pop(L, 1); // cur tab locals
                                 }
+                                lua_pop(L, 1); // cur tab
                         }
-                        lua_pop(L, 1); // cur tab
                 } else { // cur code
                         codes = lua_tostring(L, -1);
                 }
@@ -6554,23 +6560,49 @@ private:
                 bool isBlock = (usage == ExpUsage::Common) && (chainList.size() < 2 || (chainList.size() == 2 && ast_is<Invoke_t, InvokeArgs_t>(chainList.back())));
                 ParseInfo info;
                 if (type == "lua"sv) {
-                        if (!isBlock) {
-                                throw CompileError("lua macro can only be placed where block macro is allowed"sv, x);
-                        }
                         auto macroChunk = "=(macro "s + _parser.toString(x->name) + ')';
                         int top = lua_gettop(L);
                         DEFER(lua_settop(L, top));
-                        if (luaL_loadbuffer(L, codes.c_str(), codes.size(), macroChunk.c_str()) != 0) {
-                                std::string err = lua_tostring(L, -1);
-                                throw CompileError(err, x);
-                        }
-                        if (!codes.empty()) {
-                                if (_config.reserveLineNumber) {
-                                        codes.insert(0, nll(chainValue).substr(1));
+                        if (isBlock) {
+                                if (luaL_loadbuffer(L, codes.c_str(), codes.size(), macroChunk.c_str()) != 0) {
+                                        std::string err = lua_tostring(L, -1);
+                                        throw CompileError("lua macro is not expanding to valid block\n"s + err, x);
+                                } else {
+                                        Proto* f = ((LClosure*)lua_topointer(L, -1))->p;
+                                        for (int i = 0; i < f->sizelocvars; i++) {
+                                                localVars.push_back(getstr(f->locvars[i].varname));
+                                        }
+                                }
+                                if (!codes.empty()) {
+                                        if (_config.reserveLineNumber) {
+                                                codes.insert(0, nll(chainValue).substr(1));
+                                        }
+                                        codes.append(nlr(chainValue));
                                 }
-                                codes.append(nlr(chainValue));
+                                return {nullptr, nullptr, std::move(codes), std::move(localVars)};
+                        } else {
+                                auto expCode = "return ("s + codes + ')';
+                                if (luaL_loadbuffer(L, expCode.c_str(), expCode.size(), macroChunk.c_str()) != 0) {
+                                        std::string err = lua_tostring(L, -1);
+                                        throw CompileError("lua macro is not expanding to valid expression\n"s + err, x);
+                                }
+                                Utils::trim(codes);
+                                codes = '(' + codes + ')';
+                                auto plainItem = toAst<PlainItem_t>(codes, x);
+                                auto newChain = x->new_ptr<ChainValue_t>();
+                                newChain->items.push_back(plainItem);
+                                ast_ptr<false, ast_node> exp;
+                                exp.set(newExp(newChain, x));
+                                if (chainList.size() > 2 || (chainList.size() == 2 && !ast_is<Invoke_t, InvokeArgs_t>(chainList.back()))) {
+                                        auto it = chainList.begin();
+                                        it++;
+                                        if (chainList.size() > 1 && ast_is<Invoke_t, InvokeArgs_t>(*it)) it++;
+                                        for (; it != chainList.end(); ++it) {
+                                                newChain->items.push_back(*it);
+                                        }
+                                }
+                                return {exp, nullptr, Empty, std::move(localVars)};
                         }
-                        return {nullptr, nullptr, std::move(codes), std::move(localVars)};
                 } else if (type == "text"sv) {
                         if (!isBlock) {
                                 throw CompileError("text macro can only be placed where block macro is allowed"sv, x);
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index ef242ea..5993ed5 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -653,6 +653,8 @@ YueParser::YueParser() {
 
         ExistentialOp = '?' >> not_('?');
         TableAppendingOp = and_('[') >> ("[]" | brackets_expression_error);
+        PlainItem = +any_char;
+
         chain_call = (
                 Callable >> -ExistentialOp >> -chain_items
         ) | (
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index f18cc1b..02292e1 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -381,6 +381,7 @@ private:
         AST_RULE(Invoke);
         AST_RULE(ExistentialOp);
         AST_RULE(TableAppendingOp);
+        AST_RULE(PlainItem);
         AST_RULE(SpreadExp);
         AST_RULE(SpreadListExp);
         AST_RULE(TableLit);