diff --git a/src/LuaAST.ts b/src/LuaAST.ts index 424d09057..2639dbd85 100644 --- a/src/LuaAST.ts +++ b/src/LuaAST.ts @@ -813,6 +813,12 @@ export function createIdentifier( return expression; } +export function createAnnonymousIdentifier(tsOriginal?: ts.Node, parent?: Node): Identifier { + const expression = createNode(SyntaxKind.Identifier, tsOriginal, parent) as Identifier; + expression.text = "____"; + return expression; +} + export interface TableIndexExpression extends Expression { kind: SyntaxKind.TableIndexExpression; table: Expression; diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts index bd91692ef..5aa73cde1 100644 --- a/src/LuaTransformer.ts +++ b/src/LuaTransformer.ts @@ -10,7 +10,6 @@ import {TSTLErrors} from "./TSTLErrors"; export type StatementVisitResult = tstl.Statement | tstl.Statement[] | undefined; export type ExpressionVisitResult = tstl.Expression | undefined; - export enum ScopeType { File = 0x1, Function = 0x2, @@ -1031,6 +1030,144 @@ export class LuaTransformer { }); } + private transformGeneratorFunction( + parameters: ts.NodeArray, + body: ts.Block, + transformedParameters: tstl.Identifier[], + dotsLiteral: tstl.DotsLiteral, + spreadIdentifier?: tstl.Identifier + ): [tstl.Statement[], Scope] + { + this.importLuaLibFeature(LuaLibFeature.Symbol); + const [functionBody, functionScope] = this.transformFunctionBody( + parameters, + body, + spreadIdentifier + ); + + const coroutineIdentifier = tstl.createIdentifier("____co"); + const valueIdentifier = tstl.createIdentifier("____value"); + const errIdentifier = tstl.createIdentifier("____err"); + const itIdentifier = tstl.createIdentifier("____it"); + + //local ____co = coroutine.create(originalFunction) + const coroutine = + tstl.createVariableDeclarationStatement(coroutineIdentifier, + tstl.createCallExpression( + tstl.createTableIndexExpression(tstl.createIdentifier("coroutine"), + tstl.createStringLiteral("create") + ), + [tstl.createFunctionExpression( + tstl.createBlock(functionBody), + transformedParameters, + dotsLiteral, + spreadIdentifier), + ] + ) + ); + + const nextBody = []; + // coroutine.resume(__co, ...) + const resumeCall = tstl.createCallExpression( + tstl.createTableIndexExpression( + tstl.createIdentifier("coroutine"), + tstl.createStringLiteral("resume") + ), + [coroutineIdentifier, tstl.createDotsLiteral()] + ); + + // ____err, ____value = coroutine.resume(____co, ...) + nextBody.push(tstl.createVariableDeclarationStatement( + [errIdentifier, valueIdentifier], + resumeCall) + ); + + //coroutine.status(____co) ~= "dead"; + const coStatus = tstl.createCallExpression( + tstl.createTableIndexExpression( + tstl.createIdentifier("coroutine"), + tstl.createStringLiteral("status") + ), + [coroutineIdentifier] + ); + const status = tstl.createBinaryExpression( + coStatus, + tstl.createStringLiteral("dead"), + tstl.SyntaxKind.EqualityOperator + ); + nextBody.push(status); + //if(not ____err){error(____value)} + const errorCheck = tstl.createIfStatement( + tstl.createUnaryExpression( + errIdentifier, + tstl.SyntaxKind.NotOperator + ), + tstl.createBlock([ + tstl.createExpressionStatement( + tstl.createCallExpression( + tstl.createIdentifier("error"), + [valueIdentifier] + ) + ), + ]) + ); + nextBody.push(errorCheck); + //{done = coroutine.status(____co) ~= "dead"; value = ____value} + const iteratorResult = tstl.createTableExpression([ + tstl.createTableFieldExpression( + status, + tstl.createStringLiteral("done") + ), + tstl.createTableFieldExpression( + valueIdentifier, + tstl.createStringLiteral("value") + ), + ]); + nextBody.push(tstl.createReturnStatement([iteratorResult])); + + //function(____, ...) + const nextFunctionDeclaration = tstl.createFunctionExpression( + tstl.createBlock(nextBody), + [tstl.createAnnonymousIdentifier()], + tstl.createDotsLiteral()); + + //____it = {next = function(____, ...)} + const iterator = tstl.createVariableDeclarationStatement( + itIdentifier, + tstl.createTableExpression([ + tstl.createTableFieldExpression( + nextFunctionDeclaration, + tstl.createStringLiteral("next") + ), + ]) + ); + + const symbolIterator = tstl.createTableIndexExpression( + tstl.createIdentifier("Symbol"), + tstl.createStringLiteral("iterator") + ); + + const block = [ + coroutine, + iterator, + //____it[Symbol.iterator] = {return ____it} + tstl.createAssignmentStatement( + tstl.createTableIndexExpression( + itIdentifier, + symbolIterator + ), + tstl.createFunctionExpression( + tstl.createBlock( + [tstl.createReturnStatement([itIdentifier])] + ) + ) + ), + //return ____it + tstl.createReturnStatement([itIdentifier]), + ]; + return [block, functionScope]; + } + public transformFunctionDeclaration(functionDeclaration: ts.FunctionDeclaration): StatementVisitResult { // Don't transform functions without body (overload declarations) if (!functionDeclaration.body) { @@ -1044,14 +1181,21 @@ export class LuaTransformer { const [params, dotsLiteral, restParamName] = this.transformParameters(functionDeclaration.parameters, context); const name = this.transformIdentifier(functionDeclaration.name); - const [body, functionScope] = this.transformFunctionBody( - functionDeclaration.parameters, - functionDeclaration.body, - restParamName - ); + const [body, functionScope] = functionDeclaration.asteriskToken + ? this.transformGeneratorFunction( + functionDeclaration.parameters, + functionDeclaration.body, + params, + dotsLiteral, + restParamName + ) + : this.transformFunctionBody( + functionDeclaration.parameters, + functionDeclaration.body, + restParamName + ); const block = tstl.createBlock(body); const functionExpression = tstl.createFunctionExpression(block, params, dotsLiteral, restParamName); - // Remember symbols referenced in this function for hoisting later if (!this.options.noHoisting && name.symbolId !== undefined) { const scope = this.peekScope(); @@ -1059,7 +1203,6 @@ export class LuaTransformer { const functionInfo = {referencedSymbols: functionScope.referencedSymbols || new Set()}; scope.functionDefinitions.set(name.symbolId, functionInfo); } - return this.createLocalOrExportedOrGlobalDeclaration(name, functionExpression, functionDeclaration); } @@ -1219,6 +1362,12 @@ export class LuaTransformer { return tstl.createExpressionStatement(this.transformExpression(expression)); } + public transformYield(expression: ts.YieldExpression): tstl.Expression { + return tstl.createCallExpression( + tstl.createTableIndexExpression(tstl.createIdentifier("coroutine"), tstl.createStringLiteral("yield")), + expression.expression?[this.transformExpression(expression.expression)]:[], expression); + } + public transformReturn(statement: ts.ReturnStatement): tstl.Statement { if (statement.expression) { const returnType = tsHelper.getContainingFunctionReturnType(statement, this.checker); @@ -1731,6 +1880,8 @@ export class LuaTransformer { return this.transformSpreadElement(expression as ts.SpreadElement); case ts.SyntaxKind.NonNullExpression: return this.transformExpression((expression as ts.NonNullExpression).expression); + case ts.SyntaxKind.YieldExpression: + return this.transformYield(expression as ts.YieldExpression); case ts.SyntaxKind.EmptyStatement: return undefined; case ts.SyntaxKind.NotEmittedStatement: @@ -1871,7 +2022,6 @@ export class LuaTransformer { throw TSTLErrors.UnsupportedUnionAccessor(lhs); } } - return tstl.createAssignmentStatement( this.transformExpression(lhs) as tstl.IdentifierOrTableIndexExpression, right, diff --git a/test/unit/functions.spec.ts b/test/unit/functions.spec.ts index e50e9d3c5..40b49ccb7 100644 --- a/test/unit/functions.spec.ts +++ b/test/unit/functions.spec.ts @@ -389,6 +389,64 @@ export class FunctionTests { Expect(result).toBe("foobar"); } + @TestCase(1, 1) + @TestCase(2, 42) + @Test("Generator functions value") + public generatorFunctionValue(iterations: number, expectedResult: number): void { + const code = `function* seq(value: number) { + let a = yield value + 1; + return 42; + } + const gen = seq(0); + let ret: number; + for(let i = 0; i < ${iterations}; ++i) + { + ret = gen.next(i).value; + } + return ret; + `; + const result = util.transpileAndExecute(code); + Expect(result).toBe(expectedResult); + } + + @TestCase(1, false) + @TestCase(2, true) + @Test("Generator functions done") + public generatorFunctionDone(iterations: number, expectedResult: boolean): void { + const code = `function* seq(value: number) { + let a = yield value + 1; + return 42; + } + const gen = seq(0); + let ret: boolean; + for(let i = 0; i < ${iterations}; ++i) + { + ret = gen.next(i).done; + } + return ret; + `; + const result = util.transpileAndExecute(code); + Expect(result).toBe(expectedResult); + } + + @Test("Generator for..of") + public generatorFunctionForOf(): void { + const code = `function* seq() { + yield(1); + yield(2); + yield(3); + return 4; + } + let result = 0; + for(let i of seq()) + { + result = result * 10 + i; + } + return result`; + const result = util.transpileAndExecute(code); + Expect(result).toBe(123); + } + @Test("Function local overriding export") public functionLocalOverridingExport(): void { const code =