yield* support (#923)

ark120202 · web-flow · commit c76813f91ae5 · 2020-10-17T12:26:32.000+02:00
* `yield*` support

* Remove `.debug()` call
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -29,6 +29,7 @@ export enum LuaLibFeature {
     Class = "Class",
     ClassExtends = "ClassExtends",
     Decorate = "Decorate",
+    DelegatedYield = "DelegatedYield",
     Descriptors = "Descriptors",
     Error = "Error",
     FunctionBind = "FunctionBind",
diff --git a/src/lualib/DelegatedYield.ts b/src/lualib/DelegatedYield.ts
@@ -0,0 +1,32 @@
+function __TS__DelegatedYield<T>(this: void, iterable: string | GeneratorIterator | Iterable<T> | readonly T[]) {
+    if (typeof iterable === "string") {
+        for (const index of forRange(0, iterable.length - 1)) {
+            coroutine.yield(iterable[index]);
+        }
+    } else if ("____coroutine" in iterable) {
+        const co = iterable.____coroutine;
+        while (true) {
+            const [status, value] = coroutine.resume(co);
+            if (!status) throw value;
+            if (coroutine.status(co) === "dead") {
+                return value;
+            } else {
+                coroutine.yield(value);
+            }
+        }
+    } else if (iterable[Symbol.iterator]) {
+        const iterator = iterable[Symbol.iterator]();
+        while (true) {
+            const result = iterator.next();
+            if (result.done) {
+                return result.value;
+            } else {
+                coroutine.yield(result.value);
+            }
+        }
+    } else {
+        for (const value of iterable as readonly T[]) {
+            coroutine.yield(value);
+        }
+    }
+}
diff --git a/src/lualib/Iterator.ts b/src/lualib/Iterator.ts
@@ -26,15 +26,15 @@ function __TS__IteratorStringStep(this: string, index: number): [number, string]
 /** @tupleReturn */
 function __TS__Iterator<T>(
     this: void,
-    iterable: Iterable<T> | GeneratorIterator | readonly T[]
+    iterable: string | GeneratorIterator | Iterable<T> | readonly T[]
 ): [(...args: any[]) => [any, any] | [], ...any[]] {
-    if ("____coroutine" in iterable) {
+    if (typeof iterable === "string") {
+        return [__TS__IteratorStringStep, iterable, 0];
+    } else if ("____coroutine" in iterable) {
         return [__TS__IteratorGeneratorStep, iterable];
     } else if (iterable[Symbol.iterator]) {
         const iterator = iterable[Symbol.iterator]();
         return [__TS__IteratorIteratorStep, iterator];
-    } else if (typeof iterable === "string") {
-        return [__TS__IteratorStringStep, iterable, 0];
     } else {
         return ipairs(iterable as readonly T[]) as any;
     }
diff --git a/src/transformation/visitors/function.ts b/src/transformation/visitors/function.ts
@@ -279,9 +279,13 @@ 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
-    );
+export const transformYieldExpression: FunctionVisitor<ts.YieldExpression> = (expression, context) => {
+    const parameters = expression.expression ? [context.transformExpression(expression.expression)] : [];
+    return expression.asteriskToken
+        ? transformLuaLibFunction(context, LuaLibFeature.DelegatedYield, expression, ...parameters)
+        : lua.createCallExpression(
+              lua.createTableIndexExpression(lua.createIdentifier("coroutine"), lua.createStringLiteral("yield")),
+              parameters,
+              expression
+          );
+};
diff --git a/test/unit/functions/generators.spec.ts b/test/unit/functions/generators.spec.ts
@@ -51,6 +51,59 @@ test("for..of", () => {
     `.expectToMatchJsResult();
 });
 
+describe("yield*", () => {
+    test("generator", () => {
+        util.testFunction`
+            function* subGenerator() {
+                yield 1;
+                yield 2;
+                yield 3;
+            }
+
+            function* generator() {
+                yield 0;
+                return yield* subGenerator();
+            }
+
+            const it = generator();
+            return [it.next(), it.next(), it.next(), it.next(), it.next()];
+        `.expectToMatchJsResult();
+    });
+
+    test("array", () => {
+        util.testFunction`
+            function* generator() {
+                return yield* [1, 2, 3];
+            }
+
+            const it = generator();
+            return [it.next(), it.next(), it.next(), it.next()];
+        `.expectToMatchJsResult();
+    });
+
+    test("string", () => {
+        util.testFunction`
+            function* generator() {
+                return yield* "abc";
+            }
+
+            const it = generator();
+            return [it.next(), it.next(), it.next(), it.next()];
+        `.expectToMatchJsResult();
+    });
+
+    test("iterable", () => {
+        util.testFunction`
+            function* generator() {
+                return yield* new Set([1, 2, 3]);
+            }
+
+            const it = generator();
+            return [it.next(), it.next(), it.next(), it.next()];
+        `.expectToMatchJsResult();
+    });
+});
+
 test("function expression", () => {
     util.testFunction`
         const generator = function*() {
