Advertisement
cajphrase

EntityHashGrid

Sep 10th, 2024
35
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C# 7.55 KB | None | 0 0
  1. using SBP;
  2. using SBP.LINK;
  3. using SBP.Meshing;
  4. using System;
  5. using System.Collections;
  6. using System.Collections.Generic;
  7. using Unity.Collections;
  8. using Unity.Collections.LowLevel.Unsafe;
  9. using Unity.Jobs;
  10. using Unity.Mathematics;
  11. using UnityEngine;
  12.  
  13.  
  14. // wraps a pointer to a hashset so we can nest in a native container without unity complaining
  15. // could use without the pointer but would need to be careful that the internal state of the struct is stored, you can't pass it around
  16. public unsafe struct UnsafeHashSetHolder<T> : IDisposable , IEnumerable<T>
  17.         where T : unmanaged, IEquatable<T>
  18. {
  19.     // TODO disable ptr restriction
  20.     public UnsafeHashSet<T>* data;
  21.  
  22.     /// <summary>
  23.     /// Removes all values.
  24.     /// </summary>
  25.     /// <remarks>Does not change the capacity.</remarks>
  26.     public void Clear() => data->Clear();
  27.  
  28.     /// <summary>
  29.     /// Adds a new value (unless it is already present).
  30.     /// </summary>
  31.     /// <param name="item">The value to add.</param>
  32.     /// <returns>True if the value was not already present.</returns>
  33.     public bool Add(T item) => data->Add(item);
  34.  
  35.     /// <summary>
  36.     /// Removes a particular value.
  37.     /// </summary>
  38.     /// <param name="item">The value to remove.</param>
  39.     /// <returns>True if the value was present.</returns>
  40.     public bool Remove(T item) => data->Remove(item);
  41.  
  42.     /// <summary>
  43.     /// Returns true if a particular value is present.
  44.     /// </summary>
  45.     /// <param name="item">The value to check for.</param>
  46.     /// <returns>True if the value was present.</returns>
  47.     public bool Contains(T item) => data->Contains(item);
  48.  
  49.     /// <summary>
  50.     /// Sets the capacity to match what it would be if it had been originally initialized with all its entries.
  51.     /// </summary>
  52.     public void TrimExcess() => data->TrimExcess();
  53.  
  54.     /// <summary>
  55.     /// Returns an enumerator over the values of this set.
  56.     /// </summary>
  57.     /// <returns>An enumerator over the values of this set.</returns>
  58.     public UnsafeHashSet<T>.Enumerator GetEnumerator()
  59.     {
  60.         return data->GetEnumerator();
  61.     }
  62.  
  63.     /// <summary>
  64.     /// This method is not implemented. Use <see cref="GetEnumerator"/> instead.
  65.     /// </summary>
  66.     /// <returns>Throws NotImplementedException.</returns>
  67.     /// <exception cref="NotImplementedException">Method is not implemented.</exception>
  68.     IEnumerator<T> IEnumerable<T>.GetEnumerator()
  69.     {
  70.         throw new NotImplementedException();
  71.     }
  72.  
  73.     /// <summary>
  74.     /// This method is not implemented. Use <see cref="GetEnumerator"/> instead.
  75.     /// </summary>
  76.     /// <returns>Throws NotImplementedException.</returns>
  77.     /// <exception cref="NotImplementedException">Method is not implemented.</exception>
  78.     IEnumerator IEnumerable.GetEnumerator()
  79.     {
  80.         throw new NotImplementedException();
  81.     }
  82. }
  83.  
  84. public static class JobsExtensions
  85. {
  86.     public static Allocator Temp => JobsUtility.IsExecutingJob() ? Allocator.Temp : Allocator.TempJob;
  87.  
  88.     public static void RunSafe<T>(this T job) where T : IJob
  89.     {
  90.         if (JobsUtility.IsExecutingJob())
  91.             job.Execute();
  92.         else
  93.             job.Run();
  94.     }
  95.  
  96.  
  97.     public static Allocator AsJobSafe(this Allocator allocator)
  98.     {
  99.         JobsUtility.IsExecutingJob() ? Allocator.Temp : allocator == Allocator.Temp : Allocator.TempJob : allocator;
  100.     }
  101. }
  102. public struct EntityHashGrid : IDisposable
  103. {
  104.     // might be more performant to use int2 xz grid, the majority of air and ground chunks will be empty
  105.     // this will reduce iterating the grid from n^3 to n^2, saving performance on large queries and reducing the
  106.     // performance cost of smaller chunkSizes
  107.     public NativeHashMap<int2, UnsafeHashSetHolder<uint>> entitiesPerChunk;
  108.     public NativeParallelMultiHashMap<uint, int2> chunksPerEntity;
  109.     public NativeHashMap<uint, AABB> boundsPerEntity;
  110.  
  111.     public uint ecsId;
  112.     // not clear how large this should be. smaller values will result in more accurate results at the expense of
  113.     // more lookups and more duplication.
  114.     public float chunkSize;
  115.  
  116.     public EntityHashGrid(uint ecsId, float chunkSize)
  117.     {
  118.         entitiesPerChunk = new(32, Allocator.Persistent);
  119.         chunksPerEntity = new(32, Allocator.Persistent);
  120.         boundsPerEntity = new(32, Allocator.Persistent);
  121.         this.ecsId = ecsId;
  122.         this.chunkSize = chunkSize;
  123.     }
  124.  
  125.     public void AddOrUpdate(Entity entity, AABB bounds, EntityTransform transform)
  126.     {
  127.         Remove(entity);
  128.         var transformedBounds = transform * bounds;
  129.         var minIndex = Index.PositionToIndex(transformedBounds.Min, chunkSize);
  130.         var maxIndex = Index.PositionToIndex(transformedBounds.Max, chunkSize);
  131.  
  132.         for (int x = minIndex.x; x <= maxIndex.x; x++)
  133.         {
  134.             for (int z = minIndex.z; z <= maxIndex.z; z++)
  135.             {
  136.                 var chunk = new int2(x, z);
  137.                 EnsureChunk(chunk).Add(entity.id);
  138.                 chunksPerEntity.Add(entity.id, chunk);
  139.             }
  140.         }
  141.         boundsPerEntity[entity.id] = transformedBounds;
  142.     }
  143.  
  144.     private UnsafeHashSetHolder<uint> EnsureChunk(int2 index)
  145.     {
  146.         if (!entitiesPerChunk.TryGetValue(index, out var entities))
  147.         {
  148.             entitiesPerChunk[index] = entities = new(32, Allocator.Persistent);
  149.         }
  150.         return entities;
  151.     }
  152.  
  153.     public void Remove(Entity entity)
  154.     {
  155.         if (chunksPerEntity.TryGetValue(entity.id, out var chunks))
  156.         {
  157.             foreach (var chunk in chunks)
  158.             {
  159.                 EnsureChunk(chunk).Remove(entity.id);
  160.             }
  161.             chunksPerEntity.Remove(entity.id);
  162.             boundsPerEntity.Remove(entity.id);
  163.         }
  164.     }
  165.  
  166.     public NativeList<Entity> Query(AABB bounds, Allocator allocator)
  167.     {
  168.         var results = new NativeList<Entity>(128, allocator.AsJobSafe());
  169.         var job = new QueryJob()
  170.         {
  171.             hashGrid = this,
  172.             chunkSize = chunkSize,
  173.             results = results
  174.         }
  175.         job.RunSafe();
  176.         return results;
  177.     }
  178.  
  179.     public void Dispose()
  180.     {
  181.         foreach (var item in entitiesPerChunk)
  182.         {
  183.             item.Value.Dispose();
  184.         }
  185.         entitiesPerChunk.Dispose();
  186.         chunksPerEntity.Dispose();
  187.         boundsPerEntity.Dispose();
  188.     }
  189.  
  190.     [BurstCompile]
  191.     public struct QueryJob : IJob
  192.     {
  193.         [ReadOnly] public EntityHashGrid hashGrid;
  194.         public float chunkSize;
  195.         public AABB bounds;
  196.         public NativeList<Entity> results;
  197.        
  198.         public void Execute()
  199.         {
  200.             var minIndex = Index.PositionToIndex(transformedBounds.Min, chunkSize);
  201.             var maxIndex = Index.PositionToIndex(transformedBounds.Max, chunkSize);
  202.  
  203.             var hashSet = new NativeHashSet<uint>(256, Allocator.Temp);
  204.             for (int x = minIndex.x; x <= maxIndex.x; x++)
  205.             {
  206.                 for (int z = minIndex.z; z <= maxIndex.z; z++)
  207.                 {
  208.                     var chunk = new int2(x, z);
  209.                     if (!hashGrid.entitiesPerChunk.TryGetValue(chunk, out var entities))
  210.                         continue;
  211.                     foreach (var entityId in entities)
  212.                     {
  213.                         if (hashSet.Add(entityId) && hashGrid.boundsPerEntity[entityId].Intersects(bounds))
  214.                         {
  215.                             results.Add(new Entity(hashGrid.ecsId, entityId));
  216.                         }
  217.                     }
  218.                 }
  219.             }
  220.         }
  221.     }
  222. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement