diff --git a/packages/array-flatten/index.cjs b/packages/array-flatten/index.cjs index 78f053369..f81a4727f 100644 --- a/packages/array-flatten/index.cjs +++ b/packages/array-flatten/index.cjs @@ -7,20 +7,26 @@ module.exports = flatten; function flattenHelper(arr, depth) { var stack = arr.slice(); - var result = []; - - while (stack.length) { - var item = stack.pop(); - - if (Array.isArray(item) && depth > 0) { - stack.push.apply(stack, item); - depth--; - } else { - result.push(item); + var stackContainsArray = true; + + for (; depth > 0 && stackContainsArray; depth--) { + stackContainsArray = false; + var nextStack = [] + + while (stack.length) { + var item = stack.shift(); + if (Array.isArray(item)) { + stackContainsArray = true; + nextStack.push.apply(nextStack, item); + } else { + nextStack.push(item); + } } + + stack = nextStack; } - return result.reverse(); + return stack; } function flatten(arr, depth) { diff --git a/packages/array-flatten/index.mjs b/packages/array-flatten/index.mjs index f109b5d9c..4e992be88 100644 --- a/packages/array-flatten/index.mjs +++ b/packages/array-flatten/index.mjs @@ -7,20 +7,26 @@ var arrayFlatten = flatten; function flattenHelper(arr, depth) { var stack = arr.slice(); - var result = []; - - while (stack.length) { - var item = stack.pop(); - - if (Array.isArray(item) && depth > 0) { - stack.push.apply(stack, item); - depth--; - } else { - result.push(item); + var stackContainsArray = true; + + for (; depth > 0 && stackContainsArray; depth--) { + stackContainsArray = false; + var nextStack = []; + + while (stack.length) { + var item = stack.shift(); + if (Array.isArray(item)) { + stackContainsArray = true; + nextStack.push.apply(nextStack, item); + } else { + nextStack.push(item); + } } + + stack = nextStack; } - return result.reverse(); + return stack; } function flatten(arr, depth) { diff --git a/test/array-flatten/index.cjs b/test/array-flatten/index.cjs index babcf5a48..cd0ce8660 100644 --- a/test/array-flatten/index.cjs +++ b/test/array-flatten/index.cjs @@ -40,12 +40,12 @@ test('returns a new array with all sub-arrays items concatenated into it up to t function(t) { t.plan(5); - var arr = [1, [2, [3, [4, [5]]]]]; + var arr = [[1, [2, 3], 4], [5, [6, [7, [8]]]]]; t.deepEqual(flatten(arr, 0), arr); - t.deepEqual(flatten(arr, 1), [1, 2, [3, [4, [5]]]]); - t.deepEqual(flatten(arr, 2), [1, 2, 3, [4, [5]]]); - t.deepEqual(flatten(arr, 3), [1, 2, 3, 4, [5]]); - t.deepEqual(flatten(arr, 4), [1, 2, 3, 4, 5]); + t.deepEqual(flatten(arr, 1), [1, [2, 3], 4, 5, [6, [7, [8]]]]); + t.deepEqual(flatten(arr, 2), [1, 2, 3, 4, 5, 6, [7, [8]]]); + t.deepEqual(flatten(arr, 3), [1, 2, 3, 4, 5, 6, 7, [8]]); + t.deepEqual(flatten(arr, 4), [1, 2, 3, 4, 5, 6, 7, 8]); t.end(); });