From 60b5e602e95384e6f6d5ba47975b4903c12e2c59 Mon Sep 17 00:00:00 2001
From: Ali Rezvani <3788964+rzvxa@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:07:25 +0330
Subject: [PATCH] fix: session restore for nerdtree buffers. (#1405)

---
 autoload/nerdtree.vim    | 32 ++++++++++++++++++++++++++++++++
 lib/nerdtree/creator.vim | 24 ++++++++++++++++++------
 plugin/NERD_tree.vim     |  2 +-
 3 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/autoload/nerdtree.vim b/autoload/nerdtree.vim
index 1c10ec8..d7246dc 100644
--- a/autoload/nerdtree.vim
+++ b/autoload/nerdtree.vim
@@ -234,6 +234,38 @@ function! nerdtree#pathEquals(lhs, rhs) abort
     endif
 endfunction
 
+"FUNCTION: nerdtree#onBufLeave() {{{2
+" used for handling the nerdtree BufLeave/WinLeave events.
+function! nerdtree#onBufLeave() abort
+    " detect whether we are in the middle of sourcing a session.
+    " if it is a buffer from the sourced session we need to restore it.
+    if exists('g:SessionLoad') && !exists('b:NERDTree')
+        let bname = bufname('%')
+        " is the buffer for a tab tree?
+        if bname =~# '^' . g:NERDTreeCreator.BufNamePrefix() . 'tab_\d\+$'
+            " rename loaded buffer and mark it as trash to prevent this event
+            " getting fired again
+            exec 'file TRASH_' . bname
+            " delete the trash buffer
+            exec 'bwipeout!'
+            " rescue the tab tree at the current working directory
+            call g:NERDTreeCreator.CreateTabTree(getcwd())
+        " is the buffer for a window tree?
+        elseif bname =~# '^' . g:NERDTreeCreator.BufNamePrefix(). 'win_\d\+$'
+            " rescue the window tree at the current working directory
+            call g:NERDTreeCreator.CreateWindowTree(getcwd())
+        else " unknown buffer type
+            " rename buffer to mark it as broken.
+            exec 'file BROKEN_' . bname
+            call nerdtree#echoError('Failed to restore "' . bname . '" from session. Is this session created with an older version of NERDTree?')
+        endif
+    else
+        if g:NERDTree.IsOpen()
+            call b:NERDTree.ui.saveScreenState()
+        endif
+    endif
+endfunction
+
 " SECTION: View Functions {{{1
 "============================================================
 
diff --git a/lib/nerdtree/creator.vim b/lib/nerdtree/creator.vim
index 7e1940b..e794e0d 100644
--- a/lib/nerdtree/creator.vim
+++ b/lib/nerdtree/creator.vim
@@ -118,7 +118,7 @@ function! s:Creator.createWindowTree(dir)
 
     "we need a unique name for each window tree buffer to ensure they are
     "all independent
-    exec g:NERDTreeCreatePrefix . ' edit ' . self._nextBufferName()
+    exec g:NERDTreeCreatePrefix . ' edit ' . self._nextBufferName('win')
 
     call self._createNERDTree(path, 'window')
     let b:NERDTree._previousBuf = bufnr(previousBuf)
@@ -210,7 +210,7 @@ function! s:Creator._createTreeWin()
     let l:splitSize = g:NERDTreeWinSize
 
     if !g:NERDTree.ExistsForTab()
-        let t:NERDTreeBufName = self._nextBufferName()
+        let t:NERDTreeBufName = self._nextBufferName('tab')
         silent! execute l:splitLocation . l:splitDirection . ' ' . l:splitSize . ' new'
         silent! execute 'edit ' . t:NERDTreeBufName
         silent! execute l:splitDirection . ' resize '. l:splitSize
@@ -244,10 +244,22 @@ function! s:Creator.New()
     return newCreator
 endfunction
 
-" FUNCTION: s:Creator._nextBufferName() {{{1
-" returns the buffer name for the next nerd tree
-function! s:Creator._nextBufferName()
-    let name = s:Creator.BufNamePrefix() . self._nextBufferNumber()
+" FUNCTION: s:Creator._nextBufferName(type='') {{{1
+" gets an optional buffer type of either 'tab' or 'win'.
+" returns the buffer name for the next nerd tree of such type.
+function! s:Creator._nextBufferName(...)
+    if a:0 > 0
+        let type = a:1
+    else
+        let type = ''
+    end
+    let name = s:Creator.BufNamePrefix()
+    if type ==# 'tab'
+        let name = name . 'tab_'
+    elseif type ==# 'win'
+        let name = name . 'win_'
+    endif
+    let name = name . self._nextBufferNumber()
     return name
 endfunction
 
diff --git a/plugin/NERD_tree.vim b/plugin/NERD_tree.vim
index 3e0c9e0..73085d0 100644
--- a/plugin/NERD_tree.vim
+++ b/plugin/NERD_tree.vim
@@ -151,7 +151,7 @@ call nerdtree#ui_glue#setupCommands()
 "============================================================
 augroup NERDTree
     "Save the cursor position whenever we close the nerd tree
-    exec 'autocmd BufLeave,WinLeave '. g:NERDTreeCreator.BufNamePrefix() .'* if g:NERDTree.IsOpen() | call b:NERDTree.ui.saveScreenState() | endif'
+    exec 'autocmd BufLeave,WinLeave '. g:NERDTreeCreator.BufNamePrefix() .'* call nerdtree#onBufLeave()'
 
     "disallow insert mode in the NERDTree
     exec 'autocmd BufEnter,WinEnter '. g:NERDTreeCreator.BufNamePrefix() .'* stopinsert'