TypeScriptToLua · tomblind · Jun 13, 2019 · Jun 6, 2019 · Jun 7, 2019
diff --git a/src/Decorator.ts b/src/Decorator.ts
@@ -29,6 +29,8 @@ export class Decorator {
                 return DecoratorKind.NoSelf;
             case "noselfinfile":
                 return DecoratorKind.NoSelfInFile;
+            case "forrange":
+                return DecoratorKind.ForRange;
         }
 
         return undefined;
@@ -61,4 +63,5 @@ export enum DecoratorKind {
     LuaTable = "LuaTable",
     NoSelf = "NoSelf",
     NoSelfInFile = "NoSelfInFile",
+    ForRange = "ForRange",
 }
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -2379,11 +2379,59 @@ export class LuaTransformer {
         }
     }
 
+    protected transformForRangeStatement(statement: ts.ForOfStatement, body: tstl.Block): StatementVisitResult {
+        if (!ts.isCallExpression(statement.expression)) {
+            throw TSTLErrors.InvalidForRangeCall(statement.expression, "Expression must be a call expression.");
+        }
+
+        if (statement.expression.arguments.length < 2 || statement.expression.arguments.length > 3) {
+            throw TSTLErrors.InvalidForRangeCall(
+                statement.expression,
+                "@forRange function must take 2 or 3 arguments."
+            );
+        }
+
+        if (statement.expression.arguments.some(a => !tsHelper.isNumberType(this.checker.getTypeAtLocation(a)))) {
+            throw TSTLErrors.InvalidForRangeCall(statement.expression, "@forRange arguments must be number types.");
+        }
+
+        if (!ts.isVariableDeclarationList(statement.initializer)) {
+            throw TSTLErrors.InvalidForRangeCall(
+                statement.initializer,
+                "@forRange loop must declare its own control variable."
+            );
+        }
+
+        const controlDeclaration = statement.initializer.declarations[0];
+        if (!ts.isIdentifier(controlDeclaration.name)) {
+            throw TSTLErrors.InvalidForRangeCall(statement.initializer, "@forRange loop cannot use destructuring.");
+        }
+
+        const controlType = this.checker.getTypeAtLocation(controlDeclaration);
+        if (controlType && !tsHelper.isNumberType(controlType)) {
+            throw TSTLErrors.InvalidForRangeCall(
+                statement.expression,
+                "@forRange function must return Iterable<number> or Array<number>."
+            );
+        }
+
+        const control = this.transformIdentifier(controlDeclaration.name);
+        const signature = this.checker.getResolvedSignature(statement.expression);
+        const [start, limit, step] = this.transformArguments(statement.expression.arguments, signature);
+        return tstl.createForStatement(body, control, start, limit, step, statement);
+    }
+
     public transformForOfStatement(statement: ts.ForOfStatement): StatementVisitResult {
         // Transpile body
         const body = tstl.createBlock(this.transformLoopBody(statement));
 
-        if (tsHelper.isLuaIteratorType(statement.expression, this.checker)) {
+        if (
+            ts.isCallExpression(statement.expression) &&
+            tsHelper.isForRangeType(statement.expression.expression, this.checker)
+        ) {
+            // ForRange
+            return this.transformForRangeStatement(statement, body);
+        } else if (tsHelper.isLuaIteratorType(statement.expression, this.checker)) {
             // LuaIterators
             return this.transformForOfLuaIteratorStatement(statement, body);
         } else if (
@@ -4626,6 +4674,16 @@ export class LuaTransformer {
             return tstl.createIdentifier("nil");
         }
 
+        if (tsHelper.isForRangeType(identifier, this.checker)) {
+            const callExpression = tsHelper.findFirstNodeAbove(identifier, ts.isCallExpression);
+            if (!callExpression || !callExpression.parent || !ts.isForOfStatement(callExpression.parent)) {
+                throw TSTLErrors.InvalidForRangeCall(
+                    identifier,
+                    "@forRange function can only be used as an iterable in a for...of loop."
+                );
+            }
+        }
+
         const text = this.hasUnsafeIdentifierName(identifier)
             ? this.createSafeName(this.getIdentifierText(identifier))
             : this.getIdentifierText(identifier);

diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -184,6 +184,11 @@ export class TSHelper {
         return TSHelper.getCustomDecorators(type, checker).has(DecoratorKind.LuaIterator);
     }
 
+    public static isForRangeType(node: ts.Node, checker: ts.TypeChecker): boolean {
+        const type = checker.getTypeAtLocation(node);
+        return TSHelper.getCustomDecorators(type, checker).has(DecoratorKind.ForRange);
+    }
+
     public static isTupleReturnCall(node: ts.Node, checker: ts.TypeChecker): boolean {
         if (ts.isCallExpression(node)) {
             const signature = checker.getResolvedSignature(node);

diff --git a/src/TSTLErrors.ts b/src/TSTLErrors.ts
@@ -204,4 +204,8 @@ export class TSTLErrors {
             node
         );
     };
+
+    public static InvalidForRangeCall = (node: ts.Node, message: string) => {
+        return new TranspileError(`Invalid @forRange call: ${message}`, node);
+    };
 }
diff --git a/test/unit/loops.spec.ts b/test/unit/loops.spec.ts
@@ -797,3 +797,112 @@ test("while dead code after return", () => {
 
     expect(result).toBe(3);
 });
+
+test.each([
+    { args: [1, 10], expectResult: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] },
+    { args: [1, 10, 2], expectResult: [1, 3, 5, 7, 9] },
+    { args: [10, 1, -1], expectResult: [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] },
+    { args: [10, 1, -2], expectResult: [10, 8, 6, 4, 2] },
+])("@forRange loop", ({ args, expectResult }) => {
+    const tsHeader = "/** @forRange **/ declare function luaRange(i: number, j: number, k?: number): number[];";
+    const code = `
+        const results: number[] = [];
+        for (const i of luaRange(${args})) {
+            results.push(i);
+        }
+        return JSONStringify(results);`;
+
+    const result = util.transpileAndExecute(code, undefined, undefined, tsHeader);
+    expect(JSON.parse(result)).toEqual(expectResult);
+});
+
+test("invalid non-ambient @forRange function", () => {
+    const code = `
+        /** @forRange **/ function luaRange(i: number, j: number, k?: number): number[] { return []; }
+        for (const i of luaRange(1, 10, 2)) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(
+            ts.createEmptyStatement(),
+            "@forRange function can only be used as an iterable in a for...of loop."
+        ).message
+    );
+});
+
+test.each([[1], [1, 2, 3, 4]])("invalid @forRange argument count", args => {
+    const code = `
+        /** @forRange **/ declare function luaRange(...args: number[]): number[] { return []; }
+        for (const i of luaRange(${args})) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(ts.createEmptyStatement(), "@forRange function must take 2 or 3 arguments.")
+            .message
+    );
+});
+
+test("invalid @forRange control variable", () => {
+    const code = `
+        /** @forRange **/ declare function luaRange(i: number, j: number, k?: number): number[];
+        let i: number;
+        for (i of luaRange(1, 10, 2)) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(
+            ts.createEmptyStatement(),
+            "@forRange loop must declare its own control variable."
+        ).message
+    );
+});
+
+test("invalid @forRange argument type", () => {
+    const code = `
+        /** @forRange **/ declare function luaRange(i: string, j: number): number[] { return []; }
+        for (const i of luaRange("foo", 2)) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(ts.createEmptyStatement(), "@forRange arguments must be number types.").message
+    );
+});
+
+test("invalid @forRange destructuring", () => {
+    const code = `
+        /** @forRange **/ declare function luaRange(i: number, j: number, k?: number): number[][];
+        for (const [i] of luaRange(1, 10, 2)) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(ts.createEmptyStatement(), "@forRange loop cannot use destructuring.").message
+    );
+});
+
+test("invalid @forRange return type", () => {
+    const code = `
+        /** @forRange **/ declare function luaRange(i: number, j: number, k?: number): string[];
+        for (const i of luaRange(1, 10)) {}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(
+            ts.createEmptyStatement(),
+            "@forRange function must return Iterable<number> or Array<number>."
+        ).message
+    );
+});
+
+test.each([
+    "const range = luaRange(1, 10);",
+    "console.log(luaRange);",
+    "luaRange.call(null, 0, 0, 0);",
+    "let array = [0, luaRange, 1];",
+    "const call: any; call(luaRange);",
+    "for (const i of [...luaRange(1, 10)]) {}",
+])("invalid @forRange reference (%p)", statement => {
+    const code = `
+        /** @forRange **/ declare function luaRange(i: number, j: number, k?: number): number[];
+        ${statement}`;
+
+    expect(() => util.transpileString(code)).toThrow(
+        TSTLErrors.InvalidForRangeCall(
+            ts.createEmptyStatement(),
+            "@forRange function can only be used as an iterable in a for...of loop."
+        ).message
+    );
+});
-Original file line number
+Diff line change
@@ Expand Up / @@ -204,4 +204,8 @@ export class TSTLErrors { @@
                 node
             );
         };
+        public static InvalidForRangeCall = (node: ts.Node, message: string) => {
+            return new TranspileError(`Invalid @forRange call: ${message}`, node);
+        };
     }