From 344a0082b42b0b03fbed54868bd950c95d585049 Mon Sep 17 00:00:00 2001
From: cosmonaut <evan@moonside.games>
Date: Thu, 26 Oct 2023 11:02:45 -0700
Subject: [PATCH] fast snapshot restore

---
 src/Rev2/Snapshot.cs | 25 ++++++++++++++++---------
 src/Rev2/World.cs    |  7 +++----
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/Rev2/Snapshot.cs b/src/Rev2/Snapshot.cs
index d510e09..e3fb33b 100644
--- a/src/Rev2/Snapshot.cs
+++ b/src/Rev2/Snapshot.cs
@@ -87,23 +87,30 @@ public class Snapshot
 
 		public void Restore(Archetype archetype)
 		{
-			// Clear out existing entities
-			archetype.ClearAll();
-
 			// Copy all component data
 			for (int i = 0; i < ComponentColumns.Count; i += 1)
 			{
 				ComponentColumns[i].CopyAllTo(archetype.ComponentColumns[i]);
 			}
 
-			// Clear the row to entity list
-			archetype.RowToEntity.Clear();
+			var archetypeCount = archetype.Count;
 
-			// Create new entities and repopulate the row to entity list
-			for (int i = 0; i < Count; i += 1)
+			if (Count < archetypeCount)
 			{
-				var entityId = archetype.World.CreateEntityOnArchetype(archetype);
-				archetype.RowToEntity.Add(entityId);
+				// if snapshot has fewer entities than archetype, remove extra entities
+				for (int i = archetypeCount - 1; i >= Count; i -= 1)
+				{
+					archetype.World.FreeEntity(archetype.RowToEntity[i]);
+					archetype.RowToEntity.RemoveAt(i);
+				}
+			}
+			else if (Count > archetypeCount)
+			{
+				// if snapshot has more entities than archetype, add entities
+				for (int i = archetypeCount; i < Count; i += 1)
+				{
+					archetype.World.CreateEntityOnArchetype(archetype);
+				}
 			}
 		}
 	}
diff --git a/src/Rev2/World.cs b/src/Rev2/World.cs
index c981560..8e8858b 100644
--- a/src/Rev2/World.cs
+++ b/src/Rev2/World.cs
@@ -67,16 +67,15 @@ namespace MoonTools.ECS.Rev2
 			return entityId;
 		}
 
-		// used as a fast path by Archetype.Transfer
-		internal EntityId CreateEntityOnArchetype(Archetype archetype)
+		// used as a fast path by snapshot restore
+		internal void CreateEntityOnArchetype(Archetype archetype)
 		{
 			var entityId = EntityIdAssigner.Assign();
 			EntityIndex.Add(entityId, new Record(archetype, archetype.Count));
 			archetype.RowToEntity.Add(entityId);
-			return entityId;
 		}
 
-		// used as a fast path by Archetype.ClearAll
+		// used as a fast path by Archetype.ClearAll and snapshot restore
 		internal void FreeEntity(EntityId entityId)
 		{
 			EntityIndex.Remove(entityId);