From 583718a62a9661121260326e7e129e4ca38a9576 Mon Sep 17 00:00:00 2001
From: linepi <2428513224@qq.com>
Date: Sun, 28 Jul 2024 14:47:12 +0800
Subject: [PATCH] feat: support collapse a single folder under cursor

---
 doc/nvim-tree-lua.txt                         | 17 +++++++++++--
 .../actions/tree/modifiers/collapse-all.lua   | 25 +++++++++++++++----
 lua/nvim-tree/api.lua                         |  5 ++++
 lua/nvim-tree/commands.lua                    | 13 +++++++++-
 4 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/doc/nvim-tree-lua.txt b/doc/nvim-tree-lua.txt
index 8f03efab87f..9d1abec76a3 100644
--- a/doc/nvim-tree-lua.txt
+++ b/doc/nvim-tree-lua.txt
@@ -339,7 +339,15 @@ See |nvim-tree-highlight| for details.
 
     See |nvim-tree-api.tree.collapse_all()|
 
-    Calls: `api.tree.collapse_all(false)`
+    Calls: api.tree.collapse_all({ under_cursor = false, keep_buffers = false, })
+
+*:NvimTreeCollapseFolder*
+
+    Collapses the folder under cursor
+
+    See |nvim-tree-api.tree.collapse_all()|
+
+    Calls: api.tree.collapse_all({ under_cursor = true, keep_buffers = false, })
 
 *:NvimTreeCollapseKeepBuffers*
 
@@ -1764,10 +1772,14 @@ tree.find_file({opts})                        *nvim-tree-api.tree.find_file()*
 tree.search_node()                          *nvim-tree-api.tree.search_node()*
     Open the search dialogue as per the search_node action.
 
-tree.collapse_all({keep_buffers})          *nvim-tree-api.tree.collapse_all()*
+tree.collapse_all({opts})          *nvim-tree-api.tree.collapse_all()*
     Collapse the tree.
 
     Parameters: ~
+      • {opts} (table) optional parameters
+
+    Parameters: ~
+      • {under_cursor} (boolean) only collapse the node under cursor
       • {keep_buffers} (boolean) do not collapse nodes with open buffers.
 
 tree.expand_all()                            *nvim-tree-api.tree.expand_all()*
@@ -3043,6 +3055,7 @@ highlight group is not, hard linking as follows: >
 |nvim-tree-api.tree.close_in_all_tabs()|
 |nvim-tree-api.tree.close_in_this_tab()|
 |nvim-tree-api.tree.collapse_all()|
+|nvim-tree-api.tree.collapse
 |nvim-tree-api.tree.expand_all()|
 |nvim-tree-api.tree.find_file()|
 |nvim-tree-api.tree.focus()|
diff --git a/lua/nvim-tree/actions/tree/modifiers/collapse-all.lua b/lua/nvim-tree/actions/tree/modifiers/collapse-all.lua
index c6add436f1e..64cbae2068d 100644
--- a/lua/nvim-tree/actions/tree/modifiers/collapse-all.lua
+++ b/lua/nvim-tree/actions/tree/modifiers/collapse-all.lua
@@ -23,18 +23,33 @@ local function buf_match()
   end
 end
 
----@param keep_buffers boolean
-function M.fn(keep_buffers)
-  local node = lib.get_node_at_cursor()
-  local explorer = core.get_explorer()
+---@param opts ApiTreeCollapseOpts|nil
+function M.fn(opts)
+  opts = opts or {}
+  local keep_buffers = opts.keep_buffers or false
+  local under_cursor = opts.under_cursor or false
 
+  local explorer = core.get_explorer()
   if explorer == nil then
     return
   end
 
   local matches = buf_match()
 
-  Iterator.builder(explorer.nodes)
+  local node = lib.get_node_at_cursor()
+
+  local selectedNodes
+  if under_cursor then
+    if node == nil or node.nodes == nil then
+      return
+    end
+    selectedNodes = node.nodes
+    node.open = false
+  else
+    selectedNodes = explorer.nodes
+  end
+
+  Iterator.builder(selectedNodes)
     :hidden()
     :applier(function(n)
       if n.nodes ~= nil then
diff --git a/lua/nvim-tree/api.lua b/lua/nvim-tree/api.lua
index fdc1eb35808..53703138152 100644
--- a/lua/nvim-tree/api.lua
+++ b/lua/nvim-tree/api.lua
@@ -160,6 +160,11 @@ Api.tree.get_nodes = wrap(lib.get_nodes)
 
 Api.tree.find_file = wrap(actions.tree.find_file.fn)
 Api.tree.search_node = wrap(actions.finders.search_node.fn)
+
+---@class ApiTreeCollapseOpts
+---@field under_cursor boolean
+---@field keep_buffers boolean
+
 Api.tree.collapse_all = wrap(actions.tree.modifiers.collapse_all.fn)
 Api.tree.expand_all = wrap_node(actions.tree.modifiers.expand_all.fn)
 Api.tree.toggle_enable_filters = wrap(actions.tree.modifiers.toggles.enable)
diff --git a/lua/nvim-tree/commands.lua b/lua/nvim-tree/commands.lua
index f5a7ac82106..f0131d18df7 100644
--- a/lua/nvim-tree/commands.lua
+++ b/lua/nvim-tree/commands.lua
@@ -1,5 +1,6 @@
 local api = require "nvim-tree.api"
 local view = require "nvim-tree.view"
+local lib = require "nvim-tree.lib"
 
 local M = {}
 
@@ -121,7 +122,17 @@ local CMDS = {
       bar = true,
     },
     command = function()
-      api.tree.collapse_all(false)
+      api.tree.collapse_all({ under_cursor = false, keep_buffers = false})
+    end,
+  },
+  {
+    name = "NvimTreeCollapseFolder",
+    opts = {
+      desc = "nvim-tree: collapse the folder under cursor",
+      bar = true,
+    },
+    command = function()
+      api.tree.collapse_all({ under_cursor = true, keep_buffers = false})
     end,
   },
   {