2 files changed, 74 insertions, 19 deletions

diff --git a/src/nvim/ex_cmds.c b/src/nvim/ex_cmds.c
index 6987bcfd89..0190db258f 100644
--- a/src/nvim/ex_cmds.c
+++ b/src/nvim/ex_cmds.c
@@ -6140,6 +6140,8 @@ void ex_substitute(exarg_T *eap)
   char_u *save_eap = eap->arg;
   save_search_patterns();
   int save_changedtick = curbuf->b_changedtick;
+  time_t save_b_u_time_cur = curbuf->b_u_time_cur;
+  u_header_T *save_b_u_newhead = curbuf->b_u_newhead;
   long save_b_p_ul = curbuf->b_p_ul;
   curbuf->b_p_ul = LONG_MAX;  // make sure we can undo all changes
   block_autocmds();   // disable events before show_sub() opens window/buffer
@@ -6147,15 +6149,18 @@ void ex_substitute(exarg_T *eap)
 
   buf_T *preview_buf = do_sub(eap);
 
-  if (save_changedtick != curbuf->b_changedtick
-      && !u_undo_and_forget(1)) {
-    abort();
+  if (save_changedtick != curbuf->b_changedtick) {
+    if (!u_undo_and_forget(1)) { abort(); }
+    // Restore newhead. It is meaningless when curhead is valid, but we must
+    // restore it so that undotree() is identical before/after the preview.
+    curbuf->b_u_newhead = save_b_u_newhead;
+    curbuf->b_u_time_cur = save_b_u_time_cur;
+    curbuf->b_changedtick = save_changedtick;
   }
   if (buf_valid(preview_buf)) {
     // XXX: Must do this *after* u_undo_and_forget(), why?
     close_windows(preview_buf, false);
   }
-  curbuf->b_changedtick = save_changedtick;
   curbuf->b_p_ul = save_b_p_ul;
   eap->arg = save_eap;
   restore_search_patterns();
diff --git a/test/functional/ui/inccommand_spec.lua b/test/functional/ui/inccommand_spec.lua
index 520fadc8e0..f2282e3fb8 100644
--- a/test/functional/ui/inccommand_spec.lua
+++ b/test/functional/ui/inccommand_spec.lua
@@ -107,9 +107,8 @@ describe(":substitute, 'inccommand' preserves", function()
     ]])
   end)
 
-  it(':substitute with various delimiters', function()
-    for _, case in pairs{"", "split", "nosplit"} do
-      clear()
+  for _, case in pairs{"", "split", "nosplit"} do
+    it("various delimiters (inccommand="..case..")", function()
       insert(default_text)
       execute("set inccommand=" .. case)
 
@@ -122,12 +121,11 @@ describe(":substitute, 'inccommand' preserves", function()
           ]])
         execute("undo")
       end
-    end
-  end)
+    end)
+  end
 
-  it("'undolevels'", function()
-    for _, case in pairs{"", "split", "nosplit"} do
-      clear()
+  for _, case in pairs{"", "split", "nosplit"} do
+    it("'undolevels' (inccommand="..case..")", function()
       execute("set undolevels=139")
       execute("setlocal undolevels=34")
       execute("set inccommand=" .. case)
@@ -135,12 +133,64 @@ describe(":substitute, 'inccommand' preserves", function()
       feed(":%s/as/glork/<enter>")
       eq(meths.get_option('undolevels'), 139)
       eq(curbufmeths.get_option('undolevels'), 34)
-    end
-  end)
+    end)
+  end
 
-  it("b:changedtick", function()
-    for _, case in pairs{"", "split", "nosplit"} do
-      clear()
+  for _, case in ipairs({"", "split", "nosplit"}) do
+    it("empty undotree() (inccommand="..case..")", function()
+      execute("set undolevels=1000")
+      execute("set inccommand=" .. case)
+      local expected_undotree = eval("undotree()")
+
+      -- Start typing an incomplete :substitute command.
+      feed([[:%s/e/YYYY/g]])
+      wait()
+      -- Cancel the :substitute.
+      feed([[<C-\><C-N>]])
+
+      -- The undo tree should be unchanged.
+      eq(expected_undotree, eval("undotree()"))
+      eq({}, eval("undotree()")["entries"])
+    end)
+  end
+
+  for _, case in ipairs({"", "split", "nosplit"}) do
+    it("undotree() with branches (inccommand="..case..")", function()
+      execute("set undolevels=1000")
+      execute("set inccommand=" .. case)
+      -- Make some changes.
+      feed([[isome text 1<C-\><C-N>]])
+      feed([[osome text 2<C-\><C-N>]])
+      -- Add an undo branch.
+      feed([[u]])
+      -- More changes, more undo branches.
+      feed([[osome text 3<C-\><C-N>]])
+      feed([[AX<C-\><C-N>]])
+      feed([[...]])
+      feed([[uu]])
+      feed([[osome text 4<C-\><C-N>]])
+      feed([[u<C-R>u]])
+      feed([[osome text 5<C-\><C-N>]])
+      expect([[
+        some text 1
+        some text 3XX
+        some text 5]])
+      local expected_undotree = eval("undotree()")
+      eq(5, #expected_undotree["entries"])  -- sanity
+
+      -- Start typing an incomplete :substitute command.
+      feed([[:%s/e/YYYY/g]])
+      wait()
+      -- Cancel the :substitute.
+      feed([[<C-\><C-N>]])
+
+      -- The undo tree should be unchanged.
+      eq(expected_undotree, eval("undotree()"))
+    end)
+  end
+
+  for _, case in pairs{"", "split", "nosplit"} do
+    it("b:changedtick (inccommand="..case..")", function()
       execute("set inccommand=" .. case)
       feed([[isome text 1<C-\><C-N>]])
       feed([[osome text 2<C-\><C-N>]])
@@ -154,8 +204,8 @@ describe(":substitute, 'inccommand' preserves", function()
       wait()
 
       eq(expected_tick, eval("b:changedtick"))
-    end
-  end)
+    end)
+  end
 
 end)