From 66e9848bef7f163d2c9c08714f625b70a86f6525 Mon Sep 17 00:00:00 2001
From: Max Cahill <1bardesign@gmail.com>
Date: Thu, 3 Nov 2022 16:19:18 +1100
Subject: [PATCH] added pubsub:subscribe_once and fixed a bug if you subscribe
 or unsubscribe from inside a callback

---
 pubsub.lua | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 66 insertions(+), 1 deletion(-)

diff --git a/pubsub.lua b/pubsub.lua
index 4556718..f5dad81 100644
--- a/pubsub.lua
+++ b/pubsub.lua
@@ -5,6 +5,7 @@
 local path = (...):gsub("pubsub", "")
 local class = require(path .. "class")
 local set = require(path .. "set")
+local tablex = require(path .. "tablex")
 
 local pubsub = class({
 	name = "pubsub",
@@ -13,14 +14,61 @@ local pubsub = class({
 --create a new pubsub bus
 function pubsub:new()
 	self.subscriptions = {}
+	self._defer = {}
+	self._defer_stack = 0
+end
+
+--(internal; deferred area check)
+function pubsub:_deferred()
+	return self._defer_stack > 0
+end
+
+--(internal; enter deferred area)
+function pubsub:_push_defer()
+	self._defer_stack = self._defer_stack + 1
+	if self._defer_stack > 255 then
+		error("pubsub defer stack overflow; event infinite loop")
+	end
+end
+
+--(internal; enter deferred area)
+function pubsub:_defer_call(defer_f, event, callback)
+	if not self:_deferred() then
+		error("attempt to defer pubsub call when not required")
+	end
+	table.insert(self._defer, defer_f)
+	table.insert(self._defer, event)
+	table.insert(self._defer, callback)
+end
+
+--(internal; unwind deferred sub/unsub)
+function pubsub:_pop_defer()
+	self._defer_stack = self._defer_stack - 1
+	if self._defer_stack < 0 then
+		error("pubsub defer stack underflow; don't call the defer methods directly")
+	end
+	if self._defer_stack == 0 then
+		local defer_len = #self._defer
+		if defer_len then
+			for i = 1, defer_len, 3 do
+				local defer_f = self._defer[i]
+				local defer_event = self._defer[i+1]
+				local defer_cb = self._defer[i+2]
+				self[defer_f](self, defer_event, defer_cb)
+			end
+			tablex.clear(self._defer)
+		end
+	end
 end
 
 --(internal; notify a callback set of an event)
 function pubsub:_notify(callbacks, ...)
 	if callbacks then
-		for _, f in callbacks:ipairs() do
+		self:_push_defer()
+		for _, f in ipairs(callbacks:values()) do
 			f(...)
 		end
+		self:_pop_defer()
 	end
 end
 
@@ -35,6 +83,10 @@ end
 --can be a specifically named event, or "everything" to get notified for any event
 --for "everything", the callback will receive the event name as the first argument
 function pubsub:subscribe(event, callback)
+	if self:_deferred() then
+		self:_defer_call("subscribe", event, callback)
+		return
+	end
 	local callbacks = self.subscriptions[event]
 	if not callbacks then
 		callbacks = set()
@@ -43,8 +95,21 @@ function pubsub:subscribe(event, callback)
 	callbacks:add(callback)
 end
 
+--subscribe to an event, automatically unsubscribe once called
+function pubsub:subscribe_once(event, callback)
+	f = function(...)
+		callback(...)
+		self:unsubscribe(event, f)
+	end
+	self:subscribe(event, f)
+end
+
 --unsubscribe from an event
 function pubsub:unsubscribe(event, callback)
+	if self:_deferred() then
+		self:_defer_call("unsubscribe", event, callback)
+		return
+	end
 	local callbacks = self.subscriptions[event]
 	if callbacks then
 		callbacks:remove(callback)