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 |
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))))]
]],
func = [[
[(function_declaration name: (identifier) @_name)
(method_declaration name: (field_identifier) @_name)]
]],
package = [[
(package_identifier) @_name
]],
interface = [[
(type_spec
name: (type_identifier) @_name
type: (interface_type))
]],
}
---@param parent_type string[]
---@param node TSNode
---@return TSNode?
local function get_parrent_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
---@field start_line integer
---@field end_line integer
---@field is_varstruct boolean
---@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 cursor"
end
local parent_node = get_parrent_node(parent_type, node)
if not parent_node then
error "No parent node found under 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()
res["start_line"] = start_row + 1
res["end_line"] = 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 strict{} )`
--- i will be choosing always last struct in the list
---
--- var_declaration 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_declaration",
"short_var_declaration",
}, queries.struct)
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)
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
return ts
|