GetSkinnedEntityAuthoringV2.cs

// note - requires asmref trick to get internal access for using the DeformedEntity

using System;
using System.Linq;
using System.Reflection;
using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.Entities.Serialization;
using Unity.Rendering;
using UnityEngine;
using UnityEngine.Assertions;

namespace Junk.Rendering.Hybrid
{
    // Just used for testing purposes
    public class GetSkinnedEntityAuthoringV2 : MonoBehaviour
    {
        public SkinnedMeshRenderer SkinnedMeshRenderer;
        
        public class GetSkinnedEntityAuthoringBaker : Baker<GetSkinnedEntityAuthoringV2>
        {
            public override void Bake(GetSkinnedEntityAuthoringV2 authoringV2)
            {
                var entity = GetEntity(TransformUsageFlags.Dynamic);
                
                var smrEntity = GetEntity(authoringV2.SkinnedMeshRenderer, TransformUsageFlags.Dynamic);
                AddComponent(entity, new FindSkinnedEntityDataV2
                {
                    SkinnedEntity = smrEntity
                });
            }
        }
    }

    [BakingType]
    public struct FindSkinnedEntityDataV2 : IComponentData
    {
        public Entity SkinnedEntity;
    }

    /// <summary>
    /// System for finding the actual DeformedEntity from the SkinnedMeshRenderer gameobject.
    /// </summary>
    [WorldSystemFilter(WorldSystemFilterFlags.BakingSystem)]
    public partial struct GetSkinnedEntitySystemV2 : ISystem
    {
        [BurstCompile]
        public void OnUpdate(ref SystemState state)
        {
            var ecb = new EntityCommandBuffer(Allocator.Temp);
            
            foreach (var (findSkinnedEntityData, entity) in SystemAPI.Query<RefRO<FindSkinnedEntityDataV2>>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).WithEntityAccess())
            {
                var smrEntity  = findSkinnedEntityData.ValueRO.SkinnedEntity;
                var skinnedEntityList = new NativeList<Entity>(Allocator.Temp);
                
                foreach (var (deformedEntity, skinnedEntity) in SystemAPI.Query<RefRO<DeformedEntity>>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).WithEntityAccess())
                {
                    if (deformedEntity.ValueRO.Value.Equals(smrEntity))
                    {
                        skinnedEntityList.Add(skinnedEntity);
                    }
                }
                
                //Debug.Log($"Found {skinnedEntityList.Length} skinned entities");
                if (skinnedEntityList.Length > 0)
                {
                    var skinnedEntityBuffer = ecb.AddBuffer<SkinnedEntityData>(entity);
                    foreach (var skinnedEntity in skinnedEntityList)
                    {
                        skinnedEntityBuffer.Add(new SkinnedEntityData { Value = skinnedEntity });
                        //ecb.AddComponent(skinnedEntity, new SavableEntity());
                    }
                }
            }
            ecb.Playback(state.EntityManager);
        }
    }
    
    public struct SkinnedEntityData : IBufferElementData
    {
        public Entity Value;

        public EntityPrefabReference p;
    }
    

    public  static partial class DeformationBakingExtension
    {
        [BakingType]
        public struct RemapToDeformedComponent : IComponentData
        {
            /// <summary>
            /// The original baker entity
            /// </summary>
            public Entity Entity;
            
            /// <summary>
            /// The entity that has the component which we want to add to the deformed entity.
            /// This entity is temporary and will not exist after the baking process.
            /// </summary>
            public Entity TemporaryEntity;
            
            /// <summary>
            /// The skinned mesh renderer entity, not to be confused with the deformed entity.
            /// </summary>
            public Entity SmrEntity;
            
            // Stablehash of the component to be added to the deformed entity
            public ulong StableTypeHash;
            public bool  EnabledState;
        }
        
        [BakingType]
        public struct RemapToDeformedBuffer : IComponentData
        {
            /// <summary>
            /// The original baker entity
            /// </summary>
            public Entity Entity;
            
            /// <summary>
            /// The entity that has the component which we want to add to the deformed entity.
            /// This entity is temporary and will not exist after the baking process.
            /// </summary>
            public Entity TemporaryEntity;
            
            /// <summary>
            /// The skinned mesh renderer entity, not to be confused with the deformed entity.
            /// </summary>
            public Entity SmrEntity;
            
            /// <summary>
            /// If not set to zero, the buffer will be created with this capacity
            /// </summary>
            public int   BufferCapacity;
            
            // Stablehash of the component to be added to the deformed entity
            public ulong StableTypeHash;
            public bool  EnabledState;
        }

        /// <summary>
        /// Adds a component to the deformed entity. The component will be added to the deformed entity
        /// </summary>
        /// <param name="baker"></param>
        /// <param name="authoring"></param>
        /// <param name="smr"></param>
        /// <param name="component"></param>
        /// <param name="enabled"></param>
        /// <typeparam name="T"></typeparam>
        public static void AddDeformedComponent<T>(this IBaker baker, Component authoring, SkinnedMeshRenderer smr, T component, bool enabled = true) where T : unmanaged, IComponentData
        {
            // Get the original entity
            var entity           = baker.GetEntity(authoring, TransformUsageFlags.None);
            var additionalEntity = baker.CreateAdditionalEntity(TransformUsageFlags.None, true);
            
            baker.AddComponent<T>(additionalEntity, component);
            
            var typeIndex  = TypeManager.GetTypeIndex<T>();
            var stableHash = TypeManager.GetTypeInfo(typeIndex).StableTypeHash;
            
            var isEnableable = TypeManager.IsEnableableType(typeIndex);
            if (isEnableable)
            {
                
            }
            
            baker.AddComponent<RemapToDeformedComponent>(additionalEntity, new RemapToDeformedComponent
            {
                Entity          = entity,
                TemporaryEntity = additionalEntity,
                SmrEntity       = baker.GetEntity(smr, TransformUsageFlags.Dynamic),
                StableTypeHash  = TypeManager.GetTypeInfo<T>().StableTypeHash,
                EnabledState    = enabled
            });
        }

        public static void AddDeformedBuffer<T>(this IBaker baker, Component authoring, SkinnedMeshRenderer smr, int capacity = 0, bool enabled = false) where T : unmanaged, IBufferElementData
        {
            // Get the original entity
            var entity           = baker.GetEntity(authoring, TransformUsageFlags.None);
            var additionalEntity = baker.CreateAdditionalEntity(TransformUsageFlags.None, true);
            
            var typeIndex    = TypeManager.GetTypeIndex<T>();
            var stableHash   = TypeManager.GetTypeInfo(typeIndex).StableTypeHash;
            
            baker.AddComponent<RemapToDeformedBuffer>(additionalEntity, new RemapToDeformedBuffer
            {
                Entity          = entity,
                TemporaryEntity = additionalEntity,
                SmrEntity       = baker.GetEntity(smr, TransformUsageFlags.Dynamic),
                BufferCapacity  = capacity,
                StableTypeHash  = TypeManager.GetTypeInfo<T>().StableTypeHash,
                EnabledState    = enabled
            });
        }
        

        /// <summary>
        /// System for finding the actual DeformedEntity from the SkinnedMeshRenderer gameobject.
        /// </summary>
        [WorldSystemFilter(WorldSystemFilterFlags.BakingSystem)]
        public partial struct RemapDeformedComponentSystem : ISystem
        {
            public void OnUpdate(ref SystemState state)
            {
                var bufferRemap = SystemAPI.QueryBuilder().WithAll<RemapToDeformedBuffer>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).Build().ToComponentDataArray<RemapToDeformedBuffer>(Allocator.Temp);
                
                // We need to create arrays of each query so we can add components without using ecbs or the codegen system getting mad at us.
                var componentRemap    = SystemAPI.QueryBuilder().WithAll<RemapToDeformedComponent>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).Build().ToComponentDataArray<RemapToDeformedComponent>(Allocator.Temp);
                
                var deformedEntities = SystemAPI.QueryBuilder().WithAll<DeformedEntity>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).Build().ToComponentDataArray<DeformedEntity>(Allocator.Temp);
                var skinnedEntities  = SystemAPI.QueryBuilder().WithAll<DeformedEntity>().WithOptions(EntityQueryOptions.IncludePrefab | EntityQueryOptions.IncludeDisabledEntities).Build().ToEntityArray(Allocator.Temp);
                
                // Loop through the remap entities
                foreach (var remapData in componentRemap)
                {
                    // Get all pertinent data
                    var smrEntity  = remapData.SmrEntity;
                    var tempEntity = remapData.TemporaryEntity;
                    var stableHash = remapData.StableTypeHash;
                    
                    var typeIndex      = TypeManager.GetTypeIndexFromStableTypeHash(stableHash);
                    var type           = TypeManager.GetType(typeIndex);
                    var isEnableable   = TypeManager.IsEnableableType(typeIndex);
                    var enabledState   = remapData.EnabledState;
                    
                    var isBuffer = TypeManager.IsBuffer(typeIndex);
                    if (!isBuffer)
                    {
                        // Get the original component we want to transfer
                        var component = GetComponentData(ref state, tempEntity, type);
                        
                        for (var index = 0; index < deformedEntities.Length; index++)
                        {
                            var deformedEntity = deformedEntities[index];
                            var skinnedEntity  = skinnedEntities[index];
                            
                            if (deformedEntity.Value.Equals(smrEntity))
                            {
                                //Debug.Log($"Adding component {type} to deformed entity {deformedEntity.Value}");
                                AddComponentData(ref state, skinnedEntity, type, component);

                                if (isEnableable)
                                {
                                    state.EntityManager.SetComponentEnabled(skinnedEntity, type, enabledState);
                                }
                            }
                        }
                    }

                    if (isBuffer)
                    {
                        // Ok so thought process: we will need to loop through the buffer, collect all the data to an array, 
                    }
                }
                // Loop through the remap entities
                foreach (var remapData in bufferRemap)
                {
                    // Get all pertinent data
                    var smrEntity  = remapData.SmrEntity;
                    var tempEntity = remapData.TemporaryEntity;
                    var stableHash = remapData.StableTypeHash;
                    
                    // buffer capacity
                    var bufferCapacity = remapData.BufferCapacity;
                    
                    var typeIndex      = TypeManager.GetTypeIndexFromStableTypeHash(stableHash);
                    var type           = TypeManager.GetType(typeIndex);
                    var isEnableable   = TypeManager.IsEnableableType(typeIndex);
                    var enabledState   = remapData.EnabledState;
                    
                    var isBuffer = TypeManager.IsBuffer(typeIndex);
                    Assert.IsTrue(isBuffer, $"Type {type} is not a buffer element data type");

                    for (var index = 0; index < deformedEntities.Length; index++)
                    {
                        var deformedEntity = deformedEntities[index];
                        var skinnedEntity  = skinnedEntities[index];
                            
                        if (deformedEntity.Value.Equals(smrEntity))
                        {
                            //Debug.Log($"Adding component {type} to deformed entity {deformedEntity.Value}");
                            AddBufferByType(skinnedEntity, type, ref state, bufferCapacity);

                            if (isEnableable)
                            {
                                state.EntityManager.SetComponentEnabled(skinnedEntity, type, enabledState);
                            }
                        }
                    }
                }
            }
        }
        
        private static 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);
            
            Assert.IsNotNull(methodInfo, $"Method not found for {componentType}");

            return methodInfo.Invoke(state.EntityManager, new object[] { entity });
        }
        
        private static object GetBuffer(ref SystemState state, Entity entity, Type bufferType)
        {
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "GetBuffer")
                    //&& m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(bufferType);

            Assert.IsNotNull(methodInfo, $"Method not found for setting {bufferType}");
            // first parameter is the entity, second is the bool for readonly, set always false
            return methodInfo.Invoke(state.EntityManager, new object[] { entity, false });
        }
        
        private static void AddBuffer(ref SystemState state, Entity entity, Type componentType)
        {
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "AddBuffer" 
                    && m.GetParameters().Length           == 1 
                    && m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(componentType);
            
            Assert.IsNotNull(methodInfo, $"Method not found for setting {componentType}");

            methodInfo.Invoke(state.EntityManager, new object[] { entity, componentType });
        }
        
        private static void AddComponentData(ref SystemState state, Entity entity, Type componentType, object component)
        {
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name               == "AddComponentData" 
                    && m.GetParameters().Length           == 2 
                    && m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(componentType);
            
            Assert.IsNotNull(methodInfo, $"Method not found for setting {componentType}");

            methodInfo.Invoke(state.EntityManager, new object[] { entity, component });
        }
        
        private static 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);
            
            Assert.IsNotNull(methodInfo, $"Method not found for setting {componentType}");

            methodInfo.Invoke(state.EntityManager, new object[] { entity, component });
        }
        
        
        ////////////////////////////////////////
        
        public static object AddBufferByType(Entity entity, Type bufferElementType, ref SystemState state, int capacity = 0)
        {
            // Get the AddBuffer<T> method from EntityManager
            var methodInfo = typeof(EntityManager)
                .GetMethods()
                .FirstOrDefault(m => m.Name == "AddBuffer" 
                    && m.IsGenericMethodDefinition
                    && m.GetParameters().Length           == 1 
                    && m.GetParameters()[0].ParameterType == typeof(Entity))
                ?.MakeGenericMethod(bufferElementType);

            if (methodInfo == null)
            {
                throw new InvalidOperationException("AddBuffer<T> method not found on EntityManager.");
            }

            // Ensure the bufferElementType implements IBufferElementData
            if (!typeof(IBufferElementData).IsAssignableFrom(bufferElementType))
            {
                throw new ArgumentException($"The type {bufferElementType.FullName} does not implement IBufferElementData.");
            }

            // Invoke AddBuffer<T> dynamically
            var buffer = methodInfo.Invoke(state.EntityManager, new object[] { entity });
            
            // Resize the buffer if capacity is specified
            if (capacity > 0)
            {
                if (buffer == null)
                {
                    throw new InvalidOperationException("AddBuffer<T> returned null.");
                }

                // Get the Resize method on the buffer
                var resizeMethod = buffer.GetType().GetMethod("Resize", BindingFlags.Public | BindingFlags.Instance);

                if (resizeMethod == null)
                {
                    throw new InvalidOperationException($"Resize method not found on buffer of type {buffer.GetType()}.");
                }

                // Invoke Resize to set the buffer capacity
                resizeMethod.Invoke(buffer, new object[] { capacity, NativeArrayOptions.ClearMemory });
            }
            
            return buffer;
        }

        public static object AddBufferByTypeIndirect(Entity entity, Type bufferElementType)
        {
            // Search for the AddBuffer<T> method dynamically
            var        assemblies      = AppDomain.CurrentDomain.GetAssemblies();
            MethodInfo addBufferMethod = null;

            foreach (var assembly in assemblies)
            {
                foreach (var type in assembly.GetTypes())
                {
                    // Look for the AddBuffer<T> method
                    var method = type.GetMethod("AddBuffer", BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.Static);
                    if (method != null && method.IsGenericMethodDefinition)
                    {
                        addBufferMethod = method;
                        break;
                    }
                }

                if (addBufferMethod != null)
                    break;
            }

            if (addBufferMethod == null)
            {
                throw new InvalidOperationException("AddBuffer<T> method could not be located in any loaded assembly.");
            }

            // Ensure the bufferElementType implements IBufferElementData
            if (!typeof(IBufferElementData).IsAssignableFrom(bufferElementType))
            {
                throw new ArgumentException($"The type {bufferElementType.FullName} does not implement IBufferElementData.");
            }

            // Create a generic method for the specified bufferElementType
            var genericMethod = addBufferMethod.MakeGenericMethod(bufferElementType);

            // Determine if the method is static or instance-based
            if (addBufferMethod.IsStatic)
            {
                // Invoke static method
                return genericMethod.Invoke(null, new object[] { entity });
            }
            else
            {
                // For instance methods, create an instance dynamically
                var declaringType = addBufferMethod.DeclaringType;
                var instance      = Activator.CreateInstance(declaringType);
                return genericMethod.Invoke(instance, new object[] { entity });
            }
        }


    }
}