vim-patch:8.2.0935: flattening a list with existing code is slow

Problem:    Flattening a list with existing code is slow.
Solution:   Add flatten(). (Mopp, closes vim/vim#3676)
https://github.com/vim/vim/commit/077a1e670ad69ef4cefc22103ca6635bd269e764

6 files changed, 186 insertions, 0 deletions

diff --git a/runtime/doc/eval.txt b/runtime/doc/eval.txt
index 1c1baa943a..a346262b0c 100644
--- a/runtime/doc/eval.txt
+++ b/runtime/doc/eval.txt
@@ -2126,6 +2126,7 @@ finddir({name} [, {path} [, {count}]])
 				String	find directory {name} in {path}
 findfile({name} [, {path} [, {count}]])
 				String	find file {name} in {path}
+flatten({list} [, {maxdepth}])	List	flatten {list} up to {maxdepth} levels
 float2nr({expr})		Number	convert Float {expr} to a Number
 floor({expr})			Float	round {expr} down
 fmod({expr1}, {expr2})		Float	remainder of {expr1} / {expr2}
@@ -3918,6 +3919,25 @@ findfile({name} [, {path} [, {count}]])				*findfile()*
 <		Searches from the directory of the current file upwards until
 		it finds the file "tags.vim".
 
+flatten({list} [, {maxdepth}])					*flatten()*
+		Flatten {list} up to {maxdepth} levels.  Without {maxdepth}
+		the result is a |List| without nesting, as if {maxdepth} is
+		a very large number.
+		The {list} is changed in place, make a copy first if you do
+		not want that.
+							        *E964*
+		{maxdepth} means how deep in nested lists changes are made.
+		{list} is not modified when {maxdepth} is 0.
+		{maxdepth} must be positive number.
+
+		If there is an error the number zero is returned.
+
+		Example: >
+			:echo flatten([1, [2, [3, 4]], 5])
+<			[1, 2, 3, 4, 5] >
+			:echo flatten([1, [2, [3, 4]], 5], 1)
+<			[1, 2, [3, 4], 5]
+
 float2nr({expr})					*float2nr()*
 		Convert {expr} to a Number by omitting the part after the
 		decimal point.
diff --git a/runtime/doc/usr_41.txt b/runtime/doc/usr_41.txt
index 234f7801ab..31da51bfbe 100644
--- a/runtime/doc/usr_41.txt
+++ b/runtime/doc/usr_41.txt
@@ -644,6 +644,7 @@ List manipulation:					*list-functions*
 	min()			minimum value in a List
 	count()			count number of times a value appears in a List
 	repeat()		repeat a List multiple times
+	flatten()		flatten a List
 
 Dictionary manipulation:				*dict-functions*
 	get()			get an entry without an error for a wrong key
diff --git a/src/nvim/eval.lua b/src/nvim/eval.lua
index 51872a7ba8..410cce05b0 100644
--- a/src/nvim/eval.lua
+++ b/src/nvim/eval.lua
@@ -111,6 +111,7 @@ return {
     filter={args=2},
     finddir={args={1, 3}},
     findfile={args={1, 3}},
+    flatten={args={1, 2}},
     float2nr={args=1},
     floor={args=1, func="float_op_wrapper", data="&floor"},
     fmod={args=2},
diff --git a/src/nvim/eval/funcs.c b/src/nvim/eval/funcs.c
index 99014d1a09..95686b97bf 100644
--- a/src/nvim/eval/funcs.c
+++ b/src/nvim/eval/funcs.c
@@ -2160,6 +2160,42 @@ static void f_expandcmd(typval_T *argvars, typval_T *rettv, FunPtr fptr)
   rettv->vval.v_string = cmdstr;
 }
 
+
+/// "flatten(list[, {maxdepth}])" function
+static void f_flatten(typval_T *argvars, typval_T *rettv, FunPtr fptr)
+{
+  list_T *list;
+  long maxdepth;
+  bool error = false;
+
+  if (argvars[0].v_type != VAR_LIST) {
+    EMSG2(_(e_listarg), "flatten()");
+    return;
+  }
+
+  if (argvars[1].v_type == VAR_UNKNOWN) {
+    maxdepth = 999999;
+  } else {
+    maxdepth = (long)tv_get_number_chk(&argvars[1], &error);
+    if (error) {
+      return;
+    }
+    if (maxdepth < 0) {
+      EMSG(_("E900: maxdepth must be non-negative number"));
+      return;
+    }
+  }
+
+  list = argvars[0].vval.v_list;
+  if (list != NULL
+      && !tv_check_lock(tv_list_locked(list),
+                        N_("flatten() argument"),
+                        TV_TRANSLATE)
+      && tv_list_flatten(list, maxdepth) == OK) {
+    tv_copy(&argvars[0], rettv);
+  }
+}
+
 /*
  * "extend(list, list [, idx])" function
  * "extend(dict, dict [, action])" function
diff --git a/src/nvim/eval/typval.c b/src/nvim/eval/typval.c
index 89ca2db59b..576948f052 100644
--- a/src/nvim/eval/typval.c
+++ b/src/nvim/eval/typval.c
@@ -640,6 +640,53 @@ tv_list_copy_error:
   return NULL;
 }
 
+/// Flatten "list" in place to depth "maxdepth".
+/// Does nothing if "maxdepth" is 0.
+///
+/// @param[in] 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_WARN_UNUSED_RESULT
+{
+  listitem_T *item;
+  int n;
+  if (maxdepth == 0) {
+    return OK;
+  }
+
+  n = 0;
+  item = list->lv_first;
+  while (item != NULL) {
+    fast_breakcheck();
+    if (got_int) {
+      return FAIL;
+    }
+    if (item->li_tv.v_type == VAR_LIST) {
+      listitem_T *next = item->li_next;
+
+      tv_list_drop_items(list, item, item);
+      tv_list_extend(list, item->li_tv.vval.v_list, next);
+
+      if (item->li_prev == NULL) {
+        item = list->lv_first;
+      } else {
+        item = item->li_prev->li_next;
+      }
+
+      if (++n >= maxdepth) {
+        n = 0;
+        item = next;
+      }
+    } else {
+      n = 0;
+      item = item->li_next;
+    }
+  }
+  return OK;
+}
+
 /// Extend first list with the second
 ///
 /// @param[out]  l1  List to extend.
diff --git a/src/nvim/testdir/test_flatten.vim b/src/nvim/testdir/test_flatten.vim
new file mode 100644
index 0000000000..99086611e1
--- /dev/null
+++ b/src/nvim/testdir/test_flatten.vim
@@ -0,0 +1,81 @@
+" Test for flatting list.
+func Test_flatten()
+  call assert_fails('call flatten(1)', 'E686:')
+  call assert_fails('call flatten({})', 'E686:')
+  call assert_fails('call flatten("string")', 'E686:')
+  call assert_fails('call flatten([], [])', 'E745:')
+  call assert_fails('call flatten([], -1)', 'E900: maxdepth')
+
+  call assert_equal([], flatten([]))
+  call assert_equal([], flatten([[]]))
+  call assert_equal([], flatten([[[]]]))
+
+  call assert_equal([1, 2, 3], flatten([1, 2, 3]))
+  call assert_equal([1, 2, 3], flatten([[1], 2, 3]))
+  call assert_equal([1, 2, 3], flatten([1, [2], 3]))
+  call assert_equal([1, 2, 3], flatten([1, 2, [3]]))
+  call assert_equal([1, 2, 3], flatten([[1], [2], 3]))
+  call assert_equal([1, 2, 3], flatten([1, [2], [3]]))
+  call assert_equal([1, 2, 3], flatten([[1], 2, [3]]))
+  call assert_equal([1, 2, 3], flatten([[1], [2], [3]]))
+
+  call assert_equal([1, 2, 3], flatten([[1, 2, 3], []]))
+  call assert_equal([1, 2, 3], flatten([[], [1, 2, 3]]))
+  call assert_equal([1, 2, 3], flatten([[1, 2], [], [3]]))
+  call assert_equal([1, 2, 3], flatten([[], [1, 2, 3], []]))
+  call assert_equal([1, 2, 3, 4], flatten(range(1, 4)))
+
+  " example in the help
+  call assert_equal([1, 2, 3, 4, 5], flatten([1, [2, [3, 4]], 5]))
+  call assert_equal([1, 2, [3, 4], 5], flatten([1, [2, [3, 4]], 5], 1))
+
+  call assert_equal([0, [1], 2, [3], 4], flatten([[0, [1]], 2, [[3], 4]], 1))
+  call assert_equal([1, 2, 3], flatten([[[[1]]], [2], [3]], 3))
+  call assert_equal([[1], [2], [3]], flatten([[[1], [2], [3]]], 1))
+  call assert_equal([[1]], flatten([[1]], 0))
+
+  " Make it flatten if the given maxdepth is larger than actual depth.
+  call assert_equal([1, 2, 3], flatten([[1, 2, 3]], 1))
+  call assert_equal([1, 2, 3], flatten([[1, 2, 3]], 2))
+
+  let l:list = [[1], [2], [3]]
+  call assert_equal([1, 2, 3], flatten(l:list))
+  call assert_equal([1, 2, 3], l:list)
+
+  " Tests for checking reference counter works well.
+  let l:x = {'foo': 'bar'}
+  call assert_equal([1, 2, l:x, 3], flatten([1, [2, l:x], 3]))
+  call test_garbagecollect_now()
+  call assert_equal('bar', l:x.foo)
+
+  let l:list = [[1], [2], [3]]
+  call assert_equal([1, 2, 3], flatten(l:list))
+  call test_garbagecollect_now()
+  call assert_equal([1, 2, 3], l:list)
+
+  " Tests for checking circular reference list can be flatten.
+  let l:x = [1]
+  let l:y = [x]
+  let l:z = flatten(l:y)
+  call assert_equal([1], l:z)
+  call test_garbagecollect_now()
+  let l:x[0] = 2
+  call assert_equal([2], l:x)
+  call assert_equal([1], l:z) " NOTE: primitive types are copied.
+  call assert_equal([1], l:y)
+
+  let l:x = [2]
+  let l:y = [1, [l:x], 3] " [1, [[2]], 3]
+  let l:z = flatten(l:y, 1)
+  call assert_equal([1, [2], 3], l:z)
+  let l:x[0] = 9
+  call assert_equal([1, [9], 3], l:z) " Reference to l:x is kept.
+  call assert_equal([1, [9], 3], l:y)
+
+  let l:x = [1]
+  let l:y = [2]
+  call add(x, y) " l:x = [1, [2]]
+  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)
+endfunc