Add array.flat and array.flatMap (#496)

ark120202 · Perryvw · commit 8fb72f4f03a9 · 2019-03-27T19:35:16.000+01:00
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -17,6 +17,8 @@ export enum LuaLibFeature {
     ArraySlice = "ArraySlice",
     ArraySome = "ArraySome",
     ArraySplice = "ArraySplice",
+    ArrayFlat = "ArrayFlat",
+    ArrayFlatMap = "ArrayFlatMap",
     ClassIndex = "ClassIndex",
     ClassNewIndex = "ClassNewIndex",
     FunctionApply = "FunctionApply",
@@ -42,6 +44,8 @@ export enum LuaLibFeature {
 }
 
 const luaLibDependencies: {[lib in LuaLibFeature]?: LuaLibFeature[]} = {
+    ArrayFlat: [LuaLibFeature.ArrayConcat],
+    ArrayFlatMap: [LuaLibFeature.ArrayConcat],
     Iterator: [LuaLibFeature.Symbol],
     Map: [LuaLibFeature.InstanceOf, LuaLibFeature.Iterator, LuaLibFeature.Symbol],
     Set: [LuaLibFeature.InstanceOf, LuaLibFeature.Iterator, LuaLibFeature.Symbol],
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -3888,6 +3888,10 @@ export class LuaTransformer {
                     parameters,
                     node
                 );
+            case "flat":
+                return this.transformLuaLibFunction(LuaLibFeature.ArrayFlat, node, caller, ...params);
+            case "flatMap":
+                return this.transformLuaLibFunction(LuaLibFeature.ArrayFlatMap, node, caller, ...params);
             default:
                 throw TSTLErrors.UnsupportedProperty("array", expressionName as string, node);
         }
diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -25,6 +25,8 @@ const defaultArrayCallMethodNames = new Set<string>([
     "slice",
     "splice",
     "join",
+    "flat",
+    "flatMap",
 ]);
 
 const defaultArrayPropertyNames = new Set<string>([
diff --git a/src/lualib/ArrayFlat.ts b/src/lualib/ArrayFlat.ts
@@ -0,0 +1,12 @@
+function __TS__ArrayFlat(this: void, array: any[], depth = 1): any[] {
+    let result: any[] = [];
+    for (const value of array) {
+        if (depth > 0 && type(value) === "table" && 1 in value) {
+            result = result.concat(__TS__ArrayFlat(value, depth - 1));
+        } else {
+            result[result.length] = value;
+        }
+    }
+
+    return result;
+}
diff --git a/src/lualib/ArrayFlatMap.ts b/src/lualib/ArrayFlatMap.ts
@@ -0,0 +1,17 @@
+function __TS__ArrayFlatMap<T, U>(
+    this: void,
+    array: T[],
+    callback: (value: T, index: number, array: T[]) => U | ReadonlyArray<U>
+): U[] {
+    let result: U[] = [];
+    for (let i = 0; i < array.length; i++) {
+        const value = callback(array[i], i, array);
+        if (type(value) === "table" && 1 in value) {
+            result = result.concat(value);
+        } else {
+            result[result.length] = value as U;
+        }
+    }
+
+    return result;
+}
diff --git a/test/unit/lualib/lualib.spec.ts b/test/unit/lualib/lualib.spec.ts
@@ -377,6 +377,39 @@ test.each([
     expect(result).toBe(JSON.stringify(array.sort(compareFn)));
 });
 
+test.each([
+    { array: [1, [2, 3], 4], expected: [1, 2, 3, 4] },
+    { array: [1, [2, 3], 4], depth: 0, expected: [1, [2, 3], 4] },
+    { array: [1, [[2], [3]], 4], expected: [1, [2], [3], 4] },
+    { array: [1, [[[2], [3]]], 4], depth: Infinity, expected: [1, 2, 3, 4] },
+])("array.flat (%p)", ({ array, depth, expected }) => {
+    // TODO: Remove once `Infinity` would be implemented
+    const luaDepth = depth === Infinity ? "1 / 0" : depth;
+    const result = util.transpileAndExecute(`
+        return JSONStringify(${JSON.stringify(array)}.flat(${luaDepth}))
+    `);
+
+    expect(JSON.parse(result)).toEqual(expected);
+});
+
+test.each([
+    { array: [1, [2, 3], [4]], map: <T>(value: T) => value },
+    { array: [1, 2, 3], map: (v: number) => v * 2 },
+    { array: [1, 2, 3], map: (v: number) => [v, v * 2] },
+    { array: [1, 2, 3], map: (v: number) => [v, [v]] },
+    { array: [1, 2, 3], map: (v: number, i: number) => [v * 2 * i] },
+])("array.flatMap (%p)", ({ array, map }) => {
+    const result = util.transpileAndExecute(`
+        const array = ${JSON.stringify(array)};
+        const result = array.flatMap(${map.toString()});
+        return JSONStringify(result);
+    `);
+
+    // TODO(node 12): array.flatMap(map)
+    const expected = [].concat(...(array as any[]).map(map));
+    expect(JSON.parse(result)).toEqual(expected);
+});
+
 test.each([
     { condition: "true", lhs: "4", rhs: "5", expected: 4 },
     { condition: "false", lhs: "4", rhs: "5", expected: 5 },
