From 0ce1009c9218ad7983c42da5f09334b033fecab0 Mon Sep 17 00:00:00 2001
From: ray-x <rayx.cn@gmail.com>
Date: Sat, 11 Nov 2023 21:17:29 +1100
Subject: [PATCH] Update config setup for fold_ts

---
 README.md                |  6 +++++-
 lua/navigator.lua        | 16 ++++++++++++----
 lua/navigator/foldts.lua | 17 ++++++++++-------
 3 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 3dbd91d..b3668d9 100644
--- a/README.md
+++ b/README.md
@@ -264,7 +264,11 @@ require'navigator'.setup({
   -- end,
   -- The attach code will apply to all LSP clients
 
-  ts_fold = false,  -- modified version of treesitter folding
+  ts_fold = {
+    enable = false,
+    comment_fold = true, -- fold with comment string
+    max_lines_scan_comments = 20, -- only fold when the fold level higher than this value
+  },  -- modified version of treesitter folding
   default_mapping = true,  -- set to false if you will remap every key or if you using old version of nvim-
   keymaps = {{key = "gK", func = vim.lsp.declaration, desc = 'declaration'}}, -- a list of key maps
   -- this kepmap gK will override "gD" mapping function declaration()  in default kepmap
diff --git a/lua/navigator.lua b/lua/navigator.lua
index 7d860ab..395e875 100644
--- a/lua/navigator.lua
+++ b/lua/navigator.lua
@@ -30,7 +30,12 @@ _NgConfigValues = {
   on_attach = function(client, bufnr)
     -- your on_attach will be called at end of navigator on_attach
   end,
-  ts_fold = false,
+  -- ts_fold = false, -- deprecated
+  ts_fold = {
+    enable = false,
+    comment = true, -- ts fold text object
+    max_lines_scan_comments = 2000, -- maximum lines to scan for comments
+  },
   treesitter_analysis = true, -- treesitter variable context
   treesitter_navigation = true, -- bool|table
   treesitter_analysis_max_num = 100, -- how many items to run treesitter analysis
@@ -236,6 +241,9 @@ M.deprecated = function(cfg)
   if cfg.lsp and cfg.lsp.sumneko_lua then
     warn('sumneko_lua option deprecated, refer to README for more details')
   end
+  if cfg.ts_fold ~= nil and type(cfg.ts_fold) == "boolean" then
+    warn('ts_fold option changed, refer to README for more details')
+  end
 end
 
 local extend_config = function(opts)
@@ -348,9 +356,9 @@ M.setup = function(cfg)
     require('navigator.implementation')
     local ts_installed = pcall(require, 'nvim-treesitter')
     if not ts_installed then
-      if _NgConfigValues.ts_fold == true then
+      if _NgConfigValues.ts_fold.enable == true then
         warn('treesitter not installed ts_fold disabled')
-        _NgConfigValues.ts_fold = false
+        _NgConfigValues.ts_fold.enable = false
       end
       if _NgConfigValues.treesitter_analysis == true then
         warn('nvim-treesitter not installed, disable treesitter_analysis')
@@ -370,7 +378,7 @@ M.setup = function(cfg)
       _NgConfigValues.loaded = true
     end
 
-    if _NgConfigValues.ts_fold == true then
+    if _NgConfigValues.ts_fold.enable == true then
       require('navigator.foldts').on_attach()
     end
 
diff --git a/lua/navigator/foldts.lua b/lua/navigator/foldts.lua
index b2ecfbb..0165a20 100644
--- a/lua/navigator/foldts.lua
+++ b/lua/navigator/foldts.lua
@@ -2,7 +2,6 @@
 
 local log = require('navigator.util').log
 local trace = require('navigator.util').trace
-trace = log
 local api = vim.api
 local tsutils = require('nvim-treesitter.ts_utils')
 local query = require('nvim-treesitter.query')
@@ -52,10 +51,10 @@ function NG_custom_fold_text()
         spaces[2] = { '@keyword' }
       end
     end
-    local sep2 = ' ' .. string.rep(sep, 3)
+    local sep2 = ' ' .. string.rep(sep, 3) .. '  '
     table.insert(line_syntax, { sep2, { '@comment' } })
-    table.insert(line_syntax, { '  ' .. tostring(line_count), { '@number' } })
-    table.insert(line_syntax, { ' lines ', { '@comment' } })
+    table.insert(line_syntax, { tostring(line_count), { '@number' } })
+    table.insert(line_syntax, { ' lines', { '@text.title' } })
     table.insert(line_syntax, { sep2, { '@comment' } })
     return line_syntax
   end
@@ -88,19 +87,23 @@ end
 
 local function is_comment(line_number)
   local node = get_node_at_line(line_number)
-  trace(node, node:type())
+  trace(line_number, node, node:type())
   if not node then
     return false
   end
   local node_type = node:type()
-  trace(node_type)
-  return node_type == 'comment' or node_type == 'comment_block'
+  trace(line_number, node_type)
+  return node_type:find('comment')
 end
 
 local function get_comment_scopes(total_lines)
+  if not _NgConfigValues.ts_fold.comment then
+    return {}
+  end
   local comment_scopes = {}
   local comment_start = nil
 
+  total_lines = math.min(total_lines, _NgConfigValues.ts_fold.max_lines_scan_comments)
   for line = 0, total_lines - 1 do
     if is_comment(line + 1) then
       if not comment_start then