From 32e341cbc4b47fbc192128c687eea17fc981bee2 Mon Sep 17 00:00:00 2001
From: cosmonaut <evan@moonside.games>
Date: Tue, 21 Nov 2023 11:47:58 -0800
Subject: [PATCH] implement Snapshot.Dispose

---
 src/Snapshot.cs | 95 +++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 92 insertions(+), 3 deletions(-)

diff --git a/src/Snapshot.cs b/src/Snapshot.cs
index 5a64679..0421406 100644
--- a/src/Snapshot.cs
+++ b/src/Snapshot.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using MoonTools.ECS.Collections;
@@ -5,7 +6,7 @@ using MoonTools.ECS.Collections;
 namespace MoonTools.ECS;
 
 // TODO: we should implement a NativeDictionary that can be memcopied
-public class Snapshot
+public class Snapshot : IDisposable
 {
 	private Dictionary<TypeId, ComponentSnapshot> ComponentSnapshots = new Dictionary<TypeId, ComponentSnapshot>();
 
@@ -24,6 +25,8 @@ public class Snapshot
 
 	private IdAssigner EntityIdAssigner = new IdAssigner();
 
+	private bool IsDisposed;
+
 	public void Restore(World world)
 	{
 		// restore id assigner state
@@ -212,11 +215,13 @@ public class Snapshot
 		snapshot.Take(relationStorage);
 	}
 
-	private class ComponentSnapshot
+	private class ComponentSnapshot : IDisposable
 	{
 		private readonly NativeArray Components;
 		private readonly NativeArray<Entity> EntityIDs;
 
+		private bool IsDisposed;
+
 		public ComponentSnapshot(int elementSize)
 		{
 			Components = new NativeArray(elementSize);
@@ -241,13 +246,36 @@ public class Snapshot
 				componentStorage.EntityIDToStorageIndex[entityID] = i;
 			}
 		}
+
+		protected virtual void Dispose(bool disposing)
+		{
+			if (!IsDisposed)
+			{
+				if (disposing)
+				{
+					Components.Dispose();
+					EntityIDs.Dispose();
+				}
+
+				IsDisposed = true;
+			}
+		}
+
+		public void Dispose()
+		{
+			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+			Dispose(disposing: true);
+			GC.SuppressFinalize(this);
+		}
 	}
 
-	private class RelationSnapshot
+	private class RelationSnapshot : IDisposable
 	{
 		private NativeArray Relations;
 		private NativeArray RelationDatas;
 
+		private bool IsDisposed;
+
 		public RelationSnapshot(int elementSize)
 		{
 			Relations = new NativeArray(Unsafe.SizeOf<(Entity, Entity)>());
@@ -291,5 +319,66 @@ public class Snapshot
 				relationStorage.InRelationSets[relation.Item2].Add(relation.Item1);
 			}
 		}
+
+		protected virtual void Dispose(bool disposing)
+		{
+			if (!IsDisposed)
+			{
+				if (disposing)
+				{
+					Relations.Dispose();
+					RelationDatas.Dispose();
+				}
+
+				IsDisposed = true;
+			}
+		}
+
+		public void Dispose()
+		{
+			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+			Dispose(disposing: true);
+			GC.SuppressFinalize(this);
+		}
+	}
+
+	protected virtual void Dispose(bool disposing)
+	{
+		if (!IsDisposed)
+		{
+			if (disposing)
+			{
+				foreach (var componentSnapshot in ComponentSnapshots.Values)
+				{
+					componentSnapshot.Dispose();
+				}
+
+				foreach (var relationSnapshot in RelationSnapshots.Values)
+				{
+					relationSnapshot.Dispose();
+				}
+
+				foreach (var componentSet in EntityComponentIndex.Values)
+				{
+					componentSet.Dispose();
+				}
+
+				foreach (var relationSet in EntityRelationIndex.Values)
+				{
+					relationSet.Dispose();
+				}
+
+				EntityIdAssigner.Dispose();
+			}
+
+			IsDisposed = true;
+		}
+	}
+
+	public void Dispose()
+	{
+		// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+		Dispose(disposing: true);
+		GC.SuppressFinalize(this);
 	}
 }