Nerfed/Nerfed.Runtime/Systems/LocalToWorldSystem.cs

using MoonTools.ECS;
using Nerfed.Runtime.Components;
using Nerfed.Runtime.Util;
using System;
using System.Collections.Generic;
using System.Numerics;

// TODO:
// Explore if having a WorldTransform and LocalTransfom component each holding position, rotation, scale values and the matricies is useful.
// Often you need to either get or set these values.
// If so, we probably need a utility funciton to do so. Since changing these values means that we need to update all the related data + children as well.

// TODO:
// When modifying transform all the children need to be updated as well.

namespace Nerfed.Runtime.Systems
{
    public class LocalToWorldSystem : MoonTools.ECS.System
    {
        public override IReadOnlySet<Type> ReadsComponents { get; } = new HashSet<Type> { typeof(LocalTransform) };
        public override IReadOnlySet<Type> WritesComponents { get; } = new HashSet<Type> { typeof(LocalToWorld) };

        private readonly bool useParallelFor = true;
        private const int ParallelForMinCount = 32; // Below this, parallel overhead costs more than it saves.
        private static readonly System.Threading.Tasks.ParallelOptions ParallelOptions = new()
        {
            MaxDegreeOfParallelism = Environment.ProcessorCount
        };
        private readonly Filter rootEntitiesFilter;
        private readonly Filter entitiesWithoutLocalToWorldFilter;
        private readonly Action<int> updateWorldTransform;
        private ParallelWriter<LocalToWorld> _parallelWriter;

        public LocalToWorldSystem(World world) : base(world)
        {
            rootEntitiesFilter = FilterBuilder.Include<LocalTransform>().Exclude<Child>().Build();
            if (useParallelFor)
            {
                entitiesWithoutLocalToWorldFilter = FilterBuilder.Include<LocalTransform>().Exclude<LocalToWorld>().Build();
                updateWorldTransform = UpdateWorldTransformByIndex;
            }
        }

        public override void Update(TimeSpan delta)
        {
            if (rootEntitiesFilter.Empty)
            {
                return;
            }

            if (useParallelFor)
            {
                // This check is needed because some entities might not have a LocalToWorld component yet.
                // Adding this during the loop will break.
                Profiler.BeginSample("ParallelFor.LocalToWorldCheck");
                foreach (Entity entity in entitiesWithoutLocalToWorldFilter.Entities) {
                    Set(entity, new LocalToWorld(Matrix4x4.Identity));
                }
                Profiler.EndSample();

                // Acquire a ParallelWriter AFTER pre-allocation — all entities now have LocalToWorld.
                // This writer only permits updating existing values; no structural mutations allowed.
                _parallelWriter = World.GetParallelWriter<LocalToWorld>();

                Profiler.BeginSample("ParallelFor.LocalToWorldUpdate");
                if (rootEntitiesFilter.Count >= ParallelForMinCount)
                {
                    Parallel.For(0, rootEntitiesFilter.Count, ParallelOptions, updateWorldTransform);
                }
                else
                {
                    // Not enough work to justify thread overhead — run serially.
                    for (int i = 0; i < rootEntitiesFilter.Count; i++)
                    {
                        updateWorldTransform(i);
                    }
                }
                Profiler.EndSample();
            }
            else
            {
                foreach (Entity entity in rootEntitiesFilter.Entities)
                {
                    // Profiler.BeginSample("UpdateWorldTransform");
                    UpdateWorldTransform(entity, Matrix4x4.Identity);
                    // Profiler.EndSample();
                }
            }
        }

        private void UpdateWorldTransformByIndex(int entityFilterIndex)
        {
            // Profiler.BeginSample("UpdateWorldTransformByIndex");
            Entity entity = rootEntitiesFilter.NthEntity(entityFilterIndex);
            UpdateWorldTransform(entity, Matrix4x4.Identity);
            // Profiler.EndSample();
        }

        private void UpdateWorldTransform(in Entity entity, Matrix4x4 localToWorldMatrix)
        {
            if (Has<LocalTransform>(entity))
            {
                LocalTransform localTransform = Get<LocalTransform>(entity);
                localToWorldMatrix = Matrix4x4.Multiply(localToWorldMatrix, localTransform.TRS());
                LocalToWorld localToWorld = new(localToWorldMatrix);

                if (useParallelFor)
                    _parallelWriter.Set(entity, localToWorld); // thread-safe: direct write, no structural mutation
                else
                    Set(entity, localToWorld);
            }

            ReverseSpanEnumerator<Entity> childEntities = World.InRelations<ChildParentRelation>(entity);
            foreach (Entity childEntity in childEntities)
            {
                UpdateWorldTransform(childEntity, localToWorldMatrix);
            }
        }
    }
}