ColliderPartition.cs

using System.Collections.Generic;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Physics;
using Unity.Collections;
using static Unity.Physics.CompoundCollider;
using Unity.Transforms;
using Unity.Collections.LowLevel.Unsafe;

namespace AiGame.ColliderPartitions
{

    [UpdateInGroup(typeof(InitializationSystemGroup))]
    public class ColliderPartitionSystem : SystemBase
    {
        public static ColliderPartitionSystem Instance => World.DefaultGameObjectInjectionWorld.GetExistingSystem<ColliderPartitionSystem>();

        public Entity LookupEntity { get; private set; }

        private Dictionary<int2, Entity> PartitionEntities = new Dictionary<int2, Entity>();
        private Dictionary<Entity, float3> ColliderEntityToPositionMap = new Dictionary<Entity, float3>();
        private NativeList<ColliderBlobInstance> TempInstances;
        private EndInitializationEntityCommandBufferSystem EcbSystem;
        private ChildMapComponent ChildMap;
        private BlobAssetReference<Collider> DefaultCollider;

        public void ClearPartitions()
        {
            Dependency.Complete();
            ChildMap.Colliders.Clear();
            Entities
                .ForEach((ref ColliderPartition partition, ref DynamicBuffer<PartitionChangeBuffer> changeBuffer) =>
                {
                    partition.Reset = true;
                    changeBuffer.Clear();
                })
                .Run();

        }

        public void AddCollider(BlobAssetReference<Collider> collider, float3 position, quaternion rotation, Entity entity)
        {
            Entity partitionEntity = FindOrCreatePartition(position);

            ChildColliderDef child = new ChildColliderDef
            {
                Collider = collider,
                Transform = new RigidTransform { pos = position, rot = rotation },
                Entity = entity
            };

            PartitionChange change = new PartitionChange
            {
                ChangeType = PartitionChangeType.AddCollider,
                ColliderDef = child
            };

            var ecb = EcbSystem.CreateCommandBuffer();
            ecb.AppendToBuffer<PartitionChangeBuffer>(partitionEntity, change);

            if (!entity.Equals(Entity.Null))
            {
                ColliderEntityToPositionMap[entity] = position;
            }
        }

        public bool RemoveCollider(Entity entity)
        {
            if (ColliderEntityToPositionMap.TryGetValue(entity, out float3 position))
            {
                var commandBuffer = EcbSystem.CreateCommandBuffer();
                if (RemoveCollider(position, commandBuffer))
                {
                    ColliderEntityToPositionMap.Remove(entity);
                    return true;
                }
            }

            return false;
        }

        public void RemoveColliders(NativeArray<Entity> entities)
        {
            var commandBuffer = EcbSystem.CreateCommandBuffer();
            for (int i=0;i<entities.Length;i++)
            {
                Entity entity = entities[i];
                if (ColliderEntityToPositionMap.TryGetValue(entity, out float3 position))
                {
                    if (RemoveCollider(position, commandBuffer))
                    {
                        ColliderEntityToPositionMap.Remove(entity);
                    }
                }
            }
        }

        public bool RemoveCollider(float3 position, EntityCommandBuffer commandBuffer)
        {
            int2 key = ColliderPartition.GetPartitionId(position);
            if (PartitionEntities.TryGetValue(key, out Entity entity))
            {
                PartitionChange change = new PartitionChange
                {
                    ChangeType = PartitionChangeType.RemoveCollider,
                    RemovePosition = position.ToInt3()
                };

                commandBuffer.AppendToBuffer<PartitionChangeBuffer>(entity, change);
                return true;
            }
            else
            {
                return false;
            }
        }

        private Entity FindOrCreatePartition(float3 position)
        {
            int2 id = ColliderPartition.GetPartitionId(position);
            if (!PartitionEntities.TryGetValue(id, out Entity entity))
            {
                ColliderPartition partition = new ColliderPartition();
                partition.Id = id;
                partition.Colliders = new UnsafeHashMap<int3, ChildColliderDef>(4, Allocator.Persistent);

                entity = EntityManager.CreateEntity();
                EntityManager.AddComponent(entity, ComponentType.ReadOnly<Static>());
                EntityManager.AddComponentData(entity, new Translation { Value = default });
                EntityManager.AddComponentData(entity, new Rotation { Value = quaternion.identity });
                EntityManager.AddComponentData(entity, partition);
                EntityManager.AddBuffer<PartitionChangeBuffer>(entity);

                PhysicsCollider physicsCollider = new PhysicsCollider();
                physicsCollider.Value = DefaultCollider;
                EntityManager.AddComponentData(entity, physicsCollider);

                PartitionEntities[id] = entity;
            }

            return entity;
        }


        protected override void OnCreate()
        {
            base.OnCreate();
            EcbSystem = World.GetOrCreateSystem<EndInitializationEntityCommandBufferSystem>();
            TempInstances = new NativeList<ColliderBlobInstance>(Allocator.Persistent);

            ChildMap = new ChildMapComponent
            {
                Colliders = new UnsafeHashMap<int3, ChildColliderInstance>(4, Allocator.Persistent)
            };
            LookupEntity = EntityManager.CreateEntity();
            EntityManager.AddComponentData(LookupEntity, ChildMap);

            CreateDefaultCollider();
        }

        protected override void OnDestroy()
        {
            base.OnDestroy();
            TempInstances.Dispose();

            ChildMap.Colliders.Dispose();

            Entities
                .ForEach((Entity entity, ref ColliderPartition partition, ref DynamicBuffer<PartitionChangeBuffer> changeBuffer) =>
                {
                    partition.Colliders.Dispose();
                })
                .Run();

        }

        protected override void OnUpdate()
        {
            var ecb = EcbSystem.CreateCommandBuffer();

            PartitionUpdateJob updateJob = new PartitionUpdateJob
            {
                Ecb = ecb,
                ColliderLookup = ChildMap.Colliders,
                TempInstances = TempInstances,
                DefaultCollider = DefaultCollider
            };

            Entities
                .ForEach((Entity entity, ref ColliderPartition partition, ref DynamicBuffer<PartitionChangeBuffer> changeBuffer) =>
                {
                    updateJob.Execute(entity, ref partition, ref changeBuffer);

                })
                .Schedule();

            EcbSystem.AddJobHandleForProducer(Dependency);

        }

        private void CreateDefaultCollider()
        {
            CollisionFilter filter = CollisionFilter.Zero;
            SphereGeometry geom = new SphereGeometry { Center = new float3(0f, -1000f, 0f), Radius = 0.1f };
            DefaultCollider = SphereCollider.Create(geom, filter);
        }

        struct PartitionUpdateJob
        {
            public EntityCommandBuffer Ecb;
            public UnsafeHashMap<int3, ChildColliderInstance> ColliderLookup;
            public NativeList<ColliderBlobInstance> TempInstances;
            public BlobAssetReference<Collider> DefaultCollider;

            public void Execute(Entity entity, ref ColliderPartition partition, ref DynamicBuffer<PartitionChangeBuffer> changeBuffer)
            {
                if (partition.Reset)
                {
                    partition.Reset = false;
                    partition.Colliders.Clear();
                    SetDefaultCollider(entity, ref partition);
                }

                if (changeBuffer.Length == 0)
                {
                    return;
                }

                for (int i = 0; i < changeBuffer.Length; i++)
                {
                    var change = changeBuffer[i].Value;

                    switch (change.ChangeType)
                    {
                        case PartitionChangeType.RemoveCollider:
                            partition.Colliders.Remove(change.RemovePosition);
                            break;
                        case PartitionChangeType.AddCollider:
                            int3 quantized = change.ColliderDef.Transform.pos.ToInt3();
                            if (!partition.Colliders.ContainsKey(quantized))
                            {
                                partition.Colliders[quantized] = change.ColliderDef;
                            }
                            break;

                    }
                }

                changeBuffer.Clear();

                if (partition.Colliders.Count() == 0)
                {
                    SetDefaultCollider(entity, ref partition);
                }
                else
                {
                    Rebuild(entity, ref partition);
                }
            }

            private void SetDefaultCollider(Entity entity, ref ColliderPartition partition)
            {
                PhysicsCollider physicsCollider = new PhysicsCollider();
                physicsCollider.Value = DefaultCollider;
                Ecb.SetComponent(entity, physicsCollider);
                partition.ChildCount = 0;
            }

            private void Rebuild(Entity entity, ref ColliderPartition partition)
            {
                TempInstances.Clear();
                var keys = partition.Colliders.GetKeyArray(Allocator.Temp);
                for (int i = 0; i < keys.Length; i++)
                {
                    int3 key = keys[i];
                    ChildColliderDef colliderDef = partition.Colliders[key];

                    ColliderBlobInstance instance = new ColliderBlobInstance
                    {
                        Collider = colliderDef.Collider,
                        CompoundFromChild = colliderDef.Transform
                    };

                    TempInstances.Add(instance);

                    ChildColliderInstance data = new ChildColliderInstance
                    {
                        Entity = colliderDef.Entity,
                        QuantizedPosition = key
                    };
                    ColliderLookup[key] = data;
                }

                partition.ChildCount = TempInstances.Length;
                var compound = CompoundCollider.Create(TempInstances);
                PhysicsCollider physicsCollider = new PhysicsCollider();
                physicsCollider.Value = compound;
                Ecb.SetComponent(entity, physicsCollider);
            }
        }


    }
}

using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Physics;

namespace AiGame.ColliderPartitions
{
    internal struct ColliderPartition : IComponentData
    {
        public const float PartitionFactor = 64f;

        public int2 Id;
        public int ChildCount;
        public bool Reset;
        public UnsafeHashMap<int3, ChildColliderDef> Colliders;


        public static int2 GetPartitionId(float3 position)
        {
            return new int2(FloorToPartition(position.x), FloorToPartition(position.z));
        }

        private static int FloorToPartition(float pos)
        {
            float value = math.floor(pos / PartitionFactor) * PartitionFactor;
            return (int)value;
        }

    }

    internal struct ChildColliderDef
    {
        public BlobAssetReference<Collider> Collider;
        public RigidTransform Transform;
        public Entity Entity;
    }

    internal enum PartitionChangeType
    {
        AddCollider,
        RemoveCollider
    }
    internal struct PartitionChange : IComponentData
    {
        public PartitionChangeType ChangeType;
        public ChildColliderDef ColliderDef;
        public int3 RemovePosition;
    }

    internal struct PartitionChangeBuffer : IBufferElementData
    {
        public static implicit operator PartitionChange(PartitionChangeBuffer e) { return e.Value; }
        public static implicit operator PartitionChangeBuffer(PartitionChange e) { return new PartitionChangeBuffer { Value = e }; }

        public PartitionChange Value;
    }

    public struct ChildColliderInstance
    {
        public Entity Entity;
        public int3 QuantizedPosition;
    }

    public struct ChildMapComponent : IComponentData
    {
        public UnsafeHashMap<int3, ChildColliderInstance> Colliders;

        public static ChildMapComponent GetComponent()
        {
            var system = ColliderPartitionSystem.Instance;
            return system.EntityManager.GetComponentData<ChildMapComponent>(system.LookupEntity);
        }

        

        public bool TryGetChild(CollisionWorld world, int bodyIndex, ColliderKey key, out ChildColliderInstance childData)
        {
            RigidBody body = world.Bodies[bodyIndex];

            if (!key.Equals(ColliderKey.Empty))
            {
                if (body.Collider.Value.GetChild(ref key, out ChildCollider child))
                {
                    int3 quantized = child.TransformFromChild.pos.ToInt3();
                    if (Colliders.TryGetValue(quantized, out childData))
                    {
                        return true;
                    }
                }
            }

            childData = default;
            return false;
        }
    }

    public struct ChildMapComponentLookup
    {
        [NoAlias]
        [ReadOnly]
        public ComponentDataFromEntity<ChildMapComponent> Lookup;

        public Entity LookupEntity;

        public static ChildMapComponentLookup Create(SystemBase system)
        {
            return new ChildMapComponentLookup
            {
                Lookup = system.GetComponentDataFromEntity<ChildMapComponent>(true),
                LookupEntity = ColliderPartitionSystem.Instance.LookupEntity
            };
        }

        public ChildMapComponent GetComponent()
        {
            return Lookup[LookupEntity];
        }

    }

}

using System;
using Unity.Mathematics;

namespace AiGame.ColliderPartitions
{
    public static class ColliderPartitionExtensions
    {
        public static int3 ToInt3(this float3 value)
        {
            return new int3(value.x.ToInt(), value.y.ToInt(), value.z.ToInt());
        }

        public static int ToInt(this float value)
        {
            return (int)(Math.Round(value * 100, 2));
        }
    }
}