Merge pull request #64 from vigoux/feature/folding

Treesitter based code folding.

5 files changed, 47 insertions, 1 deletions

diff --git a/README.md b/README.md
index e87bc2ad..48e9e0de 100644
--- a/README.md
+++ b/README.md
@@ -124,7 +124,7 @@ The goal of `nvim-treesitter` is both to provide a simple and easy way to use th
 but also to add some functionalities to it:
 Some of these features are :
   - [x] Incremental selection
-  - [ ] Syntax based code folding
+  - [x] Syntax based code folding (`set foldmethod=expr foldexpr=nvim_treesitter#foldexr()`)
   - [x] Consistent syntax highlighting (the api is not quite stable yet)
   - [x] Statusline indicator (`require'nvim-treesitter'.statusline(size)`)
 
diff --git a/autoload/nvim_treesitter.vim b/autoload/nvim_treesitter.vim
index 715befe0..65986a8b 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/doc/nvim-treesitter.txt b/doc/nvim-treesitter.txt
index 564d9ee8..85bcbddd 100644
--- a/doc/nvim-treesitter.txt
+++ b/doc/nvim-treesitter.txt
@@ -165,4 +165,16 @@ could be used as a statusline indicator.
 Note: The `size` argument is optionnal. When specified, the string will not be
       longer than `size`.
 
+|nvim_treesitter#foldexpr()|
+						  *nvim_treesitter#foldexpr()*
+
+Functions to be used to determine the fold level at a given line number.
+To use it: >
+  set foldmethod=expr
+  set foldexpr=nvim_treesitter#foldexr()
+<
+Note: This is highly experimental, and folding can break on some types of
+      edits. If you encounter such breakage, hiting `zx` should fix folding.
+      In any case, feel free to open an issue with the reproducing steps.
+
  vim:tw=78:ts=8:noet:ft=help:norl:
diff --git a/doc/tags b/doc/tags
index e58f0d65..0a12b18d 100644
--- a/doc/tags
+++ b/doc/tags
@@ -11,6 +11,7 @@ nvim-treesitter-commands	nvim-treesitter.txt	/*nvim-treesitter-commands*
 nvim-treesitter-functions	nvim-treesitter.txt	/*nvim-treesitter-functions*
 nvim-treesitter-intro	nvim-treesitter.txt	/*nvim-treesitter-intro*
 nvim-treesitter-quickstart	nvim-treesitter.txt	/*nvim-treesitter-quickstart*
+nvim_treesitter#foldexpr()	nvim-treesitter.txt	/*nvim_treesitter#foldexpr()*
 nvim_treesitter#statusline()	nvim-treesitter.txt	/*nvim_treesitter#statusline()*
 ts_api.containing_scope	nvim-treesitter.txt	/*ts_api.containing_scope*
 ts_api.get_named_children	nvim-treesitter.txt	/*ts_api.get_named_children*
diff --git a/lua/nvim-treesitter/fold.lua b/lua/nvim-treesitter/fold.lua
new file mode 100644
index 00000000..6a39461a
--- /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