From fe3ed12860b0f76dba3d3aa2fd517824d410ce95 Mon Sep 17 00:00:00 2001
From: David Saltares <david.saltares@gmail.com>
Date: Tue, 31 May 2016 19:11:35 +0100
Subject: [PATCH] #221 - Fixes returning pooled entity to the entity pool twice

---
 .../badlogic/ashley/core/EntityManager.java   | 18 ++++----
 .../ashley/core/PooledEngineTests.java        | 45 +++++++------------
 2 files changed, 25 insertions(+), 38 deletions(-)

diff --git a/ashley/src/com/badlogic/ashley/core/EntityManager.java b/ashley/src/com/badlogic/ashley/core/EntityManager.java
index 36493a57..1d5d9252 100644
--- a/ashley/src/com/badlogic/ashley/core/EntityManager.java
+++ b/ashley/src/com/badlogic/ashley/core/EntityManager.java
@@ -101,13 +101,15 @@ public void processPendingOperations() {
 	}
 	
 	protected void removeEntityInternal(Entity entity) {
-		entity.scheduledForRemoval = false;
-		entity.removing = true;
-		entities.removeValue(entity, true);
-		entitySet.remove(entity);
-		
-		listener.entityRemoved(entity);
-		entity.removing = false;
+		boolean removed = entitySet.remove(entity);
+
+		if (removed) {
+			entity.scheduledForRemoval = false;
+			entity.removing = true;
+			entities.removeValue(entity, true);
+			listener.entityRemoved(entity);
+			entity.removing = false;
+		}
 	}
 
 	protected void addEntityInternal(Entity entity) {
@@ -120,7 +122,7 @@ protected void addEntityInternal(Entity entity) {
 
 		listener.entityAdded(entity);
 	}
-	
+
 	private static class EntityOperation implements Pool.Poolable {
 		public enum Type {
 			Add,
diff --git a/ashley/tests/com/badlogic/ashley/core/PooledEngineTests.java b/ashley/tests/com/badlogic/ashley/core/PooledEngineTests.java
index a6bf1a4a..cd2cdeeb 100644
--- a/ashley/tests/com/badlogic/ashley/core/PooledEngineTests.java
+++ b/ashley/tests/com/badlogic/ashley/core/PooledEngineTests.java
@@ -62,33 +62,6 @@ public void receive (Signal<Entity> signal, Entity object) {
 			totalCalls++;
 		}
 	}
-
-	private static class RemoveEntityTwiceSystem extends EntitySystem {
-		private ImmutableArray<Entity> entities;
-
-		@Override
-		public void addedToEngine (Engine engine) {
-			entities = engine.getEntitiesFor(Family.all(PositionComponent.class).get());
-		}
-
-		@Override
-		public void update (float deltaTime) {
-			Entity entity;
-			PooledEngine engine = (PooledEngine)getEngine();
-			for (int i = 0; i < 10; i++) {
-				entity = engine.createEntity();
-				assertEquals(0, entity.flags);
-				entity.flags = 1;
-				entity.add(engine.createComponent(PositionComponent.class));
-				engine.addEntity(entity);
-			}
-			for (int i = 0; i < entities.size(); ++i) {
-				entity = entities.get(i);
-				engine.removeEntity(entity);
-				engine.removeEntity(entity);
-			}
-		}
-	}
 	
 	private static class PooledComponentSpy implements Component, Poolable {
 		public boolean recycled = false;
@@ -210,10 +183,22 @@ public void recycleEntity () {
 	@Test
 	public void removeEntityTwice () {
 		PooledEngine engine = new PooledEngine();
-		engine.addSystem(new RemoveEntityTwiceSystem());
 
-		for (int j = 0; j < 2; j++) {
-			engine.update(0);
+		for (int i = 0; i < 100; ++i) {
+			Array<Entity> entities = new Array<Entity>();
+
+			for (int j = 0; j < 100; ++j) {
+				Entity entity = engine.createEntity();
+				engine.addEntity(entity);
+				assertEquals(0, entity.flags);
+				entity.flags = 1;
+				entities.add(entity);
+			}
+
+			for (Entity entity : entities) {
+				engine.removeEntity(entity);
+				engine.removeEntity(entity);
+			}
 		}
 	}