riversong-code-showcase/Source/Riversong/Game/World/Pathfinding/PathfindingSystem.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using Cysharp.Threading.Tasks;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;

namespace DanieleMarotta.RiversongCodeShowcase
{
    [Service(typeof(IPathfinder))]
    [BurstCompile]
    public unsafe class PathfindingSystem : GameSystem, IInitializable, IDisposable, IUpdatable, IPathfinder, IOnWorldGenerationCompletedCallback
    {
        private const int MaxConcurrentQueries = 8;

        [InjectService]
        private World _world;

        [InjectService]
        private IEntityCache _entityCache;

        [InjectService]
        private ITileSpace _tileSpace;

        [InjectService]
        private ISignalBus _signalBus;

        private Queue<PathQuery> _pendingQueries = new();

        private Dictionary<int, (TilePath, PathQueryResult)> _paths = new();

        private int _nextQueryHandle;

        private NativeArray<AStarNode> _nodeCache;

        private NativeArray<PathQuery> _runningQueries;

        private NativeArray<PathQueryResult> _results;

        private JobHandle? _currentJob;

        private NativeStream _pathStream;

        private int _version;

        private int _jobAge;

        public PathfindingSystem(IServiceLocator serviceLocator) : base(serviceLocator)
        {
        }

        public UniTask InitializeAsync()
        {
            _signalBus.Subscribe<WorldGenerationCompletedSignal>(OnWorldGenerationCompleted);

            return UniTask.CompletedTask;
        }

        public void Dispose()
        {
            _signalBus.Unsubscribe<WorldGenerationCompletedSignal>(OnWorldGenerationCompleted);

            if (_nodeCache.IsCreated) _nodeCache.Dispose();
            if (_runningQueries.IsCreated) _runningQueries.Dispose();
            if (_results.IsCreated) _results.Dispose();
        }

        public UniTask OnWorldGenerationCompletedAsync(World world)
        {
            var tileCount = _world.Size.x * _world.Size.y;
            _nodeCache = new NativeArray<AStarNode>(MaxConcurrentQueries * tileCount, Allocator.Persistent);
            for (var threadIndex = 0; threadIndex < MaxConcurrentQueries; threadIndex++)
            {
                var slice = threadIndex * tileCount;
                var tileIndex = 0;

                for (var y = 0; y < _world.Size.y; y++)
                for (var x = 0; x < _world.Size.x; x++)
                    _nodeCache[slice++] = new AStarNode
                    {
                        Index = tileIndex++,
                        Tile = new int2(x, y)
                    };
            }

            _runningQueries = new NativeArray<PathQuery>(MaxConcurrentQueries, Allocator.Persistent);
            _results = new NativeArray<PathQueryResult>(MaxConcurrentQueries, Allocator.Persistent);

            _version = 1;

            return UniTask.CompletedTask;
        }

        public int FindPath(in TilePath path, ref PathQuery query)
        {
            query.Handle = _nextQueryHandle++;

            _pendingQueries.Enqueue(query);
            _paths.Add(query.Handle, (path, PathQueryResult.InProgress));

            return query.Handle;
        }

        public void Update()
        {
            if (!CompleteCurrentJob()) return;

            var queryCount = math.min(_pendingQueries.Count, MaxConcurrentQueries);
            if (queryCount <= 0) return;

            for (var i = 0; i < queryCount; i++) _runningQueries[i] = _pendingQueries.Dequeue();

            var blockMap = new NativeArray<BlockReason>(_world.BlockMap.GetNativeArray(), Allocator.TempJob);
            var fertility = new NativeArray<FertilityMapValue>(_world.Fertility.GetNativeArray(), Allocator.TempJob);
            var entityIdMap = new NativeArray<EntityIdMapValue>(_world.EntityIdMap.GetNativeArray(), Allocator.TempJob);
            CopyCritterIds(entityIdMap);

            _pathStream = new NativeStream(queryCount, Allocator.TempJob);

            _currentJob = new PathJob
            {
                Version = _version++,
                Queries = _runningQueries,
                Results = _results,
                NodeCache = _nodeCache,
                PathStream = _pathStream,
                BlockMap = blockMap,
                Fertility = fertility,
                EntityIdMap = entityIdMap,
                WorldSize = _world.Size
            }.Schedule(queryCount, 1);

            blockMap.Dispose(_currentJob.Value);
            entityIdMap.Dispose(_currentJob.Value);
            fertility.Dispose(_currentJob.Value);

            _jobAge = 1;
        }

        private void OnWorldGenerationCompleted(WorldGenerationCompletedSignal worldGenerationCompleted)
        {
            worldGenerationCompleted.Callbacks.Add(this);
        }

        private bool CompleteCurrentJob()
        {
            if (_currentJob == null) return true;

            if (!_currentJob.Value.IsCompleted && ++_jobAge < 4) return false;

            _currentJob.Value.Complete();
            _currentJob = null;

            var reader = _pathStream.AsReader();

            for (var i = 0; i < reader.ForEachCount; i++)
            {
                if (_results[i] != PathQueryResult.Success) continue;

                var queryHandle = _runningQueries[i].Handle;
                var (path, _) = _paths[queryHandle];

                reader.BeginForEachIndex(i);
                while (reader.RemainingItemCount > 0) path.Steps.Add(reader.Read<int2>());
                reader.EndForEachIndex();
            }

            _pathStream.Dispose();
            _pathStream = default;

            for (var i = 0; i < _runningQueries.Length; i++) _paths[_runningQueries[i].Handle] = (default, _results[i]);

            return true;
        }

        private void CopyCritterIds(NativeArray<EntityIdMapValue> entityIdMap)
        {
            foreach (var critter in _entityCache.GetCritterAgents())
            {
                ref var critterState = ref critter.GetCritterStateRW();
                if (critterState.LockId != Entity.InvalidId) continue;

                var tile = _tileSpace.WorldToTile(critter.Position);
                var index = math.mad(tile.y, _world.Size.x, tile.x);

                ref var value = ref UnsafeUtility.ArrayElementAsRef<EntityIdMapValue>(entityIdMap.GetUnsafePtr(), index);
                value.CritterHerdId = critterState.CritterDefinitionId;
            }
        }

        public PathQueryResult CompletePath(int queryHandle)
        {
            if (!_paths.TryGetValue(queryHandle, out var result)) return PathQueryResult.InvalidQuery;

            if (result.Item2 != PathQueryResult.InProgress) _paths.Remove(queryHandle);

            return result.Item2;
        }

        private struct AStarNode
        {
            public int Index;

            public int2 Tile;

            public int Version;

            public int Parent;

            public int G;

            public bool IsClosed;

            public void Clear(int version)
            {
                Version = version;
                Parent = -1;
                G = int.MaxValue;
                IsClosed = false;
            }
        }

        [BurstCompile]
        private struct PathJob : IJobParallelFor
        {
            private const int StraightStepCost = 10;

            private const int DiagonalStepCost = 14;

            public int Version;

            [ReadOnly]
            public NativeArray<PathQuery> Queries;

            [WriteOnly]
            public NativeArray<PathQueryResult> Results;

            [NativeDisableParallelForRestriction]
            public NativeArray<AStarNode> NodeCache;

            [NativeDisableParallelForRestriction]
            public NativeStream PathStream;

            [ReadOnly]
            public NativeArray<BlockReason> BlockMap;

            [ReadOnly]
            public NativeArray<FertilityMapValue> Fertility;

            [ReadOnly]
            public NativeArray<EntityIdMapValue> EntityIdMap;

            public int2 WorldSize;

            [BurstCompile]
            public void Execute(int index)
            {
                var query = Queries[index];

                var tileCount = WorldSize.x * WorldSize.y;
                var nodes = new NativeSlice<AStarNode>(NodeCache, tileCount * index, tileCount);

                var sourceIndex = GetNodeIndex(query.SourceTile);
                ref var source = ref GetNodeRW(nodes, sourceIndex);
                if (IsGoal(query, source))
                {
                    Results[index] = PathQueryResult.Success;
                    return;
                }

                source.Clear(Version);
                source.G = 0;
                var open = new NativePriorityQueue((int)(0.2f * tileCount), Allocator.Temp);
                open.Add(sourceIndex, GetHScore(query.SourceTile, query.DestinationRect, query.SearchType));

                var directions = (query.TraversalRules & PathTraversalRules.CardinalDirectionsOnly) != 0 ? DirectionVectors.Directions4 : DirectionVectors.Directions8;
                var pathFound = false;

                while (open.TryPop(out var currentIndex))
                {
                    ref var current = ref GetNodeRW(nodes, currentIndex);

                    current.IsClosed = true;

                    foreach (var direction in directions)
                    {
                        var neighborTile = current.Tile + direction;
                        if (!ValidateTile(query.SourceTile, neighborTile, query.MaxDistance)) continue;

                        ref var neighbor = ref GetNodeRW(nodes, neighborTile);
                        if (neighbor.Version < Version) neighbor.Clear(Version);

                        if (neighbor.IsClosed) continue;

                        var moveScore = EvaluateMove(currentIndex, neighbor.Index, query.TargetId, query.TraversalRules);
                        if (moveScore == int.MaxValue) continue;

                        if (IsGoal(query, neighbor))
                        {
                            neighbor.Parent = currentIndex;
                            WritePath(index, query, nodes, neighbor);
                            pathFound = true;
                            break;
                        }

                        var tentativeG = current.G + moveScore;
                        if (tentativeG >= neighbor.G) continue;

                        neighbor.Parent = currentIndex;
                        neighbor.G = tentativeG;

                        var f = tentativeG + GetHScore(neighborTile, query.DestinationRect, query.SearchType);
                        open.Add(neighbor.Index, f);
                    }

                    if (pathFound) break;
                }

                Results[index] = pathFound ? PathQueryResult.Success : PathQueryResult.Failure;
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private bool IsGoal(in PathQuery query, in AStarNode node)
            {
                if (query.SearchType == PathSearchType.FindPath && query.DestinationRect.Contains(node.Tile)) return true;

                var entityIdValue = EntityIdMap[node.Index];
                switch (query.SearchType)
                {
                    case PathSearchType.SearchResourceNode:
                        return entityIdValue.ResourceNodeId == query.TargetId;

                    case PathSearchType.SearchFertileTile:
                        var fertility = Fertility[node.Index];
                        return fertility.MaxFertility > 0 && fertility.CurrentFertility >= fertility.MaxFertility && fertility.LockId == Entity.InvalidId;

                    case PathSearchType.SearchCritterHerd:
                        return entityIdValue.CritterHerdId == query.TargetId;
                }

                return false;
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private bool ValidateTile(int2 pathSource, int2 tile, int maxDistance)
            {
                return math.all(tile >= 0) && math.all(tile < WorldSize) && (maxDistance <= 0 || TileMath.StepCount(pathSource, tile) <= maxDistance);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private int GetNodeIndex(int2 tile)
            {
                return math.mad(tile.y, WorldSize.x, tile.x);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private ref AStarNode GetNodeRW(in NativeSlice<AStarNode> nodes, int index)
            {
                return ref UnsafeUtility.ArrayElementAsRef<AStarNode>(nodes.GetUnsafePtr(), index);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private ref AStarNode GetNodeRW(in NativeSlice<AStarNode> nodes, int2 tile)
            {
                return ref GetNodeRW(nodes, GetNodeIndex(tile));
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private int EvaluateMove(int leavingIndex, int enteringIndex, int targetId, PathTraversalRules rules)
            {
                var enteringBuildingId = EntityIdMap[enteringIndex].BuildingId;
                if (enteringBuildingId != Entity.InvalidId)
                {
                    var leavingBuildingId = EntityIdMap[leavingIndex].BuildingId;
                    if (enteringBuildingId == leavingBuildingId || (enteringBuildingId == targetId && leavingBuildingId == Entity.InvalidId)) return 0;
                }

                var blockReason = BlockMap[enteringIndex];

                var blocksAgents = (blockReason & BlockReason.BlocksAgents) != 0;
                if (blocksAgents) return int.MaxValue;

                var roadsOnly = (rules & PathTraversalRules.RoadsOnly) != 0;
                var isRoad = (blockReason & BlockReason.Road) != 0;
                if (roadsOnly && !isRoad) return int.MaxValue;

                var fertileTilesOnly = (rules & PathTraversalRules.FertileTilesOnly) != 0;
                if (fertileTilesOnly && Fertility[enteringIndex].MaxFertility <= 0) return int.MaxValue;

                return StepCost(leavingIndex, enteringIndex);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private int StepCost(int leavingIndex, int enteringIndex)
            {
                var delta = math.abs(leavingIndex - enteringIndex);
                return delta == 1 || delta == WorldSize.x ? StraightStepCost : DiagonalStepCost;
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private static int GetHScore(int2 tile, TileRect rect, PathSearchType searchType)
            {
                return searchType switch
                {
                    PathSearchType.FindPath => StraightStepCost * TileMath.StepCount(tile, rect),
                    _ => 0
                };
            }

            private void WritePath(int threadIndex, in PathQuery query, in NativeSlice<AStarNode> nodes, in AStarNode current)
            {
                var temp = new NativeList<int2>(Allocator.Temp);
                while (true)
                {
                    if ((BlockMap[current.Index] & BlockReason.BlocksAgents) == 0 || EntityIdMap[current.Index].BuildingId == query.TargetId) temp.Add(current.Tile);

                    if (math.all(current.Tile == query.SourceTile)) break;

                    current = ref GetNodeRW(nodes, current.Parent);
                }

                var pathWriter = PathStream.AsWriter();
                pathWriter.BeginForEachIndex(threadIndex);
                for (var i = temp.Length - 1; i >= 0; i--) pathWriter.Write(temp[i]);
                pathWriter.EndForEachIndex();
            }
        }

        [BurstCompile]
        private struct NativePriorityQueue : IDisposable
        {
            private NativeList<Item> _items;

            public NativePriorityQueue(int capacity, Allocator allocator)
            {
                _items = new NativeList<Item>(capacity, allocator);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            public void Dispose()
            {
                _items.Dispose();
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            public void Clear()
            {
                _items.Clear();
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            public void Add(int index, int score)
            {
                var i = _items.Length;
                _items.Add(new Item(index, score));
                BubbleUp(i);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            public bool TryPop(out int index)
            {
                var length = _items.Length;
                if (length == 0)
                {
                    index = 0;
                    return false;
                }

                index = _items[0].Index;

                var last = length - 1;
                _items[0] = _items[last];
                _items.RemoveAt(last);

                if (_items.Length > 0) BubbleDown(0);

                return true;
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private void BubbleUp(int i)
            {
                while (i > 0)
                {
                    var parent = (i - 1) >> 1;
                    if (_items[i].Score >= _items[parent].Score) break;
                    (_items[i], _items[parent]) = (_items[parent], _items[i]);
                    i = parent;
                }
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            private void BubbleDown(int i)
            {
                var count = _items.Length;

                while (true)
                {
                    var left = (i << 1) + 1;
                    var right = left + 1;
                    var smallest = i;

                    if (left < count && _items[left].Score < _items[smallest].Score) smallest = left;
                    if (right < count && _items[right].Score < _items[smallest].Score) smallest = right;

                    if (smallest == i) break;

                    (_items[i], _items[smallest]) = (_items[smallest], _items[i]);
                    i = smallest;
                }
            }

            private struct Item
            {
                public int Index;

                public int Score;

                public Item(int index, int score)
                {
                    Index = index;
                    Score = score;
                }
            }
        }
    }
}