diff --git a/src/LuaLib.ts b/src/LuaLib.ts index 3db60196c..4dc023de0 100644 --- 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 index 97b789989..2e5b04d51 100644 --- 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 index b0326fe3d..717ba023c 100644 --- a/src/TSHelper.ts +++ b/src/TSHelper.ts @@ -25,6 +25,8 @@ const defaultArrayCallMethodNames = new Set([ "slice", "splice", "join", + "flat", + "flatMap", ]); const defaultArrayPropertyNames = new Set([ diff --git a/src/lualib/ArrayFlat.ts b/src/lualib/ArrayFlat.ts new file mode 100644 index 000000000..4006b51d5 --- /dev/null +++ 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 new file mode 100644 index 000000000..9c73db579 --- /dev/null +++ b/src/lualib/ArrayFlatMap.ts @@ -0,0 +1,17 @@ +function __TS__ArrayFlatMap( + this: void, + array: T[], + callback: (value: T, index: number, array: T[]) => U | ReadonlyArray +): 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 index 8c4052562..803ed4858 100644 --- 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: (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 },