vim-patch:8.2.4627: flatten() does not use maxdepth correctly

Problem:    flatten() does not use maxdepth correctly.
Solution:   Use a recursive implementation. (closes vim/vim#10020)

https://github.com/vim/vim/commit/acf7d73a7f5cdd63b34de777a4ce5eb3e2ba0ab3

Co-authored-by: Bram Moolenaar <Bram@vim.org>

3 files changed, 39 insertions, 21 deletions

diff --git a/src/nvim/eval/funcs.c b/src/nvim/eval/funcs.c
index 3a58b90f96..6ac94706c7 100644
--- a/src/nvim/eval/funcs.c
+++ b/src/nvim/eval/funcs.c
@@ -1914,7 +1914,7 @@ static void flatten_common(typval_T *argvars, typval_T *rettv, bool make_copy)
     tv_list_ref(list);
   }
 
-  tv_list_flatten(list, maxdepth);
+  tv_list_flatten(list, NULL, tv_list_len(list), maxdepth);
 }
 
 /// "flatten(list[, {maxdepth}])" function
diff --git a/src/nvim/eval/typval.c b/src/nvim/eval/typval.c
index ab4245b59a..125a4310d8 100644
--- a/src/nvim/eval/typval.c
+++ b/src/nvim/eval/typval.c
@@ -647,53 +647,55 @@ tv_list_copy_error:
   return NULL;
 }
 
-/// Flatten "list" in place to depth "maxdepth".
+/// Flatten up to "maxitems" in "list", starting at "first" to depth "maxdepth".
+/// When "first" is NULL use the first item.
 /// Does nothing if "maxdepth" is 0.
 ///
 /// @param[in,out] list   List to flatten
 /// @param[in] maxdepth   Maximum depth that will be flattened
 ///
 /// @return OK or FAIL
-void tv_list_flatten(list_T *list, long maxdepth)
+void tv_list_flatten(list_T *list, listitem_T *first, long maxitems, long maxdepth)
   FUNC_ATTR_NONNULL_ARG(1)
 {
   listitem_T *item;
   listitem_T *to_free;
-  int n;
+  int done = 0;
   if (maxdepth == 0) {
     return;
   }
 
-  n = 0;
-  item = list->lv_first;
-  while (item != NULL) {
+  if (first == NULL) {
+    item = list->lv_first;
+  } else {
+    item = first;
+  }
+
+  while (item != NULL && done < maxitems) {
+    listitem_T *next = item->li_next;
+
     fast_breakcheck();
     if (got_int) {
       return;
     }
     if (item->li_tv.v_type == VAR_LIST) {
-      listitem_T *next = item->li_next;
+      list_T *itemlist = item->li_tv.vval.v_list;
 
       tv_list_drop_items(list, item, item);
-      tv_list_extend(list, item->li_tv.vval.v_list, next);
+      tv_list_extend(list, itemlist, next);
       tv_clear(&item->li_tv);
       to_free = item;
 
-      if (item->li_prev == NULL) {
-        item = list->lv_first;
-      } else {
-        item = item->li_prev->li_next;
+      if (maxdepth > 0) {
+        tv_list_flatten(list,
+                        item->li_prev == NULL ? list->lv_first : item->li_prev->li_next,
+                        itemlist->lv_len, maxdepth - 1);
       }
       xfree(to_free);
-
-      if (++n >= maxdepth) {
-        n = 0;
-        item = next;
-      }
-    } else {
-      n = 0;
-      item = item->li_next;
     }
+
+    done++;
+    item = next;
   }
 }
 
diff --git a/src/nvim/testdir/test_flatten.vim b/src/nvim/testdir/test_flatten.vim
index d8abc1ad68..aa91060313 100644
--- a/src/nvim/testdir/test_flatten.vim
+++ b/src/nvim/testdir/test_flatten.vim
@@ -78,6 +78,14 @@ func Test_flatten()
   call add(y, x) " l:y = [2, [1, [...]]]
   call assert_equal([1, 2, 1, 2], flatten(l:x, 2))
   call assert_equal([2, l:x], l:y)
+
+  let l4 = [ 1, [ 11, [ 101, [ 1001 ] ] ] ]
+  call assert_equal(l4, flatten(deepcopy(l4), 0))
+  call assert_equal([1, 11, [101, [1001]]], flatten(deepcopy(l4), 1))
+  call assert_equal([1, 11, 101, [1001]], flatten(deepcopy(l4), 2))
+  call assert_equal([1, 11, 101, 1001], flatten(deepcopy(l4), 3))
+  call assert_equal([1, 11, 101, 1001], flatten(deepcopy(l4), 4))
+  call assert_equal([1, 11, 101, 1001], flatten(deepcopy(l4)))
 endfunc
 
 func Test_flattennew()
@@ -87,6 +95,14 @@ func Test_flattennew()
 
   call assert_equal([1, 2, [3, 4], 5], flattennew(l, 1))
   call assert_equal([1, [2, [3, 4]], 5], l)
+
+  let l4 = [ 1, [ 11, [ 101, [ 1001 ] ] ] ]
+  call assert_equal(l4, flatten(deepcopy(l4), 0))
+  call assert_equal([1, 11, [101, [1001]]], flattennew(l4, 1))
+  call assert_equal([1, 11, 101, [1001]], flattennew(l4, 2))
+  call assert_equal([1, 11, 101, 1001], flattennew(l4, 3))
+  call assert_equal([1, 11, 101, 1001], flattennew(l4, 4))
+  call assert_equal([1, 11, 101, 1001], flattennew(l4))
 endfunc
 
 " vim: shiftwidth=2 sts=2 expandtab