Merge pull request #5 from wlindley/recursive-event-handler-removal

Recursive event handler removal
wlindley · Nov 24, 2016 · 36b14ea · 36b14ea
2 parents 2b548f5 + f716dee
commit 36b14ea
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 3 deletions.
diff --git a/Bantam/EventBus.cs b/Bantam/EventBus.cs
@@ -12,11 +12,14 @@ public class EventBus
 		private Dictionary<Type, ArrayList> listeners = new Dictionary<Type, ArrayList>();
 		private Dictionary<Type, ArrayList> onceListeners = new Dictionary<Type, ArrayList>();
 		private ArrayList allListeners = new ArrayList();
+		private ArrayList toRemove = new ArrayList();
+		private ArrayList currentlyDispatching;
 		private ObjectPool pool;
 
 		public EventBus(ObjectPool pool)
 		{
 			this.pool = pool;
+			currentlyDispatching = null;
 		}
 
 		public void AddListener<T>(EventListener<T> listener) where T : class, Event, new()
@@ -40,9 +43,7 @@ public void RemoveListener<T>(EventListener<T> listener) where T : class, Event
 		{
 			EnsureKeyExists<T>();
 			var type = typeof(T);
-			listeners[type].Remove(listener);
-			onceListeners[type].Remove(listener);
-			allListeners.Remove(listener);
+			SafeRemoveFromAll(listener, type);
 		}
 
 		public void Dispatch<T>(EventInitializer<T> initializer = null) where T : class, Event, new()
@@ -60,6 +61,7 @@ public void RemoveListener<T>(EventListener<T> listener) where T : class, Event
 		private void EnsureKeyExists<T>() where T : Event
 		{
 			var type = typeof(T);
+
 			if (!listeners.ContainsKey(type))
 				listeners[type] = new ArrayList();
 			if (!onceListeners.ContainsKey(type))
@@ -79,9 +81,36 @@ private void DispatchEvent<T>(T ev) where T : Event
 
 		private void DispatchEventToListeners<T>(T ev, ArrayList eventListeners) where T : Event
 		{
+			currentlyDispatching = eventListeners;
 			var numListeners = eventListeners.Count;
 			for (var i = 0; i < numListeners; i++)
 				(eventListeners[i] as EventListener<T>)(ev);
+			currentlyDispatching = null;
+
+			ProcessRemovals(eventListeners);
+		}
+
+		private void ProcessRemovals(ArrayList eventListeners)
+		{
+			var numToRemove = toRemove.Count;
+			for (var i = 0; i < numToRemove; i++)
+				eventListeners.Remove(toRemove[i]);
+			toRemove.Clear();
+		}
+
+		private void SafeRemoveFromAll(object listener, Type type)
+		{
+			SafeRemoveFromList(listener, listeners[type]);
+			SafeRemoveFromList(listener, onceListeners[type]);
+			SafeRemoveFromList(listener, allListeners);
+		}
+
+		private void SafeRemoveFromList(object listener, ArrayList list)
+		{
+			if (list != currentlyDispatching)
+				list.Remove(listener);
+			else
+				toRemove.Add(listener);
 		}
 	}
 }
diff --git a/BantamTest/EventBusTest.cs b/BantamTest/EventBusTest.cs
@@ -122,6 +122,40 @@ public void RemoveListenerPreventsEventsFromGoingToAnAllListener()
 			testObj.Dispatch<DummyEvent>();
 			Assert.IsFalse(wasCalled);
 		}
+
+		[Test]
+		public void RemoveListenerSucceedsWhenCalledFromHandlerForEvent()
+		{
+			var wasSecondListenerCalled = false;
+			Action action = () => {};
+			var listener = new EventListener<DummyEvent>(evt => action());
+			action = () => {
+				testObj.RemoveListener<DummyEvent>(listener);
+			};
+			testObj.AddListener<DummyEvent>(listener);
+			testObj.AddListener<DummyEvent>(evt => wasSecondListenerCalled = true);
+
+			testObj.Dispatch<DummyEvent>();
+
+			Assert.IsTrue(wasSecondListenerCalled, "Second listener was not executed.");
+		}
+
+		[Test]
+		public void RemoveListenerSucceedsWhenCalledFromOnceHandlerForEvent()
+		{
+			var wasSecondListenerCalled = false;
+			Action action = () => {};
+			var listener = new EventListener<DummyEvent>(evt => action());
+			action = () => {
+				testObj.RemoveListener<DummyEvent>(listener);
+			};
+			testObj.AddOnce<DummyEvent>(listener);
+			testObj.AddOnce<DummyEvent>(evt => wasSecondListenerCalled = true);
+
+			testObj.Dispatch<DummyEvent>();
+
+			Assert.IsTrue(wasSecondListenerCalled, "Second listener was not executed.");
+		}
 	}
 
 	public class DummyEvent : Event