Nerfed/Nerfed.Runtime/Systems/LocalToWorldSystem.cs

using MoonTools.ECS;
using Nerfed.Runtime.Components;
using Nerfed.Runtime.Util;
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
    {
        private readonly JobSystem jobs;
        private readonly Filter rootEntitiesFilter;
        private readonly Filter entitiesWithoutLocalToWorldFilter;
        private readonly Action<int> updateWorldTransformByIndex;

        public LocalToWorldSystem(World world, JobSystem jobs = null) : base(world)
        {
            this.jobs = jobs ?? JobSystem.Default;
            rootEntitiesFilter = FilterBuilder.Include<LocalTransform>().Exclude<Child>().Build();
            entitiesWithoutLocalToWorldFilter = FilterBuilder.Include<LocalTransform>().Exclude<LocalToWorld>().Build();
            updateWorldTransformByIndex = UpdateWorldTransformByIndex;
        }

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

            if (this.jobs.WorkerCount > 0)
            {
                Profiler.BeginSample("LocalToWorldCheck");
                // Structural pre-pass: ensure LocalToWorld exists on all entities before parallel writes.
                foreach (Entity entity in entitiesWithoutLocalToWorldFilter.Entities)
                {
                    Set(entity, new LocalToWorld(Matrix4x4.Identity));
                }
                Profiler.EndSample();

                Profiler.BeginSample("LocalToWorldUpdate");
                this.jobs.Dispatch(rootEntitiesFilter.Count, updateWorldTransformByIndex);
                Profiler.EndSample();
            }
            else
            {
                foreach (Entity entity in rootEntitiesFilter.Entities)
                {
                    Profiler.BeginSample("UpdateWorldTransform");
                    UpdateWorldTransform(entity, Matrix4x4.Identity);
                    Profiler.EndSample();
                }
            }
        }

        private void UpdateWorldTransformByIndex(int entityFilterIndex)
        {
            using ProfilerScope scope = new("UpdateWorldTransformByIndex");
            Entity entity = rootEntitiesFilter.NthEntity(entityFilterIndex);
            UpdateWorldTransform(entity, Matrix4x4.Identity);
        }

        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 DEBUG
                if (!Has<LocalToWorld>(entity))
                {
                    throw new InvalidOperationException(
                        $"Entity {entity} is missing LocalToWorld. Ensure the structural pre-pass runs before parallel dispatch.");
                }
#endif
                Set(entity, localToWorld);
            }

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