From 67fa09d3098c567a4ebaa771167f87e08ba88db3 Mon Sep 17 00:00:00 2001
From: Perryvw <perry.wesel@gmail.com>
Date: Sat, 6 Aug 2022 21:32:52 +0200
Subject: [PATCH] Support spreading lua iterables

---
 src/LuaLib.ts                                 |  1 +
 src/lualib/LuaIteratorSpread.ts               | 13 ++++++
 .../utils/language-extensions.ts              |  4 ++
 src/transformation/visitors/spread.ts         |  5 +++
 test/unit/spread.spec.ts                      | 40 +++++++++++++++++++
 5 files changed, 63 insertions(+)
 create mode 100644 src/lualib/LuaIteratorSpread.ts

diff --git a/src/LuaLib.ts b/src/LuaLib.ts
index 60ffcf8c7..7780a35eb 100644
--- a/src/LuaLib.ts
+++ b/src/LuaLib.ts
@@ -44,6 +44,7 @@ export enum LuaLibFeature {
     InstanceOf = "InstanceOf",
     InstanceOfObject = "InstanceOfObject",
     Iterator = "Iterator",
+    LuaIteratorSpread = "LuaIteratorSpread",
     Map = "Map",
     MathAtan2 = "MathAtan2",
     MathSign = "MathSign",
diff --git a/src/lualib/LuaIteratorSpread.ts b/src/lualib/LuaIteratorSpread.ts
new file mode 100644
index 000000000..abe18e454
--- /dev/null
+++ b/src/lualib/LuaIteratorSpread.ts
@@ -0,0 +1,13 @@
+export function __TS__LuaIteratorSpread<TKey, TValue, TState>(
+    this: (this: void, state: TState, key: TKey) => LuaMultiReturn<[TKey, TValue]>,
+    state: TState,
+    firstKey: TKey
+): LuaMultiReturn<Array<[TKey, TValue]>> {
+    const results = [];
+    let [key, value] = this(state, firstKey);
+    while (key) {
+        results.push([key, value]);
+        [key, value] = this(state, key);
+    }
+    return $multi(...results);
+}
diff --git a/src/transformation/utils/language-extensions.ts b/src/transformation/utils/language-extensions.ts
index bb904471f..881da105b 100644
--- a/src/transformation/utils/language-extensions.ts
+++ b/src/transformation/utils/language-extensions.ts
@@ -98,6 +98,10 @@ export enum IterableExtensionKind {
     PairsKey = "PairsKey",
 }
 
+export function isLuaIterable(context: TransformationContext, type: ts.Type): boolean {
+    return getPropertyValue(context, type, "__tstlIterable") !== undefined;
+}
+
 export function getIterableExtensionTypeForType(
     context: TransformationContext,
     type: ts.Type
diff --git a/src/transformation/visitors/spread.ts b/src/transformation/visitors/spread.ts
index 270ee7d36..84a13f0fd 100644
--- a/src/transformation/visitors/spread.ts
+++ b/src/transformation/visitors/spread.ts
@@ -1,6 +1,7 @@
 import * as ts from "typescript";
 import * as lua from "../../LuaAST";
 import { FunctionVisitor, TransformationContext } from "../context";
+import { isLuaIterable } from "../utils/language-extensions";
 import { createUnpackCall } from "../utils/lua-ast";
 import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import {
@@ -73,6 +74,10 @@ export const transformSpreadElement: FunctionVisitor<ts.SpreadElement> = (node,
     if (isMultiReturnCall(context, tsInnerExpression)) return innerExpression;
 
     const type = context.checker.getTypeAtLocation(node.expression); // not ts-inner expression, in case of casts
+    if (isLuaIterable(context, type)) {
+        return transformLuaLibFunction(context, LuaLibFeature.LuaIteratorSpread, node, innerExpression);
+    }
+
     if (isArrayType(context, type)) {
         return createUnpackCall(context, innerExpression, node);
     }
diff --git a/test/unit/spread.spec.ts b/test/unit/spread.spec.ts
index c3fe72b07..29f7110d3 100644
--- a/test/unit/spread.spec.ts
+++ b/test/unit/spread.spec.ts
@@ -478,3 +478,43 @@ describe("vararg spread in IIFE", () => {
         `.expectToMatchJsResult();
     });
 });
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1244
+test.each(["pairs", "ipairs"])("can spread %s (#1244)", func => {
+    util.testFunction`
+        const arr = ["a", "b", "c"];
+        return [...${func}(arr)];
+    `
+        .withLanguageExtensions()
+        .setTsHeader(
+            `
+            declare function ipairs<T>(this: void, t: T): LuaIterable<LuaMultiReturn<[number, NonNullable<T[keyof T]>]>>;
+            declare function pairs<T>(this: void, t: T): LuaIterable<LuaMultiReturn<[keyof T, NonNullable<T[keyof T]>]>>;
+            `
+        )
+        .expectToEqual([
+            [1, "a"],
+            [2, "b"],
+            [3, "c"],
+        ]);
+});
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1244
+test.each(["LuaTable", "LuaMap"])("can spread %s (#1244)", type => {
+    const result: Array<[string, string]> = util.testFunction`
+        const tbl = new ${type}();
+        tbl.set("foo", "bar");
+        tbl.set("fizz", "buzz");
+        return [...pairs(tbl)];
+    `
+        .withLanguageExtensions()
+        .setTsHeader(
+            "declare function pairs<T>(this: void, t: T): LuaIterable<LuaMultiReturn<[keyof T, NonNullable<T[keyof T]>]>>;"
+        )
+        .getLuaExecutionResult();
+
+    // We don't know the order so match like this
+    expect(result).toHaveLength(2);
+    expect(result.some(([k, v]) => k === "foo" && v === "bar")).toBe(true);
+    expect(result.some(([k, v]) => k === "fizz" && v === "buzz")).toBe(true);
+});