From ced4e476cf56bc882e014bc6ff059802c4408210 Mon Sep 17 00:00:00 2001
From: charburgx <charburgx@gmail.com>
Date: Wed, 17 Aug 2022 16:14:16 -0500
Subject: [PATCH] feat: Add folding

---
 lua/symbols-outline.lua            | 128 +++++++++++++++++++++++------
 lua/symbols-outline/config.lua     |   8 ++
 lua/symbols-outline/folding.lua    |  28 +++++++
 lua/symbols-outline/parser.lua     |  93 ++++++++++++++-------
 lua/symbols-outline/utils/init.lua |  53 ++++++++++++
 lua/symbols-outline/writer.lua     |  25 +++++-
 6 files changed, 281 insertions(+), 54 deletions(-)
 create mode 100644 lua/symbols-outline/folding.lua

diff --git a/lua/symbols-outline.lua b/lua/symbols-outline.lua
index 816ee1b..fedeaa1 100644
--- a/lua/symbols-outline.lua
+++ b/lua/symbols-outline.lua
@@ -5,11 +5,12 @@ local writer = require 'symbols-outline.writer'
 local config = require 'symbols-outline.config'
 local utils = require 'symbols-outline.utils.init'
 local View = require 'symbols-outline.view'
+local folding = require 'symbols-outline.folding'
 
 local M = {}
 
 local function setup_global_autocmd()
-  if config.options.highlight_hovered_item then
+  if config.options.highlight_hovered_item or config.options.auto_unfold_hover then
     vim.api.nvim_create_autocmd('CursorHold', {
       pattern = '*',
       callback = function()
@@ -63,6 +64,15 @@ local function wipe_state()
   M.state = { outline_items = {}, flattened_outline_items = {}, code_win = 0 }
 end
 
+local function _update_lines()
+  M.state.flattened_outline_items = parser.flatten(M.state.outline_items)
+  writer.parse_and_write(M.view.bufnr, M.state.flattened_outline_items)
+end
+
+local function _merge_items(items)
+  utils.merge_items_rec({ children = items }, { children = M.state.outline_items })
+end
+
 local function __refresh()
   if M.view:is_open() then
     local function refresh_handler(response)
@@ -71,12 +81,11 @@ local function __refresh()
       end
 
       local items = parser.parse(response)
+      _merge_items(items)
 
       M.state.code_win = vim.api.nvim_get_current_win()
-      M.state.outline_items = items
-      M.state.flattened_outline_items = parser.flatten(items)
 
-      writer.parse_and_write(M.view.bufnr, M.state.flattened_outline_items)
+      _update_lines()
     end
 
     providers.request_symbols(refresh_handler)
@@ -85,9 +94,13 @@ end
 
 M._refresh = utils.debounce(__refresh, 100)
 
-local function goto_location(change_focus)
+function M._current_node()
   local current_line = vim.api.nvim_win_get_cursor(M.view.winnr)[1]
-  local node = M.state.flattened_outline_items[current_line]
+  return M.state.flattened_outline_items[current_line]
+end
+
+local function goto_location(change_focus)
+  local node = M._current_node()
   vim.api.nvim_win_set_cursor(
     M.state.code_win,
     { node.line + 1, node.character }
@@ -100,6 +113,53 @@ local function goto_location(change_focus)
   end
 end
 
+function M._set_folded(folded, move_cursor, node_index)
+  local node = M.state.flattened_outline_items[node_index] or M._current_node()
+  local changed = (folded ~= folding.is_folded(node))
+
+  if folding.is_foldable(node) and changed then
+    node.folded = folded
+
+    if move_cursor then
+      vim.api.nvim_win_set_cursor(M.view.winnr, { node_index, 0 })
+    end
+
+    _update_lines()
+  elseif node.parent then
+    for i, n in ipairs(M.state.flattened_outline_items) do
+      if n == node.parent then
+        M._set_folded(folded, not node.parent.folded and folded, i)
+      end
+    end
+  end
+end
+
+function M._set_all_folded(folded, nodes)
+  local is_root_exec = not nodes
+  nodes = nodes or M.state.outline_items
+
+  for _, node in ipairs(nodes) do
+    node.folded = folded
+    if node.children then
+      M._set_all_folded(folded, node.children)
+    end
+  end
+
+  _update_lines()
+end
+
+local function _items_dfs(callback, children)
+  children = children or M.state.outline_items
+
+  for _, val in ipairs(children) do
+    callback(val)
+
+    if val.children then
+      _items_dfs(callback, val.children)
+    end
+  end
+end
+
 function M._highlight_current_item(winnr)
   local has_provider = providers.has_provider()
 
@@ -127,26 +187,28 @@ function M._highlight_current_item(winnr)
 
   local hovered_line = vim.api.nvim_win_get_cursor(win)[1] - 1
 
-  local nodes = {}
-  for index, value in ipairs(M.state.flattened_outline_items) do
-    if
-      value.line == hovered_line
-      or (hovered_line > value.range_start and hovered_line < value.range_end)
-    then
-      value.line_in_outline = index
-      table.insert(nodes, value)
+  local leaf_node = nil
+
+  local cb = function(value)
+    value.hovered = nil
+
+    if value.line == hovered_line or (hovered_line > value.range_start and hovered_line < value.range_end) then
+      value.hovered = true
+      leaf_node = value
     end
   end
 
-  -- clear old highlight
-  ui.clear_hover_highlight(M.view.bufnr)
-  for _, value in ipairs(nodes) do
-    ui.add_hover_highlight(
-      M.view.bufnr,
-      value.line_in_outline - 1,
-      value.depth * 2
-    )
-    vim.api.nvim_win_set_cursor(M.view.winnr, { value.line_in_outline, 1 })
+  _items_dfs(cb)
+
+  _update_lines()
+
+  if leaf_node then
+    for index, node in ipairs(M.state.flattened_outline_items) do
+      if node == leaf_node then
+          vim.api.nvim_win_set_cursor(M.view.winnr, { index, 1 })
+          break
+      end
+    end
   end
 end
 
@@ -188,6 +250,26 @@ local function setup_keymaps(bufnr)
   map(config.options.keymaps.close, function()
     M.view:close()
   end)
+  -- fold selection
+  map( config.options.keymaps.fold, function ()
+    M._set_folded(true)
+  end)
+  -- unfold selection
+  map(config.options.keymaps.unfold, function ()
+    M._set_folded(false)
+  end)
+  -- fold all
+  map(config.options.keymaps.fold_all, function ()
+    M._set_all_folded(true)
+  end)
+  -- unfold all
+  map(config.options.keymaps.unfold_all, function ()
+    M._set_all_folded(false)
+  end)
+  -- fold reset
+  map(config.options.keymaps.fold_reset,function ()
+    M._set_all_folded(nil)
+  end)
 end
 
 local function handler(response)
diff --git a/lua/symbols-outline/config.lua b/lua/symbols-outline/config.lua
index 5376b91..e5ed735 100644
--- a/lua/symbols-outline/config.lua
+++ b/lua/symbols-outline/config.lua
@@ -16,6 +16,9 @@ M.defaults = {
   show_symbol_details = true,
   preview_bg_highlight = 'Pmenu',
   winblend = 0,
+  autofold_depth = nil,
+  auto_unfold_hover = true,
+  fold_markers = { '', '' },
   keymaps = { -- These keymaps can be a string or a table for multiple keys
     close = { '<Esc>', 'q' },
     goto_location = '<Cr>',
@@ -25,6 +28,11 @@ M.defaults = {
     rename_symbol = 'r',
     code_actions = 'a',
     show_help = '?',
+    fold = 'h',
+    unfold = 'l',
+    fold_all = 'W',
+    unfold_all = 'E',
+    fold_reset = 'R',
   },
   lsp_blacklist = {},
   symbol_blacklist = {},
diff --git a/lua/symbols-outline/folding.lua b/lua/symbols-outline/folding.lua
new file mode 100644
index 0000000..d2ce275
--- /dev/null
+++ b/lua/symbols-outline/folding.lua
@@ -0,0 +1,28 @@
+local M = {}
+local config = require 'symbols-outline.config'
+
+M.is_foldable = function(node)
+  return node.children and #node.children > 0
+end
+
+local get_default_folded = function(depth)
+  local fold_past = config.options.autofold_depth
+  if not fold_past then
+    return false
+  else
+    return depth >= fold_past
+  end
+end
+
+M.is_folded = function(node)
+  if node.folded ~= nil then
+    return node.folded
+  elseif node.hovered and config.options.auto_unfold_hover then
+    return false
+  else
+    return get_default_folded(node.depth)
+  end
+end
+
+return M
+
diff --git a/lua/symbols-outline/parser.lua b/lua/symbols-outline/parser.lua
index 6be870a..59cf924 100644
--- a/lua/symbols-outline/parser.lua
+++ b/lua/symbols-outline/parser.lua
@@ -2,6 +2,7 @@ local symbols = require 'symbols-outline.symbols'
 local ui = require 'symbols-outline.ui'
 local config = require 'symbols-outline.config'
 local t_utils = require 'symbols-outline.utils.table'
+local folding = require 'symbols-outline.folding'
 
 local M = {}
 
@@ -9,8 +10,9 @@ local M = {}
 ---@param result table The result from a language server.
 ---@param depth number? The current depth of the symbol in the hierarchy.
 ---@param hierarchy table? A table of booleans which tells if a symbols parent was the last in its group.
+---@param parent table? A reference to the current symbol's parent in the function's recursion
 ---@return table
-local function parse_result(result, depth, hierarchy)
+local function parse_result(result, depth, hierarchy, parent)
   local ret = {}
 
   for index, value in pairs(result) do
@@ -24,14 +26,6 @@ local function parse_result(result, depth, hierarchy)
       -- whether this node is the last in its group
       local isLast = index == #result
 
-      local children = nil
-      if value.children ~= nil then
-        -- copy by value because we dont want it messing with the hir table
-        local child_hir = t_utils.array_copy(hir)
-        table.insert(child_hir, isLast)
-        children = parse_result(value.children, level + 1, child_hir)
-      end
-
       -- support SymbolInformation[]
       -- https://microsoft.github.io/language-server-protocol/specification#textDocument_documentSymbol
       local selectionRange = value.selectionRange
@@ -44,7 +38,7 @@ local function parse_result(result, depth, hierarchy)
         range = value.location.range
       end
 
-      table.insert(ret, {
+      local node = {
         deprecated = value.deprecated,
         kind = value.kind,
         icon = symbols.icon_from_kind(value.kind),
@@ -54,11 +48,24 @@ local function parse_result(result, depth, hierarchy)
         character = selectionRange.start.character,
         range_start = range.start.line,
         range_end = range['end'].line,
-        children = children,
         depth = level,
         isLast = isLast,
         hierarchy = hir,
-      })
+        parent = parent,
+      }
+
+      table.insert(ret, node)
+
+      local children = nil
+      if value.children ~= nil then
+        -- copy by value because we dont want it messing with the hir table
+        local child_hir = t_utils.array_copy(hir)
+        table.insert(child_hir, isLast)
+        children = parse_result(value.children, level + 1, child_hir, node)
+      end
+
+      node.children = children
+
     end
   end
   return ret
@@ -139,14 +146,23 @@ function M.parse(response)
   return parse_result(sorted, nil, nil)
 end
 
-function M.flatten(outline_items, ret)
+function M.flatten(outline_items, ret, depth)
+  depth = depth or 1
   ret = ret or {}
   for _, value in ipairs(outline_items) do
     table.insert(ret, value)
-    if value.children ~= nil then
-      M.flatten(value.children, ret)
+    value.line_in_outline = #ret
+    if value.children ~= nil and not folding.is_folded(value) then
+      M.flatten(value.children, ret, depth + 1)
     end
   end
+
+  -- if depth == 1 then
+  --   for index, value in ipairs(ret) do
+  --     value.line_in_outline = index
+  --   end
+  -- end
+
   return ret
 end
 
@@ -155,7 +171,10 @@ function M.get_lines(flattened_outline_items)
   local hl_info = {}
 
   for node_line, node in ipairs(flattened_outline_items) do
-    local line = t_utils.str_to_table(string.rep(' ', node.depth))
+    local depth = node.depth
+    local marker_space = (config.options.fold_markers and 1) or 0
+
+    local line = t_utils.str_to_table(string.rep(' ', depth + marker_space))
     local running_length = 1
 
     local function add_guide_hl(from, to)
@@ -176,27 +195,43 @@ function M.get_lines(flattened_outline_items)
           -- i f index is last, add a bottom marker if current item is last,
           -- else add a middle marker
         elseif index == #line then
-          if node.isLast then
-            line[index] = ui.markers.bottom
+          -- add fold markers
+          if config.options.fold_markers and folding.is_foldable(node) then
+            if folding.is_folded(node) then
+              line[index] = config.options.fold_markers[1]
+            else
+              line[index] = config.options.fold_markers[2]
+            end
+
             add_guide_hl(
               running_length,
-              running_length + vim.fn.strlen(ui.markers.bottom) - 1
-            )
-          else
-            line[index] = ui.markers.middle
-            add_guide_hl(
-              running_length,
-              running_length + vim.fn.strlen(ui.markers.middle) - 1
+              running_length + vim.fn.strlen(line[index]) - 1
             )
+
+            -- the root level has no vertical markers
+          elseif depth > 1 then
+            if node.isLast then
+              line[index] = ui.markers.bottom
+              add_guide_hl(
+                running_length,
+                running_length + vim.fn.strlen(ui.markers.bottom) - 1
+              )
+            else
+              line[index] = ui.markers.middle
+              add_guide_hl(
+                running_length,
+                running_length + vim.fn.strlen(ui.markers.middle) - 1
+              )
+            end
           end
           -- else if the parent was not the last in its group, add a
           -- vertical marker because there are items under us and we need
           -- to point to those
-        elseif not node.hierarchy[index] then
-          line[index] = ui.markers.vertical
+        elseif not node.hierarchy[index] and depth > 1 then
+          line[index + marker_space] = ui.markers.vertical
           add_guide_hl(
-            running_length,
-            running_length + vim.fn.strlen(ui.markers.vertical) - 1
+            running_length - 1 + 2*marker_space,
+            running_length + vim.fn.strlen(ui.markers.vertical) - 1 + 2 * marker_space
           )
         end
       end
diff --git a/lua/symbols-outline/utils/init.lua b/lua/symbols-outline/utils/init.lua
index 29465dd..879618c 100644
--- a/lua/symbols-outline/utils/init.lua
+++ b/lua/symbols-outline/utils/init.lua
@@ -38,4 +38,57 @@ function M.debounce(f, delay)
   end
 end
 
+---Merges a symbol tree recursively, only replacing nodes
+---which have changed. This will maintain the folding
+---status of any unchanged nodes.
+---@param new_node table New node
+---@param old_node table Old node
+---@param index? number Index of old_item in parent
+---@param parent? table Parent of old_item
+M.merge_items_rec = function(new_node, old_node, index, parent)
+  local failed = false
+
+  if not new_node or not old_node then
+    failed = true
+  else
+    for key, _ in pairs(new_node) do
+      if vim.tbl_contains({ 'parent', 'children', 'folded', 'hovered', 'line_in_outline', 'hierarchy' }, key) then
+        goto continue
+      end
+
+      if key == 'name' then
+        -- in the case of a rename, just rename the existing node
+        old_node['name'] = new_node['name']
+      else
+        if not vim.deep_equal(new_node[key], old_node[key]) then
+          failed = true
+          break
+        end
+      end
+
+      ::continue::
+    end
+  end
+
+  if failed then
+    if parent and index then
+      parent[index] = new_node
+    end
+  else
+    local next_new_item = new_node.children or {}
+
+    -- in case new children are created on a node which
+    -- previously had no children
+    if #next_new_item > 0 and not old_node.children then
+      old_node.children = {}
+    end
+
+    local next_old_item = old_node.children or {}
+
+    for i = 1, math.max(#next_new_item, #next_old_item) do
+      M.merge_items_rec(next_new_item[i], next_old_item[i], i, next_old_item)
+    end
+  end
+end
+
 return M
diff --git a/lua/symbols-outline/writer.lua b/lua/symbols-outline/writer.lua
index 1849a89..13df583 100644
--- a/lua/symbols-outline/writer.lua
+++ b/lua/symbols-outline/writer.lua
@@ -1,5 +1,6 @@
 local parser = require 'symbols-outline.parser'
 local config = require 'symbols-outline.config'
+local ui = require 'symbols-outline.ui'
 
 local M = {}
 
@@ -21,7 +22,7 @@ function M.write_outline(bufnr, lines)
   vim.api.nvim_buf_set_option(bufnr, 'modifiable', false)
 end
 
-function M.add_highlights(bufnr, hl_info)
+function M.add_highlights(bufnr, hl_info, nodes)
   for _, line_hl in ipairs(hl_info) do
     local line, hl_start, hl_end, hl_type = unpack(line_hl)
     vim.api.nvim_buf_add_highlight(
@@ -33,6 +34,8 @@ function M.add_highlights(bufnr, hl_info)
       hl_end
     )
   end
+
+  M.add_hover_highlights(bufnr, nodes)
 end
 
 local ns = vim.api.nvim_create_namespace 'symbols-outline-virt-text'
@@ -58,6 +61,24 @@ local function clear_virt_text(bufnr)
   vim.api.nvim_buf_clear_namespace(bufnr, -1, 0, -1)
 end
 
+M.add_hover_highlights = function(bufnr, nodes)
+  if not config.options.highlight_hovered_item then
+    return
+  end
+
+  -- clear old highlight
+  ui.clear_hover_highlight(bufnr)
+  for _, node in ipairs(nodes) do
+    if not node.hovered then
+      goto continue
+    end
+    
+    local marker_fac = (config.options.fold_markers and 1) or 0
+    ui.add_hover_highlight(bufnr, node.line_in_outline - 1, (node.depth + marker_fac) * 2)
+    ::continue::
+  end
+end
+
 -- runs the whole writing routine where the text is cleared, new data is parsed
 -- and then written
 function M.parse_and_write(bufnr, flattened_outline_items)
@@ -66,7 +87,7 @@ function M.parse_and_write(bufnr, flattened_outline_items)
 
   clear_virt_text(bufnr)
   local details = parser.get_details(flattened_outline_items)
-  M.add_highlights(bufnr, hl_info)
+  M.add_highlights(bufnr, hl_info, flattened_outline_items)
   M.write_details(bufnr, details)
 end