1 files changed, 34 insertions, 26 deletions

diff --git a/src/nvim/window.c b/src/nvim/window.c
index 90f7145f51..6bae92c909 100644
--- a/src/nvim/window.c
+++ b/src/nvim/window.c
@@ -2813,10 +2813,11 @@ int win_close(win_T *win, bool free_buf, bool force)
   wp = win_free_mem(win, &dir, NULL);
 
   if (help_window) {
-    // Closing the help window moves the cursor back to the original window.
-    win_T *tmpwp = get_snapshot_focus(SNAP_HELP_IDX);
-    if (tmpwp != NULL) {
-      wp = tmpwp;
+    // Closing the help window moves the cursor back to the current window
+    // of the snapshot.
+    win_T *prev_win = get_snapshot_curwin(SNAP_HELP_IDX);
+    if (win_valid(prev_win)) {
+      wp = prev_win;
     }
   }
 
@@ -6827,6 +6828,35 @@ static void clear_snapshot_rec(frame_T *fr)
   }
 }
 
+/// Traverse a snapshot to find the previous curwin.
+static win_T *get_snapshot_curwin_rec(frame_T *ft)
+{
+  win_T *wp;
+
+  if (ft->fr_next != NULL) {
+    if ((wp = get_snapshot_curwin_rec(ft->fr_next)) != NULL) {
+      return wp;
+    }
+  }
+  if (ft->fr_child != NULL) {
+    if ((wp = get_snapshot_curwin_rec(ft->fr_child)) != NULL) {
+      return wp;
+    }
+  }
+
+  return ft->fr_win;
+}
+
+/// @return  the current window stored in the snapshot or NULL.
+static win_T *get_snapshot_curwin(int idx)
+{
+  if (curtab->tp_snapshot[idx] == NULL) {
+    return NULL;
+  }
+
+  return get_snapshot_curwin_rec(curtab->tp_snapshot[idx]);
+}
+
 /// Restore a previously created snapshot, if there is any.
 /// This is only done if the screen size didn't change and the window layout is
 /// still the same.
@@ -6899,28 +6929,6 @@ static win_T *restore_snapshot_rec(frame_T *sn, frame_T *fr)
   return wp;
 }
 
-/// Gets the focused window (the one holding the cursor) of the snapshot.
-static win_T *get_snapshot_focus(int idx)
-{
-  if (curtab->tp_snapshot[idx] == NULL) {
-    return NULL;
-  }
-
-  frame_T *sn = curtab->tp_snapshot[idx];
-  // This should be equivalent to the recursive algorithm found in
-  // restore_snapshot as far as traveling nodes go.
-  while (sn->fr_child != NULL || sn->fr_next != NULL) {
-    while (sn->fr_child != NULL) {
-      sn = sn->fr_child;
-    }
-    if (sn->fr_next != NULL) {
-      sn = sn->fr_next;
-    }
-  }
-
-  return win_valid(sn->fr_win) ? sn->fr_win : NULL;
-}
-
 /// Set "win" to be the curwin and "tp" to be the current tab page.
 /// restore_win() MUST be called to undo, also when FAIL is returned.
 /// No autocommands will be executed until restore_win() is called.