Optimize array push array

GlassBricks · GlassBricks · commit 4a68d9ceb690 · 2022-03-17T13:58:31.000-07:00
diff --git a/src/LuaAST.ts b/src/LuaAST.ts
@@ -129,6 +129,7 @@ export enum NodeFlags {
     None = 0,
     Inline = 1 << 0, // Keep function body on same line
     Declaration = 1 << 1, // Prefer declaration syntax `function foo()` over assignment syntax `foo = function()`
+    TableUnpackCall = 1 << 2, // This is a table.unpack call
 }
 
 export interface TextRange {
@@ -197,6 +198,11 @@ export function getOriginalPos(node: Node): TextRange {
     return { line: node.line, column: node.column };
 }
 
+export function setNodeFlags<T extends Node>(node: T, flags: NodeFlags): T {
+    node.flags = flags;
+    return node;
+}
+
 export interface File extends Node {
     kind: SyntaxKind.File;
     statements: Statement[];
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -16,6 +16,7 @@ export enum LuaLibFeature {
     ArrayJoin = "ArrayJoin",
     ArrayMap = "ArrayMap",
     ArrayPush = "ArrayPush",
+    ArrayPushArray = "ArrayPushArray",
     ArrayReduce = "ArrayReduce",
     ArrayReduceRight = "ArrayReduceRight",
     ArrayReverse = "ArrayReverse",
diff --git a/src/lualib/ArrayPushArray.ts b/src/lualib/ArrayPushArray.ts
@@ -0,0 +1,8 @@
+export function __TS__ArrayPushArray<T>(this: T[], items: T[]): number {
+    let len = this.length;
+    for (const i of $range(1, items.length)) {
+        len++;
+        this[len - 1] = items[i - 1];
+    }
+    return len;
+}
diff --git a/src/transformation/builtins/array.ts b/src/transformation/builtins/array.ts
@@ -6,6 +6,7 @@ import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import { PropertyCallExpression, transformArguments, transformCallAndArguments } from "../visitors/call";
 import { isStringType, isNumberType, findFirstNonOuterParent } from "../utils/typescript";
 import { moveToPrecedingTemp } from "../visitors/expression-list";
+import { isUnpackCall } from "../utils/lua-ast";
 
 export function transformArrayConstructorCall(
     context: TransformationContext,
@@ -76,8 +77,20 @@ export function transformArrayPrototypeCall(
         case "entries":
             return transformLuaLibFunction(context, LuaLibFeature.ArrayEntries, node, caller);
         case "push":
-            if (node.arguments.length === 1 && !ts.isSpreadElement(node.arguments[0])) {
-                return transformSingleElementArrayPush(context, node, caller, params[0]);
+            if (node.arguments.length === 1) {
+                const param = params[0];
+                if (isUnpackCall(param)) {
+                    return transformLuaLibFunction(
+                        context,
+                        LuaLibFeature.ArrayPushArray,
+                        node,
+                        caller,
+                        (param as lua.CallExpression).params[0]
+                    );
+                }
+                if (!lua.isDotsLiteral(param)) {
+                    return transformSingleElementArrayPush(context, node, caller, param);
+                }
             }
 
             return transformLuaLibFunction(context, LuaLibFeature.ArrayPush, node, caller, ...params);
diff --git a/src/transformation/utils/lua-ast.ts b/src/transformation/utils/lua-ast.ts
@@ -76,7 +76,11 @@ export function createUnpackCall(
             ? lua.createIdentifier("unpack")
             : lua.createTableIndexExpression(lua.createIdentifier("table"), lua.createStringLiteral("unpack"));
 
-    return lua.createCallExpression(unpack, [expression], tsOriginal);
+    return lua.setNodeFlags(lua.createCallExpression(unpack, [expression], tsOriginal), lua.NodeFlags.TableUnpackCall);
+}
+
+export function isUnpackCall(node: lua.Node): boolean {
+    return lua.isCallExpression(node) && (node.flags & lua.NodeFlags.TableUnpackCall) !== 0;
 }
 
 export function wrapInTable(...expressions: lua.Expression[]): lua.TableExpression {
diff --git a/test/unit/builtins/array.spec.ts b/test/unit/builtins/array.spec.ts
@@ -472,12 +472,33 @@ test.each([
     util.testExpression`${util.formatCode(array)}.indexOf(${util.formatCode(...args)})`.expectToMatchJsResult();
 });
 
-test.each([{ args: "1" }, { args: "1, 2, 3" }, { args: "...[1, 2, 3]" }])("array.push (%p)", ({ args }) => {
+test.each([
+    { args: "1" },
+    { args: "1, 2, 3" },
+    { args: "...[1, 2, 3]" },
+    { args: "1, 2, ...[3, 4]" },
+    { args: "1, 2, ...[3, 4], 5" },
+])("array.push (%p)", ({ args }) => {
     util.testFunction`
         const array = [0];
         const value = array.push(${args});
         return { array, value };
-    `.expectToMatchJsResult();
+    `
+        .debug()
+        .expectToMatchJsResult();
+});
+
+test("array.push(...)", () => {
+    util.testFunction`
+        const array = [0];
+        function push(...args: any[]) {
+            return array.push(...args);
+        }
+        const value = push(1, 2, 3);
+        return { array, value };
+    `
+        .debug()
+        .expectToMatchJsResult();
 });
 
 test.each([
