Feature: code folding

README.md

@@ -103,6 +103,16 @@ vim.g.aerial = {
   -- Set to 0 or false to disable
   highlight_on_jump = 300,
 
+  -- Fold code when folding the tree. Only works when manage_folds is enabled
+  link_tree_to_folds = true,
+
+  -- Fold the tree when folding code. Only works when manage_folds is enabled
+  link_folds_to_tree = false,
+
+  -- Use symbol tree for folding. Set to true or false to enable/disable
+  -- 'auto' will manage folds if your previous foldmethod was 'manual'
+  manage_folds = 'auto',
+
   -- The maximum width of the aerial window
   max_width = 40,
 
@@ -124,13 +134,13 @@ vim.g.aerial = {
   -- If open_automatic is true, only open aerial if there are at least this many symbols
   open_automatic_min_symbols = 0,
 
-  -- Run this command after jumping to a symbol ('' will disable)
-  post_jump_cmd = 'normal! zvzz',
+  -- Run this command after jumping to a symbol (false will disable)
+  post_jump_cmd = 'normal! zz',
 
   -- Set to false to not update the symbols when there are LSP errors
   update_when_errors = true,
 
-  -- A list of all symbols to display. Set to false to show all symbols.
+  -- A list of all symbols to display. Set to false to display all symbols.
   filter_kind = {
     "Class",
     "Constructor",
@@ -199,8 +209,15 @@ Key     | Command
 `[[`    | Jump to the previous symbol at the same tree level
 `]]`    | Jump to the next symbol at the same tree level
 `q`     | Close the aerial window
-`o`     | Toggle the tree
-`O`     | Toggle the tree recursively
+`o`     | Toggle the symbol under the cursor open/closed
+`O`     | Recursive toggle the symbol under the cursor open/closed
+`zx`    | Sync code folding to the tree (useful if they get out of sync)
+`za`    | Toggle the symbol under the cursor open/closed
+`zA`    | Recursive toggle the symbol under the cursor open/closed
+`zo`    | Open the symbol under the cursor
+`zO`    | Recursive open the symbol under the cursor
+`zc`    | Close the symbol under the cursor
+`zC`    | Recursive close the symbol under the cursor
 
 ## Highlight
 

autoload/aerial.vim

@@ -0,0 +1,4 @@
+
+function! aerial#foldexpr() abort
+	return luaeval('require"aerial.fold".foldexpr(_A)', v:lnum)
+endfunction

doc/aerial.txt

@@ -109,6 +109,22 @@ g:aerial_highlight_on_jump                           *g:aerial_highlight_on_jump
     300). It can also be set to `v:true` for the default, or `v:false` to
     disable the highlight.
 
+g:aerial_link_folds_to_tree                         *g:aerial_link_folds_to_tree*
+    When you fold code with |za|, |zo|, or |zc|, update the tree as well.
+    Requires |g:aerial_manage_folds| to be enabled. Default `false`.
+
+g:aerial_link_tree_to_folds                         *g:aerial_link_tree_to_folds*
+    Update your code folds when you open/collapse symbols in the tree.
+    Requires |g:aerial_manage_folds| to be enabled. Default `true`.
+
+g:aerial_manage_folds                                     *g:aerial_manage_folds*
+    If `true`, will automatically configure your windows to use the symbols
+    tree for code folding. This is equivalent to setting 'foldmethod'=expr and
+    'foldexpr'=aerial#foldexpr() (which you can set manually if you prefer).
+        `true`    Use aerial's foldexpr
+        `false`   Do not modify fold settings
+        "auto"  Manage folds if your previous 'foldmethod' was "manual" (default)
+
 g:aerial_max_width                                           *g:aerial_max_width*
     The maximum width of the aerial window. Default 40.
 
@@ -252,6 +268,11 @@ aerial.tree_cmd({action}, [{opts}])                           *aerial.tree_cmd()
 
     {opts} can be a table with the following optional keys:
 
+        index    The index of the tree to perform the action on. Defaults to
+                 cursor location.
+        fold     If `false`, do not modify folds regardless of
+                 |g:aerial_link_tree_to_folds| setting. If "other", will only update
+                 folds in other windows. (default `true`)
         recurse  If `true`, perform the action recursively on all children
                  (default `false`)
         bubble   If `true` and current symbol has no children, perform the

doc/tags

@@ -32,6 +32,9 @@ g:aerial_filter_kind	aerial.txt	/*g:aerial_filter_kind*
 g:aerial_highlight_mode	aerial.txt	/*g:aerial_highlight_mode*
 g:aerial_highlight_on_jump	aerial.txt	/*g:aerial_highlight_on_jump*
 g:aerial_icons	aerial.txt	/*g:aerial_icons*
+g:aerial_link_folds_to_tree	aerial.txt	/*g:aerial_link_folds_to_tree*
+g:aerial_link_tree_to_folds	aerial.txt	/*g:aerial_link_tree_to_folds*
+g:aerial_manage_folds	aerial.txt	/*g:aerial_manage_folds*
 g:aerial_max_width	aerial.txt	/*g:aerial_max_width*
 g:aerial_min_width	aerial.txt	/*g:aerial_min_width*
 g:aerial_nerd_font	aerial.txt	/*g:aerial_nerd_font*

ftplugin/aerial.vim

@@ -26,6 +26,14 @@ nnoremap <buffer> q <cmd>AerialClose<CR>
 nnoremap <buffer> o <cmd>lua require"aerial".tree_cmd('toggle')<CR>
 " O toggles the tree recursively
 nnoremap <buffer> O <cmd>lua require"aerial".tree_cmd('toggle', {recurse=true})<CR>
-
 " l opens the tree and moves the cursor
 nnoremap <buffer> l l<cmd>lua require"aerial".tree_cmd('open', {bubble=false})<CR>
+" zx will sync the folds for all open windows
+nnoremap <buffer> zx <cmd>lua require"aerial".sync_folds()<CR>
+" make za/o/c function the same as for code folds
+nnoremap <buffer> za <cmd>lua require"aerial".tree_cmd('toggle')<CR>
+nnoremap <buffer> zA <cmd>lua require"aerial".tree_cmd('toggle', {recurse=true})<CR>
+nnoremap <buffer> zo <cmd>lua require"aerial".tree_cmd('open')<CR>
+nnoremap <buffer> zO <cmd>lua require"aerial".tree_cmd('open', {recurse=true})<CR>
+nnoremap <buffer> zc <cmd>lua require"aerial".tree_cmd('close')<CR>
+nnoremap <buffer> zC <cmd>lua require"aerial".tree_cmd('close', {recurse=true})<CR>

lua/aerial.lua

@@ -1,8 +1,10 @@
 local callbacks = require 'aerial.callbacks'
 local config = require 'aerial.config'
 local data = require 'aerial.data'
+local fold = require 'aerial.fold'
 local nav = require 'aerial.navigation'
 local render = require 'aerial.render'
+local util = require 'aerial.util'
 local window = require 'aerial.window'
 
 local M = {}
@@ -66,6 +68,20 @@ M.on_attach = function(client, opts)
   end
   vim.lsp.handlers['textDocument/documentSymbol'] = new_callback
 
+  if config.link_folds_to_tree then
+    local function map(key, cmd)
+      vim.api.nvim_buf_set_keymap(0, 'n', key, cmd, {noremap=true})
+    end
+
+    map('za', [[<cmd>lua require'aerial'.tree_cmd('toggle', {fold='other'})<CR>za]])
+    map('zA', [[<cmd>lua require'aerial'.tree_cmd('toggle', {fold='other'})<CR>zA]])
+    map('zo', [[<cmd>lua require'aerial'.tree_cmd('open', {fold='other'})<CR>zo]])
+    map('zO', [[<cmd>lua require'aerial'.tree_cmd('open', {fold='other'})<CR>zO]])
+    map('zc', [[<cmd>lua require'aerial'.tree_cmd('close', {fold='other'})<CR>zc]])
+    map('zC', [[<cmd>lua require'aerial'.tree_cmd('close', {fold='other'})<CR>zC]])
+  end
+
+
   if config.diagnostics_trigger_update then
     vim.cmd("autocmd User LspDiagnosticsChanged lua require'aerial.autocommands'.on_diagnostics_changed()")
   end
@@ -84,16 +100,52 @@ M.on_attach = function(client, opts)
 end
 
 M.tree_cmd = function(action, opts)
-  local did_update, row = data[0]:action(action, opts)
+  opts = vim.tbl_extend('keep', opts or {}, {
+    index = nil,
+    fold = true,
+  })
+  local index
+  if opts.index then
+    index = opts.index
+  elseif util.is_aerial_buffer() then
+    index = vim.api.nvim_win_get_cursor(0)[1]
+  else
+    index = window.get_position_in_win().lnum
+  end
+  local lnum = data[0]:item(index).lnum
+  local did_update, row = data[0]:action(index, action, opts)
   if did_update then
+    if config.link_tree_to_folds and opts.fold then
+      fold.fold_action(action, lnum, {
+        recurse = opts.recurse,
+        exclude_self = opts.fold == 'other',
+      })
+    end
     render.update_aerial_buffer()
     window.update_all_positions()
-    if row then
-      vim.api.nvim_win_set_cursor(0, {row, 0})
+    if util.is_aerial_buffer() then
+      if row then
+        vim.api.nvim_win_set_cursor(0, {row, 0})
+      end
+    else
+      window.update_position(0, true)
     end
   end
 end
 
+M.sync_folds = function()
+  local mywin = vim.api.nvim_get_current_win()
+  if util.is_aerial_buffer() then
+    local bufnr = util.get_source_buffer()
+    for _,winid in ipairs(util.get_fixed_wins(bufnr)) do
+      fold.sync_tree_folds(winid)
+    end
+  else
+    fold.sync_tree_folds(mywin)
+  end
+  util.go_win_no_au(mywin)
+end
+
 -- @deprecated
 M.set_open_automatic = function(ft_or_mapping, bool)
   local opts = vim.g.aerial or {}

lua/aerial/autocommands.lua

@@ -1,6 +1,7 @@
 -- Functions that are called in response to autocommands
 local config = require 'aerial.config'
 local data = require 'aerial.data'
+local fold = require 'aerial.fold'
 local util = require 'aerial.util'
 local render = require 'aerial.render'
 local window = require 'aerial.window'
@@ -27,10 +28,12 @@ M.on_enter_buffer = function()
 
   -- We only care if we enter an LSP-enabled buffer or an aerial buffer
   if vim.tbl_isempty(vim.lsp.buf_get_clients()) and not util.is_aerial_buffer(mybuf) then
+    fold.restore_foldmethod()
     close_orphans()
     return
   end
 
+  fold.maybe_set_foldmethod()
   if util.is_aerial_buffer(mybuf) then
     if (config.close_behavior ~= 'persist' and util.is_aerial_buffer_orphaned(mybuf))
       or vim.tbl_count(vim.api.nvim_list_wins()) == 1 then

lua/aerial/callbacks.lua

@@ -1,6 +1,6 @@
 local config = require 'aerial.config'
 local data = require 'aerial.data'
-local nav = require 'aerial.navigation'
+local fold = require 'aerial.fold'
 local protocol = require 'vim.lsp.protocol'
 local render = require 'aerial.render'
 local window = require 'aerial.window'
@@ -60,8 +60,11 @@ local function handle_symbols(result, bufnr)
 
   render.update_aerial_buffer(bufnr)
   window.update_all_positions(bufnr)
-  if not had_symbols and bufnr == vim.api.nvim_get_current_buf() then
-    window.maybe_open_automatic()
+  if not had_symbols then
+    fold.maybe_set_foldmethod(bufnr)
+    if bufnr == vim.api.nvim_get_current_buf() then
+      window.maybe_open_automatic()
+    end
   end
 end
 

lua/aerial/config.lua

@@ -5,7 +5,7 @@ local has_devicons = pcall(require, 'nvim-web-devicons')
 local default_options = {
   -- Enum: persist, close, auto
   --   persist - aerial window will stay open until closed
-  --   close   - aerial window will close when original file is left
+  --   close   - aerial window will close when original file is no longer visible
   --   auto    - aerial window will stay open as long as there is a visible
   --             buffer to attach to
   close_behavior = 'auto',
@@ -31,6 +31,16 @@ local default_options = {
   -- Set to 0 or false to disable
   highlight_on_jump = 300,
 
+  -- Fold code when folding the tree. Only works when manage_folds is enabled
+  link_tree_to_folds = true,
+
+  -- Fold the tree when folding code. Only works when manage_folds is enabled
+  link_folds_to_tree = false,
+
+  -- Use symbol tree for folding. Set to true or false to enable/disable
+  -- 'auto' will manage folds if your previous foldmethod was 'manual'
+  manage_folds = 'auto',
+
   -- The maximum width of the aerial window
   max_width = 40,
 
@@ -52,13 +62,13 @@ local default_options = {
   -- If open_automatic is true, only open aerial if there are at least this many symbols
   open_automatic_min_symbols = 0,
 
-  -- Run this command after jumping to a symbol ('' will disable)
-  post_jump_cmd = 'normal! zvzz',
+  -- Run this command after jumping to a symbol (false will disable)
+  post_jump_cmd = 'normal! zz',
 
   -- Set to false to not update the symbols when there are LSP errors
   update_when_errors = true,
 
-  -- A list of all symbols to display
+  -- A list of all symbols to display. Set to false to display all symbols.
   filter_kind = {
     "Class",
     "Constructor",
@@ -70,13 +80,23 @@ local default_options = {
   },
 }
 
+-- config options that are valid as bools, but don't have bools as the default
+local addl_bool_opts = {
+  highlight_mode = true,
+  highlight_on_jump = true,
+  manage_folds = true,
+  nerd_font = true,
+  post_jump_cmd = true,
+  filter_kind = true,
+}
+
 local function get_option(opt)
   local ret = vim.g[string.format('aerial_%s', opt)]
   if ret == nil then
     ret = (vim.g.aerial or {})[opt]
   end
   -- People are used to using 0 for v:false in vimscript
-  if ret == 0 and (default_options[opt] == true or default_options[opt] == false) then
+  if ret == 0 and (type(default_options[opt]) == 'boolean' or addl_bool_opts[opt]) then
     ret = false
   end
   if ret == nil then

lua/aerial/data.lua

@@ -44,35 +44,43 @@ local BufData = {
     return key
   end,
 
-  _is_collapsed = function(self, item)
+  is_collapsed = function(self, item)
     local key = self:_get_item_key(item)
     return self.collapsed[key]
   end,
 
-  _is_collapsable = function(_, item)
-      return item.children and not vim.tbl_isempty(item.children)
+  is_collapsable = function(_, item)
+      return item.level == 0 or (item.children and not vim.tbl_isempty(item.children))
   end,
 
   _get_config = function(self, item)
     return {
-      collapsed = self:_is_collapsed(item),
-      has_children = self:_is_collapsable(item),
+      collapsed = self:is_collapsed(item),
+      has_children = self:is_collapsable(item),
     }
   end,
 
-  action = function(self, action, opts)
+  _get_target = function(self, action, item, bubble)
+    if not bubble then
+      return item
+    end
+    while item and
+      (not self:is_collapsable(item)
+       or (action == 'close' and self:is_collapsed(item))) do
+      item = item.parent
+    end
+    return item
+  end,
+
+  action = function(self, index, action, opts)
     opts = vim.tbl_extend('keep', opts or {}, {
-      recurse = false,
       bubble = true,
+      recurse = false,
     })
     local did_update = false
     local function do_action(item, bubble)
-      if bubble then
-        while item and not self:_is_collapsable(item) do
-          item = item.parent
-        end
-      end
-      if not item or not self:_is_collapsable(item) then
+      item = self:_get_target(action, item, bubble)
+      if not item or not self:is_collapsable(item) then
         return
       end
       local key = self:_get_item_key(item)
@@ -98,8 +106,7 @@ local BufData = {
         error(string.format("Unknown action '%s'", action))
       end
     end
-    local cursor = vim.api.nvim_win_get_cursor(0)
-    local current_item = self:item(cursor[1])
+    local current_item = self:item(index)
     local item = do_action(current_item, opts.bubble)
     return did_update, self:indexof(item)
   end,
@@ -112,7 +119,7 @@ local BufData = {
       local ret = callback(item, self:_get_config(item))
       if ret then return ret end
       if item.children
-        and (opts.incl_hidden or not self:_is_collapsed(item)) then
+        and (opts.incl_hidden or not self:is_collapsed(item)) then
         for _,child in ipairs(item.children) do
           ret = visit_item(child)
           if ret then return ret end
@@ -128,7 +135,7 @@ local BufData = {
   flatten = function(self, filter, opts)
     local items = {}
     self:visit(function(item)
-      if filter(item) then
+      if not filter or filter(item) then
         table.insert(items, item)
       end
     end, opts)
@@ -160,7 +167,8 @@ function Data:has_symbols(bufnr)
   if not bufnr or bufnr == 0 then
     bufnr = vim.api.nvim_get_current_buf()
   end
-  return rawget(self, bufnr) ~= nil
+  local bufdata = rawget(self, bufnr)
+  return bufdata ~= nil and not vim.tbl_isempty(bufdata.items)
 end
 
 return Data

lua/aerial/fold.lua

@@ -0,0 +1,146 @@
+local config = require 'aerial.config'
+local data = require 'aerial.data'
+local util = require 'aerial.util'
+local M = {}
+
+M.foldexpr = function(lnum, debug)
+  if util.is_aerial_buffer() then
+    return '0'
+  end
+  if not data:has_symbols(0) then
+    return '0'
+  end
+  local bufdata = data[0]
+  local lastItem = {}
+  local foldItem = {level = -1}
+  bufdata:visit(function(item)
+    lastItem = item
+    if item.lnum > lnum then
+      return true
+    elseif bufdata:is_collapsable(item) then
+      foldItem = item
+    end
+  end, {incl_hidden = true})
+  local levelstr = string.format("%d", foldItem.level + 1)
+  if lnum == foldItem.lnum then
+    levelstr = '>' .. levelstr
+  elseif vim.api.nvim_buf_get_lines(0, lnum-1, lnum, true)[1] == '' then
+    levelstr = '-1'
+  end
+
+  if debug then
+    levelstr = string.format("%s %s:%d:%d", levelstr, lastItem.name, lastItem.level, lastItem.lnum)
+  end
+  return levelstr
+end
+
+local prev_fdm = '_aerial_prev_foldmethod'
+local prev_fde = '_aerial_prev_foldexpr'
+M.restore_foldmethod = function()
+  local ok, prev_foldmethod = pcall(vim.api.nvim_win_get_var, 0, prev_fdm)
+  if ok and prev_foldmethod then
+    vim.api.nvim_win_del_var(0, prev_fdm)
+    vim.wo.foldmethod = prev_foldmethod
+  end
+  local ok2, prev_foldexpr = pcall(vim.api.nvim_win_get_var, 0, prev_fde)
+  if ok2 and prev_foldexpr then
+    vim.api.nvim_win_del_var(0, prev_fde)
+    vim.wo.foldexpr = prev_foldexpr
+  end
+end
+
+M.maybe_set_foldmethod = function(bufnr)
+  local manage_folds = config.manage_folds
+  if not manage_folds then
+    return
+  end
+  if not data:has_symbols(bufnr) then
+    return
+  end
+  local winids
+  if bufnr then
+    winids = util.get_fixed_wins(bufnr)
+  else
+    winids = {vim.api.nvim_get_current_win()}
+  end
+  for _,winid in ipairs(winids) do
+    local fdm = vim.api.nvim_win_get_option(winid, 'foldmethod')
+    local fde = vim.api.nvim_win_get_option(winid, 'foldexpr')
+    if manage_folds == true or manage_folds == 1
+      or (manage_folds == 'auto' and fdm == 'manual') then
+      vim.api.nvim_win_set_var(winid, prev_fdm, fdm)
+      vim.api.nvim_win_set_var(winid, prev_fde, fde)
+      vim.api.nvim_win_set_option(winid, 'foldmethod', 'expr')
+      vim.api.nvim_win_set_option(winid, 'foldexpr', 'aerial#foldexpr()')
+    end
+  end
+end
+
+M.sync_tree_folds = function(winid)
+  if not util.is_managing_folds(winid) then
+    return
+  end
+  util.go_win_no_au(winid)
+  local view = vim.fn.winsaveview()
+  vim.cmd('normal! zx')
+  local bufdata = data[0]
+  local items = bufdata:flatten(nil, {incl_hidden = true})
+  table.sort(items, function(a, b)
+    return a.level > b.level
+  end)
+  for _,item in ipairs(items) do
+    if bufdata:is_collapsed(item) then
+      vim.api.nvim_win_set_cursor(0, {item.lnum, 0})
+      vim.cmd('normal! zc')
+    end
+  end
+
+  vim.fn.winrestview(view)
+end
+
+local function win_do_action(winid, action, lnum, recurse)
+  util.go_win_no_au(winid)
+  if vim.fn.foldlevel(lnum) == 0 then
+    M.sync_tree_folds(winid)
+  end
+  if vim.fn.foldlevel(lnum) == 0 then
+    return
+  end
+  local view = vim.fn.winsaveview()
+  vim.api.nvim_win_set_cursor(0, {lnum, 0})
+  local key
+  if action == 'open' then
+    key = 'o'
+  elseif action == 'close' then
+    key = 'c'
+  elseif action == 'toggle' then
+    key = 'a'
+  end
+  if key and recurse then
+    key = string.upper(key)
+  end
+  if key then
+    vim.cmd('normal! z' .. key)
+  end
+  vim.fn.winrestview(view)
+end
+
+M.fold_action = function(action, lnum, opts)
+  opts = vim.tbl_extend('keep', opts or {}, {
+    exclude_self = false,
+    recurse = false,
+  })
+  local my_winid = vim.api.nvim_get_current_win()
+  local wins
+  local bufnr, _ = util.get_buffers()
+  wins = util.get_fixed_wins(bufnr)
+  for _,winid in ipairs(wins) do
+    if util.is_managing_folds(winid)
+      and (not opts.exclude_self or winid ~= my_winid) then
+      win_do_action(winid, action, lnum, opts.recurse)
+    end
+  end
+  util.go_win_no_au(my_winid)
+end
+
+return M

lua/aerial/navigation.lua

@@ -138,8 +138,9 @@ M.select = function(opts)
   local bufnr, _ = util.get_buffers()
   vim.api.nvim_win_set_buf(winid, bufnr)
   vim.api.nvim_win_set_cursor(winid, {item.lnum, item.col})
-  if config.post_jump_cmd ~= '' then
-    vim.fn.win_execute(winid, config.post_jump_cmd, true)
+  local cmd = config.post_jump_cmd
+  if cmd and cmd ~= '' then
+    vim.fn.win_execute(winid, cmd, true)
   end
 
   if opts.jump then

lua/aerial/render.lua

@@ -71,7 +71,8 @@ end
 
 -- Update the highlighted lines in the aerial buffer
 M.update_highlights = function(buf)
-  if config.highlight_mode == 'none' then
+  local hl_mode = config.highlight_mode
+  if not hl_mode or hl_mode == 'none' then
     return
   end
   local bufnr, aer_bufnr = util.get_buffers(buf)
@@ -86,7 +87,6 @@ M.update_highlights = function(buf)
     return
   end
   local hl_width = math.floor(util.get_width(aer_bufnr) / #winids)
-  local hl_mode = config.highlight_mode
 
   if hl_mode == 'last' then
     local row = data[bufnr].last_position

lua/aerial/util.lua

@@ -135,6 +135,10 @@ M.is_floating_win = function(winid)
   return vim.api.nvim_win_get_config(winid).relative ~= ""
 end
 
+M.is_managing_folds = function(winid)
+  return vim.api.nvim_win_get_option(winid or 0, 'foldexpr') == 'aerial#foldexpr()'
+end
+
 M.detect_split_direction = function(bufnr)
   local default = config.default_direction
   if default == 'left' then

lua/aerial/window.lua

@@ -196,7 +196,7 @@ M.update_all_positions = function(bufnr)
 end
 
 M.update_position = function(winid, update_last)
-  if config.highlight_mode == 'none' then
+  if not config.highlight_mode or config.highlight_mode == 'none' then
     return
   end
   if winid == 0 then