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
|