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 #3676)

runtime/doc/eval.txt

@@ -2451,6 +2451,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}
@@ -4514,6 +4515,25 @@ findfile({name} [, {path} [, {count}]])				*findfile()*
 		Can also be used as a |method|: >
 			GetName()->findfile()
 
+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.

runtime/doc/usr_41.txt

@@ -650,6 +650,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

src/evalfunc.c

@@ -541,6 +541,7 @@ static funcentry_T global_functions[] =
     {"filter",		2, 2, FEARG_1,	  ret_any,	f_filter},
     {"finddir",		1, 3, FEARG_1,	  ret_string,	f_finddir},
     {"findfile",	1, 3, FEARG_1,	  ret_string,	f_findfile},
+    {"flatten",		1, 2, FEARG_1,	  ret_list_any,	f_flatten},
     {"float2nr",	1, 1, FEARG_1,	  ret_number,	FLOAT_FUNC(f_float2nr)},
     {"floor",		1, 1, FEARG_1,	  ret_float,	FLOAT_FUNC(f_floor)},
     {"fmod",		2, 2, FEARG_1,	  ret_float,	FLOAT_FUNC(f_fmod)},

src/list.c

@@ -730,6 +730,95 @@ list_insert(list_T *l, listitem_T *ni, listitem_T *item)
     }
 }
 
+/*
+ * Flatten "list" to depth "maxdepth".
+ * It does nothing if "maxdepth" is 0.
+ * Returns FAIL when out of memory.
+ */
+    static int
+list_flatten(list_T *list, long maxdepth)
+{
+    listitem_T	*item;
+    int		n;
+
+    if (maxdepth == 0)
+	return OK;
+    CHECK_LIST_MATERIALIZE(list);
+
+    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;
+
+	    vimlist_remove(list, item, item);
+	    if (list_extend(list, item->li_tv.vval.v_list, next) == FAIL)
+		return FAIL;
+
+	    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;
+}
+
+/*
+ * "flatten(list[, {maxdepth}])" function
+ */
+    void
+f_flatten(typval_T *argvars, typval_T *rettv)
+{
+    list_T  *l;
+    long    maxdepth;
+    int	    error = FALSE;
+
+    if (argvars[0].v_type != VAR_LIST)
+    {
+	semsg(_(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;
+	}
+    }
+
+    l = argvars[0].vval.v_list;
+    if (l != NULL && !var_check_lock(l->lv_lock,
+				      (char_u *)N_("flatten() argument"), TRUE)
+		 && list_flatten(l, maxdepth) == OK)
+	copy_tv(&argvars[0], rettv);
+}
+
 /*
  * Extend "l1" with "l2".
  * If "bef" is NULL append at the end, otherwise insert before this item.

src/proto/list.pro

@@ -30,6 +30,7 @@ int list_append_string(list_T *l, char_u *str, int len);
 int list_append_number(list_T *l, varnumber_T n);
 int list_insert_tv(list_T *l, typval_T *tv, listitem_T *item);
 void list_insert(list_T *l, listitem_T *ni, listitem_T *item);
+void f_flatten(typval_T *argvars, typval_T *rettv);
 int list_extend(list_T *l1, list_T *l2, listitem_T *bef);
 int list_concat(list_T *l1, list_T *l2, typval_T *tv);
 list_T *list_copy(list_T *orig, int deep, int copyID);
@@ -37,7 +38,7 @@ void vimlist_remove(list_T *l, listitem_T *item, listitem_T *item2);
 char_u *list2string(typval_T *tv, int copyID, int restore_copyID);
 int list_join(garray_T *gap, list_T *l, char_u *sep, int echo_style, int restore_copyID, int copyID);
 void f_join(typval_T *argvars, typval_T *rettv);
-int get_list_tv(char_u **arg, typval_T *rettv, int evaluate, int do_error);
+int get_list_tv(char_u **arg, typval_T *rettv, int flags, int do_error);
 int write_list(FILE *fd, list_T *list, int binary);
 void init_static_list(staticList10_T *sl);
 void f_list2str(typval_T *argvars, typval_T *rettv);

src/testdir/Make_all.mak

@@ -135,6 +135,7 @@ NEW_TESTS = \
 	test_find_complete \
 	test_findfile \
 	test_fixeol \
+	test_flatten \
 	test_float_func \
 	test_fnameescape \
 	test_fnamemodify \
@@ -373,6 +374,7 @@ NEW_TESTS_RES = \
 	test_find_complete.res \
 	test_findfile.res \
 	test_fixeol.res \
+	test_flatten.res \
 	test_float_func.res \
 	test_fnameescape.res \
 	test_fold.res \

src/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

src/version.c

@@ -754,6 +754,8 @@ static char *(features[]) =
 
 static int included_patches[] =
 {   /* Add new patch number below this line */
+/**/
+    935,
 /**/
     934,
 /**/