NB_FX

Packages/NB_FX/XuanXuanRenderUtility/Runtime/NBParticleLocalTransformHelper.cs

@@ -0,0 +1,192 @@
+using UnityEngine;
+
+namespace NBShader
+{
+    [ExecuteAlways]
+    [DisallowMultipleComponent]
+    [RequireComponent(typeof(ParticleSystem))]
+    public sealed class NBParticleLocalTransformHelper : MonoBehaviour
+    {
+        private const string NBShaderName = "Effects/NBShader";
+        private const string LegacyShaderName = "Effects/NBShader(Legacy)";
+        private const string CustomLocalTransformKeyword = "_CUSTOM_LOCAL_TRANSFORM";
+
+        private static readonly int CustomLocalTransformLocalToWorldId =
+            Shader.PropertyToID("_CustomLocalTransformLocalToWorld");
+
+        private static readonly int CustomLocalTransformWorldToLocalId =
+            Shader.PropertyToID("_CustomLocalTransformWorldToLocal");
+
+        private ParticleSystemRenderer _particleRenderer;
+        private Material _runtimeMaterial;
+        private Material _lastAppliedMaterial;
+        private string _lastWarning;
+
+        private void OnEnable()
+        {
+            CacheRenderer();
+            ApplyCustomTransform();
+        }
+
+        private void LateUpdate()
+        {
+            ApplyCustomTransform();
+        }
+
+        private void OnValidate()
+        {
+            CacheRenderer();
+            ApplyCustomTransform();
+        }
+
+        private void OnDisable()
+        {
+            SetCustomTransformKeyword(false);
+        }
+
+        private void OnDestroy()
+        {
+            SetCustomTransformKeyword(false);
+        }
+
+        private void CacheRenderer()
+        {
+            if (_particleRenderer == null)
+            {
+                TryGetComponent(out _particleRenderer);
+            }
+        }
+
+        private void ApplyCustomTransform()
+        {
+            if (!TryGetWritableMaterial(out Material material))
+            {
+                DisableLastAppliedKeyword();
+                return;
+            }
+
+            if (_lastAppliedMaterial != null && _lastAppliedMaterial != material)
+            {
+                _lastAppliedMaterial.DisableKeyword(CustomLocalTransformKeyword);
+            }
+
+            Matrix4x4 localToWorld = transform.localToWorldMatrix;
+            Matrix4x4 worldToLocal = transform.worldToLocalMatrix;
+
+            material.SetMatrix(CustomLocalTransformLocalToWorldId, localToWorld);
+            material.SetMatrix(CustomLocalTransformWorldToLocalId, worldToLocal);
+            material.EnableKeyword(CustomLocalTransformKeyword);
+
+            _lastAppliedMaterial = material;
+            _lastWarning = null;
+        }
+
+        private void SetCustomTransformKeyword(bool enabled)
+        {
+            Material material = GetExistingMaterial();
+            if (material == null)
+            {
+                return;
+            }
+
+            if (enabled)
+            {
+                material.EnableKeyword(CustomLocalTransformKeyword);
+            }
+            else
+            {
+                material.DisableKeyword(CustomLocalTransformKeyword);
+                if (material == _lastAppliedMaterial)
+                {
+                    _lastAppliedMaterial = null;
+                }
+            }
+        }
+
+        private Material GetExistingMaterial()
+        {
+            if (_lastAppliedMaterial != null)
+            {
+                return _lastAppliedMaterial;
+            }
+
+            CacheRenderer();
+            if (_particleRenderer == null)
+            {
+                return null;
+            }
+
+            if (Application.isPlaying)
+            {
+                return _runtimeMaterial;
+            }
+
+            return _particleRenderer.sharedMaterial;
+        }
+
+        private bool TryGetWritableMaterial(out Material material)
+        {
+            CacheRenderer();
+            if (_particleRenderer == null)
+            {
+                material = null;
+                LogWarningOnce("NBParticleLocalTransformHelper requires a ParticleSystemRenderer on the same GameObject.");
+                return false;
+            }
+
+            material = Application.isPlaying ? GetRuntimeMaterial() : _particleRenderer.sharedMaterial;
+            if (material == null)
+            {
+                LogWarningOnce("NBParticleLocalTransformHelper could not find a material on the ParticleSystemRenderer.");
+                return false;
+            }
+
+            if (material.shader == null || !IsSupportedShader(material.shader))
+            {
+                material = null;
+                LogWarningOnce("NBParticleLocalTransformHelper only supports NBShader materials using shader '" +
+                               NBShaderName + "' or '" + LegacyShaderName + "'.");
+                return false;
+            }
+
+            return true;
+        }
+
+        private static bool IsSupportedShader(Shader shader)
+        {
+            return shader.name == NBShaderName || shader.name == LegacyShaderName;
+        }
+
+        private Material GetRuntimeMaterial()
+        {
+            if (_runtimeMaterial == null && _particleRenderer != null)
+            {
+                _runtimeMaterial = _particleRenderer.material;
+            }
+
+            return _runtimeMaterial;
+        }
+
+        private void DisableLastAppliedKeyword()
+        {
+            if (_lastAppliedMaterial == null)
+            {
+                return;
+            }
+
+            _lastAppliedMaterial.DisableKeyword(CustomLocalTransformKeyword);
+            _lastAppliedMaterial = null;
+        }
+
+        private void LogWarningOnce(string message)
+        {
+            if (_lastWarning == message)
+            {
+                return;
+            }
+
+            _lastWarning = message;
+            Debug.LogWarning(message, this);
+        }
+    }
+}