all repos

gopher.nvim @ 27ba078f14dd39fcd5e8a57d7009cdd77450ebbb

Minimalistic plugin for Go development

gopher.nvim/lua/gopher/_utils/ts.lua(view raw)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
local ts = {}
local queries = {
  struct = [[
    [(type_spec name: (type_identifier) @_name
                type: (struct_type))
     (var_declaration (var_spec
                        name: (identifier) @_name @_var
                        type: (struct_type)))
     (short_var_declaration
       left:  (expression_list (identifier) @_name @_var)
       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_elem          name: (field_identifier) @_name)]
  ]],
  package = [[
    (package_identifier) @_name
  ]],
  interface = [[
    (type_spec
      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_parent_node(parent_type, node)
  ---@type TSNode?
  local current = node
  while current do
    if vim.tbl_contains(parent_type, current:type()) then
      break
    end

    current = current:parent()
    if current == nil then
      return nil
    end
  end
  return current
end

---@param query vim.treesitter.Query
---@param node TSNode
---@param bufnr integer
---@return {name:string, is_varstruct:boolean}
local function get_captures(query, node, bufnr)
  local res = {}
  for id, _node in query:iter_captures(node, bufnr) do
    if query.captures[id] == "_name" then
      res["name"] = vim.treesitter.get_node_text(_node, bufnr)
    end

    if query.captures[id] == "_var" then
      res["is_varstruct"] = true
    end
  end

  return res
end

---@class gopher.TsResult
---@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[]
---@param query string
---@return gopher.TsResult
local function do_stuff(bufnr, parent_type, query)
  if not vim.treesitter.get_parser(bufnr, "go") then
    error "No treesitter parser found for go"
  end

  local node = vim.treesitter.get_node { bufnr = bufnr }
  if not node then
    error "No nodes found under the cursor"
  end

  local parent_node = get_parent_node(parent_type, node)
  if not parent_node then
    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, start_col, end_row, _ = parent_node:range()
  res["indent"] = start_col
  res["start"] = start_row + 1
  res["end_"] = end_row + 1

  return res
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 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",
    "var_declaration",
    "short_var_declaration",
  }, 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, {
    "method_elem",
    "function_declaration",
    "method_declaration",
  }, queries.func)
end

---@param bufnr integer
function ts.get_package_under_cursor(bufnr)
  return do_stuff(bufnr, { "package_clause" }, queries.package)
end

---@param bufnr integer
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