feat: add syntax-based folding

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