Coroutine-like scheduler-friendly interface for Roblox Lua

_Examples.lua

---

-- basics
local yield = Yield:new(function()
	Yield:yield(5)
	Yield:yield(7)
	Yield:yield(9)
end)

print(yield:resume())  -- 5
print(yield:resume())  -- 7
print(yield:resume())  -- 9

---

-- parameters
local yield = Yield:new(function(i)
	Yield:yield(i)
	Yield:yield(i + 2)
	Yield:yield(i + 4)
end)

print(yield:resume(10))  -- 10
print(yield:resume())  -- 12
print(yield:resume())  -- 14

---

-- passing things into `:resume`
local yield = Yield:new(function(num1)
	local num2 = Yield:yield(num1)
	local num3 = Yield:yield(num1 + num2)
	return num1 + num2 + num3
end)

-- pass `2` in as `num1`: we get `num1`
print(yield:resume(2))  -- 2
-- pass `3` in as `num2`: we get `num1 + num2`
print(yield:resume(3))  -- 5
-- pass `5` in as `num3`: we get `num1 + num2 + num3`
print(yield:resume(5))  -- 10

---

-- infinite loops and resume
-- this accumulates numbers: every time you give it a number,
--  it adds your number to its value and returns its value
local yield = Yield:new(function(num)
	while true do
		local nextNum = Yield:yield(num)
		num = num + nextNum
	end
end)

print(yield:resume(2))  -- 2
print(yield:resume(3))  -- 5
print(yield:resume(5))  -- 10

---

-- syntax sugar: we can make things look nicer!
local accumulate = Yield(function(num)
	while true do
		num = num + Yield:yield(num)
	end
end)

print(accumulate(2))  -- 2
print(accumulate(3))  -- 5
print(accumulate(5))  -- 10

---

-- yields as loop handlers
local yield = Yield(function()
	for i = 1, 10 do
		Yield:yield(i)
	end
end)

for num in yield do
	print(num)  -- will print 1 through 10
end

---

-- yields as loop handlers 2
local getNums = function(start, count)
	return Yield:new(function()
		for i = start, count do
			Yield:yield(i)
		end
	end)
end

for num in getNums(10, 20) do
	print(num)  -- wil print 1 through 20
end

---

Yield.lua

--[[ API Reference
	
	Class Yield
		METHODS
			static Yield(f: function) --> yield: Yield
			static :new(f: function) --> yield: Yield
				Creates a new Yield that will run `f` when resumed
			static :running() --> Yield currentYield
				Returns the Yield that is currently running, or nil if none.
			static :yield(... [1]) --> ... [2]
				Same as non-static `:yield`, but acts on whatever the currently-running Yield is.
				Errors if there is no current Yield, or if this Yield is not running.
			:yield(... [a]) --> ... [b]
				Passes ... [a] to whatever `:resume`d this Yield, then waits to be `:resume`d.
				Returns whatever ... [b] in `:resume(... [b])` will be, once resumed.
				Errors if this Yield is not running.
			Yield(... [b]) --> ... [a]
			:resume(... [b]) --> ... [a]
				Passes ... [b] into this Yield as the result of the earlier `:yield` call,
				 then waits for the `:yield`
				Returns whatever the next `:yield(... [a])` will be, once yield is called.
				 If the Yield returns and finished, this returns whatever yield returned
				If this Yield errors, it will return `nil` and the `error` propert will be set.
				Errors if this Yield is running or already finished.
			:getResumeCaller() --> resumeCaller: function
				Returns a function that calls `:resume` on this yield and returns the result
			:getYieldCaller() --> yieldCaller: function
				Returns a function that calls `:yield` on this yield and returns the result
			:finished() --> isFinished: bool
				Returns whether or not the Yield has finished. (`.state < Yield.FINISHED`)
		PROPERTIES
			state: YieldState
			error: Variant
		CONSTANTS
			STOPPED:  YieldState (0)
			RUNNING:  YieldState (1)
			PAUSED:   YieldState (2)
			FINISHED: YieldState (3)
			ERROR:    YieldState (4)
--]]

local function runYield(this)
	this.outArguments = {this.func(unpack(this.inArguments))}
end

local function assertMetatable(tbl, meta, err)
	return assert(type(tbl) == "table" and getmetatable(tbl) == meta, err)
end

local YieldWrap, YieldWrapMeta
YieldWrapMeta = {
	__index = {
		STOPPED  = 0,
		RUNNING  = 1,
		PAUSED   = 2,
		FINISHED = 3,
		ERROR    = 4,
		globalIndex = {},
		globalStack = {},
		pushStack = function(this, yieldWrap)
			this.globalStack[#this.globalStack + 1] = yieldWrap
		end,
		popStack = function(this)
			local yieldWrap = this.globalStack[#this.globalStack]
			this.globalStack[#this.globalStack] = nil
			return yieldWrap
		end,
		running = function(this)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function running")
			local current = this.globalIndex[coroutine.running()]
			if current then
				-- easy! current Yield is whichever the current coroutine is!
				return current
			else
				-- looks like it shifted to another coroutine, but hasn't waited at all yet
				-- in that case, we need to find which coroutine is active, but not suspended or dead
				-- we want to find the most recent one like this: it's possible for a Yield to
				--  resume another Yield, so we need to get the most recent one resumed!
				local globalStack = this.globalStack
				for i = #globalStack, 1, -1 do
					local v = globalStack[i]
					if coroutine.status(v.coroutine) == "normal" then
						return v
					end
				end
			end
		end,
		new = function(this, ...)
			assert(this == YieldWrap, "Expected ':' not '.' calling constructor Yield")
			local newYield = setmetatable({}, YieldWrapMeta)
			newYield:construct(...)
			return newYield
		end,
		construct = function(this, func)
			assert(type(func) == "function" or type(func) == "table", "`f` should be a function or table")
			this.func = func
			this.inEvent = Instance.new("BindableEvent")
			this.outEvent = Instance.new("BindableEvent")
			this.inArguments = {}
			this.outArguments = {}
			this.state = this.STOPPED
			local conn
			conn = this.inEvent.Event:connect(function()
				conn:disconnect()
				this.state = this.RUNNING
				this.coroutine = coroutine.running()
				this.globalIndex[this.coroutine] = this
				this:pushStack(this)
				local success, err = pcall(runYield, this)
				this:popStack()
				this.globalIndex[this.coroutine] = nil
				if success then
					this.state = this.FINISHED
				else
					-- warn("Error in Yield: "..tostring(err))
					this.outArguments = {}
					this.state = this.ERROR
					this.error = err
				end
				this.outEvent:Fire()
			end)
		end,
		yield = function(this, ...)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function yield")
			if this == YieldWrap then
				return this.globalStack[#this.globalStack]:yield(...)
			end
			if this.state >= this.FINISHED then
				error("Cannot yield when already finished")
			elseif this.state == this.STOPPED then
				error("Cannot yield before started")
			elseif this.state == this.PAUSED then
				error("Cannot yield while paused")
			end
			this.outArguments = {...}
			local inArgs
			local conn
			conn = this.inEvent.Event:connect(function()
				conn:disconnect()
				inArgs = this.inArguments
			end)
			this.state = this.PAUSED
			this:popStack()
			this.outEvent:Fire()
			if not inArgs then
				this.inEvent.Event:wait()
				inArgs = this.inArguments
			end
			this:pushStack(this)
			this.state = this.RUNNING
			return unpack(inArgs)
		end,
		resume = function(this, ...)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function resume")
			if this.state >= this.FINISHED then
				error("Cannot resume when already finished")
			elseif this.state == this.RUNNING then
				error("Cannot resume while running")
			end
			this.inArguments = {...}
			local outArgs
			local conn
			conn = this.outEvent.Event:connect(function()
				conn:disconnect()
				outArgs = this.outArguments
			end)
			this.inEvent:Fire()
			if not outArgs then
				this.outEvent.Event:wait()
				outArgs = this.outArguments
			end
			return unpack(outArgs)
		end,
		getYieldCaller = function(this)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function getYieldCaller")
			if not this.yieldCaller then
				this.yieldCaller = function(...)
					return this:yield(...)
				end
			end
			return this.yieldCaller
		end,
		getResumeCaller = function(this)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function getResumeCaller")
			if not this.resumeCaller then
				this.resumeCaller = function(...)
					return this:resume(...)
				end
			end
			return this.resumeCaller
		end,
		finished = function(this)
			assertMetatable(this, YieldWrapMeta, "Expected ':' not '.' calling member function finished")
			return this.state < this.FINISHED
		end
	},
	__call = function(this, ...)
		if this == YieldWrap then
			return this:new(...)
		else
			return this:resume(...)
		end
	end
}

YieldWrap = setmetatable({}, YieldWrapMeta)

return YieldWrap