5 files changed, 198 insertions, 6 deletions

diff --git a/spec/inputs/test/ast_spec.yue b/spec/inputs/test/ast_spec.yue
new file mode 100644
index 0000000..11efc94
--- /dev/null
+++ b/spec/inputs/test/ast_spec.yue
@@ -0,0 +1,91 @@
+yue = require "yue"
+
+describe "yue.to_ast", ->
+        it "should return AST with end position for simple expression", ->
+                ast = yue.to_ast "x = 1"
+                assert.is_not_nil ast
+                -- ast format: [name, begin_line, begin_col, end_line, end_col, ...children]
+                assert.same ast[1], "File"
+                assert.same ast[2], 1  -- begin line
+                assert.same ast[3], 1  -- begin col
+                assert.is_number ast[4]  -- end line
+                assert.is_number ast[5]  -- end col
+
+        it "should have correct end position for leaf nodes", ->
+                ast = yue.to_ast "1"
+                assert.is_not_nil ast
+                -- Leaf node with no children should have format: [name, begin_line, begin_col, end_line, end_col, value]
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                assert.is_number ast[4]
+                assert.is_number ast[5]
+
+        it "should have end position for multi-line code", ->
+                code = [[
+x = 1
+y = 2]]
+                ast = yue.to_ast code
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                -- End position should be on line 2
+                assert.is_true ast[4] >= 2
+
+        it "should have end position for function definition", ->
+                code = [[
+add = (a, b) ->
+        a + b]]
+                ast = yue.to_ast code
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                assert.is_number ast[4]
+                assert.is_number ast[5]
+
+        it "should have end position for table literal", ->
+                ast = yue.to_ast "{a: 1, b: 2}"
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                assert.is_number ast[4]
+                assert.is_number ast[5]
+
+        it "should have end position for class definition", ->
+                code = [[
+class Person
+        new: (@name) =>
+        getName: => @name]]
+                ast = yue.to_ast code
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                assert.is_true ast[4] >= 3
+
+        it "should return nil and error message for invalid syntax", ->
+                ast, err = yue.to_ast "if then else"
+                assert.is_nil ast
+                assert.is_string err
+
+        it "should support flatten level parameter", ->
+                ast = yue.to_ast "x = 1", 0
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.is_number ast[4]
+                assert.is_number ast[5]
+
+        it "should have end position in nested structures", ->
+                code = "x = [i for i = 1, 10]"
+                ast = yue.to_ast code
+                assert.is_not_nil ast
+                assert.same ast[1], "File"
+                assert.same ast[2], 1
+                assert.same ast[3], 1
+                assert.is_number ast[4]
+                assert.is_number ast[5]
+                -- End column should reflect the end of the expression
+                assert.is_true ast[5] > 1
diff --git a/spec/inputs/test/format_spec.yue b/spec/inputs/test/format_spec.yue
index 8c6096a..310b610 100644
--- a/spec/inputs/test/format_spec.yue
+++ b/spec/inputs/test/format_spec.yue
@@ -165,7 +165,9 @@ import "yue"
 rewriteLineCol = (item)->
         item[2] = 0
         item[3] = 0
-        for i = 4, #item
+        item[4] = 0
+        item[5] = 0
+        for i = 6, #item
                 switch type item[i] when "table"
                         if item[i][1] == "comment"
                                 table.remove item, i
diff --git a/spec/outputs/test/ast_spec.lua b/spec/outputs/test/ast_spec.lua
new file mode 100644
index 0000000..bab5b9e
--- /dev/null
+++ b/spec/outputs/test/ast_spec.lua
@@ -0,0 +1,85 @@
+local yue = require("yue")
+return describe("yue.to_ast", function()
+        it("should return AST with end position for simple expression", function()
+                local ast = yue.to_ast("x = 1")
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                assert.is_number(ast[4])
+                return assert.is_number(ast[5])
+        end)
+        it("should have correct end position for leaf nodes", function()
+                local ast = yue.to_ast("1")
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                assert.is_number(ast[4])
+                return assert.is_number(ast[5])
+        end)
+        it("should have end position for multi-line code", function()
+                local code = [[x = 1
+y = 2]]
+                local ast = yue.to_ast(code)
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                return assert.is_true(ast[4] >= 2)
+        end)
+        it("should have end position for function definition", function()
+                local code = [[add = (a, b) ->
+        a + b]]
+                local ast = yue.to_ast(code)
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                assert.is_number(ast[4])
+                return assert.is_number(ast[5])
+        end)
+        it("should have end position for table literal", function()
+                local ast = yue.to_ast("{a: 1, b: 2}")
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                assert.is_number(ast[4])
+                return assert.is_number(ast[5])
+        end)
+        it("should have end position for class definition", function()
+                local code = [[class Person
+        new: (@name) =>
+        getName: => @name]]
+                local ast = yue.to_ast(code)
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                return assert.is_true(ast[4] >= 3)
+        end)
+        it("should return nil and error message for invalid syntax", function()
+                local ast, err = yue.to_ast("if then else")
+                assert.is_nil(ast)
+                return assert.is_string(err)
+        end)
+        it("should support flatten level parameter", function()
+                local ast = yue.to_ast("x = 1", 0)
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.is_number(ast[4])
+                return assert.is_number(ast[5])
+        end)
+        return it("should have end position in nested structures", function()
+                local code = "x = [i for i = 1, 10]"
+                local ast = yue.to_ast(code)
+                assert.is_not_nil(ast)
+                assert.same(ast[1], "File")
+                assert.same(ast[2], 1)
+                assert.same(ast[3], 1)
+                assert.is_number(ast[4])
+                assert.is_number(ast[5])
+                return assert.is_true(ast[5] > 1)
+        end)
+end)
diff --git a/spec/outputs/test/format_spec.lua b/spec/outputs/test/format_spec.lua
index c9ea3c2..d38a0ad 100644
--- a/spec/outputs/test/format_spec.lua
+++ b/spec/outputs/test/format_spec.lua
@@ -164,7 +164,9 @@ local rewriteLineCol
 rewriteLineCol = function(item)
         item[2] = 0
         item[3] = 0
-        for i = 4, #item do
+        item[4] = 0
+        item[5] = 0
+        for i = 6, #item do
                 local _exp_0 = type(item[i])
                 if "table" == _exp_0 then
                         if item[i][1] == "comment" then
diff --git a/src/yuescript/yuescript.cpp b/src/yuescript/yuescript.cpp
index 0dbfe2f..7645ef4 100644
--- a/src/yuescript/yuescript.cpp
+++ b/src/yuescript/yuescript.cpp
@@ -380,18 +380,22 @@ static int yuetoast(lua_State* L) {
                                         int count = current.children ? static_cast<int>(current.children->size()) : 0;
                                         switch (count) {
                                                 case 0: {
-                                                        lua_createtable(L, 4, 0);
+                                                        lua_createtable(L, 6, 0);
                                                         getName(node);
                                                         lua_rawseti(L, -2, 1);
                                                         lua_pushinteger(L, node->m_begin.m_line);
                                                         lua_rawseti(L, -2, 2);
                                                         lua_pushinteger(L, node->m_begin.m_col);
                                                         lua_rawseti(L, -2, 3);
+                                                        lua_pushinteger(L, node->m_end.m_line);
+                                                        lua_rawseti(L, -2, 4);
+                                                        lua_pushinteger(L, node->m_end.m_col);
+                                                        lua_rawseti(L, -2, 5);
                                                         formatter.indent = 0;
                                                         auto str = node->to_string(&formatter);
                                                         yue::Utils::trim(str);
                                                         lua_pushlstring(L, str.c_str(), str.length());
-                                                        lua_rawseti(L, -2, 4);
+                                                        lua_rawseti(L, -2, 6);
                                                         lua_rawseti(L, tableIndex, static_cast<int>(lua_objlen(L, tableIndex)) + 1);
                                                         break;
                                                 }
@@ -404,6 +408,10 @@ static int yuetoast(lua_State* L) {
                                                                 lua_rawseti(L, -2, 2);
                                                                 lua_pushinteger(L, node->m_begin.m_col);
                                                                 lua_rawseti(L, -2, 3);
+                                                                lua_pushinteger(L, node->m_end.m_line);
+                                                                lua_rawseti(L, -2, 4);
+                                                                lua_pushinteger(L, node->m_end.m_col);
+                                                                lua_rawseti(L, -2, 5);
                                                                 lua_pop(L, 1);
                                                                 break;
                                                         }
@@ -411,14 +419,18 @@ static int yuetoast(lua_State* L) {
                                                 }
                                                 default: {
                                                         auto len = static_cast<int>(lua_objlen(L, tableIndex));
-                                                        lua_createtable(L, count + 3, 0);
+                                                        lua_createtable(L, count + 5, 0);
                                                         getName(node);
                                                         lua_rawseti(L, -2, 1);
                                                         lua_pushinteger(L, node->m_begin.m_line);
                                                         lua_rawseti(L, -2, 2);
                                                         lua_pushinteger(L, node->m_begin.m_col);
                                                         lua_rawseti(L, -2, 3);
-                                                        for (int i = count, j = 4; i >= 1; i--, j++) {
+                                                        lua_pushinteger(L, node->m_end.m_line);
+                                                        lua_rawseti(L, -2, 4);
+                                                        lua_pushinteger(L, node->m_end.m_col);
+                                                        lua_rawseti(L, -2, 5);
+                                                        for (int i = count, j = 6; i >= 1; i--, j++) {
                                                                 lua_rawgeti(L, tableIndex, len - i + 1);
                                                                 lua_rawseti(L, -2, j);
                                                         }