util.iterators: Add join() method and tests

1 files changed, 46 insertions, 0 deletions

diff --git a/util/iterators.lua b/util/iterators.lua
index a152e7be..5d16d8c1 100644
--- a/util/iterators.lua
+++ b/util/iterators.lua
@@ -14,6 +14,11 @@ local t_insert = table.insert;
 local select, next = select, next;
 local unpack = table.unpack or unpack; --luacheck: ignore 113 143
 local pack = table.pack or function (...) return { n = select("#", ...), ... }; end -- luacheck: ignore 143
+local type = type;
+local table, setmetatable = table, setmetatable;
+
+local _ENV = nil;
+--luacheck: std none
 
 -- Reverse an iterator
 function it.reverse(f, s, var)
@@ -184,4 +189,45 @@ function it.to_table(f, s, var)
 	return t;
 end
 
+local function _join_iter(j_s, j_var)
+	local iterators, current_idx = j_s[1], j_s[2];
+	local f, s, var = unpack(iterators[current_idx], 1, 3);
+	if j_var ~= nil then
+		var = j_var;
+	end
+	local ret = pack(f(s, var));
+	local var1 = ret[1];
+	if var1 == nil then
+		-- End of this iterator, advance to next
+		if current_idx == #iterators then
+			-- No more iterators, return nil
+			return;
+		end
+		j_s[2] = current_idx + 1;
+		return _join_iter(j_s);
+	end
+	return unpack(ret, 1, ret.n);
+end
+local join_methods = {};
+local join_mt = {
+	__index = join_methods;
+	__call = function (t, s, var) --luacheck: ignore 212/t
+		return _join_iter(s, var);
+	end;
+};
+
+function join_methods:append(f, s, var)
+	table.insert(self, { f, s, var });
+	return self, { self, 1 };
+end
+
+function join_methods:prepend(f, s, var)
+	table.insert(self, { f, s, var }, 1);
+	return self, { self, 1 };
+end
+
+function it.join(f, s, var)
+	return setmetatable({ {f, s, var} }, join_mt);
+end
+
 return it;