using System; using System.Runtime.CompilerServices; using Unity.Burst; using Unity.Jobs; using Unity.Jobs.LowLevel.Unsafe; using UnityEngine.Assertions; namespace Unity.Collections.LowLevel.Unsafe { [GenerateTestsForBurstCompatibility] internal unsafe struct UnsafeStreamBlock { internal UnsafeStreamBlock* Next; internal fixed byte Data[1]; } [GenerateTestsForBurstCompatibility] internal unsafe struct UnsafeStreamRange { internal UnsafeStreamBlock* Block; internal int OffsetInFirstBlock; internal int ElementCount; /// One byte past the end of the last byte written internal int LastOffset; internal int NumberOfBlocks; } [GenerateTestsForBurstCompatibility] internal unsafe struct UnsafeStreamBlockData { internal const int AllocationSize = 4 * 1024; internal AllocatorManager.AllocatorHandle Allocator; internal UnsafeStreamBlock** Blocks; internal int BlockCount; internal AllocatorManager.Block Ranges; internal int RangeCount; internal UnsafeStreamBlock* Allocate(UnsafeStreamBlock* oldBlock, int threadIndex) { Assert.IsTrue(threadIndex < BlockCount && threadIndex >= 0); UnsafeStreamBlock* block = (UnsafeStreamBlock*)Memory.Unmanaged.Array.Resize(null, 0, AllocationSize, Allocator, 1, 16); block->Next = null; if (oldBlock == null) { // Append our new block in front of the previous head. block->Next = Blocks[threadIndex]; Blocks[threadIndex] = block; } else { block->Next = oldBlock->Next; oldBlock->Next = block; } return block; } internal void Free(UnsafeStreamBlock* oldBlock) { Memory.Unmanaged.Array.Resize(oldBlock, AllocationSize, 0, Allocator, 1, 16); } } /// /// A set of untyped, append-only buffers. Allows for concurrent reading and concurrent writing without synchronization. /// /// /// As long as each individual buffer is written in one thread and read in one thread, multiple /// threads can read and write the stream concurrently, *e.g.* /// while thread *A* reads from buffer *X* of a stream, thread *B* can read from /// buffer *Y* of the same stream. /// /// Each buffer is stored as a chain of blocks. When a write exceeds a buffer's current capacity, another block /// is allocated and added to the end of the chain. Effectively, expanding the buffer never requires copying the existing /// data (unlike, for example, with ). /// /// **All writing to a stream should be completed before the stream is first read. Do not write to a stream after the first read.** /// /// Writing is done with , and reading is done with . /// An individual reader or writer cannot be used concurrently across threads. Each thread must use its own. /// /// The data written to an individual buffer can be heterogeneous in type, and the data written /// to different buffers of a stream can be entirely different in type, number, and order. Just make sure /// that the code reading from a particular buffer knows what to expect to read from it. /// [GenerateTestsForBurstCompatibility] public unsafe struct UnsafeStream : INativeDisposable { [NativeDisableUnsafePtrRestriction] internal AllocatorManager.Block m_BlockData; /// /// Initializes and returns an instance of UnsafeStream. /// /// The number of buffers to give the stream. You usually want /// one buffer for each thread that will read or write the stream. /// The allocator to use. public UnsafeStream(int bufferCount, AllocatorManager.AllocatorHandle allocator) { AllocateBlock(out this, allocator); AllocateForEach(bufferCount); } /// /// Creates and schedules a job to allocate a new stream. /// /// The stream can be used on the main thread after completing the returned job or used in other jobs that depend upon the returned job. /// /// Using a job to allocate the buffers can be more efficient, particularly for a stream with many buffers. /// /// Ignored. /// Outputs the new stream. /// A list whose length determines the number of buffers in the stream. /// A job handle. The new job will depend upon this handle. /// The allocator to use. /// The handle of the new job. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public static JobHandle ScheduleConstruct(out UnsafeStream stream, NativeList bufferCount, JobHandle dependency, AllocatorManager.AllocatorHandle allocator) where T : unmanaged { AllocateBlock(out stream, allocator); var jobData = new ConstructJobList { List = (UntypedUnsafeList*)bufferCount.GetUnsafeList(), Container = stream }; return jobData.Schedule(dependency); } /// /// Creates and schedules a job to allocate a new stream. /// /// The stream can be used on the main thread after completing the returned job or used in other jobs that depend upon the returned job. /// /// Allocating the buffers in a job can be more efficient, particularly for a stream with many buffers. /// /// Outputs the new stream. /// An array whose value at index 0 determines the number of buffers in the stream. /// A job handle. The new job will depend upon this handle. /// The allocator to use. /// The handle of the new job. public static JobHandle ScheduleConstruct(out UnsafeStream stream, NativeArray bufferCount, JobHandle dependency, AllocatorManager.AllocatorHandle allocator) { AllocateBlock(out stream, allocator); var jobData = new ConstructJob { Length = bufferCount, Container = stream }; return jobData.Schedule(dependency); } internal static void AllocateBlock(out UnsafeStream stream, AllocatorManager.AllocatorHandle allocator) { #if UNITY_2022_2_14F1_OR_NEWER int maxThreadCount = JobsUtility.ThreadIndexCount; #else int maxThreadCount = JobsUtility.MaxJobThreadCount; #endif int blockCount = maxThreadCount; int allocationSize = sizeof(UnsafeStreamBlockData) + sizeof(UnsafeStreamBlock*) * blockCount; AllocatorManager.Block blk = AllocatorManager.AllocateBlock(ref allocator, allocationSize, 16, 1); UnsafeUtility.MemClear( (void*)blk.Range.Pointer, blk.AllocatedBytes); stream.m_BlockData = blk; var blockData = (UnsafeStreamBlockData*)blk.Range.Pointer; blockData->Allocator = allocator; blockData->BlockCount = blockCount; blockData->Blocks = (UnsafeStreamBlock**)(blk.Range.Pointer + sizeof(UnsafeStreamBlockData)); blockData->Ranges = default; blockData->RangeCount = 0; } internal void AllocateForEach(int forEachCount) { long allocationSize = sizeof(UnsafeStreamRange) * forEachCount; var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; blockData->Ranges = AllocatorManager.AllocateBlock(ref m_BlockData.Range.Allocator, sizeof(UnsafeStreamRange), 16, forEachCount); blockData->RangeCount = forEachCount; UnsafeUtility.MemClear((void*)blockData->Ranges.Range.Pointer, blockData->Ranges.AllocatedBytes); } /// /// Returns true if this stream is empty. /// /// True if this stream is empty or the stream has not been constructed. public readonly bool IsEmpty() { if (!IsCreated) { return true; } var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; var ranges = (UnsafeStreamRange*)blockData->Ranges.Range.Pointer; for (int i = 0; i != blockData->RangeCount; i++) { if (ranges[i].ElementCount > 0) { return false; } } return true; } /// /// Whether this stream has been allocated (and not yet deallocated). /// /// Does not necessarily reflect whether the buffers of the stream have themselves been allocated. /// True if this stream has been allocated (and not yet deallocated). public readonly bool IsCreated { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => m_BlockData.Range.Pointer != IntPtr.Zero; } /// /// The number of buffers in this stream. /// /// The number of buffers in this stream. public readonly int ForEachCount => ((UnsafeStreamBlockData*)m_BlockData.Range.Pointer)->RangeCount; /// /// Returns a reader of this stream. /// /// A reader of this stream. public Reader AsReader() { return new Reader(ref this); } /// /// Returns a writer of this stream. /// /// A writer of this stream. public Writer AsWriter() { return new Writer(ref this); } /// /// Returns the total number of items in the buffers of this stream. /// /// Each and call increments this number. /// The total number of items in the buffers of this stream. public int Count() { int itemCount = 0; var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; var ranges = (UnsafeStreamRange*)blockData->Ranges.Range.Pointer; for (int i = 0; i != blockData->RangeCount; i++) { itemCount += ranges[i].ElementCount; } return itemCount; } /// /// Returns a new NativeArray copy of this stream's data. /// /// The length of the array will equal the count of this stream. /// /// Each buffer of this stream is copied to the array, one after the other. /// /// The type of values in the array. /// The allocator to use. /// A new NativeArray copy of this stream's data. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public NativeArray ToNativeArray(AllocatorManager.AllocatorHandle allocator) where T : unmanaged { var array = CollectionHelper.CreateNativeArray(Count(), allocator, NativeArrayOptions.UninitializedMemory); var reader = AsReader(); int offset = 0; for (int i = 0; i != reader.ForEachCount; i++) { reader.BeginForEachIndex(i); int rangeItemCount = reader.RemainingItemCount; for (int j = 0; j < rangeItemCount; ++j) { array[offset] = reader.Read(); offset++; } reader.EndForEachIndex(); } return array; } void Deallocate() { if (!IsCreated) { return; } var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; for (int i = 0; i != blockData->BlockCount; i++) { UnsafeStreamBlock* block = blockData->Blocks[i]; while (block != null) { UnsafeStreamBlock* next = block->Next; blockData->Free(block); block = next; } } blockData->Ranges.Dispose(); m_BlockData.Dispose(); m_BlockData = default; } /// /// Releases all resources (memory). /// public void Dispose() { if (!IsCreated) { return; } Deallocate(); } /// /// Creates and schedules a job that will release all resources (memory and safety handles) of this stream. /// /// A job handle which the newly scheduled job will depend upon. /// The handle of a new job that will release all resources (memory and safety handles) of this stream. public JobHandle Dispose(JobHandle inputDeps) { if (!IsCreated) { return inputDeps; } var jobHandle = new DisposeJob { Container = this }.Schedule(inputDeps); m_BlockData = default; return jobHandle; } [BurstCompile] struct DisposeJob : IJob { public UnsafeStream Container; public void Execute() { Container.Deallocate(); } } [BurstCompile] struct ConstructJobList : IJob { public UnsafeStream Container; [ReadOnly] [NativeDisableUnsafePtrRestriction] public UntypedUnsafeList* List; public void Execute() { Container.AllocateForEach(List->m_length); } } [BurstCompile] struct ConstructJob : IJob { public UnsafeStream Container; [ReadOnly] public NativeArray Length; public void Execute() { Container.AllocateForEach(Length[0]); } } /// /// Writes data into a buffer of an . /// /// An individual writer can only be used for one buffer of one stream. /// Do not create more than one writer for an individual buffer. [GenerateTestsForBurstCompatibility] public unsafe struct Writer { [NativeDisableUnsafePtrRestriction] internal AllocatorManager.Block m_BlockData; [NativeDisableUnsafePtrRestriction] UnsafeStreamBlock* m_CurrentBlock; [NativeDisableUnsafePtrRestriction] byte* m_CurrentPtr; [NativeDisableUnsafePtrRestriction] byte* m_CurrentBlockEnd; internal int m_ForeachIndex; int m_ElementCount; [NativeDisableUnsafePtrRestriction] UnsafeStreamBlock* m_FirstBlock; int m_FirstOffset; int m_NumberOfBlocks; [NativeSetThreadIndex] int m_ThreadIndex; internal Writer(ref UnsafeStream stream) { m_BlockData = stream.m_BlockData; m_ForeachIndex = int.MinValue; m_ElementCount = -1; m_CurrentBlock = null; m_CurrentBlockEnd = null; m_CurrentPtr = null; m_FirstBlock = null; m_NumberOfBlocks = 0; m_FirstOffset = 0; m_ThreadIndex = 0; } /// /// The number of buffers in the stream of this writer. /// /// The number of buffers in the stream of this writer. public int ForEachCount => ((UnsafeStreamBlockData*)m_BlockData.Range.Pointer)->RangeCount; /// /// Readies this writer to write to a particular buffer of the stream. /// /// Must be called before using this writer. For an individual writer, call this method only once. /// /// When done using this writer, you must call . /// The index of the buffer to write. public void BeginForEachIndex(int foreachIndex) { m_ForeachIndex = foreachIndex; m_ElementCount = 0; m_NumberOfBlocks = 0; m_FirstBlock = m_CurrentBlock; m_FirstOffset = (int)(m_CurrentPtr - (byte*)m_CurrentBlock); } /// /// Readies the buffer written by this writer for reading. /// /// Must be called before reading the buffer written by this writer. public void EndForEachIndex() { var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; var ranges = (UnsafeStreamRange*)blockData->Ranges.Range.Pointer; ranges[m_ForeachIndex].ElementCount = m_ElementCount; ranges[m_ForeachIndex].OffsetInFirstBlock = m_FirstOffset; ranges[m_ForeachIndex].Block = m_FirstBlock; ranges[m_ForeachIndex].LastOffset = (int)(m_CurrentPtr - (byte*)m_CurrentBlock); ranges[m_ForeachIndex].NumberOfBlocks = m_NumberOfBlocks; } /// /// Write a value to a buffer. /// /// The value is written to the buffer which was specified /// with . /// The type of value to write. /// The value to write. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public void Write(T value) where T : unmanaged { ref T dst = ref Allocate(); dst = value; } /// /// Allocate space in a buffer. /// /// The space is allocated in the buffer which was specified /// with . /// The type of value to allocate space for. /// A reference to the allocation. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public ref T Allocate() where T : unmanaged { int size = UnsafeUtility.SizeOf(); return ref UnsafeUtility.AsRef(Allocate(size)); } /// /// Allocate space in a buffer. /// /// The space is allocated in the buffer which was specified /// with . /// The number of bytes to allocate. /// The allocation. public byte* Allocate(int size) { byte* ptr = m_CurrentPtr; m_CurrentPtr += size; if (m_CurrentPtr > m_CurrentBlockEnd) { UnsafeStreamBlock* oldBlock = m_CurrentBlock; var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; m_CurrentBlock = blockData->Allocate(oldBlock, m_ThreadIndex); m_CurrentPtr = m_CurrentBlock->Data; if (m_FirstBlock == null) { m_FirstOffset = (int)(m_CurrentPtr - (byte*)m_CurrentBlock); m_FirstBlock = m_CurrentBlock; } else { m_NumberOfBlocks++; } m_CurrentBlockEnd = (byte*)m_CurrentBlock + UnsafeStreamBlockData.AllocationSize; ptr = m_CurrentPtr; m_CurrentPtr += size; } m_ElementCount++; return ptr; } } /// /// Reads data from a buffer of an . /// /// An individual reader can only be used for one buffer of one stream. /// Do not create more than one reader for an individual buffer. [GenerateTestsForBurstCompatibility] public unsafe struct Reader { [NativeDisableUnsafePtrRestriction] internal AllocatorManager.Block m_BlockData; [NativeDisableUnsafePtrRestriction] internal UnsafeStreamBlock* m_CurrentBlock; [NativeDisableUnsafePtrRestriction] internal byte* m_CurrentPtr; [NativeDisableUnsafePtrRestriction] internal byte* m_CurrentBlockEnd; internal int m_RemainingItemCount; internal int m_LastBlockSize; internal Reader(ref UnsafeStream stream) { m_BlockData = stream.m_BlockData; m_CurrentBlock = null; m_CurrentPtr = null; m_CurrentBlockEnd = null; m_RemainingItemCount = 0; m_LastBlockSize = 0; } /// /// Readies this reader to read a particular buffer of the stream. /// /// Must be called before using this reader. For an individual reader, call this method only once. /// /// When done using this reader, you must call . /// The index of the buffer to read. /// The number of remaining elements to read from the buffer. public int BeginForEachIndex(int foreachIndex) { var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; var ranges = (UnsafeStreamRange*)blockData->Ranges.Range.Pointer; m_RemainingItemCount = ranges[foreachIndex].ElementCount; m_LastBlockSize = ranges[foreachIndex].LastOffset; m_CurrentBlock = ranges[foreachIndex].Block; m_CurrentPtr = (byte*)m_CurrentBlock + ranges[foreachIndex].OffsetInFirstBlock; m_CurrentBlockEnd = (byte*)m_CurrentBlock + UnsafeStreamBlockData.AllocationSize; return m_RemainingItemCount; } /// /// Does nothing. /// /// Included only for consistency with . public void EndForEachIndex() { } /// /// The number of buffers in the stream of this reader. /// /// The number of buffers in the stream of this reader. public int ForEachCount => ((UnsafeStreamBlockData*)m_BlockData.Range.Pointer)->RangeCount; /// /// The number of items not yet read from the buffer. /// /// The number of items not yet read from the buffer. public int RemainingItemCount => m_RemainingItemCount; /// /// Returns a pointer to the next position to read from the buffer. Advances the reader some number of bytes. /// /// The number of bytes to advance the reader. /// A pointer to the next position to read from the buffer. /// Thrown if the reader has been advanced past the end of the buffer. public byte* ReadUnsafePtr(int size) { m_RemainingItemCount--; byte* ptr = m_CurrentPtr; m_CurrentPtr += size; if (m_CurrentPtr > m_CurrentBlockEnd) { m_CurrentBlock = m_CurrentBlock->Next; m_CurrentPtr = m_CurrentBlock->Data; m_CurrentBlockEnd = (byte*)m_CurrentBlock + UnsafeStreamBlockData.AllocationSize; ptr = m_CurrentPtr; m_CurrentPtr += size; } return ptr; } /// /// Reads the next value from the buffer. /// /// Each read advances the reader to the next item in the buffer. /// The type of value to read. /// A reference to the next value from the buffer. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public ref T Read() where T : unmanaged { int size = UnsafeUtility.SizeOf(); return ref UnsafeUtility.AsRef(ReadUnsafePtr(size)); } /// /// Reads the next value from the buffer. Does not advance the reader. /// /// The type of value to read. /// A reference to the next value from the buffer. [GenerateTestsForBurstCompatibility(GenericTypeArguments = new[] { typeof(int) })] public ref T Peek() where T : unmanaged { int size = UnsafeUtility.SizeOf(); byte* ptr = m_CurrentPtr; if (ptr + size > m_CurrentBlockEnd) { ptr = m_CurrentBlock->Next->Data; } return ref UnsafeUtility.AsRef(ptr); } /// /// Returns the total number of items in the buffers of the stream. /// /// The total number of items in the buffers of the stream. public int Count() { var blockData = (UnsafeStreamBlockData*)m_BlockData.Range.Pointer; var ranges = (UnsafeStreamRange*)blockData->Ranges.Range.Pointer; int itemCount = 0; for (int i = 0; i != blockData->RangeCount; i++) { itemCount += ranges[i].ElementCount; } return itemCount; } } } }