From a9f2e9c1d91e7683b1cb2e041b0ac8b4250aaec3 Mon Sep 17 00:00:00 2001
From: Hisham Muhammad <hisham@gobolinux.org>
Date: Fri, 5 Jan 2024 16:26:48 -0300
Subject: [PATCH] fix: hashbang, if exists and wanted, stays on line 1

Fixes: #729
Closes: #730

Co-authored-by: JLPLabs LLC <jlplabs@man.com>
---
 spec/cli/gen_spec.lua        | 62 +++++++++++++++++++++++-------------
 spec/lexer/hashbang_spec.lua |  2 +-
 tl.lua                       |  7 ++--
 tl.tl                        |  7 ++--
 4 files changed, 51 insertions(+), 27 deletions(-)

diff --git a/spec/cli/gen_spec.lua b/spec/cli/gen_spec.lua
index d3132871..fd3d67f5 100644
--- a/spec/cli/gen_spec.lua
+++ b/spec/cli/gen_spec.lua
@@ -72,15 +72,31 @@ end
 local c = 100
 ]]
 
-local script_with_hashbang = [[
+local hashbang_cases = {
+[1] = {
+with_hashbang = [[
 #!/usr/bin/env lua
+
 print("hello world")
-]]
+]],
+without_hashbang = [[
 
-local script_without_hashbang = [[
 
 print("hello world")
-]]
+]],
+},
+
+[2] = {
+with_hashbang = [[
+#!/usr/bin/env lua
+print("hello world")
+]],
+without_hashbang = [[
+
+print("hello world")
+]],
+},
+}
 
 local function tl_to_lua(name)
    return (name:gsub("%.tl$", ".lua"):gsub("^" .. util.os_tmp .. util.os_sep, ""))
@@ -195,25 +211,27 @@ describe("tl gen", function()
       end)
    end)
 
-   it("preserves hashbang with --keep-hashbang", function()
-      local name = util.write_tmp_file(finally, script_with_hashbang)
-      local pd = io.popen(util.tl_cmd("gen", "--keep-hashbang", name), "r")
-      local output = pd:read("*a")
-      util.assert_popen_close(0, pd:close())
-      local lua_name = tl_to_lua(name)
-      assert.match("Wrote: " .. lua_name, output, 1, true)
-      util.assert_line_by_line(script_with_hashbang, util.read_file(lua_name))
-   end)
+   for i, case in ipairs(hashbang_cases) do
+      it("[" .. i .. "] preserves hashbang with --keep-hashbang", function()
+         local name = util.write_tmp_file(finally, case.with_hashbang)
+         local pd = io.popen(util.tl_cmd("gen", "--keep-hashbang", name), "r")
+         local output = pd:read("*a")
+         util.assert_popen_close(0, pd:close())
+         local lua_name = tl_to_lua(name)
+         assert.match("Wrote: " .. lua_name, output, 1, true)
+         assert.equal(case.with_hashbang, util.read_file(lua_name))
+      end)
 
-   it("drops hashbang when not using --keep-hashbang", function()
-      local name = util.write_tmp_file(finally, script_with_hashbang)
-      local pd = io.popen(util.tl_cmd("gen", name), "r")
-      local output = pd:read("*a")
-      util.assert_popen_close(0, pd:close())
-      local lua_name = tl_to_lua(name)
-      assert.match("Wrote: " .. lua_name, output, 1, true)
-      util.assert_line_by_line(script_without_hashbang, util.read_file(lua_name))
-   end)
+      it("[" .. i .. "] drops hashbang when not using --keep-hashbang", function()
+         local name = util.write_tmp_file(finally, case.with_hashbang)
+         local pd = io.popen(util.tl_cmd("gen", name), "r")
+         local output = pd:read("*a")
+         util.assert_popen_close(0, pd:close())
+         local lua_name = tl_to_lua(name)
+         assert.match("Wrote: " .. lua_name, output, 1, true)
+         util.assert_line_by_line(case.without_hashbang, util.read_file(lua_name))
+      end)
+   end
 
    describe("with --gen-target=5.1", function()
       it("targets generated code to Lua 5.1+", function()
diff --git a/spec/lexer/hashbang_spec.lua b/spec/lexer/hashbang_spec.lua
index 7d69f1cc..2b3272dd 100644
--- a/spec/lexer/hashbang_spec.lua
+++ b/spec/lexer/hashbang_spec.lua
@@ -12,7 +12,7 @@ describe("lexer", function()
    it("skips hashbang at the beginning of a file", function()
       local syntax_errors = {}
       local tokens = tl.lex("#!/usr/bin/env lua\nlocal x = 1")
-      assert.same({"#!/usr/bin/env lua\n", "local", "x", "=", "1", "$EOF$"}, map(function(x) return x.tk end, tokens))
+      assert.same({"#!/usr/bin/env lua", "local", "x", "=", "1", "$EOF$"}, map(function(x) return x.tk end, tokens))
 
       tl.parse_program(tokens, syntax_errors)
       assert.same({}, syntax_errors)
diff --git a/tl.lua b/tl.lua
index b6cb0e26..3d1f2a76 100644
--- a/tl.lua
+++ b/tl.lua
@@ -605,7 +605,7 @@ do
          if not i then
             i = len + 1
          end
-         end_token_here("hashbang")
+         end_token_prev("hashbang")
          y = 2
          x = 0
       end
@@ -3854,9 +3854,12 @@ function tl.pretty_print_ast(ast, gen_target, mode)
    visit_node.cbs = {
       ["statements"] = {
          after = function(node, children)
-            local out = { y = node.y, h = 0 }
+            local out
             if opts.preserve_hashbang and node.hashbang then
+               out = { y = 1, h = 0 }
                table.insert(out, node.hashbang)
+            else
+               out = { y = node.y, h = 0 }
             end
             local space
             for i, child in ipairs(children) do
diff --git a/tl.tl b/tl.tl
index 421a6b70..8291157c 100644
--- a/tl.tl
+++ b/tl.tl
@@ -605,7 +605,7 @@ do
          if not i then
             i = len + 1
          end
-         end_token_here("hashbang")
+         end_token_prev("hashbang")
          y = 2
          x = 0
       end
@@ -3854,9 +3854,12 @@ function tl.pretty_print_ast(ast: Node, gen_target: TargetMode, mode: boolean |
    visit_node.cbs = {
       ["statements"] = {
          after = function(node: Node, children: {Output}): Output
-            local out: Output = { y = node.y, h = 0 }
+            local out: Output
             if opts.preserve_hashbang and node.hashbang then
+               out = { y = 1, h = 0 }
                table.insert(out, node.hashbang)
+            else
+               out = { y = node.y, h = 0 }
             end
             local space: string
             for i, child in ipairs(children) do