Access wrapper duplicate type

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/JobConfig/AbstractJobConfig.cs

@@ -124,12 +124,14 @@ private void ResolveDuplicateAccessWrappers(AbstractAccessWrapper accessWrapper)
             // If the existing access wrapper facilitates the needs of the new one then just keep the existing one.
             if (existingAccessWrapper.AccessType.IsCompatibleWith(accessWrapper.AccessType))
             {
+                existingAccessWrapper.MergeStateFrom(accessWrapper);
                 accessWrapper.Dispose();
             }
             // If the new access wrapper facilitates the needs of the existing one then dispose the existing wrapper and
             // use the new access wrapper.
             else if (accessWrapper.AccessType.IsCompatibleWith(existingAccessWrapper.AccessType))
             {
+                accessWrapper.MergeStateFrom(existingAccessWrapper);
                 existingAccessWrapper.Dispose();
                 m_AccessWrappers[accessWrapper.ID] = accessWrapper;
             }
@@ -438,21 +440,23 @@ CancelProgressLookupAccessWrapper cancelProgressLookupAccessWrapper
             return cancelProgressLookupAccessWrapper.ProgressLookup;
         }
 
-        internal CancelRequestsDataStream GetCancelRequestsDataStream()
+        internal CancelRequestsDataStream GetCancelRequestsDataStream(IAbstractCancelRequestDataStream explicitSource = null)
         {
             CancelRequestsPendingAccessWrapper cancelRequestsPendingAccessWrapper
                 = GetAccessWrapper<CancelRequestsPendingAccessWrapper>(Usage.RequestCancel);
 
-            return cancelRequestsPendingAccessWrapper.CancelRequestsDataStream;
+            return cancelRequestsPendingAccessWrapper.GetInstance(explicitSource);
         }
 
-        internal EntityProxyDataStream<TInstance> GetPendingDataStream<TInstance>(Usage usage)
+        internal EntityProxyDataStream<TInstance> GetPendingDataStream<TInstance>(
+            Usage usage,
+            IAbstractDataStream<TInstance> explicitSource = null)
             where TInstance : unmanaged, IEntityKeyedTask
         {
             DataStreamPendingAccessWrapper<TInstance> dataStreamAccessWrapper
                 = GetAccessWrapper<DataStreamPendingAccessWrapper<TInstance>>(usage);
 
-            return dataStreamAccessWrapper.DataStream;
+            return dataStreamAccessWrapper.GetInstance(explicitSource);
         }
 
         internal EntityProxyDataStream<TInstance> GetActiveDataStream<TInstance>(Usage usage)
@@ -617,7 +621,7 @@ private void Debug_EnsureWrapperExists(JobConfigDataID id)
         [Conditional("ANVIL_DEBUG_SAFETY")]
         private void Debug_EnsureWrapperUsage(AbstractAccessWrapper wrapper)
         {
-            if (wrapper.Debug_WrapperType != typeof(AbstractDataStreamAccessWrapper<>))
+            if (wrapper.Debug_WrapperType != typeof(AbstractDataStreamActiveAccessWrapper<>))
             {
                 return;
             }

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/JobData/AbstractJobData.cs

@@ -34,19 +34,19 @@ protected AbstractJobData(IJobConfig jobConfig)
         /// <summary>
         /// Fulfills an instance of the provided type for the job.
         /// </summary>
-        public void Fulfill(out CancelRequestsWriter instance)
+        public void Fulfill(out CancelRequestsWriter instance, IAbstractCancelRequestDataStream explicitSource = null)
         {
-            CancelRequestsDataStream cancelRequestDataStream = m_JobConfig.GetCancelRequestsDataStream();
+            CancelRequestsDataStream cancelRequestDataStream = m_JobConfig.GetCancelRequestsDataStream(explicitSource);
             instance = cancelRequestDataStream.CreateCancelRequestsWriter();
         }
 
         /// <summary>
         /// Fulfills an instance of the provided type for the job.
         /// </summary>
-        public void Fulfill<TInstance>(out DataStreamPendingWriter<TInstance> instance)
+        public void Fulfill<TInstance>(out DataStreamPendingWriter<TInstance> instance, IAbstractDataStream<TInstance> explicitSource = null)
             where TInstance : unmanaged, IEntityKeyedTask
         {
-            EntityProxyDataStream<TInstance> dataStream = m_JobConfig.GetPendingDataStream<TInstance>(AbstractJobConfig.Usage.Default);
+            EntityProxyDataStream<TInstance> dataStream = m_JobConfig.GetPendingDataStream(AbstractJobConfig.Usage.Default, explicitSource);
             instance = dataStream.CreateDataStreamPendingWriter();
         }
 
@@ -60,11 +60,11 @@ public void Fulfill<TInstance>(out DataStreamActiveReader<TInstance> instance)
             instance = dataStream.CreateDataStreamActiveReader();
         }
 
-        
+
         //*************************************************************************************************************
         // ENTITY SPAWNER
         //*************************************************************************************************************
-        
+
         public void Fulfill(out EntitySpawner entitySpawner)
         {
             entitySpawner = m_JobConfig.GetEntitySpawner();

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/AbstractAccessWrapper.cs

@@ -1,4 +1,5 @@
 using Anvil.CSharp.Core;
+using Anvil.CSharp.Logging;
 using Anvil.Unity.DOTS.Jobs;
 using System;
 using System.Diagnostics;
@@ -23,6 +24,16 @@ protected AbstractAccessWrapper(AccessType accessType, AbstractJobConfig.Usage u
         public abstract JobHandle AcquireAsync();
         public abstract void ReleaseAsync(JobHandle dependsOn);
 
+        /// <summary>
+        /// Merge the state from another wrapper instance of the same type
+        /// </summary>
+        /// <param name="other">The wrapper instance to read state from.</param>
+        public virtual void MergeStateFrom(AbstractAccessWrapper other)
+        {
+            Debug_EnsureSameID(other);
+            Debug_EnsureSameType(other);
+        }
+
         //*************************************************************************************************************
         // SAFETY
         //*************************************************************************************************************
@@ -38,5 +49,25 @@ private void Debug_SetWrapperType()
                 ? type.GetGenericTypeDefinition()
                 : type;
         }
+
+        [Conditional("ANVIL_DEBUG_SAFETY")]
+        private void Debug_EnsureSameID(AbstractAccessWrapper wrapper)
+        {
+            if (ID != wrapper.ID)
+            {
+                throw new Exception($"Cannot merge two wrappers with different IDs. This:{ID}, Other:{wrapper.ID}");
+            }
+        }
+
+        [Conditional("ANVIL_DEBUG_SAFETY")]
+        private void Debug_EnsureSameType(AbstractAccessWrapper wrapper)
+        {
+            Type thisType = GetType();
+            Type otherType = wrapper.GetType();
+            if (thisType != otherType)
+            {
+                throw new Exception($"Cannot merge two wrappers of different types. This:{thisType.GetReadableName()}, Other:{otherType.GetReadableName()}");
+            }
+        }
     }
 }

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/AbstractDataStreamAccessWrapper.cs → Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/AbstractDataStreamActiveAccessWrapper.cs

@@ -2,12 +2,12 @@
 
 namespace Anvil.Unity.DOTS.Entities.TaskDriver
 {
-    internal abstract class AbstractDataStreamAccessWrapper<T> : AbstractAccessWrapper
+    internal abstract class AbstractDataStreamActiveAccessWrapper<T> : AbstractAccessWrapper
         where T : unmanaged, IEntityKeyedTask
     {
         public EntityProxyDataStream<T> DataStream { get; }
 
-        protected AbstractDataStreamAccessWrapper(EntityProxyDataStream<T> dataStream, AccessType accessType, AbstractJobConfig.Usage usage)
+        protected AbstractDataStreamActiveAccessWrapper(EntityProxyDataStream<T> dataStream, AccessType accessType, AbstractJobConfig.Usage usage)
             : base(accessType, usage)
         {
             DataStream = dataStream;

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/AbstractDataStreamPendingAccessWrapper.cs

@@ -0,0 +1,79 @@
+using Anvil.CSharp.Logging;
+using Anvil.Unity.DOTS.Jobs;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+
+namespace Anvil.Unity.DOTS.Entities.TaskDriver
+{
+    internal abstract class AbstractDataStreamPendingAccessWrapper<T> : AbstractAccessWrapper
+        where T : AbstractDataStream, IAbstractDataStream
+    {
+        protected readonly T m_DefaultStream;
+
+        protected AbstractDataStreamPendingAccessWrapper(T defaultStream, AccessType accessType, AbstractJobConfig.Usage usage)
+            : base(accessType, usage)
+        {
+            m_DefaultStream = defaultStream;
+            DEBUG_TrackRequiredStream(defaultStream);
+        }
+
+        public T GetInstance(IAbstractDataStream explicitStream = null)
+        {
+            DEBUG_EnforceExplicitStream(explicitStream);
+            DEBUG_EnsureStreamWasRequired(explicitStream);
+
+            return (T)explicitStream ?? m_DefaultStream;
+        }
+
+        public override void MergeStateFrom(AbstractAccessWrapper other)
+        {
+            base.MergeStateFrom(other);
+
+            AbstractDataStreamPendingAccessWrapper<T> otherTyped = (AbstractDataStreamPendingAccessWrapper<T>)other;
+            m_DEBUG_RequiredStreams.UnionWith(otherTyped.m_DEBUG_RequiredStreams);
+            DEBUG_TrackRequiredStream(otherTyped.m_DefaultStream);
+        }
+
+        //*************************************************************************************************************
+        // SAFETY
+        //*************************************************************************************************************
+
+        // Used only with #if ANVIL_DEBUG_SAFETY
+        private HashSet<T> m_DEBUG_RequiredStreams;
+
+        [Conditional("ANVIL_DEBUG_SAFETY")]
+        private void DEBUG_TrackRequiredStream(T stream)
+        {
+            m_DEBUG_RequiredStreams ??= new HashSet<T>(1);
+            m_DEBUG_RequiredStreams.Add(stream);
+        }
+
+        [Conditional("ANVIL_DEBUG_SAFETY")]
+        private void DEBUG_EnsureStreamWasRequired(IAbstractDataStream stream)
+        {
+            T typedStream = (T)stream;
+            if (stream == null || m_DEBUG_RequiredStreams.Contains(typedStream))
+            {
+                return;
+            }
+
+            throw new Exception($"The explicit stream instance requested was not set as required. DataTargetID:{typedStream.DataTargetID}");
+        }
+
+        [Conditional("ANVIL_DEBUG_SAFETY")]
+        private void DEBUG_EnforceExplicitStream(IAbstractDataStream stream)
+        {
+            int requiredStreamCount = m_DEBUG_RequiredStreams.Count;
+            if (stream == null && requiredStreamCount > 1)
+            {
+                throw new Exception($"More than one stream has set this type as a requirement. The exact stream must be provided on retrieval. Type:{typeof(T).GetReadableName()}");
+            }
+
+            if (stream != null && requiredStreamCount == 1)
+            {
+                Logger.Warning($"An explicit stream was provided when not required. Consider using default fulfillment. Type:{typeof(T).GetReadableName()}");
+            }
+        }
+    }
+}

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/CancelProgressLookupAccessWrapper.cs

@@ -11,7 +11,11 @@ internal class CancelProgressLookupAccessWrapper : AbstractAccessWrapper
         public UnsafeParallelHashMap<EntityKeyedTaskID, bool> ProgressLookup { get; }
 
 
-        public CancelProgressLookupAccessWrapper(ActiveLookupData<EntityKeyedTaskID> cancelProgressLookupData, AccessType accessType, AbstractJobConfig.Usage usage) : base(accessType, usage)
+        public CancelProgressLookupAccessWrapper(
+            ActiveLookupData<EntityKeyedTaskID> cancelProgressLookupData,
+            AccessType accessType,
+            AbstractJobConfig.Usage usage)
+            : base(accessType, usage)
         {
             m_CancelProgressLookupData = cancelProgressLookupData;
             ProgressLookup = m_CancelProgressLookupData.Lookup;
@@ -27,4 +31,4 @@ public override void ReleaseAsync(JobHandle dependsOn)
             m_CancelProgressLookupData.ReleaseAsync(dependsOn);
         }
     }
-}
+}

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/CancelRequestsPendingAccessWrapper.cs

@@ -3,27 +3,22 @@
 
 namespace Anvil.Unity.DOTS.Entities.TaskDriver
 {
-    internal class CancelRequestsPendingAccessWrapper : AbstractAccessWrapper
+    internal class CancelRequestsPendingAccessWrapper : AbstractDataStreamPendingAccessWrapper<CancelRequestsDataStream>
     {
-        public CancelRequestsDataStream CancelRequestsDataStream { get; }
-
         public CancelRequestsPendingAccessWrapper(
-            CancelRequestsDataStream cancelRequestsDataStream,
+            CancelRequestsDataStream defaultStream,
             AccessType accessType,
             AbstractJobConfig.Usage usage)
-            : base(accessType, usage)
-        {
-            CancelRequestsDataStream = cancelRequestsDataStream;
-        }
+            : base(defaultStream, accessType, usage) { }
 
         public override JobHandle AcquireAsync()
         {
-            return CancelRequestsDataStream.AcquirePendingAsync(AccessType);
+            return m_DefaultStream.AcquirePendingAsync(AccessType);
         }
 
         public override void ReleaseAsync(JobHandle dependsOn)
         {
-            CancelRequestsDataStream.ReleasePendingAsync(dependsOn);
+            m_DefaultStream.ReleasePendingAsync(dependsOn);
         }
     }
-}
+}

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/DataStreamActiveAccessWrapper.cs

@@ -5,7 +5,7 @@ namespace Anvil.Unity.DOTS.Entities.TaskDriver
 {
     //TODO: Can we simplify this a lot and get rid of a bunch of special Access Wrapper types?
     //TODO: https://github.com/decline-cookies/anvil-unity-dots/pull/105#discussion_r1043593841
-    internal class DataStreamActiveAccessWrapper<T> : AbstractDataStreamAccessWrapper<T>
+    internal class DataStreamActiveAccessWrapper<T> : AbstractDataStreamActiveAccessWrapper<T>
         where T : unmanaged, IEntityKeyedTask
     {
         public DataStreamActiveAccessWrapper(

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/DataStreamActiveCancelAccessWrapper.cs

@@ -3,7 +3,7 @@
 
 namespace Anvil.Unity.DOTS.Entities.TaskDriver
 {
-    internal class DataStreamActiveCancelAccessWrapper<T> : AbstractDataStreamAccessWrapper<T>,
+    internal class DataStreamActiveCancelAccessWrapper<T> : AbstractDataStreamActiveAccessWrapper<T>,
                                                             IDataStreamPendingAccessWrapper
         where T : unmanaged, IEntityKeyedTask
     {
@@ -19,4 +19,4 @@ public override void ReleaseAsync(JobHandle dependsOn)
             DataStream.ReleaseActiveCancelAsync(dependsOn);
         }
     }
-}
+}

Scripts/Runtime/Entities/TaskDriver/TaskSet/Job/Wrapper/DataStreamPendingAccessWrapper.cs

@@ -3,20 +3,20 @@
 
 namespace Anvil.Unity.DOTS.Entities.TaskDriver
 {
-    internal class DataStreamPendingAccessWrapper<T> : AbstractDataStreamAccessWrapper<T>, IDataStreamPendingAccessWrapper
+    internal class DataStreamPendingAccessWrapper<T> : AbstractDataStreamPendingAccessWrapper<EntityProxyDataStream<T>>, IDataStreamPendingAccessWrapper
         where T : unmanaged, IEntityKeyedTask
     {
         public DataStreamPendingAccessWrapper(EntityProxyDataStream<T> dataStream, AccessType accessType, AbstractJobConfig.Usage usage)
             : base(dataStream, accessType, usage) { }
 
         public override JobHandle AcquireAsync()
         {
-            return DataStream.AcquirePendingAsync(AccessType);
+            return m_DefaultStream.AcquirePendingAsync(AccessType);
         }
 
         public override void ReleaseAsync(JobHandle dependsOn)
         {
-            DataStream.ReleasePendingAsync(dependsOn);
+            m_DefaultStream.ReleasePendingAsync(dependsOn);
         }
     }
 }

Scripts/Runtime/Entities/TaskDriver/TaskSet/TaskData/DataStream/Cancellation/CancelRequestsDataStream.cs

@@ -9,7 +9,7 @@ internal class CancelRequestsDataStream : AbstractDataStream,
                                               ISystemCancelRequestDataStream
     {
         private const string UNIQUE_CONTEXT_IDENTIFIER = "CANCEL_REQUEST";
-        
+
         private readonly CancelRequestsDataSource m_DataSource;
         public ActiveLookupData<EntityKeyedTaskID> ActiveLookupData { get; }
 
@@ -22,7 +22,7 @@ public override IDataSource DataSource
         {
             get => m_DataSource;
         }
-        
+
         /// <inheritdoc cref="IAbstractDataStream.ActiveDataVersion"/>
         public uint ActiveDataVersion
         {
@@ -88,11 +88,11 @@ public void ReleaseCancelRequestsWriter()
         {
             ReleasePending();
         }
-        
+
         /// <inheritdoc cref="IAbstractDataStream.IsActiveDataInvalidated"/>
         public bool IsActiveDataInvalidated(uint lastVersion)
         {
             return ActiveLookupData.IsDataInvalidated(lastVersion);
         }
     }
-}
+}