From 1e7af1b2123de40baf81071388601fdae02573a5 Mon Sep 17 00:00:00 2001
From: Oleksandr Smirnov <olexsmir@gmail.com>
Date: Thu, 4 Sep 2025 16:52:18 +0300
Subject: [PATCH] feat(comment): add support for: interface methods, struct
 fields, variables (#123)

* refactor(comment): dont use unnecessary function

* chore: quick way to open vim in dev mode

* feat(comment): add comment on on interface method

* feat(comment): add comment on a struct field

* feat(comment): add comment on a variable

* docs: add note about the generate function

* docs: gopher.TsResult

* fix(utils): handle case when indentation is wrong
---
 Taskfile.yml                                  |  2 +
 lua/gopher/_utils/init.lua                    | 14 ++++-
 lua/gopher/_utils/ts.lua                      | 56 +++++++++++++------
 lua/gopher/comment.lua                        | 44 +++++++++------
 .../comment/interface_many_method_input.go    |  6 ++
 .../comment/interface_many_method_output.go   |  7 +++
 .../comment/interface_method_input.go         |  5 ++
 .../comment/interface_method_output.go        |  6 ++
 .../comment/many_structs_fields_input.go      | 18 ++++++
 .../comment/many_structs_fields_output.go     | 19 +++++++
 spec/fixtures/comment/struct_fields_input.go  |  7 +++
 spec/fixtures/comment/struct_fields_output.go |  8 +++
 spec/fixtures/comment/svar_input.go           |  5 ++
 spec/fixtures/comment/svar_output.go          |  6 ++
 spec/fixtures/comment/var_input.go            |  5 ++
 spec/fixtures/comment/var_output.go           |  6 ++
 .../comment/var_struct_fields_input.go        |  8 +++
 .../comment/var_struct_fields_output.go       |  9 +++
 spec/integration/comment_test.lua             | 28 ++++++++++
 spec/unit/utils_test.lua                      | 24 ++++++++
 20 files changed, 250 insertions(+), 33 deletions(-)
 create mode 100644 spec/fixtures/comment/interface_many_method_input.go
 create mode 100644 spec/fixtures/comment/interface_many_method_output.go
 create mode 100644 spec/fixtures/comment/interface_method_input.go
 create mode 100644 spec/fixtures/comment/interface_method_output.go
 create mode 100644 spec/fixtures/comment/many_structs_fields_input.go
 create mode 100644 spec/fixtures/comment/many_structs_fields_output.go
 create mode 100644 spec/fixtures/comment/struct_fields_input.go
 create mode 100644 spec/fixtures/comment/struct_fields_output.go
 create mode 100644 spec/fixtures/comment/svar_input.go
 create mode 100644 spec/fixtures/comment/svar_output.go
 create mode 100644 spec/fixtures/comment/var_input.go
 create mode 100644 spec/fixtures/comment/var_output.go
 create mode 100644 spec/fixtures/comment/var_struct_fields_input.go
 create mode 100644 spec/fixtures/comment/var_struct_fields_output.go

diff --git a/Taskfile.yml b/Taskfile.yml
index 62ccd0c..2905634 100644
--- a/Taskfile.yml
+++ b/Taskfile.yml
@@ -21,6 +21,8 @@ tasks:
              -u ./scripts/minimal_init.lua \
              -c "lua MiniTest.run()" \
              -c ":qa!"
+  nvim:
+    cmd: nvim --clean -u "./scripts/minimal_init.lua" {{ .CLI_ARGS }}
 
   docgen:
     desc: generate vimhelp
diff --git a/lua/gopher/_utils/init.lua b/lua/gopher/_utils/init.lua
index e176345..345c670 100644
--- a/lua/gopher/_utils/init.lua
+++ b/lua/gopher/_utils/init.lua
@@ -3,7 +3,7 @@ local log = require "gopher._utils.log"
 local utils = {}
 
 ---@param msg string
----@param lvl? number by default `vim.log.levels.INFO`
+---@param lvl? integer by default `vim.log.levels.INFO`
 function utils.notify(msg, lvl)
   lvl = lvl or vim.log.levels.INFO
   vim.notify(msg, lvl, {
@@ -38,4 +38,16 @@ function utils.trimend(s)
   return r
 end
 
+-- Since indentation can be spaces or tabs, that's my hack around it
+---@param line string
+---@param indent integer
+---@return string
+function utils.indent(line, indent)
+  local char = string.sub(line, 1, 1)
+  if char ~= " " and char ~= "\t" then
+    char = " "
+  end
+  return string.rep(char, indent)
+end
+
 return utils
diff --git a/lua/gopher/_utils/ts.lua b/lua/gopher/_utils/ts.lua
index a0623ad..d24fb09 100644
--- a/lua/gopher/_utils/ts.lua
+++ b/lua/gopher/_utils/ts.lua
@@ -11,9 +11,13 @@ local queries = {
        right: (expression_list (composite_literal
                                  type: (struct_type))))]
   ]],
+  struct_field = [[
+    (field_declaration name: (field_identifier) @_name)
+  ]],
   func = [[
     [(function_declaration name: (identifier)       @_name)
-     (method_declaration   name: (field_identifier) @_name)]
+     (method_declaration   name: (field_identifier) @_name)
+     (method_elem          name: (field_identifier) @_name)]
   ]],
   package = [[
     (package_identifier) @_name
@@ -23,12 +27,17 @@ local queries = {
       name: (type_identifier) @_name
       type: (interface_type))
   ]],
+  var = [[
+    [(var_declaration (var_spec name: (identifier) @_name))
+     (short_var_declaration
+       left: (expression_list (identifier) @_name @_var))]
+  ]],
 }
 
 ---@param parent_type string[]
 ---@param node TSNode
 ---@return TSNode?
-local function get_parrent_node(parent_type, node)
+local function get_parent_node(parent_type, node)
   ---@type TSNode?
   local current = node
   while current do
@@ -64,10 +73,11 @@ local function get_captures(query, node, bufnr)
 end
 
 ---@class gopher.TsResult
----@field name string
----@field start integer
----@field end_ integer
----@field is_varstruct boolean
+---@field name string Name of the struct, function, etc
+---@field start integer Line number where the declaration starts
+---@field end_ integer Line number where the declaration ends
+---@field indent integer Number of spaces/tabs in the current cursor line
+---@field is_varstruct boolean Is struct declared as `var S struct{}` or `s := struct{}{}`
 
 ---@param bufnr integer
 ---@param parent_type string[]
@@ -78,23 +88,22 @@ local function do_stuff(bufnr, parent_type, query)
     error "No treesitter parser found for go"
   end
 
-  local node = vim.treesitter.get_node {
-    bufnr = bufnr,
-  }
+  local node = vim.treesitter.get_node { bufnr = bufnr }
   if not node then
-    error "No nodes found under cursor"
+    error "No nodes found under the cursor"
   end
 
-  local parent_node = get_parrent_node(parent_type, node)
+  local parent_node = get_parent_node(parent_type, node)
   if not parent_node then
-    error "No parent node found under cursor"
+    error "No parent node found under the cursor"
   end
 
   local q = vim.treesitter.query.parse("go", query)
   local res = get_captures(q, parent_node, bufnr)
   assert(res.name ~= nil, "No capture name found")
 
-  local start_row, _, end_row, _ = parent_node:range()
+  local start_row, start_col, end_row, _ = parent_node:range()
+  res["indent"] = start_col
   res["start"] = start_row + 1
   res["end_"] = end_row + 1
 
@@ -104,11 +113,12 @@ end
 ---@param bufnr integer
 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
+  --- because in cases like `type ( T struct{}, U struct{} )`
   ---
   --- var_declaration is for cases like `var x struct{}`
   --- short_var_declaration is for cases like `x := struct{}{}`
+  ---
+  --- it always chooses last struct type in the list
   return do_stuff(bufnr, {
     "type_spec",
     "type_declaration",
@@ -117,10 +127,19 @@ function ts.get_struct_under_cursor(bufnr)
   }, queries.struct)
 end
 
+---@param bufnr integer
+function ts.get_struct_field_under_cursor(bufnr)
+  return do_stuff(bufnr, { "field_declaration" }, queries.struct_field)
+end
+
 ---@param bufnr integer
 function ts.get_func_under_cursor(bufnr)
   --- since this handles both and funcs and methods we should check for both parent nodes
-  return do_stuff(bufnr, { "function_declaration", "method_declaration" }, queries.func)
+  return do_stuff(bufnr, {
+    "method_elem",
+    "function_declaration",
+    "method_declaration",
+  }, queries.func)
 end
 
 ---@param bufnr integer
@@ -133,4 +152,9 @@ function ts.get_interface_under_cursor(bufnr)
   return do_stuff(bufnr, { "type_declaration" }, queries.interface)
 end
 
+---@param bufnr integer
+function ts.get_variable_under_cursor(bufnr)
+  return do_stuff(bufnr, { "var_declaration", "short_var_declaration" }, queries.var)
+end
+
 return ts
diff --git a/lua/gopher/comment.lua b/lua/gopher/comment.lua
index 5a91ed1..16aef1f 100644
--- a/lua/gopher/comment.lua
+++ b/lua/gopher/comment.lua
@@ -7,32 +7,40 @@
 
 local ts = require "gopher._utils.ts"
 local log = require "gopher._utils.log"
+local u = require "gopher._utils"
 local comment = {}
 
----@param name string
----@return string
----@dochide
-local function template(name)
-  return "// " .. name .. " "
-end
-
+--- NOTE: The order of functions executed inside this function is IMPORTANT.
+--- This function is extremely fragile; run tests after making any changes.
+---
 ---@param bufnr integer
+---@param line string
 ---@return string
 ---@dochide
-local function generate(bufnr)
+local function generate(bufnr, line)
+  local sf_ok, sf_res = pcall(ts.get_struct_field_under_cursor, bufnr)
+  if sf_ok then
+    return u.indent(line, sf_res.indent) .. "// " .. sf_res.name .. " "
+  end
+
   local s_ok, s_res = pcall(ts.get_struct_under_cursor, bufnr)
   if s_ok then
-    return template(s_res.name)
+    return u.indent(line, s_res.indent) .. "// " .. s_res.name .. " "
+  end
+
+  local v_ok, v_res = pcall(ts.get_variable_under_cursor, bufnr)
+  if v_ok then
+    return u.indent(line, v_res.indent) .. "// " .. v_res.name .. " "
   end
 
   local f_ok, f_res = pcall(ts.get_func_under_cursor, bufnr)
   if f_ok then
-    return template(f_res.name)
+    return u.indent(line, f_res.indent) .. "// " .. f_res.name .. " "
   end
 
   local i_ok, i_res = pcall(ts.get_interface_under_cursor, bufnr)
   if i_ok then
-    return template(i_res.name)
+    return u.indent(line, i_res.indent) .. "// " .. i_res.name .. " "
   end
 
   local p_ok, p_res = pcall(ts.get_package_under_cursor, bufnr)
@@ -45,12 +53,16 @@ end
 
 function comment.comment()
   local bufnr = vim.api.nvim_get_current_buf()
-  local cmt = generate(bufnr)
-  log.debug("generated comment: " .. cmt)
+  local lnum = vim.fn.getcurpos()[2]
+  local line = vim.fn.getline(lnum)
+  local cmt = generate(bufnr, line)
+  log.debug("generated comment:", {
+    comment = cmt,
+    line = line,
+  })
 
-  local pos = vim.fn.getcurpos()[2]
-  vim.fn.append(pos - 1, cmt)
-  vim.fn.setpos(".", { 0, pos, #cmt })
+  vim.fn.append(lnum - 1, cmt)
+  vim.fn.setpos(".", { bufnr, lnum, #cmt })
   vim.cmd "startinsert!"
 end
 
diff --git a/spec/fixtures/comment/interface_many_method_input.go b/spec/fixtures/comment/interface_many_method_input.go
new file mode 100644
index 0000000..e5dca33
--- /dev/null
+++ b/spec/fixtures/comment/interface_many_method_input.go
@@ -0,0 +1,6 @@
+package main
+
+type Testinger interface {
+	Get(id string) int
+	Set(id string, val int)
+}
diff --git a/spec/fixtures/comment/interface_many_method_output.go b/spec/fixtures/comment/interface_many_method_output.go
new file mode 100644
index 0000000..cdeb97c
--- /dev/null
+++ b/spec/fixtures/comment/interface_many_method_output.go
@@ -0,0 +1,7 @@
+package main
+
+type Testinger interface {
+	Get(id string) int
+	// Set 
+	Set(id string, val int)
+}
diff --git a/spec/fixtures/comment/interface_method_input.go b/spec/fixtures/comment/interface_method_input.go
new file mode 100644
index 0000000..7e39a39
--- /dev/null
+++ b/spec/fixtures/comment/interface_method_input.go
@@ -0,0 +1,5 @@
+package main
+
+type Testinger interface {
+	Method(input string) error
+}
diff --git a/spec/fixtures/comment/interface_method_output.go b/spec/fixtures/comment/interface_method_output.go
new file mode 100644
index 0000000..5fa2388
--- /dev/null
+++ b/spec/fixtures/comment/interface_method_output.go
@@ -0,0 +1,6 @@
+package main
+
+type Testinger interface {
+	// Method 
+	Method(input string) error
+}
diff --git a/spec/fixtures/comment/many_structs_fields_input.go b/spec/fixtures/comment/many_structs_fields_input.go
new file mode 100644
index 0000000..f5c6367
--- /dev/null
+++ b/spec/fixtures/comment/many_structs_fields_input.go
@@ -0,0 +1,18 @@
+package main
+
+type (
+	TestOne struct {
+		Asdf string
+		ID   int
+	}
+
+	TestTwo struct {
+		Fesa int
+		A    bool
+	}
+
+	TestThree struct {
+		Asufj int
+		Fs    string
+	}
+)
diff --git a/spec/fixtures/comment/many_structs_fields_output.go b/spec/fixtures/comment/many_structs_fields_output.go
new file mode 100644
index 0000000..648f39f
--- /dev/null
+++ b/spec/fixtures/comment/many_structs_fields_output.go
@@ -0,0 +1,19 @@
+package main
+
+type (
+	TestOne struct {
+		Asdf string
+		ID   int
+	}
+
+	TestTwo struct {
+		// Fesa 
+		Fesa int
+		A    bool
+	}
+
+	TestThree struct {
+		Asufj int
+		Fs    string
+	}
+)
diff --git a/spec/fixtures/comment/struct_fields_input.go b/spec/fixtures/comment/struct_fields_input.go
new file mode 100644
index 0000000..6d6da38
--- /dev/null
+++ b/spec/fixtures/comment/struct_fields_input.go
@@ -0,0 +1,7 @@
+package main
+
+type CommentStruct struct {
+	Name    string
+	Address string
+	Aliases []string
+}
diff --git a/spec/fixtures/comment/struct_fields_output.go b/spec/fixtures/comment/struct_fields_output.go
new file mode 100644
index 0000000..06e6483
--- /dev/null
+++ b/spec/fixtures/comment/struct_fields_output.go
@@ -0,0 +1,8 @@
+package main
+
+type CommentStruct struct {
+	Name    string
+	// Address 
+	Address string
+	Aliases []string
+}
diff --git a/spec/fixtures/comment/svar_input.go b/spec/fixtures/comment/svar_input.go
new file mode 100644
index 0000000..fb75949
--- /dev/null
+++ b/spec/fixtures/comment/svar_input.go
@@ -0,0 +1,5 @@
+package main
+
+func varTest() {
+	s := "something"
+}
diff --git a/spec/fixtures/comment/svar_output.go b/spec/fixtures/comment/svar_output.go
new file mode 100644
index 0000000..e8c917a
--- /dev/null
+++ b/spec/fixtures/comment/svar_output.go
@@ -0,0 +1,6 @@
+package main
+
+func varTest() {
+	// s 
+	s := "something"
+}
diff --git a/spec/fixtures/comment/var_input.go b/spec/fixtures/comment/var_input.go
new file mode 100644
index 0000000..47f8257
--- /dev/null
+++ b/spec/fixtures/comment/var_input.go
@@ -0,0 +1,5 @@
+package main
+
+func test() {
+	var imAVar string
+}
diff --git a/spec/fixtures/comment/var_output.go b/spec/fixtures/comment/var_output.go
new file mode 100644
index 0000000..e38aba1
--- /dev/null
+++ b/spec/fixtures/comment/var_output.go
@@ -0,0 +1,6 @@
+package main
+
+func test() {
+	// imAVar 
+	var imAVar string
+}
diff --git a/spec/fixtures/comment/var_struct_fields_input.go b/spec/fixtures/comment/var_struct_fields_input.go
new file mode 100644
index 0000000..c686212
--- /dev/null
+++ b/spec/fixtures/comment/var_struct_fields_input.go
@@ -0,0 +1,8 @@
+package main
+
+func main() {
+	var s struct {
+		API string
+		Key string
+	}
+}
diff --git a/spec/fixtures/comment/var_struct_fields_output.go b/spec/fixtures/comment/var_struct_fields_output.go
new file mode 100644
index 0000000..9676923
--- /dev/null
+++ b/spec/fixtures/comment/var_struct_fields_output.go
@@ -0,0 +1,9 @@
+package main
+
+func main() {
+	var s struct {
+		API string
+		// Key 
+		Key string
+	}
+}
diff --git a/spec/integration/comment_test.lua b/spec/integration/comment_test.lua
index fdf39fe..e8f54b1 100644
--- a/spec/integration/comment_test.lua
+++ b/spec/integration/comment_test.lua
@@ -18,6 +18,18 @@ comment["should add comment to struct"] = function()
   do_the_test("struct", { 4, 1 })
 end
 
+comment["should add a comment on struct field"] = function()
+  do_the_test("struct_fields", { 5, 8 })
+end
+
+comment["should add a comment on var struct field"] = function()
+  do_the_test("var_struct_fields", { 6, 4 })
+end
+
+comment["should add a comment on one field of many structs"] = function()
+  do_the_test("many_structs_fields", { 10, 4 })
+end
+
 comment["should add comment to function"] = function()
   do_the_test("func", { 3, 1 })
 end
@@ -30,6 +42,22 @@ comment["should add comment to interface"] = function()
   do_the_test("interface", { 3, 6 })
 end
 
+comment["should add comment on interface method"] = function()
+  do_the_test("interface_method", { 4, 2 })
+end
+
+comment["should add a comment on interface with many method"] = function()
+  do_the_test("interface_many_method", { 5, 2 })
+end
+
+comment["should add a comment on a var"] = function()
+  do_the_test("var", { 4, 2 })
+end
+
+comment["should add a comment on a short declared var"] = function()
+  do_the_test("svar", { 4, 8 })
+end
+
 comment["otherwise should add // above cursor"] = function()
   do_the_test("empty", { 1, 1 })
 end
diff --git a/spec/unit/utils_test.lua b/spec/unit/utils_test.lua
index d4898c2..2d8f0bc 100644
--- a/spec/unit/utils_test.lua
+++ b/spec/unit/utils_test.lua
@@ -22,4 +22,28 @@ utils["should .trimend()"] = function()
   t.eq(u.trimend "  hi   ", "  hi")
 end
 
+utils["should add .indent() spaces"] = function()
+  local u = require "gopher._utils"
+  local line = "    func Test() error {"
+  local indent = 4
+
+  t.eq("    ", u.indent(line, indent))
+end
+
+utils["should add .indent() a tab"] = function()
+  local u = require "gopher._utils"
+  local line = "\tfunc Test() error {"
+  local indent = 1
+
+  t.eq("\t", u.indent(line, indent))
+end
+
+utils["should add .indent() 2 tabs"] = function()
+  local u = require "gopher._utils"
+  local line = "\t\tfunc Test() error {"
+  local indent = 2
+
+  t.eq("\t\t", u.indent(line, indent))
+end
+
 return T