From 46c093dbae34bf44fb2fccca619bacd0eea43a35 Mon Sep 17 00:00:00 2001
From: Oleksandr Smirnov <olexsmir@gmail.com>
Date: Sat, 22 Mar 2025 12:46:13 +0200
Subject: [PATCH] fix(struct_tags): edge case with structs declared as var

---
 lua/gopher/_utils/ts.lua   | 25 ++++++++++++++++++++++---
 lua/gopher/struct_tags.lua | 10 ++++++++--
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/lua/gopher/_utils/ts.lua b/lua/gopher/_utils/ts.lua
index 9c74f47..bf59ab5 100644
--- a/lua/gopher/_utils/ts.lua
+++ b/lua/gopher/_utils/ts.lua
@@ -1,8 +1,14 @@
 local ts = {}
 local queries = {
   struct = [[
-    (type_spec name: (type_identifier) @_name
-               type: (struct_type))
+    [(type_spec name: (type_identifier) @_name
+                type: (struct_type))
+     (var_spec  name: (identifier) @_name
+                type: (struct_type))
+     (short_var_declaration
+       left:  (expression_list (identifier) @_name @_var)
+       right: (expression_list (composite_literal
+                                 type: (struct_type))))]
   ]],
   func = [[
     [(function_declaration name: (identifier)       @_name)
@@ -49,6 +55,10 @@ local function get_captures(query, node, bufnr)
       if capture_name == "_name" then
         res["name"] = vim.treesitter.get_node_text(captured_node, bufnr)
       end
+
+      if capture_name == "_var" then
+        res["is_varstruct"] = true
+      end
     end
   end
 
@@ -59,6 +69,7 @@ end
 ---@field name string
 ---@field start_line integer
 ---@field end_line integer
+---@field is_varstruct boolean
 
 ---@param bufnr integer
 ---@param parent_type string[]
@@ -93,7 +104,15 @@ function ts.get_struct_under_cursor(bufnr)
   --- should be both type_spec and type_declaration
   --- because in cases like `type ( T struct{}, U strict{} )`
   --- i will be choosing always last struct in the list
-  return do_stuff(bufnr, { "type_spec", "type_declaration" }, queries.struct)
+  ---
+  --- var_spec is for cases like `var x struct{}`
+  --- short_var_declaration is for cases like `x := struct{}{}`
+  return do_stuff(bufnr, {
+    "type_spec",
+    "type_declaration",
+    "var_spec",
+    "short_var_declaration",
+  }, queries.struct)
 end
 
 ---@param bufnr integer
diff --git a/lua/gopher/struct_tags.lua b/lua/gopher/struct_tags.lua
index 4545640..86689e4 100644
--- a/lua/gopher/struct_tags.lua
+++ b/lua/gopher/struct_tags.lua
@@ -44,11 +44,18 @@ local function handle_tags(fpath, bufnr, user_args)
     c.commands.gomodifytags,
     "-transform", c.gotag.transform,
     "-format", "json",
-    "-struct", st.name,
     "-file", fpath,
     "-w",
   }
 
+  if st.is_varstruct then
+    table.insert(cmd, "-line")
+    table.insert(cmd, string.format("%d,%d", st.start_line, st.end_line))
+  else
+    table.insert(cmd, "-struct")
+    table.insert(cmd, st.name)
+  end
+
   for _, v in ipairs(user_args) do
     table.insert(cmd, v)
   end
@@ -60,7 +67,6 @@ local function handle_tags(fpath, bufnr, user_args)
   end
 
   local res = vim.json.decode(rs.stdout)
-
   if res["errors"] then
     log.error("tags: got an error " .. vim.inspect(res))
     error("failed to set tags " .. vim.inspect(res["errors"]))