Support for spread operator on iterables (#569)

tomblind · Perryvw · commit 1fc5f2607d38 · 2019-05-05T16:22:03.000+02:00
* Support for spread operator on iterables fixes #534 * adressing feedback on test
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -44,6 +44,7 @@ export enum LuaLibFeature {
     WeakMap = "WeakMap",
     WeakSet = "WeakSet",
     SourceMapTraceBack = "SourceMapTraceBack",
+    Spread = "Spread",
     StringConcat = "StringConcat",
     StringEndsWith = "StringEndsWith",
     StringReplace = "StringReplace",
@@ -63,6 +64,7 @@ const luaLibDependencies: {[lib in LuaLibFeature]?: LuaLibFeature[]} = {
     Set: [LuaLibFeature.InstanceOf, LuaLibFeature.Iterator, LuaLibFeature.Symbol],
     WeakMap: [LuaLibFeature.InstanceOf, LuaLibFeature.Iterator, LuaLibFeature.Symbol],
     WeakSet: [LuaLibFeature.InstanceOf, LuaLibFeature.Iterator, LuaLibFeature.Symbol],
+    Spread: [LuaLibFeature.Iterator],
     SymbolRegistry: [LuaLibFeature.Symbol],
 };
 
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -4566,9 +4566,14 @@ export class LuaTransformer {
         const innerExpression = this.expectExpression(this.transformExpression(expression.expression));
         if (tsHelper.isTupleReturnCall(expression.expression, this.checker)) {
             return innerExpression;
-        } else {
+        }
+
+        const type = this.checker.getTypeAtLocation(expression.expression);
+        if (tsHelper.isArrayType(type, this.checker, this.program)) {
             return this.createUnpackCall(innerExpression, expression);
         }
+
+        return this.transformLuaLibFunction(LuaLibFeature.Spread, expression, innerExpression);
     }
 
     public transformStringLiteral(literal: ts.StringLiteralLike): ExpressionVisitResult {
diff --git a/src/lualib/Spread.ts b/src/lualib/Spread.ts
@@ -0,0 +1,13 @@
+declare function unpack<T>(this: void, list: T[], i?: number, j?: number): T[];
+
+declare namespace table {
+    export function unpack<T>(this: void, list: T[], i?: number, j?: number): T[];
+}
+
+function __TS__Spread<T>(this: void, iterable: Iterable<T>): T[] {
+    const arr: T[] = [];
+    for (const item of iterable) {
+        arr[arr.length] = item;
+    }
+    return (table.unpack || unpack)(arr);
+}
diff --git a/test/unit/spreadElement.spec.ts b/test/unit/spreadElement.spec.ts
@@ -47,3 +47,23 @@ test("Spread Element Lua JIT", () => {
     const lua = util.transpileString(`[...[0, 1, 2]]`, options);
     expect(lua).toBe("local ____ = {unpack({\n    0,\n    1,\n    2,\n})}");
 });
+
+test("Spread Element Iterable", () => {
+    const code = `
+        const it = {
+            i: -1,
+            [Symbol.iterator]() {
+                return this;
+            },
+            next() {
+                ++this.i;
+                return {
+                    value: 2 ** this.i,
+                    done: this.i == 9,
+                }
+            }
+        };
+        const arr = [...it];
+        return JSONStringify(arr)`;
+    expect(JSON.parse(util.transpileAndExecute(code))).toEqual([1, 2, 4, 8, 16, 32, 64, 128, 256]);
+});
