vim-patch:9.0.1686: undotree() only works for the current buffer (#24665)

Problem:    undotree() only works for the current buffer
Solution:   Add an optional "buffer number" parameter to undotree().  If
            omitted, use the current buffer for backwards compatibility.

closes: vim/vim#4001
closes: vim/vim#12292

https://github.com/vim/vim/commit/5fee11114975b7405b7ccd3ee8758e54bf559760

Co-authored-by: Devin J. Pohly <djpohly@gmail.com>

6 files changed, 80 insertions, 23 deletions

diff --git a/runtime/doc/builtin.txt b/runtime/doc/builtin.txt
index e678a8d804..1cb995fd6e 100644
--- a/runtime/doc/builtin.txt
+++ b/runtime/doc/builtin.txt
@@ -8171,9 +8171,10 @@ undofile({name})                                                    *undofile()*
 		buffer without a file name will not write an undo file.
 		Useful in combination with |:wundo| and |:rundo|.
 
-undotree()                                                          *undotree()*
-		Return the current state of the undo tree in a dictionary with
-		the following items:
+undotree([{buf}])                                                   *undotree()*
+		Return the current state of the undo tree for the current
+		buffer, or for a specific buffer if {buf} is given.  The
+		result is a dictionary with the following items:
 		  "seq_last"	The highest undo sequence number used.
 		  "seq_cur"	The sequence number of the current position in
 				the undo tree.  This differs from "seq_last"
diff --git a/runtime/doc/usr_41.txt b/runtime/doc/usr_41.txt
index e1bb6c4af1..e6fcfa2fc6 100644
--- a/runtime/doc/usr_41.txt
+++ b/runtime/doc/usr_41.txt
@@ -1092,7 +1092,7 @@ Various:					*various-functions*
 	libcallnr()		idem, returning a number
 
 	undofile()		get the name of the undo file
-	undotree()		return the state of the undo tree
+	undotree()		return the state of the undo tree for a buffer
 
 	shiftwidth()		effective value of 'shiftwidth'
 
diff --git a/runtime/lua/vim/_meta/vimfn.lua b/runtime/lua/vim/_meta/vimfn.lua
index 14aae99971..bf80ccab45 100644
--- a/runtime/lua/vim/_meta/vimfn.lua
+++ b/runtime/lua/vim/_meta/vimfn.lua
@@ -9733,8 +9733,9 @@ function vim.fn.type(expr) end
 --- @return string
 function vim.fn.undofile(name) end
 
---- Return the current state of the undo tree in a dictionary with
---- the following items:
+--- Return the current state of the undo tree for the current
+--- buffer, or for a specific buffer if {buf} is given.  The
+--- result is a dictionary with the following items:
 ---   "seq_last"  The highest undo sequence number used.
 ---   "seq_cur"  The sequence number of the current position in
 ---     the undo tree.  This differs from "seq_last"
@@ -9775,8 +9776,9 @@ function vim.fn.undofile(name) end
 ---     blocks.  Each item may again have an "alt"
 ---     item.
 ---
+--- @param buf? any
 --- @return any
-function vim.fn.undotree() end
+function vim.fn.undotree(buf) end
 
 --- Remove second and succeeding copies of repeated adjacent
 --- {list} items in-place.  Returns {list}.  If you want a list
diff --git a/src/nvim/eval.lua b/src/nvim/eval.lua
index 4e6688a4a1..a09580cd75 100644
--- a/src/nvim/eval.lua
+++ b/src/nvim/eval.lua
@@ -11620,9 +11620,12 @@ M.funcs = {
     signature = 'undofile({name})',
   },
   undotree = {
+    args = { 0, 1 },
+    base = 1,
     desc = [=[
-      Return the current state of the undo tree in a dictionary with
-      the following items:
+      Return the current state of the undo tree for the current
+      buffer, or for a specific buffer if {buf} is given.  The
+      result is a dictionary with the following items:
         "seq_last"	The highest undo sequence number used.
         "seq_cur"	The sequence number of the current position in
       		the undo tree.  This differs from "seq_last"
@@ -11664,8 +11667,8 @@ M.funcs = {
       		item.
     ]=],
     name = 'undotree',
-    params = {},
-    signature = 'undotree()',
+    params = { { 'buf', 'any' } },
+    signature = 'undotree([{buf}])',
   },
   uniq = {
     args = { 1, 3 },
diff --git a/src/nvim/undo.c b/src/nvim/undo.c
index b324b777a6..695cf81f73 100644
--- a/src/nvim/undo.c
+++ b/src/nvim/undo.c
@@ -93,6 +93,7 @@
 #include "nvim/cursor.h"
 #include "nvim/drawscreen.h"
 #include "nvim/edit.h"
+#include "nvim/eval/funcs.h"
 #include "nvim/eval/typval.h"
 #include "nvim/eval/typval_defs.h"
 #include "nvim/ex_cmds_defs.h"
@@ -3118,7 +3119,7 @@ bool curbufIsChanged(void)
 /// @param[in]  first_uhp  Undo blocks list to start with.
 ///
 /// @return [allocated] List with a representation of undo blocks.
-static list_T *u_eval_tree(const u_header_T *const first_uhp)
+static list_T *u_eval_tree(buf_T *const buf, const u_header_T *const first_uhp)
   FUNC_ATTR_WARN_UNUSED_RESULT FUNC_ATTR_NONNULL_RET
 {
   list_T *const list = tv_list_alloc(kListLenMayKnow);
@@ -3127,10 +3128,10 @@ static list_T *u_eval_tree(const u_header_T *const first_uhp)
     dict_T *const dict = tv_dict_alloc();
     tv_dict_add_nr(dict, S_LEN("seq"), (varnumber_T)uhp->uh_seq);
     tv_dict_add_nr(dict, S_LEN("time"), (varnumber_T)uhp->uh_time);
-    if (uhp == curbuf->b_u_newhead) {
+    if (uhp == buf->b_u_newhead) {
       tv_dict_add_nr(dict, S_LEN("newhead"), 1);
     }
-    if (uhp == curbuf->b_u_curhead) {
+    if (uhp == buf->b_u_curhead) {
       tv_dict_add_nr(dict, S_LEN("curhead"), 1);
     }
     if (uhp->uh_save_nr > 0) {
@@ -3139,7 +3140,7 @@ static list_T *u_eval_tree(const u_header_T *const first_uhp)
 
     if (uhp->uh_alt_next.ptr != NULL) {
       // Recursive call to add alternate undo tree.
-      tv_dict_add_list(dict, S_LEN("alt"), u_eval_tree(uhp->uh_alt_next.ptr));
+      tv_dict_add_list(dict, S_LEN("alt"), u_eval_tree(buf, uhp->uh_alt_next.ptr));
     }
 
     tv_list_append_dict(list, dict);
@@ -3167,21 +3168,24 @@ void f_undofile(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
   }
 }
 
-/// "undotree()" function
+/// "undotree(expr)" function
 void f_undotree(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
 {
+  typval_T *const tv = &argvars[0];
+  buf_T *const buf = tv->v_type == VAR_UNKNOWN ? curbuf : tv_get_buf_from_arg(tv);
+
   tv_dict_alloc_ret(rettv);
 
   dict_T *dict = rettv->vval.v_dict;
 
-  tv_dict_add_nr(dict, S_LEN("synced"), (varnumber_T)curbuf->b_u_synced);
-  tv_dict_add_nr(dict, S_LEN("seq_last"), (varnumber_T)curbuf->b_u_seq_last);
-  tv_dict_add_nr(dict, S_LEN("save_last"), (varnumber_T)curbuf->b_u_save_nr_last);
-  tv_dict_add_nr(dict, S_LEN("seq_cur"), (varnumber_T)curbuf->b_u_seq_cur);
-  tv_dict_add_nr(dict, S_LEN("time_cur"), (varnumber_T)curbuf->b_u_time_cur);
-  tv_dict_add_nr(dict, S_LEN("save_cur"), (varnumber_T)curbuf->b_u_save_nr_cur);
+  tv_dict_add_nr(dict, S_LEN("synced"), (varnumber_T)buf->b_u_synced);
+  tv_dict_add_nr(dict, S_LEN("seq_last"), (varnumber_T)buf->b_u_seq_last);
+  tv_dict_add_nr(dict, S_LEN("save_last"), (varnumber_T)buf->b_u_save_nr_last);
+  tv_dict_add_nr(dict, S_LEN("seq_cur"), (varnumber_T)buf->b_u_seq_cur);
+  tv_dict_add_nr(dict, S_LEN("time_cur"), (varnumber_T)buf->b_u_time_cur);
+  tv_dict_add_nr(dict, S_LEN("save_cur"), (varnumber_T)buf->b_u_save_nr_cur);
 
-  tv_dict_add_list(dict, S_LEN("entries"), u_eval_tree(curbuf->b_u_oldhead));
+  tv_dict_add_list(dict, S_LEN("entries"), u_eval_tree(buf, buf->b_u_oldhead));
 }
 
 // Given the buffer, Return the undo header. If none is set, set one first.
diff --git a/test/old/testdir/test_undo.vim b/test/old/testdir/test_undo.vim
index 4678a51d60..08a0ba4c39 100644
--- a/test/old/testdir/test_undo.vim
+++ b/test/old/testdir/test_undo.vim
@@ -93,6 +93,53 @@ func FillBuffer()
   endfor
 endfunc
 
+func Test_undotree_bufnr()
+  new
+  let buf1 = bufnr()
+
+  normal! Aabc
+  set ul=100
+
+  " Save undo tree without bufnr as ground truth for buffer 1
+  let d1 = undotree()
+
+  new
+  let buf2 = bufnr()
+
+  normal! Adef
+  set ul=100
+
+  normal! Aghi
+  set ul=100
+
+  " Save undo tree without bufnr as ground truth for buffer 2
+  let d2 = undotree()
+
+  " Check undotree() with bufnr argument
+  let d = undotree(buf1)
+  call assert_equal(d1, d)
+  call assert_notequal(d2, d)
+
+  let d = undotree(buf2)
+  call assert_notequal(d1, d)
+  call assert_equal(d2, d)
+
+  " Switch buffers and check again
+  wincmd p
+
+  let d = undotree(buf1)
+  call assert_equal(d1, d)
+
+  let d = undotree(buf2)
+  call assert_notequal(d1, d)
+  call assert_equal(d2, d)
+
+  " Drop created windows
+  set ul&
+  new
+  only!
+endfunc
+
 func Test_global_local_undolevels()
   new one
   set undolevels=5