TypeScriptToLua · Perryvw · May 23, 2020 · May 16, 2020 · May 19, 2020 · May 19, 2020
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
 <!-- TODO: https://www.typescriptlang.org/docs/handbook/release-notes/typescript-3-9.html doesn't seem to work now -->
 
 - TypeScript has been updated to 3.9. See [release notes](https://devblogs.microsoft.com/typescript/announcing-typescript-3-9/) for details. This update includes some fixes specific to our API usage:
+
   - Importing a non-module using `import "./file"` produced a TS2307 error [#35973](https://github.com/microsoft/TypeScript/issues/35973)
   - TypeScript now tries to find a call signature even in presence of type errors (#36665)(https://github.com/microsoft/TypeScript/pull/36665):
     ```ts
@@ -18,6 +19,12 @@
     foo(1)
     ```
 
+- Reduced memory consumption and optimized performance of generators and iterators
+
+- Fixed generator syntax being ignored on methods (`*foo() {}`) and function expressions (`function*() {}`)
+
+- Fixed iteration over generators stopping at first yielded `nil` value
+
 ## 0.33.0
 
 - Added support for nullish coalescing `A ?? B`.

diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -33,6 +33,7 @@ export enum LuaLibFeature {
     FunctionApply = "FunctionApply",
     FunctionBind = "FunctionBind",
     FunctionCall = "FunctionCall",
+    Generator = "Generator",
     InstanceOf = "InstanceOf",
     InstanceOfObject = "InstanceOfObject",
     Iterator = "Iterator",
@@ -72,6 +73,7 @@ const luaLibDependencies: Partial<Record<LuaLibFeature, LuaLibFeature[]>> = {
     ArrayFlat: [LuaLibFeature.ArrayConcat],
     ArrayFlatMap: [LuaLibFeature.ArrayConcat],
     Error: [LuaLibFeature.New, LuaLibFeature.Class, LuaLibFeature.FunctionCall],
+    Generator: [LuaLibFeature.Symbol],
     InstanceOf: [LuaLibFeature.Symbol],
     Iterator: [LuaLibFeature.Symbol],
     ObjectFromEntries: [LuaLibFeature.Iterator, LuaLibFeature.Symbol],

diff --git a/src/lualib/Generator.ts b/src/lualib/Generator.ts
@@ -0,0 +1,31 @@
+interface GeneratorIterator {
+    ____coroutine: LuaThread;
+    [Symbol.iterator](): GeneratorIterator;
+    next: typeof __TS__GeneratorNext;
+}
+
+function __TS__GeneratorIterator(this: GeneratorIterator) {
+    return this;
+}
+
+function __TS__GeneratorNext(this: GeneratorIterator, ...args: Vararg<any[]>) {
+    const co = this.____coroutine;
+    if (coroutine.status(co) === "dead") return { done: true };
+
+    const [status, value] = coroutine.resume(co, ...args);
+    if (!status) throw value;
+
+    return { value, done: coroutine.status(co) === "dead" };
+}
+
+function __TS__Generator(this: void, fn: (this: void, ...args: any[]) => any) {
+    return function(this: void, ...args: Vararg<any[]>): GeneratorIterator {
+        const argsLength = select("#", ...args);
+        return {
+            // Using explicit this there, since we don't pass arguments after the first nil and context is likely to be nil
+            ____coroutine: coroutine.create(() => fn((unpack || table.unpack)(args, 1, argsLength))),
+            [Symbol.iterator]: __TS__GeneratorIterator,
+            next: __TS__GeneratorNext,
+        };
+    };
+}
diff --git a/src/lualib/Iterator.ts b/src/lualib/Iterator.ts
@@ -1,19 +1,32 @@
-function __TS__Iterator<T>(this: void, iterable: Iterable<T>): (this: void) => T {
-    if (iterable[Symbol.iterator]) {
+/** @tupleReturn */
+function __TS__IteratorGeneratorStep(this: GeneratorIterator): [true, any] | [] {
+    const co = this.____coroutine;
+
+    const [status, value] = coroutine.resume(co);
+    if (!status) throw value;
+
+    if (coroutine.status(co) === "dead") return [];
+    return [true, value];
+}
+
+/** @tupleReturn */
+function __TS__IteratorIteratorStep<T>(this: Iterator<T>): [true, T] | [] {
+    const result = this.next();
+    if (result.done) return [];
+    return [true, result.value];
+}
+
+/** @tupleReturn */
+function __TS__Iterator<T>(
+    this: void,
+    iterable: Iterable<T> | GeneratorIterator | readonly T[]
+): [(...args: any[]) => [any, T] | [], ...any[]] {
+    if ("____coroutine" in iterable) {
+        return [__TS__IteratorGeneratorStep, iterable];
+    } else if (iterable[Symbol.iterator]) {
         const iterator = iterable[Symbol.iterator]();
-        return () => {
-            const result = iterator.next();
-            if (!result.done) {
-                return result.value;
-            } else {
-                return undefined;
-            }
-        };
+        return [__TS__IteratorIteratorStep, iterator];
     } else {
-        let i = 0;
-        return () => {
-            i += 1;
-            return iterable[i];
-        };
+        return ipairs(iterable as readonly T[]);
     }
 }
diff --git a/src/lualib/declarations/coroutine.d.ts b/src/lualib/declarations/coroutine.d.ts
@@ -0,0 +1,19 @@
+/** @noSelfInFile */
+
+interface LuaThread {
+    readonly __internal__: unique symbol;
+}
+
+declare namespace coroutine {
+    function create(f: (...args: any[]) => any): LuaThread;
+
+    /** @tupleReturn */
+    function resume(co: LuaThread, ...val: any[]): [true, ...any[]] | [false, string];
+
+    function status(co: LuaThread): "running" | "suspended" | "normal" | "dead";
+
+    function wrap(f: (...args: any[]) => any): /** @tupleReturn */ (...args: any[]) => any[];
+
+    /** @tupleReturn */
+    function yield(...args: any[]): any[];
+}
diff --git a/src/lualib/declarations/global.d.ts b/src/lualib/declarations/global.d.ts
@@ -24,3 +24,9 @@ declare function unpack<T>(list: T[], i?: number, j?: number): T[];
 
 declare function select<T>(index: number, ...args: T[]): T;
 declare function select<T>(index: "#", ...args: T[]): number;
+
+/**
+ * @luaIterator
+ * @tupleReturn
+ */
+declare function ipairs<T>(t: Record<number, T>): LuaTupleIterable<[number, T]>;
diff --git a/src/transformation/visitors/class/members/method.ts b/src/transformation/visitors/class/members/method.ts
@@ -1,9 +1,7 @@
 import * as ts from "typescript";
 import * as lua from "../../../../LuaAST";
 import { TransformationContext } from "../../../context";
-import { ContextType, getFunctionContextType } from "../../../utils/function-context";
-import { createSelfIdentifier } from "../../../utils/lua-ast";
-import { transformFunctionBody, transformParameters } from "../../function";
+import { transformFunctionToExpression } from "../../function";
 import { transformPropertyName } from "../../literal";
 import { isStaticNode } from "../utils";
 
@@ -18,29 +16,17 @@ export function transformMethodDeclaration(
         return undefined;
     }
 
+    const methodTable =
+        isStaticNode(node) || noPrototype
+            ? lua.cloneIdentifier(className)
+            : lua.createTableIndexExpression(lua.cloneIdentifier(className), lua.createStringLiteral("prototype"));
+
     let methodName = transformPropertyName(context, node.name);
     if (lua.isStringLiteral(methodName) && methodName.value === "toString") {
         methodName = lua.createStringLiteral("__tostring", node.name);
     }
 
-    const type = context.checker.getTypeAtLocation(node);
-    const functionContext =
-        getFunctionContextType(context, type) !== ContextType.Void ? createSelfIdentifier() : undefined;
-    const [paramNames, dots, restParamName] = transformParameters(context, node.parameters, functionContext);
-
-    const [body] = transformFunctionBody(context, node.parameters, node.body, restParamName);
-    const functionExpression = lua.createFunctionExpression(
-        lua.createBlock(body),
-        paramNames,
-        dots,
-        lua.FunctionExpressionFlags.Declaration,
-        node.body
-    );
-
-    const methodTable =
-        isStaticNode(node) || noPrototype
-            ? lua.cloneIdentifier(className)
-            : lua.createTableIndexExpression(lua.cloneIdentifier(className), lua.createStringLiteral("prototype"));
+    const [functionExpression] = transformFunctionToExpression(context, node);
 
     return lua.createAssignmentStatement(
         lua.createTableIndexExpression(methodTable, methodName),

diff --git a/src/transformation/visitors/function.ts b/src/transformation/visitors/function.ts
@@ -1,5 +1,6 @@
 import * as ts from "typescript";
 import * as lua from "../../LuaAST";
+import { assert } from "../../utils";
 import { FunctionVisitor, TransformationContext } from "../context";
 import { isVarargType } from "../utils/annotations";
 import { createDefaultExportStringLiteral, hasDefaultExportModifier } from "../utils/export";
@@ -11,8 +12,8 @@ import {
     createSelfIdentifier,
     wrapInTable,
 } from "../utils/lua-ast";
+import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import { peekScope, performHoisting, popScope, pushScope, Scope, ScopeType } from "../utils/scope";
-import { transformGeneratorFunctionBody } from "./generator";
 import { transformIdentifier } from "./identifier";
 import { transformBindingPattern } from "./variable-declaration";
 
@@ -157,12 +158,13 @@ export function transformParameters(
     return [paramNames, dotsLiteral, restParamName];
 }
 
-export function transformFunctionLikeDeclaration(
-    node: ts.FunctionLikeDeclaration,
-    context: TransformationContext
-): lua.Expression {
-    const type = context.checker.getTypeAtLocation(node);
+export function transformFunctionToExpression(
+    context: TransformationContext,
+    node: ts.FunctionLikeDeclaration
+): [lua.Expression, Scope] {
+    assert(node.body);
 
+    const type = context.checker.getTypeAtLocation(node);
     let functionContext: lua.Identifier | undefined;
     if (getFunctionContextType(context, type) !== ContextType.Void) {
         if (ts.isArrowFunction(node)) {
@@ -176,14 +178,10 @@ export function transformFunctionLikeDeclaration(
         }
     }
 
-    // Build parameter string
-    const [paramNames, dotsLiteral, spreadIdentifier] = transformParameters(context, node.parameters, functionContext);
-
     let flags = lua.FunctionExpressionFlags.None;
-
-    if (node.body === undefined) {
-        // This code can be reached only from object methods, which is TypeScript error
-        return lua.createNilLiteral();
+    if (!ts.isBlock(node.body)) flags |= lua.FunctionExpressionFlags.Inline;
+    if (ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node)) {
+        flags |= lua.FunctionExpressionFlags.Declaration;
     }
 
     let body: ts.Block;
@@ -193,14 +191,11 @@ export function transformFunctionLikeDeclaration(
         const returnExpression = ts.createReturn(node.body);
         body = ts.createBlock([returnExpression]);
         returnExpression.parent = body;
-        if (node.body) {
-            body.parent = node.body.parent;
-        }
-        flags |= lua.FunctionExpressionFlags.Inline;
+        if (node.body) body.parent = node.body.parent;
     }
 
-    const [transformedBody, scope] = transformFunctionBody(context, node.parameters, body, spreadIdentifier);
-
+    const [paramNames, dotsLiteral, spreadIdentifier] = transformParameters(context, node.parameters, functionContext);
+    const [transformedBody, functionScope] = transformFunctionBody(context, node.parameters, body, spreadIdentifier);
     const functionExpression = lua.createFunctionExpression(
         lua.createBlock(transformedBody),
         paramNames,
@@ -209,12 +204,31 @@ export function transformFunctionLikeDeclaration(
         node
     );
 
+    return [
+        node.asteriskToken
+            ? transformLuaLibFunction(context, LuaLibFeature.Generator, undefined, functionExpression)
+            : functionExpression,
+        functionScope,
+    ];
+}
+
+export function transformFunctionLikeDeclaration(
+    node: ts.FunctionLikeDeclaration,
+    context: TransformationContext
+): lua.Expression {
+    if (node.body === undefined) {
+        // This code can be reached only from object methods, which is TypeScript error
+        return lua.createNilLiteral();
+    }
+
+    const [functionExpression, functionScope] = transformFunctionToExpression(context, node);
+
     // Handle named function expressions which reference themselves
-    if (ts.isFunctionExpression(node) && node.name && scope.referencedSymbols) {
+    if (ts.isFunctionExpression(node) && node.name && functionScope.referencedSymbols) {
         const symbol = context.checker.getSymbolAtLocation(node.name);
         if (symbol) {
             // TODO: Not using symbol ids because of https://github.com/microsoft/TypeScript/issues/37131
-            const isReferenced = [...scope.referencedSymbols].some(([, nodes]) =>
+            const isReferenced = [...functionScope.referencedSymbols].some(([, nodes]) =>
                 nodes.some(n => context.checker.getSymbolAtLocation(n)?.valueDeclaration === symbol.valueDeclaration)
             );
 
@@ -234,34 +248,19 @@ export function transformFunctionLikeDeclaration(
 
 export const transformFunctionDeclaration: FunctionVisitor<ts.FunctionDeclaration> = (node, context) => {
     // Don't transform functions without body (overload declarations)
-    if (!node.body) {
+    if (node.body === undefined) {
         return undefined;
     }
 
-    const type = context.checker.getTypeAtLocation(node);
-    const functionContext =
-        getFunctionContextType(context, type) !== ContextType.Void ? createSelfIdentifier() : undefined;
-    const [params, dotsLiteral, restParamName] = transformParameters(context, node.parameters, functionContext);
-
-    const [body, functionScope] = node.asteriskToken
-        ? transformGeneratorFunctionBody(context, node.parameters, node.body, restParamName)
-        : transformFunctionBody(context, node.parameters, node.body, restParamName);
-
-    const block = lua.createBlock(body);
-    const functionExpression = lua.createFunctionExpression(
-        block,
-        params,
-        dotsLiteral,
-        lua.FunctionExpressionFlags.Declaration
-    );
-
     if (hasDefaultExportModifier(node)) {
         return lua.createAssignmentStatement(
             lua.createTableIndexExpression(createExportsIdentifier(), createDefaultExportStringLiteral(node)),
             transformFunctionLikeDeclaration(node, context)
         );
     }
 
+    const [functionExpression, functionScope] = transformFunctionToExpression(context, node);
+
     // Name being undefined without default export is a TypeScript error
     const name = node.name ? transformIdentifier(context, node.name) : lua.createAnonymousIdentifier();
 
@@ -278,3 +277,10 @@ export const transformFunctionDeclaration: FunctionVisitor<ts.FunctionDeclaratio
 
     return createLocalOrExportedOrGlobalDeclaration(context, name, functionExpression, node);
 };
+
+export const transformYieldExpression: FunctionVisitor<ts.YieldExpression> = (expression, context) =>
+    lua.createCallExpression(
+        lua.createTableIndexExpression(lua.createIdentifier("coroutine"), lua.createStringLiteral("yield")),
+        expression.expression ? [context.transformExpression(expression.expression)] : [],
+        expression
+    );