diff --git a/src/LuaLib.ts b/src/LuaLib.ts index 1898be2cf..1b4c2aa2b 100644 --- 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 new file mode 100644 index 000000000..6b524365a --- /dev/null +++ b/src/lualib/DelegatedYield.ts @@ -0,0 +1,32 @@ +function __TS__DelegatedYield(this: void, iterable: string | GeneratorIterator | Iterable | 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 index e55de6404..809ca69ee 100644 --- 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( this: void, - iterable: Iterable | GeneratorIterator | readonly T[] + iterable: string | GeneratorIterator | Iterable | 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 index b62a115be..d3fb88a5a 100644 --- a/src/transformation/visitors/function.ts +++ b/src/transformation/visitors/function.ts @@ -279,9 +279,13 @@ export const transformFunctionDeclaration: FunctionVisitor = (expression, context) => - lua.createCallExpression( - lua.createTableIndexExpression(lua.createIdentifier("coroutine"), lua.createStringLiteral("yield")), - expression.expression ? [context.transformExpression(expression.expression)] : [], - expression - ); +export const transformYieldExpression: FunctionVisitor = (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 index 90fb67aa4..7065aecf6 100644 --- 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*() {