From ca3993c5ad029e668530fa39d606a3659c0603b4 Mon Sep 17 00:00:00 2001
From: WhiteBlackGoose <wbg@angouri.org>
Date: Sat, 11 Mar 2023 17:04:12 +0600
Subject: [PATCH] Use Treesitter for folding (#523)

---
 ftplugin/org.vim                 | 6 +-----
 lua/orgmode/utils/treesitter.lua | 6 ++++--
 queries/org/folds.scm            | 8 ++++++++
 3 files changed, 13 insertions(+), 7 deletions(-)
 create mode 100644 queries/org/folds.scm

diff --git a/ftplugin/org.vim b/ftplugin/org.vim
index 660669886..805e903f8 100644
--- a/ftplugin/org.vim
+++ b/ftplugin/org.vim
@@ -1,10 +1,6 @@
 lua require('orgmode.config'):setup_mappings('org')
 lua require('orgmode.config'):setup_mappings('text_objects')
 
-function! OrgmodeFoldExpr()
-  return luaeval('require("orgmode.org.indent").foldexpr()')
-endfunction
-
 function! OrgmodeFoldText()
   return luaeval('require("orgmode.org.indent").foldtext()')
 endfunction
@@ -20,7 +16,7 @@ endfunction
 setlocal nomodeline
 setlocal fillchars+=fold:\ 
 setlocal foldmethod=expr
-setlocal foldexpr=OrgmodeFoldExpr()
+setlocal foldexpr=nvim_treesitter#foldexpr()
 setlocal foldtext=OrgmodeFoldText()
 setlocal formatexpr=OrgmodeFormatExpr()
 setlocal foldlevel=0
diff --git a/lua/orgmode/utils/treesitter.lua b/lua/orgmode/utils/treesitter.lua
index 098026535..4d620ae85 100644
--- a/lua/orgmode/utils/treesitter.lua
+++ b/lua/orgmode/utils/treesitter.lua
@@ -39,8 +39,10 @@ function M.get_node_at_cursor(cursor, winnr, ignore_injected_langs)
   winnr = winnr or 0
   local buf = vim.api.nvim_win_get_buf(winnr)
   -- TODO: Use only this function when 0.8 is released
-  if vim.treesitter.get_node_at_pos then
-    return vim.treesitter.get_node_at_pos(buf, cursor[1] - 1, cursor[2], {
+  if vim.treesitter.get_node then
+    return vim.treesitter.get_node({
+      bufrn = buf,
+      pos = { cursor[1] - 1, cursor[2] },
       ignore_injections = ignore_injected_langs,
     })
   end
diff --git a/queries/org/folds.scm b/queries/org/folds.scm
new file mode 100644
index 000000000..ebd61aa87
--- /dev/null
+++ b/queries/org/folds.scm
@@ -0,0 +1,8 @@
+([
+  (section)
+  (table)
+  (drawer)
+  (property_drawer)
+  (block)
+  ] @fold
+  (#trim! @fold))