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
|
local api = vim.api
local tsutils = require "nvim-treesitter.ts_utils"
local query = require "nvim-treesitter.query"
local parsers = require "nvim-treesitter.parsers"
local M = {}
-- This is cached on buf tick to avoid computing that multiple times
-- Especially not for every line in the file when `zx` is hit
local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
local max_fold_level = api.nvim_win_get_option(0, "foldnestmax")
local trim_level = function(level)
if level > max_fold_level then
return max_fold_level
end
return level
end
local parser = parsers.get_parser(bufnr)
if not parser then
return {}
end
local matches = query.get_capture_matches_recursively(bufnr, function(lang)
if query.has_folds(lang) then
return "@fold", "folds"
elseif query.has_locals(lang) then
return "@scope", "locals"
end
end)
-- start..stop is an inclusive range
local start_counts = {}
local stop_counts = {}
local prev_start = -1
local prev_stop = -1
local min_fold_lines = api.nvim_win_get_option(0, "foldminlines")
for _, match in ipairs(matches) do
local start, stop, stop_col
if match.metadata and match.metadata.range then
start, _, stop, stop_col = unpack(match.metadata.range)
else
start, _, stop, stop_col = match.node:range()
end
if stop_col == 0 then
stop = stop - 1
end
local fold_length = stop - start + 1
local should_fold = fold_length > min_fold_lines
-- Fold only multiline nodes that are not exactly the same as previously met folds
-- Checking against just the previously found fold is sufficient if nodes
-- are returned in preorder or postorder when traversing tree
if should_fold and not (start == prev_start and stop == prev_stop) then
start_counts[start] = (start_counts[start] or 0) + 1
stop_counts[stop] = (stop_counts[stop] or 0) + 1
prev_start = start
prev_stop = stop
end
end
local levels = {}
local current_level = 0
-- We now have the list of fold opening and closing, fill the gaps and mark where fold start
for lnum = 0, api.nvim_buf_line_count(bufnr) do
local prefix = ""
local last_trimmed_level = trim_level(current_level)
current_level = current_level + (start_counts[lnum] or 0)
local trimmed_level = trim_level(current_level)
current_level = current_level - (stop_counts[lnum] or 0)
local next_trimmed_level = trim_level(current_level)
-- Determine if it's the start/end of a fold
-- NB: vim's fold-expr interface does not have a mechanism to indicate that
-- two (or more) folds start at this line, so it cannot distinguish between
-- ( \n ( \n )) \n (( \n ) \n )
-- versus
-- ( \n ( \n ) \n ( \n ) \n )
-- If it did have such a mechanism, (trimmed_level - last_trimmed_level)
-- would be the correct number of starts to pass on.
if trimmed_level - last_trimmed_level > 0 then
prefix = ">"
elseif trimmed_level - next_trimmed_level > 0 then
-- Ending marks tend to confuse vim more than it helps, particularly when
-- the fold level changes by at least 2; we can uncomment this if
-- vim's behavior gets fixed.
-- prefix = "<"
prefix = ""
end
levels[lnum + 1] = prefix .. tostring(trimmed_level)
end
return levels
end)
function M.get_fold_indic(lnum)
if not parsers.has_parser() or not lnum then
return "0"
end
local buf = api.nvim_get_current_buf()
local levels = folds_levels(buf) or {}
return levels[lnum] or "0"
end
return M
|
