Add annotation statements and expand annotation tests

14 files changed, 584 insertions, 27 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 42d84d3..ea6a0f3 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -28,6 +28,11 @@ if (LUA_EXEC_NAME MATCHES "luajit")
                 NAMES luajit libluajit
                 PATHS ${LUA_LIBDIR}
                 NO_DEFAULT_PATH)
+elseif (LUA_VERSION_STRING MATCHES "Lua 5.5")
+        find_library(LUA_LIBRARIES
+                NAMES lua55 lua5.5 liblua55 liblua5.5 lua liblua
+                PATHS ${LUA_LIBDIR}
+                NO_DEFAULT_PATH)
 elseif (LUA_VERSION_STRING MATCHES "Lua 5.4")
         find_library(LUA_LIBRARIES
                 NAMES lua54 lua5.4 liblua54 liblua5.4 lua liblua
diff --git a/spec/inputs/annotation.yue b/spec/inputs/annotation.yue
new file mode 100644
index 0000000..2124ad3
--- /dev/null
+++ b/spec/inputs/annotation.yue
@@ -0,0 +1,27 @@
+macro ClsDef = (code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        lines = table.concat [item\gsub "%-%-%-", "---@" for item in code\gmatch "(%-%-%-.-)\n"], "\n"
+        return
+                type: "text"
+                before: false
+                code: |
+                        ---@class #{className}
+                        #{lines}
+                        ---@class #{className}Class
+                        ---@operator call:#{className}
+                        ---@cast #{className} #{className}Class
+
+$[ClsDef]
+class A
+        ---field x number
+        ---field y number
+        new: (@x = 0, @y = 0) =>
+        ---field setAdd fun(self: A, x: number, y: number): number Set fields and add number values.
+        setAdd: (@x, @y) => @x + @y
+
+a = A!
+res = a::setAdd 1, 2
+print(a.x, a.y, a.y, res)
+
+return
+
diff --git a/spec/inputs/annotation_before.yue b/spec/inputs/annotation_before.yue
new file mode 100644
index 0000000..9b11aad
--- /dev/null
+++ b/spec/inputs/annotation_before.yue
@@ -0,0 +1,16 @@
+macro Tag = (tag, code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: tag == "before"
+                code: "-- #{tag}:#{className}"
+
+$[Tag("before")]
+class B
+        getTag: => "before"
+
+$[Tag("after")]
+class C
+        getTag: => "after"
+
+return B!\getTag!, C!\getTag!
diff --git a/spec/inputs/test/annotation_spec.yue b/spec/inputs/test/annotation_spec.yue
new file mode 100644
index 0000000..3cd1d5a
--- /dev/null
+++ b/spec/inputs/test/annotation_spec.yue
@@ -0,0 +1,160 @@
+import to_lua from require "yue"
+
+compile_and_run = (code, config = {}) ->
+        lua_code, err = to_lua code, config
+        assert.is_nil err
+        assert.is_not_nil lua_code
+        chunk, load_err = load lua_code
+        assert.is_nil load_err
+        assert.is_not_nil chunk
+        chunk!
+
+describe "annotation", ->
+        it "should append generated text after annotated class by default", ->
+                code = [[
+macro ClsDef = (code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "-- after:" .. className
+
+$[ClsDef]
+class A
+        getName: => "A"
+
+return
+]]
+                result, err = to_lua code
+                assert.is_nil err
+                assert.is_not_nil result
+                assert.is_true result\find("__name = \"A\"") != nil
+                assert.is_true result\find("%-%- after:A") != nil
+                assert.is_true result\find("__name = \"A\"") < result\find("%-%- after:A")
+
+        it "should place generated text before the annotated statement when before is true", ->
+                code = [[
+macro Before = (code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: true
+                code: "-- before:" .. className
+
+$[Before]
+class B
+        getName: => "B"
+
+return
+]]
+                result, err = to_lua code
+                assert.is_nil err
+                assert.is_not_nil result
+                assert.is_true result\find("%-%- before:B") != nil
+                assert.is_true result\find("local B") != nil
+                assert.is_true result\find("%-%- before:B") < result\find("local B")
+
+        it "should support annotation invocation arguments", ->
+                code = [[
+macro Tag = (tag, code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "-- " .. tag .. ":" .. className
+
+$[Tag("entity")]
+class C
+        getName: => "C"
+
+return
+]]
+                result, err = to_lua code
+                assert.is_nil err
+                assert.is_not_nil result
+                assert.is_true result\find("%-%- \"entity\":C") != nil
+
+        it "should report an error when annotation is not followed by a statement", ->
+                code = [[
+macro Invalid = (code) -> ""
+$[Invalid]
+]]
+                result, err = to_lua code
+                assert.is_nil result
+                assert.is_true err\match("annotation must be followed by a statement") != nil
+
+        it "should wrap annotated function to validate numeric arguments", ->
+                code = [[
+macro ValidateNumberArgs = (code) ->
+        funcName = code\match "^(%w+)%s*="
+        return
+                type: "text"
+                before: false
+                code: table.concat {
+                        "local __orig_#{funcName} = #{funcName}"
+                        "#{funcName} = function(a, b)"
+                        "\tassert(type(a) == \"number\", \"expected number for a\")"
+                        "\tassert(type(b) == \"number\", \"expected number for b\")"
+                        "\treturn __orig_#{funcName}(a, b)"
+                        "end"
+                }, "\n"
+
+$[ValidateNumberArgs]
+add = (a, b) -> a + b
+
+ok, value = pcall -> add 3, 4
+bad_ok, bad_err = pcall -> add "3", 4
+return ok, value, bad_ok, bad_err
+]]
+                ok, value, bad_ok, bad_err = compile_and_run code
+                assert.is_true ok
+                assert.same value, 7
+                assert.is_false bad_ok
+                assert.is_true bad_err\match("expected number for a") != nil
+
+        it "should wrap annotated function to validate return value", ->
+                code = [[
+macro ValidateNumberReturn = (code) ->
+        funcName = code\match "^(%w+)%s*="
+        return
+                type: "text"
+                before: false
+                code: table.concat {
+                        "local __orig_#{funcName} = #{funcName}"
+                        "#{funcName} = function(...)"
+                        "\tlocal result = __orig_#{funcName}(...)"
+                        "\tassert(type(result) == \"number\", \"expected numeric return\")"
+                        "\treturn result"
+                        "end"
+                }, "\n"
+
+$[ValidateNumberReturn]
+toText = (value) -> tostring value
+
+ok, err = pcall -> toText 42
+return ok, err
+]]
+                ok, err = compile_and_run code
+                assert.is_false ok
+                assert.is_true err\match("expected numeric return") != nil
+
+        it "should use annotation arguments to register annotated classes", ->
+                code = [[
+macro Register = (registry, code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "#{registry}[\"#{className}\"] = #{className}"
+
+registry = {}
+
+$[Register(registry)]
+class Worker
+        run: => "ok"
+
+return registry.Worker != nil, registry.Worker!\run!
+]]
+                exists, result = compile_and_run code
+                assert.is_true exists
+                assert.same result, "ok"
diff --git a/spec/inputs/test/format_spec.yue b/spec/inputs/test/format_spec.yue
index 310b610..a76a5dd 100644
--- a/spec/inputs/test/format_spec.yue
+++ b/spec/inputs/test/format_spec.yue
@@ -26,6 +26,8 @@ files = [
         "spec/inputs/export_default.yue"
         "spec/inputs/with_scope_shadow.yue"
         "spec/inputs/assign.yue"
+        "spec/inputs/annotation.yue"
+        "spec/inputs/annotation_before.yue"
         "spec/inputs/literals.yue"
         "spec/inputs/luarocks_upload.yue"
         "spec/inputs/comprehension_nested.yue"
@@ -64,6 +66,7 @@ files = [
         "spec/inputs/test/continue_spec.yue"
         "spec/inputs/test/varargs_assignment_spec.yue"
         "spec/inputs/test/advanced_macro_spec.yue"
+        "spec/inputs/test/annotation_spec.yue"
         "spec/inputs/test/pipe_spec.yue"
         "spec/inputs/test/export_spec.yue"
         "spec/inputs/test/existential_spec.yue"
@@ -192,4 +195,3 @@ for file in *files
         assert.is_not_nil ast
         rewriteLineCol ast
         assert.same original_ast, ast
-
diff --git a/spec/outputs/annotation.lua b/spec/outputs/annotation.lua
new file mode 100644
index 0000000..261bd7b
--- /dev/null
+++ b/spec/outputs/annotation.lua
@@ -0,0 +1,48 @@
+local A
+do
+        local _class_0
+        local _base_0 = {
+                setAdd = function(self, x, y)
+                        self.x = x
+                        self.y = y
+                        return self.x + self.y
+                end
+        }
+        if _base_0.__index == nil then
+                _base_0.__index = _base_0
+        end
+        _class_0 = setmetatable({
+                __init = function(self, x, y)
+                        if x == nil then
+                                x = 0
+                        end
+                        if y == nil then
+                                y = 0
+                        end
+                        self.x = x
+                        self.y = y
+                end,
+                __base = _base_0,
+                __name = "A"
+        }, {
+                __index = _base_0,
+                __call = function(cls, ...)
+                        local _self_0 = setmetatable({ }, _base_0)
+                        cls.__init(_self_0, ...)
+                        return _self_0
+                end
+        })
+        _base_0.__class = _class_0
+        A = _class_0
+end
+---@class A
+---@field x number
+---@field y number
+---@field setAdd fun(self: A, x: number, y: number): number Set fields and add number values.
+---@class AClass
+---@operator call:A
+---@cast A AClass
+local a = A()
+local res = a:setAdd(1, 2)
+print(a.x, a.y, a.y, res)
+return
diff --git a/spec/outputs/annotation_before.lua b/spec/outputs/annotation_before.lua
new file mode 100644
index 0000000..874ef37
--- /dev/null
+++ b/spec/outputs/annotation_before.lua
@@ -0,0 +1,55 @@
+local B
+do
+        local _class_0
+        local _base_0 = {
+                getTag = function(self)
+                        return "before"
+                end
+        }
+        if _base_0.__index == nil then
+                _base_0.__index = _base_0
+        end
+        _class_0 = setmetatable({
+                __init = function() end,
+                __base = _base_0,
+                __name = "B"
+        }, {
+                __index = _base_0,
+                __call = function(cls, ...)
+                        local _self_0 = setmetatable({ }, _base_0)
+                        cls.__init(_self_0, ...)
+                        return _self_0
+                end
+        })
+        _base_0.__class = _class_0
+        B = _class_0
+end
+-- "before":B
+local C
+do
+        local _class_0
+        local _base_0 = {
+                getTag = function(self)
+                        return "after"
+                end
+        }
+        if _base_0.__index == nil then
+                _base_0.__index = _base_0
+        end
+        _class_0 = setmetatable({
+                __init = function() end,
+                __base = _base_0,
+                __name = "C"
+        }, {
+                __index = _base_0,
+                __call = function(cls, ...)
+                        local _self_0 = setmetatable({ }, _base_0)
+                        cls.__init(_self_0, ...)
+                        return _self_0
+                end
+        })
+        _base_0.__class = _class_0
+        C = _class_0
+end
+-- "after":C
+return B():getTag(), C():getTag()
diff --git a/spec/outputs/test/annotation_spec.lua b/spec/outputs/test/annotation_spec.lua
new file mode 100644
index 0000000..866b9dc
--- /dev/null
+++ b/spec/outputs/test/annotation_spec.lua
@@ -0,0 +1,162 @@
+local to_lua
+do
+        local _obj_0 = require("yue")
+        to_lua = _obj_0.to_lua
+end
+local compile_and_run
+compile_and_run = function(code, config)
+        if config == nil then
+                config = { }
+        end
+        local lua_code, err = to_lua(code, config)
+        assert.is_nil(err)
+        assert.is_not_nil(lua_code)
+        local chunk, load_err = load(lua_code)
+        assert.is_nil(load_err)
+        assert.is_not_nil(chunk)
+        return chunk()
+end
+return describe("annotation", function()
+        it("should append generated text after annotated class by default", function()
+                local code = [[macro ClsDef = (code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "-- after:" .. className
+
+$[ClsDef]
+class A
+        getName: => "A"
+
+return
+]]
+                local result, err = to_lua(code)
+                assert.is_nil(err)
+                assert.is_not_nil(result)
+                assert.is_true(result:find("__name = \"A\"") ~= nil)
+                assert.is_true(result:find("%-%- after:A") ~= nil)
+                return assert.is_true(result:find("__name = \"A\"") < result:find("%-%- after:A"))
+        end)
+        it("should place generated text before the annotated statement when before is true", function()
+                local code = [[macro Before = (code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: true
+                code: "-- before:" .. className
+
+$[Before]
+class B
+        getName: => "B"
+
+return
+]]
+                local result, err = to_lua(code)
+                assert.is_nil(err)
+                assert.is_not_nil(result)
+                assert.is_true(result:find("%-%- before:B") ~= nil)
+                assert.is_true(result:find("local B") ~= nil)
+                return assert.is_true(result:find("%-%- before:B") < result:find("local B"))
+        end)
+        it("should support annotation invocation arguments", function()
+                local code = [[macro Tag = (tag, code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "-- " .. tag .. ":" .. className
+
+$[Tag("entity")]
+class C
+        getName: => "C"
+
+return
+]]
+                local result, err = to_lua(code)
+                assert.is_nil(err)
+                assert.is_not_nil(result)
+                return assert.is_true(result:find("%-%- \"entity\":C") ~= nil)
+        end)
+        it("should report an error when annotation is not followed by a statement", function()
+                local code = [[macro Invalid = (code) -> ""
+$[Invalid]
+]]
+                local result, err = to_lua(code)
+                assert.is_nil(result)
+                return assert.is_true(err:match("annotation must be followed by a statement") ~= nil)
+        end)
+        it("should wrap annotated function to validate numeric arguments", function()
+                local code = [[macro ValidateNumberArgs = (code) ->
+        funcName = code\match "^(%w+)%s*="
+        return
+                type: "text"
+                before: false
+                code: table.concat {
+                        "local __orig_#{funcName} = #{funcName}"
+                        "#{funcName} = function(a, b)"
+                        "\tassert(type(a) == \"number\", \"expected number for a\")"
+                        "\tassert(type(b) == \"number\", \"expected number for b\")"
+                        "\treturn __orig_#{funcName}(a, b)"
+                        "end"
+                }, "\n"
+
+$[ValidateNumberArgs]
+add = (a, b) -> a + b
+
+ok, value = pcall -> add 3, 4
+bad_ok, bad_err = pcall -> add "3", 4
+return ok, value, bad_ok, bad_err
+]]
+                local ok, value, bad_ok, bad_err = compile_and_run(code)
+                assert.is_true(ok)
+                assert.same(value, 7)
+                assert.is_false(bad_ok)
+                return assert.is_true(bad_err:match("expected number for a") ~= nil)
+        end)
+        it("should wrap annotated function to validate return value", function()
+                local code = [[macro ValidateNumberReturn = (code) ->
+        funcName = code\match "^(%w+)%s*="
+        return
+                type: "text"
+                before: false
+                code: table.concat {
+                        "local __orig_#{funcName} = #{funcName}"
+                        "#{funcName} = function(...)"
+                        "\tlocal result = __orig_#{funcName}(...)"
+                        "\tassert(type(result) == \"number\", \"expected numeric return\")"
+                        "\treturn result"
+                        "end"
+                }, "\n"
+
+$[ValidateNumberReturn]
+toText = (value) -> tostring value
+
+ok, err = pcall -> toText 42
+return ok, err
+]]
+                local ok, err = compile_and_run(code)
+                assert.is_false(ok)
+                return assert.is_true(err:match("expected numeric return") ~= nil)
+        end)
+        return it("should use annotation arguments to register annotated classes", function()
+                local code = [[macro Register = (registry, code`ClassDecl) ->
+        className = code\match "^class%s+(%w+)"
+        return
+                type: "text"
+                before: false
+                code: "#{registry}[\"#{className}\"] = #{className}"
+
+registry = {}
+
+$[Register(registry)]
+class Worker
+        run: => "ok"
+
+return registry.Worker != nil, registry.Worker!\run!
+]]
+                local exists, result = compile_and_run(code)
+                assert.is_true(exists)
+                return assert.same(result, "ok")
+        end)
+end)
diff --git a/spec/outputs/test/format_spec.lua b/spec/outputs/test/format_spec.lua
index d38a0ad..1eb2fbb 100644
--- a/spec/outputs/test/format_spec.lua
+++ b/spec/outputs/test/format_spec.lua
@@ -26,6 +26,8 @@ local files = {
         "spec/inputs/export_default.yue",
         "spec/inputs/with_scope_shadow.yue",
         "spec/inputs/assign.yue",
+        "spec/inputs/annotation.yue",
+        "spec/inputs/annotation_before.yue",
         "spec/inputs/literals.yue",
         "spec/inputs/luarocks_upload.yue",
         "spec/inputs/comprehension_nested.yue",
@@ -64,6 +66,7 @@ local files = {
         "spec/inputs/test/continue_spec.yue",
         "spec/inputs/test/varargs_assignment_spec.yue",
         "spec/inputs/test/advanced_macro_spec.yue",
+        "spec/inputs/test/annotation_spec.yue",
         "spec/inputs/test/pipe_spec.yue",
         "spec/inputs/test/export_spec.yue",
         "spec/inputs/test/existential_spec.yue",
diff --git a/src/yuescript/yue_ast.cpp b/src/yuescript/yue_ast.cpp
index e5f8d12..23f3820 100644
--- a/src/yuescript/yue_ast.cpp
+++ b/src/yuescript/yue_ast.cpp
@@ -1530,6 +1530,9 @@ std::string MacroFunc_t::to_string(void* ud) const {
 std::string Macro_t::to_string(void* ud) const {
         return "macro "s + name->to_string(ud) + " = "s + decl->to_string(ud);
 }
+std::string Annotation_t::to_string(void* ud) const {
+        return "$["s + name->to_string(ud) + (invoke ? invoke->to_string(ud) : ""s) + "]"s;
+}
 std::string MacroInPlace_t::to_string(void* ud) const {
         auto info = reinterpret_cast<YueFormat*>(ud);
         auto line = "$ ->"s;
diff --git a/src/yuescript/yue_ast.h b/src/yuescript/yue_ast.h
index 7d9a4a1..65568af 100644
--- a/src/yuescript/yue_ast.h
+++ b/src/yuescript/yue_ast.h
@@ -853,6 +853,12 @@ AST_NODE(Macro)
         AST_MEMBER(Macro, &name, &decl)
 AST_END(Macro)
 
+AST_NODE(Annotation)
+        ast_ptr<true, UnicodeName_t> name;
+        ast_sel<false, Invoke_t, InvokeArgs_t> invoke;
+        AST_MEMBER(Annotation, &name, &invoke)
+AST_END(Annotation)
+
 AST_NODE(NameOrDestructure)
         ast_sel<true, Variable_t, SimpleTable_t, TableLit_t, Comprehension_t> item;
         AST_MEMBER(NameOrDestructure, &item)
@@ -961,7 +967,8 @@ AST_END(ChainAssign)
 AST_NODE(Statement)
         ast_sel<true,
                 Import_t, While_t, Repeat_t, For_t,
-                Return_t, Local_t, Global_t, Export_t, Macro_t, MacroInPlace_t,
+                Return_t, Local_t, Global_t, Export_t,
+                Macro_t, MacroInPlace_t, Annotation_t,
                 BreakLoop_t, Label_t, Goto_t, ShortTabAppending_t,
                 Backcall_t, LocalAttrib_t, PipeBody_t, ExpListAssign_t, ChainAssign_t
         > content;
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index 76d6b0f..844d61e 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -78,7 +78,7 @@ static std::unordered_set<std::string> Metamethods = {
         "close"s // Lua 5.4
 };
 
-const std::string_view version = "0.33.10"sv;
+const std::string_view version = "0.34.0"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -5386,7 +5386,8 @@ private:
                 if (!nodes.empty()) {
                         str_list temp;
                         auto lastStmt = lastStatementFrom(nodes);
-                        for (auto node : nodes) {
+                        for (auto nodeIt = nodes.begin(); nodeIt != nodes.end(); ++nodeIt) {
+                                auto node = *nodeIt;
                                 if (auto comment = ast_cast<YueComment_t>(node)) {
                                         transformComment(comment, temp);
                                         continue;
@@ -5398,7 +5399,8 @@ private:
                                 } else if (!ast_is<Statement_t>(node)) {
                                         continue;
                                 }
-                                auto transformNode = [&]() {
+                                std::function<void()> transformNode;
+                                transformNode = [&]() {
                                         currentScope().lastStatement = (node == lastStmt) && currentScope().mode == GlobalMode::None;
                                         auto stmt = static_cast<Statement_t*>(node);
                                         if (auto importNode = stmt->content.as<Import_t>();
@@ -5418,6 +5420,56 @@ private:
                                                                 _importedGlobal->importContent = importNode->content.get();
                                                         }
                                                 }
+                                        } else if (auto annotation = stmt->content.as<Annotation_t>()) {
+#ifndef YUE_NO_MACRO
+                                                auto next = nodeIt;
+                                                ++next;
+                                                if (next == nodes.end() || !ast_is<Statement_t>(*next)) {
+                                                        throw CompileError("annotation must be followed by a statement"sv, node);
+                                                }
+                                                if (static_cast<Statement_t*>(*next)->content.is<Return_t>()) {
+                                                        throw CompileError("annotation can not be applied to a return statement"sv, node);
+                                                }
+                                                auto callable = annotation->new_ptr<Callable_t>();
+                                                auto macroName = annotation->new_ptr<MacroName_t>();
+                                                macroName->name.set(annotation->name);
+                                                callable->item.set(macroName);
+                                                auto chainValue = annotation->new_ptr<ChainValue_t>();
+                                                chainValue->items.push_back(callable);
+                                                if (annotation->invoke) {
+                                                        chainValue->items.push_back(annotation->invoke);
+                                                }
+                                                auto stmtCode = YueFormat{}.toString(*next);
+                                                ast_ptr<false, ast_node> macroNode;
+                                                std::unique_ptr<input> codes;
+                                                std::string luaCodes;
+                                                str_list localVars;
+                                                bool before = false;
+                                                std::tie(macroNode, codes, luaCodes, localVars, before) = expandMacro(chainValue, ExpUsage::Common, false, stmtCode);
+                                                if (!before) {
+                                                        ++nodeIt;
+                                                        node = *nodeIt;
+                                                        transformNode();
+                                                }
+                                                if (!macroNode) {
+                                                        temp.push_back(luaCodes);
+                                                        if (!localVars.empty()) {
+                                                                for (const auto& var : localVars) {
+                                                                        addToScope(var);
+                                                                }
+                                                        }
+                                                } else {
+                                                        if (!macroNode.to<Block_t>()->statementOrComments.empty()) {
+                                                                auto doBody = macroNode->new_ptr<Body_t>();
+                                                                doBody->content.set(macroNode);
+                                                                auto doNode = macroNode->new_ptr<Do_t>();
+                                                                doNode->body.set(doBody);
+                                                                transformDo(doNode, temp, ExpUsage::Common);
+                                                        }
+                                                }
+#else // YUE_NO_MACRO
+                                                throw CompileError("macro feature not supported"sv, annotation);
+#endif // YUE_NO_MACRO
                                         } else {
                                                 transformStatement(stmt, temp);
                                         }
@@ -5659,7 +5711,7 @@ private:
                         pushCurrentModule(); // cur
                         int top = lua_gettop(L) - 1;
                         DEFER(lua_settop(L, top));
-                        if (auto builtinCode = expandMacroChain(chainValue)) {
+                        if (auto builtinCode = expandMacroChain(chainValue, Empty)) {
                                 throw CompileError("macro generating function must return a function"sv, chainValue);
                         } // cur res
                         if (lua_isfunction(L, -1) == 0) {
@@ -7039,7 +7091,7 @@ private:
                 return Empty;
         }
 
-        std::optional<std::string> expandMacroChain(ChainValue_t* chainValue) {
+        std::optional<std::string> expandMacroChain(ChainValue_t* chainValue, const std::string& extraCode) {
                 const auto& chainList = chainValue->items.objects();
                 auto x = ast_to<Callable_t>(chainList.front())->item.to<MacroName_t>();
                 auto macroName = _parser.toString(x->name);
@@ -7088,7 +7140,7 @@ private:
                                                         auto chainValue = exp->get_by_path<UnaryExp_t, Value_t, ChainValue_t>();
                                                         BREAK_IF(!chainValue);
                                                         BREAK_IF(!isMacroChain(chainValue));
-                                                        str = std::get<1>(expandMacroStr(chainValue));
+                                                        str = std::get<1>(expandMacroStr(chainValue, Empty));
                                                         BLOCK_END
                                                 }
                                         }
@@ -7184,7 +7236,11 @@ private:
                 if (args) {
                         argIt = args->begin();
                 }
+                if (!extraCode.empty()) {
+                        argStrs.push_back(extraCode);
+                }
                 for (const auto& arg : argStrs) {
+                        ast_node* currentArg = args && argIt != args->end() ? *argIt : chainValue;
                         if (checkIt != checks.end()) {
                                 if (checkIt->empty()) {
                                         ++checkIt;
@@ -7192,18 +7248,20 @@ private:
                                         if ((*checkIt)[0] == '.') {
                                                 auto astName = checkIt->substr(3);
                                                 if (!_parser.match(astName, arg)) {
-                                                        throw CompileError("expecting \""s + astName + "\", AST mismatch"s, *argIt);
+                                                        throw CompileError("expecting \""s + astName + "\", AST mismatch"s, currentArg);
                                                 }
                                         } else {
                                                 if (!_parser.match(*checkIt, arg)) {
-                                                        throw CompileError("expecting \""s + *checkIt + "\", AST mismatch"s, *argIt);
+                                                        throw CompileError("expecting \""s + *checkIt + "\", AST mismatch"s, currentArg);
                                                 }
                                                 ++checkIt;
                                         }
                                 }
                         }
                         lua_pushlstring(L, arg.c_str(), arg.size());
-                        ++argIt;
+                        if (args && argIt != args->end()) {
+                                ++argIt;
+                        }
                 } // cur pcall macroFunc args...
                 bool success = lua_pcall(L, static_cast<int>(argStrs.size()), 1, 0) == 0;
                 if (!success) { // cur err
@@ -7218,21 +7276,22 @@ private:
                 return std::nullopt;
         }
 
-        std::tuple<std::string, std::string, str_list> expandMacroStr(ChainValue_t* chainValue) {
+        std::tuple<std::string, std::string, str_list, bool> expandMacroStr(ChainValue_t* chainValue, const std::string& extraCode) {
                 auto x = chainValue->items.front();
                 pushCurrentModule(); // cur
                 int top = lua_gettop(L) - 1;
                 DEFER(lua_settop(L, top));
-                auto builtinCode = expandMacroChain(chainValue);
+                auto builtinCode = expandMacroChain(chainValue, extraCode);
                 if (builtinCode) {
-                        return {Empty, builtinCode.value(), {}};
+                        return {Empty, builtinCode.value(), {}, false};
                 } // cur res
                 if (lua_isstring(L, -1) == 0 && lua_istable(L, -1) == 0) {
                         throw CompileError("macro function must return a string or a table"sv, x);
                 } // cur res
                 std::string codes;
-                std::string type;
+                std::string type = "yue"s;
                 str_list localVars;
+                bool before = false;
                 if (lua_istable(L, -1) != 0) { // cur tab
                         lua_getfield(L, -1, "code"); // cur tab code
                         if (lua_isstring(L, -1) != 0) {
@@ -7245,8 +7304,8 @@ private:
                         if (lua_isstring(L, -1) != 0) {
                                 type = lua_tostring(L, -1);
                         }
-                        if (type != "lua"sv && type != "text"sv) {
-                                throw CompileError("macro table must contain field \"type\" of value \"lua\" or \"text\""sv, x);
+                        if (type != "yue"sv && type != "lua"sv && type != "text"sv) {
+                                throw CompileError("macro table must contain field \"type\" of value \"yue\", \"lua\" or \"text\""sv, x);
                         }
                         lua_pop(L, 1); // cur tab
                         if (type == "text"sv) {
@@ -7269,20 +7328,26 @@ private:
                                 }
                                 lua_pop(L, 1); // cur tab
                         }
+                        lua_getfield(L, -1, "before"); // cur tab before
+                        if (lua_toboolean(L, -1) != 0) {
+                                before = true;
+                        }
+                        lua_pop(L, 1);
                 } else { // cur code
                         codes = lua_tostring(L, -1);
                 }
                 Utils::trim(codes);
                 Utils::replace(codes, "\r\n"sv, "\n"sv);
-                return {type, codes, std::move(localVars)};
+                return {type, codes, std::move(localVars), before};
         }
 
-        std::tuple<ast_ptr<false, ast_node>, std::unique_ptr<input>, std::string, str_list> expandMacro(ChainValue_t* chainValue, ExpUsage usage, bool allowBlockMacroReturn) {
+        std::tuple<ast_ptr<false, ast_node>, std::unique_ptr<input>, std::string, str_list, bool> expandMacro(ChainValue_t* chainValue, ExpUsage usage, bool allowBlockMacroReturn, const std::string& extraCode) {
                 auto x = ast_to<Callable_t>(chainValue->items.front())->item.to<MacroName_t>();
                 const auto& chainList = chainValue->items.objects();
                 std::string type, codes;
                 str_list localVars;
-                std::tie(type, codes, localVars) = expandMacroStr(chainValue);
+                bool before = false;
+                std::tie(type, codes, localVars, before) = expandMacroStr(chainValue, extraCode);
                 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) {
@@ -7298,7 +7363,7 @@ private:
                                         codes.insert(0, indent() + "do"s + nl(chainValue));
                                         codes.append(_newLine + indent() + "end"s + nl(chainValue));
                                 }
-                                return {nullptr, nullptr, std::move(codes), std::move(localVars)};
+                                return {nullptr, nullptr, std::move(codes), std::move(localVars), before};
                         } else {
                                 auto expCode = "return ("s + codes + ')';
                                 if (luaL_loadbuffer(L, expCode.c_str(), expCode.size(), macroChunk.c_str()) != 0) {
@@ -7320,7 +7385,7 @@ private:
                                                 newChain->items.push_back(*it);
                                         }
                                 }
-                                return {exp, nullptr, Empty, std::move(localVars)};
+                                return {exp, nullptr, Empty, std::move(localVars), before};
                         }
                 } else if (type == "text"sv) {
                         if (!isBlock) {
@@ -7329,7 +7394,7 @@ private:
                         if (!codes.empty()) {
                                 codes.append(_newLine);
                         }
-                        return {nullptr, nullptr, std::move(codes), std::move(localVars)};
+                        return {nullptr, nullptr, std::move(codes), std::move(localVars), before};
                 } else {
                         if (!codes.empty()) {
                                 if (isBlock) {
@@ -7395,10 +7460,10 @@ private:
                                         auto block = blockEnd->block.get();
                                         info.node.set(block);
                                 }
-                                return {info.node, std::move(info.codes), Empty, std::move(localVars)};
+                                return {info.node, std::move(info.codes), Empty, std::move(localVars), before};
                         } else {
                                 if (!isBlock) throw CompileError("failed to expand empty macro as expr"sv, x);
-                                return {x->new_ptr<Block_t>().get(), std::move(info.codes), Empty, std::move(localVars)};
+                                return {x->new_ptr<Block_t>().get(), std::move(info.codes), Empty, std::move(localVars), before};
                         }
                 }
         }
@@ -7411,7 +7476,8 @@ private:
                         std::unique_ptr<input> codes;
                         std::string luaCodes;
                         str_list localVars;
-                        std::tie(node, codes, luaCodes, localVars) = expandMacro(chainValue, usage, allowBlockMacroReturn);
+                        bool before = false;
+                        std::tie(node, codes, luaCodes, localVars, before) = expandMacro(chainValue, usage, allowBlockMacroReturn, Empty);
                         if (!node) {
                                 out.push_back(luaCodes);
                                 if (!localVars.empty()) {
diff --git a/src/yuescript/yue_parser.cpp b/src/yuescript/yue_parser.cpp
index 7ea2c97..8ab667d 100644
--- a/src/yuescript/yue_parser.cpp
+++ b/src/yuescript/yue_parser.cpp
@@ -1072,6 +1072,8 @@ YueParser::YueParser() {
         );
         MacroInPlace = '$' >> space >> "->" >> space >> Body;
 
+        Annotation = "$[" >> space >> UnicodeName >> -(Invoke | InvokeArgs) >> space >> ']';
+
         must_variable = Variable | and_(LuaKeyword >> not_alpha_num) >> keyword_as_identifier_syntax_error | expected_indentifier_error;
 
         NameList = Seperator >> must_variable >> *(space >> ',' >> space >> must_variable);
@@ -1128,7 +1130,7 @@ YueParser::YueParser() {
         StatementAppendix = IfLine | WhileLine | CompFor;
         Statement = (
                 (
-                        Import | Export | Global | Macro | MacroInPlace | Label
+                        Import | Export | Global | Macro | MacroInPlace | Annotation | Label
                 ) | (
                         Local | While | Repeat | For | Return |
                         BreakLoop | Goto | ShortTabAppending |
diff --git a/src/yuescript/yue_parser.h b/src/yuescript/yue_parser.h
index 07153fb..cfcbb48 100644
--- a/src/yuescript/yue_parser.h
+++ b/src/yuescript/yue_parser.h
@@ -464,6 +464,7 @@ private:
         AST_RULE(MacroFunc);
         AST_RULE(Macro);
         AST_RULE(MacroInPlace);
+        AST_RULE(Annotation);
         AST_RULE(NameOrDestructure);
         AST_RULE(AssignableNameList);
         AST_RULE(InvokeArgs);