Using ipairs for for...of loops (#595)

tomblind · web-flow · commit 908db3699a4c · 2019-05-26T05:59:41.000-06:00
Extracted from #585 Also includes: - Throwing error on attempts to use object destructuring in for...of loops - Fixed issue with sourcemaps in lua for...in loops
diff --git a/src/LuaPrinter.ts b/src/LuaPrinter.ts
@@ -390,12 +390,12 @@ export class LuaPrinter {
     }
 
     public printForInStatement(statement: tstl.ForInStatement): SourceNode {
-        const names = statement.names.map(i => this.printIdentifier(i)).join(", ");
-        const expressions = statement.expressions.map(e => this.printExpression(e)).join(", ");
+        const names = this.joinChunks(", ", statement.names.map(i => this.printIdentifier(i)));
+        const expressions = this.joinChunks(", ", statement.expressions.map(e => this.printExpression(e)));
 
         const chunks: SourceChunk[] = [];
 
-        chunks.push(this.indent("for "), names, " in ", expressions, " do\n");
+        chunks.push(this.indent("for "), ...names, " in ", ...expressions, " do\n");
 
         this.pushIndent();
         chunks.push(this.printBlock(statement.body));
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -2284,6 +2284,9 @@ export class LuaTransformer {
             const variableDeclarations = this.transformVariableDeclaration(initializer.declarations[0]);
             if (ts.isArrayBindingPattern(initializer.declarations[0].name)) {
                 expression = this.createUnpackCall(expression, initializer);
+
+            } else if (ts.isObjectBindingPattern(initializer.declarations[0].name)) {
+                throw TSTLErrors.UnsupportedObjectDestructuringInForOf(initializer);
             }
 
             const variableStatements = this.statementVisitResultToArray(variableDeclarations);
@@ -2302,6 +2305,10 @@ export class LuaTransformer {
                 expression = this.createUnpackCall(expression, initializer);
                 variables = initializer.elements
                     .map(e => this.transformExpression(e)) as tstl.AssignmentLeftHandSideExpression[];
+
+            } else if (ts.isObjectLiteralExpression(initializer)) {
+                throw TSTLErrors.UnsupportedObjectDestructuringInForOf(initializer);
+
             } else {
                 variables = this.transformExpression(initializer) as tstl.AssignmentLeftHandSideExpression;
             }
@@ -2336,43 +2343,35 @@ export class LuaTransformer {
     }
 
     private transformForOfArrayStatement(statement: ts.ForOfStatement, block: tstl.Block): StatementVisitResult {
-        const arrayExpression = this.transformExpression(statement.expression);
-
-        // Arrays use numeric for loop (performs better than ipairs)
-        const indexVariable = tstl.createIdentifier("____TS_index");
-        if (!ts.isIdentifier(statement.expression)) {
-            // Cache iterable expression if it's not a simple identifier
-            // local ____TS_array = ${iterable};
-            // for ____TS_index = 1, #____TS_array do
-            //     local ${initializer} = ____TS_array[____TS_index]
-            const arrayVariable = tstl.createIdentifier("____TS_array", statement.expression);
-            const arrayAccess = tstl.createTableIndexExpression(arrayVariable, indexVariable);
-            const initializer = this.transformForOfInitializer(statement.initializer, arrayAccess);
-            block.statements.splice(0, 0, initializer);
-            return [
-                tstl.createVariableDeclarationStatement(arrayVariable, arrayExpression),
-                tstl.createForStatement(
-                    block,
-                    indexVariable,
-                    tstl.createNumericLiteral(1),
-                    tstl.createUnaryExpression(arrayVariable, tstl.SyntaxKind.LengthOperator)
-                ),
-            ];
+        let valueVariable: tstl.Identifier;
+        if (ts.isVariableDeclarationList(statement.initializer)) {
+            // Declaration of new variable
+            const variables = statement.initializer.declarations[0].name;
+            if (ts.isArrayBindingPattern(variables) || ts.isObjectBindingPattern(variables)) {
+                valueVariable = tstl.createIdentifier("____TS_values");
+                block.statements.unshift(this.transformForOfInitializer(statement.initializer, valueVariable));
+
+            } else {
+                valueVariable = this.transformIdentifier(variables);
+            }
 
         } else {
-            // Simple identifier version
-            // for ____TS_index = 1, #${iterable} do
-            //     local ${initializer} = ${iterable}[____TS_index]
-            const iterableAccess = tstl.createTableIndexExpression(arrayExpression, indexVariable);
-            const initializer = this.transformForOfInitializer(statement.initializer, iterableAccess);
-            block.statements.splice(0, 0, initializer);
-            return tstl.createForStatement(
-                block,
-                indexVariable,
-                tstl.createNumericLiteral(1),
-                tstl.createUnaryExpression(arrayExpression, tstl.SyntaxKind.LengthOperator)
-            );
+            // Assignment to existing variable
+            valueVariable = tstl.createIdentifier("____TS_value");
+            block.statements.unshift(this.transformForOfInitializer(statement.initializer, valueVariable));
         }
+
+        const ipairsCall = tstl.createCallExpression(
+            tstl.createIdentifier("ipairs"),
+            [this.transformExpression(statement.expression)]
+        );
+
+        return tstl.createForInStatement(
+            block,
+            [tstl.createAnonymousIdentifier(), valueVariable],
+            [ipairsCall],
+            statement
+        );
     }
 
     private transformForOfLuaIteratorStatement(statement: ts.ForOfStatement, block: tstl.Block): StatementVisitResult {
diff --git a/src/TSTLErrors.ts b/src/TSTLErrors.ts
@@ -191,6 +191,10 @@ export class TSTLErrors {
         );
     };
 
+    public static UnsupportedObjectDestructuringInForOf = (node: ts.Node) => {
+        return new TranspileError(`Unsupported object destructuring in for...of statement.`, node);
+    };
+
     public static InvalidAmbientLuaKeywordIdentifier = (node: ts.Identifier) => {
         return new TranspileError(
             `Invalid use of lua keyword "${node.text}" as ambient identifier name.`,
diff --git a/test/translation/__snapshots__/transformation.spec.ts.snap b/test/translation/__snapshots__/transformation.spec.ts.snap
@@ -238,7 +238,7 @@ end"
 `;
 
 exports[`Transformation (forOf) 1`] = `
-"local ____TS_array = {
+"for ____, i in ipairs({
     1,
     2,
     3,
@@ -249,9 +249,7 @@ exports[`Transformation (forOf) 1`] = `
     8,
     9,
     10,
-}
-for ____TS_index = 1, #____TS_array do
-    local i = ____TS_array[____TS_index]
+}) do
 end"
 `;
 
diff --git a/test/unit/loops.spec.ts b/test/unit/loops.spec.ts
@@ -504,6 +504,37 @@ test("forof destructuring with iterator and existing variables", () => {
     expect(result).toBe("0a1b2c");
 });
 
+test("forof array which modifies length", () => {
+    const code = `
+        const arr = ["a", "b", "c"];
+        let result = "";
+        for (const v of arr) {
+            if (v === "a") {
+                arr.push("d");
+            }
+            result += v;
+        }
+        return result;`;
+
+    expect(util.transpileAndExecute(code)).toBe("abcd");
+});
+
+test.each([
+    { initializer: "const {a, b}", vars: "" },
+    { initializer: "const {a: x, b: y}", vars: "" },
+    { initializer: "{a, b}", vars: "let a: string, b: string;" },
+    { initializer: "{a: c, b: d}", vars: "let c: string, d: string;" },
+])("forof object destructuring (%p)", ({ initializer, vars }) => {
+    const code = `
+        declare const arr: {a: string, b: string}[];
+        ${vars}
+        for (${initializer} of arr) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.UnsupportedObjectDestructuringInForOf(ts.createEmptyStatement()).message,
+    );
+});
+
 test("forof with array typed as iterable", () => {
     const code = `
         function foo(): Iterable<string> {
diff --git a/test/unit/sourcemaps.spec.ts b/test/unit/sourcemaps.spec.ts
@@ -102,8 +102,8 @@ test.each([
         `,
 
         assertPatterns: [
+            { luaPattern: "for", typeScriptPattern: "for" },
             { luaPattern: "getArr()", typeScriptPattern: "getArr()" },
-            { luaPattern: "____TS_array", typeScriptPattern: "getArr()" },
             { luaPattern: "element", typeScriptPattern: "element" },
         ],
     },
