Merge pull request #22386 from zeertzjq/vim-8.2.2449

vim-patch:8.2.{2449,4627,4629,4632}: flattennew(), flatten() fixes

2 files changed, 60 insertions, 36 deletions

diff --git a/src/nvim/eval/funcs.c b/src/nvim/eval/funcs.c
index 6f983d3208..f1852f1b6d 100644
--- a/src/nvim/eval/funcs.c
+++ b/src/nvim/eval/funcs.c
@@ -1870,8 +1870,8 @@ static void f_expandcmd(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
   rettv->vval.v_string = cmdstr;
 }
 
-/// "flatten(list[, {maxdepth}])" function
-static void f_flatten(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
+/// "flatten()" and "flattennew()" functions
+static void flatten_common(typval_T *argvars, typval_T *rettv, bool make_copy)
 {
   bool error = false;
 
@@ -1895,13 +1895,38 @@ static void f_flatten(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
   }
 
   list_T *list = argvars[0].vval.v_list;
-  if (list != NULL
-      && !value_check_lock(tv_list_locked(list),
-                           N_("flatten() argument"),
-                           TV_TRANSLATE)
-      && tv_list_flatten(list, maxdepth) == OK) {
-    tv_copy(&argvars[0], rettv);
+  rettv->v_type = VAR_LIST;
+  rettv->vval.v_list = list;
+  if (list == NULL) {
+    return;
   }
+
+  if (make_copy) {
+    list = tv_list_copy(NULL, list, false, get_copyID());
+    rettv->vval.v_list = list;
+    if (list == NULL) {
+      return;
+    }
+  } else {
+    if (value_check_lock(tv_list_locked(list), N_("flatten() argument"), TV_TRANSLATE)) {
+      return;
+    }
+    tv_list_ref(list);
+  }
+
+  tv_list_flatten(list, NULL, tv_list_len(list), maxdepth);
+}
+
+/// "flatten(list[, {maxdepth}])" function
+static void f_flatten(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
+{
+  flatten_common(argvars, rettv, false);
+}
+
+/// "flattennew(list[, {maxdepth}])" function
+static void f_flattennew(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
+{
+  flatten_common(argvars, rettv, true);
 }
 
 /// "extend(list, list [, idx])" function
diff --git a/src/nvim/eval/typval.c b/src/nvim/eval/typval.c
index 9a3a1c3c0f..17499480ed 100644
--- a/src/nvim/eval/typval.c
+++ b/src/nvim/eval/typval.c
@@ -647,55 +647,54 @@ 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
-int tv_list_flatten(list_T *list, long maxdepth)
-  FUNC_ATTR_NONNULL_ARG(1) FUNC_ATTR_WARN_UNUSED_RESULT
+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 OK;
+    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 FAIL;
+      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_clear(&item->li_tv);
-      to_free = item;
+      tv_list_extend(list, itemlist, next);
 
-      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;
+      tv_clear(&item->li_tv);
+      xfree(item);
     }
+
+    done++;
+    item = next;
   }
-  return OK;
 }
 
 /// Extend first list with the second