diff options
| -rw-r--r-- | autoload/nvim_treesitter.vim | 4 | ||||
| -rw-r--r-- | lua/nvim-treesitter/fold.lua | 29 |
2 files changed, 33 insertions, 0 deletions
diff --git a/autoload/nvim_treesitter.vim b/autoload/nvim_treesitter.vim index 715befe0f..65986a8bb 100644 --- a/autoload/nvim_treesitter.vim +++ b/autoload/nvim_treesitter.vim @@ -1,3 +1,7 @@ function! nvim_treesitter#statusline(len) return luaeval("require'nvim-treesitter'.statusline(_A)", a:len) endfunction + +function! nvim_treesitter#foldexpr() + return luaeval(printf('require"nvim-treesitter.fold".get_fold_indic(%d)', v:lnum)) +endfunction diff --git a/lua/nvim-treesitter/fold.lua b/lua/nvim-treesitter/fold.lua new file mode 100644 index 000000000..6a39461a7 --- /dev/null +++ b/lua/nvim-treesitter/fold.lua @@ -0,0 +1,29 @@ +local api = vim.api +local parsers = require'nvim-treesitter.parsers' + +local M = {} + +function M.get_fold_indic(lnum) + if not parsers.has_parser() or not lnum then return '0' end + + local function smallest_multiline_containing(node, level) + for index = 0,(node:named_child_count() -1) do + local child = node:named_child(index) + local start, _, stop, _ = child:range() + + if start ~= stop and start <= (lnum -1) and stop >= (lnum -1) then + return smallest_multiline_containing(child, level + 1) + end + end + + return node, level + end + + local parser = parsers.get_parser() + + local multiline_here, level = smallest_multiline_containing(parser:parse():root(), 0) + + return tostring(level) +end + +return M |