MySkinnedBakingTestAuthoring.cs

using System;
using System.Linq;
using Unity.Collections;
using Unity.Entities;
using Unity.Entities.Hybrid.Baking;
using Unity.Transforms;
using UnityEngine;
// ReSharper disable Unity.Entities.MustBeSurroundedWithRefRwRo

namespace Junk.Rendering.Hybrid
{
    /// <summary>
    /// Tests the remapping utility for skinned mesh entities
    /// </summary>
    public class MySkinnedBakingTestAuthoring : MonoBehaviour
    {
        public SkinnedMeshRenderer SkinnedMeshRenderer;

        public class SkinnedBakingTestBaker : Baker<MySkinnedBakingTestAuthoring>
        {
            public override void Bake(MySkinnedBakingTestAuthoring authoring)
            {
                //Debug.Log("Baker GetEntity" + GetEntity(TransformUsageFlags.Dynamic));
                AddComponent(GetEntity(TransformUsageFlags.Dynamic), new MyComponent
                {
                    MyEntityWithDeformations =  this.GetSkinnedEntity(authoring, authoring.SkinnedMeshRenderer, typeof(MyComponent))
                });
            }
        }

        /// <summary>
        /// Test Component containing an entity that has deformations.
        /// </summary>
        public struct MyComponent : IComponentData
        {
            public Entity MyEntityWithDeformations;
        }
    }


    /// <summary>
    /// Remaps entities that were created by the skinned mesh baker to the correct entity with deformations.
    /// Note: currently does not support multiple entities created by the skinned mesh baker, if there are submeshes or Keep Quads is enabled.
    /// If a user only has one Entity field they want to remap, this system will work.
    /// Investigate: Could try to add support for multiple entities created by the skinned mesh baker, using an int to specify a submesh index for a mesh with multiple submeshes.
    /// </summary>
    [WorldSystemFilter(WorldSystemFilterFlags.BakingSystem)]
    public partial struct RemapSkinnedEntityBaking : ISystem
    {
        public object GetComponentData(ref SystemState state, Entity entity, Type componentType)
        {
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "GetComponentData" 
                    && m.GetParameters().Length           == 1 
                    && m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(componentType);

            if (methodInfo != null)
            {
                return methodInfo.Invoke(state.EntityManager, new object[] { entity });
            }
            else
            {
                // Handle the case where the method is not found
                Debug.LogWarning($"ErrMethod not found for {componentType}");
                return null; // or throw an exception
            }
        }
        
        public object GetBuffer(ref SystemState state, Entity entity, Type bufferType)
        {
            //var buffer = state.EntityManager.GetBuffer<Child>(entity);

            /*var methods = typeof(EntityManager)
                .GetMethods();
            foreach (var method in methods)
            {
                if (method.Name.Contains("GetBuffer"))
                {
                    Debug.Log($"Method name = {method.Name} + method parameters = {method.GetParameters()} + method length = {method.GetParameters().Length}");
                    
                    // find which parameter is the entity
                    for (var index = 0; index < method.GetParameters().Length; index++)
                    {
                        var parameter = method.GetParameters()[index];
                        Debug.Log($"Parameter type = {parameter.ParameterType} and the index = {index}");
                    }
                }
            }*/
            
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "GetBuffer")
                    //&& m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(bufferType);

            if (methodInfo != null)
            {
                // first parameter is the entity, second is the bool for readonly, set always false
                return methodInfo.Invoke(state.EntityManager, new object[] { entity, false });
            }
            else
            {
                // Handle the case where the method is not found
                Debug.LogWarning($"ErrMethod not found for {bufferType}");
                return null; // or throw an exception
            }
        }

        public void SetComponentData(ref SystemState state, Entity entity, Type componentType, object component)
        {
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "SetComponentData" 
                    && m.GetParameters().Length           == 2 
                    && m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(componentType);
            if (methodInfo != null)
            {
                methodInfo.Invoke(state.EntityManager, new object[] { entity, component });
            }
            else
            {
                // Handle the case where the method is not found
                Debug.LogWarning($"ErrMethod not found for {componentType}");
            }
        }

        /// <summary>
        /// using reflection, a method that finds all the Entity fields in a generic ComponentData, and remaps them to the correct entity if the entity matches
        /// </summary>
        /// <param name="state">The System State</param>
        /// <param name="entity">The main entity. This is the entity that the original baker converted.</param>
        /// <param name="componentType">The component type that was added in the original baker, where GetSkinnedEntity was used.</param>
        /// <param name="entityToRemap">This is the entity original resulting from baking the SkinnedMeshRenderer (which then creates additional entities with actual deformations).
        /// By default Unity does not add a Deformations component to this entity, rather it is added to another entity and parented under the main bone(usually the hips).</param>
        /// <param name="deformedEntity">The actual entity which is parented to the root bone of the rig.</param>
        public void RemapEntities(ref SystemState state, Entity entity, Type componentType, Entity entityToRemap, Entity deformedEntity, bool isBuffer = false, int bufferIndex = 0)
        {
            if (!isBuffer)
            {
                Debug.Log($"Remapping entities for {entity} with component type {componentType} and entity to remap {entityToRemap} and deformed entity {deformedEntity}");
                var component = GetComponentData(ref state, entity, componentType);
            
                //Debug.Log($"component = {component} and type = {component.GetType()}");
                var fields    = componentType.GetFields();
            
                foreach (var field in fields)
                {
                    if (field.FieldType == typeof(Entity))
                    {
                        var value = (Entity)field.GetValue(component);
                        //Debug.Log($"Found entity field {field.Name} with value {value} and original entity {originalEntity}");
                    
                        if (value == entityToRemap)
                        {
                            //Debug.Log($"Remapping entity {value} to {newEntity}");
                            field.SetValue(component, deformedEntity);
                        }
                    }
                }
                SetComponentData(ref state, entity, componentType, component);
            }

            if (isBuffer)
            {
                var bufferObject = GetBuffer(ref state, entity, componentType);
                var isNull       = bufferObject == null;
                //Debug.Log("Remapping entities for buffer" + bufferObject + " is null " + isNull + " and bufferIndex " + bufferIndex);
                var bufferType   = bufferObject.GetType();
                //Debug.Log($"bufferType = {bufferType}");
                var methodInfo   = bufferType.GetMethod("AsNativeArray");
                var bufferArray  = methodInfo.Invoke(bufferObject, null);
                var bufferLength = (int)bufferType.GetProperty("Length").GetValue(bufferObject, null);
                //Debug.Log($"bufferArray = {bufferArray} and bufferLength = {bufferLength}");
                //var arr = new NativeArray<float>(2, Allocator.Temp);
                //arr.GetType()
                for (var i = 0; i < bufferLength; i++)
                {
                    if(i != bufferIndex) 
                        continue;
                    
                    var bufferElement = bufferArray.GetType().GetMethod("get_Item").Invoke(bufferArray, new object[] { i });
                    var fields        = bufferElement.GetType().GetFields();
                    foreach (var field in fields)
                    {
                        if (field.FieldType == typeof(Entity))
                        {
                            var value = (Entity)field.GetValue(bufferElement);
                            //Debug.Log($"Found entity field {field.Name} with value {value} and original entity {originalEntity}");
                    
                            if (value == entityToRemap)
                            {
                                //Debug.Log($"Remapping entity {value} to {newEntity}");
                                field.SetValue(bufferElement, deformedEntity);
                            }
                        }
                    }
                    bufferArray.GetType().GetMethod("set_Item").Invoke(bufferArray, new object[] { i, bufferElement });
                }
            }
        }

        /// <summary>
        /// Collect data from the skinning bakers that creates an additional entity per skinnedMeshRenderer
        /// </summary>
        /// <param name="state"></param>
        private void RemapToEntityWithDeformations(ref SystemState state)
        {
            foreach (var (remapBuffer, entity) in SystemAPI.Query<DynamicBuffer<GetSkinnedMeshUtility.RemapSmrEntityBuffer>>()
                         .WithEntityAccess().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities))
            {
                var bufferAsArray = remapBuffer.AsNativeArray();

                foreach (var element in bufferAsArray)
                {
                    var bakerEntity    = element.BakerEntityWithComponent;
                    var remapEntity    = element.EntityToRemap;
                    var typeStableHash = element.ComponentTypeStableHash;
                    var bufferIndex    = element.BufferIndex;

                    if (!SystemAPI.HasBuffer<AdditionalEntitiesBakingData>(remapEntity))
                    {
                        Debug.LogWarning($"AdditionalEntitiesBakingData buffer not found on {remapEntity} when trying to remap skinned mesh entities, " +
                            $"are you sure this has a SkinnedMeshRenderer?");
                        continue;
                    }

                    var additionalEntitiesBakingDataArray = SystemAPI.GetBuffer<AdditionalEntitiesBakingData>(remapEntity).AsNativeArray();
                    var length                            = additionalEntitiesBakingDataArray.Length;
                    //Debug.Log($"additionalEntitiesBakingDataArray buffer length  = {length} ");
                    
                    if (length != 1)
                    {
                        Debug.LogError($"Expected 1 additional entity, got {length} when reading skinnedmesh baker entities. " +
                            $" Are there submeshes or is Keep Quads enabled(as this will create a submesh)? "                  +
                            $"If you are creating entities on the Skinned Mesh Renderer, the remapping system will fail to account for this.");
                    }

                    var deformedEntity = additionalEntitiesBakingDataArray[0].Value;
                    //Debug.Log($"deformedEntity = {deformedEntity}");

                    var componentIndex = TypeManager.GetTypeIndexFromStableTypeHash(typeStableHash);
                    var type           = TypeManager.GetType(componentIndex);
                    var isBuffer       = typeof(IBufferElementData).IsAssignableFrom(type);
                    //Debug.Log($"type = {type} and isBuffer = {isBuffer}");
                    RemapEntities(ref state, bakerEntity, type, remapEntity, deformedEntity, isBuffer, bufferIndex);

                }

            }
        }

        public void OnUpdate(ref SystemState state)
        {
            /*foreach (var (component, entity) in SystemAPI.Query<RefRO<MySkinnedBakingTestAuthoring.MyComponent>>()
                         .WithEntityAccess().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities))
            {
                Debug.Log($"OnUpdate Loop: Found MyComponent on entity with data {component.ValueRO.MyEntityWithDeformations}");
            }*/

            RemapToEntityWithDeformations(ref state);
        }
    }
    
    public static class GetSkinnedMeshUtility
    {
        /// <summary>
        /// The entity that is the original entity resulting from the baking process. This entity does not
        /// contain any renderable components, they are created as additional entities.
        /// </summary>
        [TemporaryBakingType]
        public struct RemapSmrEntityBuffer : IBufferElementData
        {
            public ulong  ComponentTypeStableHash;
            public int    BufferIndex;
            public Entity EntityToRemap;
            public Entity BakerEntityWithComponent;
        }

        /// <summary>
        /// First add the entity as the original system works. This original entity will be used by a baking system to get the list of additional entities created
        /// by the render baker. Then patch in the correct entity after ensuring it matches with the original entity and component stablehash.
        /// </summary>
        /// <param name="baker">The Baker</param>
        /// <param name="authoring">The authoring component</param>
        /// <param name="smr">The SkinnedMeshRenderer</param>
        /// <param name="componentType">The component type</param>
        public static Entity GetSkinnedEntity<T>(this Baker<T> baker, Component authoring, SkinnedMeshRenderer smr, Type componentType, int bufferIndex = -1) where T : Component
        {
            // Get the original entity
            var bakerEntity = baker.GetEntity(authoring, TransformUsageFlags.Dynamic);
            // Get the entity that will be remapped
            var remapEntity = baker.GetEntity(smr, TransformUsageFlags.Dynamic);
            var remapBuffer = baker.AddBuffer<RemapSmrEntityBuffer>(bakerEntity);
            
            //Debug.Log("Baker GetSkinnedMeshUtility.GetSkinnedEntity Baker entity:" + bakerEntity + " RemapEntity:" + remapEntity);

            //var typeIndex  = TypeManager.GetTypeIndexFromStableTypeHash(1234);
            var typeIndex  = TypeManager.GetTypeIndex(componentType);
            var stableHash = TypeManager.GetTypeInfo(typeIndex).StableTypeHash;

            remapBuffer.Add(new RemapSmrEntityBuffer
            {
                ComponentTypeStableHash = stableHash,
                BufferIndex             = bufferIndex,
                EntityToRemap            = remapEntity,
                BakerEntityWithComponent = bakerEntity
            });

            return remapEntity;
        }
    }
}