Add annotation statements and expand annotation tests

4 files changed, 268 insertions, 0 deletions

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",