diff options
-rw-r--r-- | src/nvim/eval/funcs.c | 2 | ||||
-rw-r--r-- | src/nvim/eval/typval.c | 42 | ||||
-rw-r--r-- | src/nvim/testdir/test_flatten.vim | 16 |
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 |