UTY-2661: Improve snapshot api

CHANGELOG.md

@@ -7,12 +7,18 @@
 - The minimum supported version of Unity is now 2020.1. [#1459](https://github.com/spatialos/gdk-for-unity/pull/1459)
 - Upgraded Unity Entities to 0.14.0-preview.18. [#1463](https://github.com/spatialos/gdk-for-unity/pull/1463)
   - Projects now require the `Unity Web Request` built-in package to compile for iOS and Android.
+- Adding an entity to the `Snapshot` class with a duplicate entity ID now throws an exception. [#1479](https://github.com/spatialos/gdk-for-unity/pull/1479)
 
 ### Added
 
 - Added `MeansImplicitUse` attribute to `RequireAttribute` to reduce warnings in Rider IDE. [#1462](https://github.com/spatialos/gdk-for-unity/pull/1462)
 - Added Event Tracing API. [#1452](https://github.com/spatialos/gdk-for-unity/pull/1452)
 - Added tooltips to the SpatialOS Project Settings. [#1470](https://github.com/spatialos/gdk-for-unity/pull/1470)
+- Added new features to the `Snapshot` class [#1479](https://github.com/spatialos/gdk-for-unity/pull/1479)
+  - Added the `Contains(EntityId entityId)` method to check if snapshot already contains an `EntityId`
+  - Improved the search for next available `EntityId`, it will no longer return already used `EntityId`s.
+  - `Snapshot` now auto-adds `Persistence` when adding entities
+  - `Snapshot` now implements `IDisposable`
 
 ### Changed
 

workers/unity/Assets/Playground/Editor/SnapshotGenerator/SnapshotGenerator.cs

@@ -19,10 +19,11 @@ public struct Arguments
         public static void Generate(Arguments arguments)
         {
             Debug.Log("Generating snapshot.");
-            var snapshot = CreateSnapshot(arguments.NumberEntities);
-
-            Debug.Log($"Writing snapshot to: {arguments.OutputPath}");
-            snapshot.WriteToFile(arguments.OutputPath);
+            using (var snapshot = CreateSnapshot(arguments.NumberEntities))
+            {
+                Debug.Log($"Writing snapshot to: {arguments.OutputPath}");
+                snapshot.WriteToFile(arguments.OutputPath);
+            }
         }
 
         private static Snapshot CreateSnapshot(int cubeCount)

workers/unity/Assets/Tests/EditmodeTests/Correctness/Utility/SnapshotTests.cs

@@ -0,0 +1,90 @@
+using System;
+using System.Linq;
+using Improbable.Gdk.Core;
+using NUnit.Framework;
+
+namespace Improbable.Gdk.EditmodeTests.Utility
+{
+    [TestFixture]
+    public class SnapshotTests
+    {
+        private Snapshot snapshot;
+
+        [SetUp]
+        public void Initialize()
+        {
+            snapshot = new Snapshot();
+        }
+
+        [Test]
+        public void Contains_returns_true_if_entityId_exists()
+        {
+            var entityId = snapshot.GetNextEntityId();
+            var template = GetTemplate();
+            snapshot.AddEntity(entityId, template);
+            Assert.IsTrue(snapshot.Contains(entityId));
+        }
+
+        [Test]
+        public void Contains_returns_false_if_entityId_does_not_exists()
+        {
+            var entityId = snapshot.GetNextEntityId();
+            Assert.IsFalse(snapshot.Contains(entityId));
+        }
+
+        [TestCase(10)]
+        [TestCase(20)]
+        [TestCase(40)]
+        public void GetNextEntityId_will_return_valid_EntityIds(int size)
+        {
+            var entities = Enumerable.Range(1, size).Select(i => new EntityId(i));
+            var template = new EntityTemplate();
+            template.AddComponent(new Position.Snapshot());
+            foreach (var entity in entities)
+            {
+                snapshot.AddEntity(entity, template);
+            }
+
+            Assert.IsFalse(snapshot.Contains(snapshot.GetNextEntityId()));
+        }
+
+        [TestCase(10)]
+        [TestCase(20)]
+        [TestCase(40)]
+        public void Count_increases_on_adding_new_entity(int size)
+        {
+            var entityIds = Enumerable.Range(1, size).Select(i => new EntityId(i));
+            var template = new EntityTemplate();
+            template.AddComponent(new Position.Snapshot());
+            foreach (var entityId in entityIds)
+            {
+                snapshot.AddEntity(entityId, template);
+            }
+
+            Assert.AreEqual(size, snapshot.Count);
+        }
+
+        [Test]
+        public void AddEntity_throws_if_EntityId_exists()
+        {
+            var entityId = new EntityId(1);
+            var template = GetTemplate();
+            var template2 = GetTemplate(new Coordinates(1, 2, 3));
+            snapshot.AddEntity(entityId, template);
+            Assert.Throws<ArgumentException>(() => snapshot.AddEntity(entityId, template2));
+        }
+
+        [TearDown]
+        public void TearDown()
+        {
+            snapshot.Dispose();
+        }
+
+        private static EntityTemplate GetTemplate(Coordinates pos = default)
+        {
+            var template = new EntityTemplate();
+            template.AddComponent(new Position.Snapshot(pos));
+            return template;
+        }
+    }
+}

workers/unity/Packages/io.improbable.gdk.core/Utility/Snapshot.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using Improbable.Worker.CInterop;
 using UnityEngine;
@@ -7,12 +8,11 @@ namespace Improbable.Gdk.Core
     /// <summary>
     ///     Convenience wrapper around the WorkerSDK Snapshot API.
     /// </summary>
-    public class Snapshot
+    public sealed class Snapshot : IDisposable
     {
+        private const int PersistenceComponentId = 55;
         private readonly Dictionary<EntityId, Entity> entities = new Dictionary<EntityId, Entity>();
-
         public int Count => entities.Count;
-
         private long nextEntityId = 1;
 
         /// <summary>
@@ -21,9 +21,24 @@ public class Snapshot
         /// <returns>The next available entity ID.</returns>
         public EntityId GetNextEntityId()
         {
+            while (Contains(new EntityId(nextEntityId)))
+            {
+                nextEntityId++;
+            }
+
             return new EntityId(nextEntityId++);
         }
 
+        /// <summary>
+        ///     Checks if the snapshot contains an entity ID
+        /// </summary>
+        /// <param name="entityId">The entity ID to check for</param>
+        /// <returns> true if the snapshot contains the input entity ID, false otherwise </returns>
+        public bool Contains(EntityId entityId)
+        {
+            return entities.ContainsKey(entityId);
+        }
+
         /// <summary>
         ///     Adds an entity to the snapshot
         /// </summary>
@@ -50,7 +65,15 @@ public EntityId AddEntity(EntityTemplate entityTemplate)
         /// </remarks>
         public void AddEntity(EntityId entityId, EntityTemplate entityTemplate)
         {
-            entities[entityId] = entityTemplate.GetEntity();
+            if (entities.ContainsKey(entityId))
+            {
+                throw new ArgumentException($"EntityId {entityId} already exists in the snapshot");
+            }
+
+            var entity = entityTemplate.GetEntity();
+            // This is a no-op if the entity already has persistence.
+            entity.Add(new ComponentData(PersistenceComponentId, SchemaComponentData.Create()));
+            entities[entityId] = entity;
         }
 
         /// <summary>
@@ -79,5 +102,19 @@ public void WriteToFile(string path)
                 }
             }
         }
+
+        public void Dispose()
+        {
+            foreach (var entity in entities.Values)
+            {
+                foreach (var id in entity.GetComponentIds())
+                {
+                    var componentData = entity.Get(id).Value;
+                    componentData.SchemaData?.Destroy();
+                }
+            }
+        }
+
+        internal Entity this[EntityId entityId] => entities[entityId];
     }
 }