riversong-code-showcase/Source/Riversong/Game/Rendering/MaterialReplacementCache.cs

using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Pool;

namespace DanieleMarotta.RiversongCodeShowcase
{
    public class MaterialReplacementCache
    {
        private static readonly int BaseMapPropertyId = Shader.PropertyToID("_BaseMap");

        private static readonly int MainTexPropertyId = Shader.PropertyToID("_MainTex");

        private static readonly int MainTexturePropertyId = Shader.PropertyToID("_Main_Texture");

        private HashSet<int> _gameObjects = new();

        private List<(Renderer, Material, int, MaterialPropertyBlock)> _cache = new();

        private List<MaterialPropertyBlock> _propertyBlockPool = new();

        private int _propertyBlockPoolIndex;

        private MaterialPropertyBlock _propertyBlock;

        public static void ReplaceRendererMaterials(Renderer renderer, Material material, MaterialPropertyBlock propertyBlock)
        {
            var materials = renderer.sharedMaterials;

            for (var i = 0; i < materials.Length; i++)
            {
                var originalMaterial = materials[i];
                materials[i] = material;
                ReplaceTexture(renderer, originalMaterial, i, propertyBlock, null);
            }

            renderer.sharedMaterials = materials;
        }

        public void ReplaceMaterials(GameObject gameObject, Material material)
        {
            if (!_gameObjects.Add(gameObject.GetInstanceID())) return;

            using var renderersScope = ListPool<Renderer>.Get(out var renderers);
            gameObject.GetComponentsInChildren(renderers);

            foreach (var renderer in renderers) ReplaceMaterials(renderer, material);
        }

        public void ReplaceMaterials(List<GameObject> gameObjects, Material material)
        {
            if (gameObjects.Count == 0) return;

            using var renderersScope = ListPool<Renderer>.Get(out var renderers);

            foreach (var gameObject in gameObjects)
            {
                if (!_gameObjects.Add(gameObject.GetInstanceID())) continue;

                renderers.Clear();
                gameObject.GetComponentsInChildren(renderers);

                foreach (var renderer in renderers) ReplaceMaterials(renderer, material);
            }
        }

        public void RestoreMaterials()
        {
            foreach (var (renderer, material, i, propertyBlock) in _cache)
            {
                if (!renderer) continue;

                var materials = renderer.sharedMaterials;
                materials[i] = material;
                renderer.sharedMaterials = materials;
                renderer.SetPropertyBlock(propertyBlock, i);
            }

            _gameObjects.Clear();
            _cache.Clear();
            _propertyBlockPoolIndex = 0;
        }

        private void ReplaceMaterials(Renderer renderer, Material material)
        {
            var materials = renderer.sharedMaterials;

            _propertyBlock ??= new MaterialPropertyBlock();
            for (var i = 0; i < materials.Length; i++)
            {
                var currentPropertyBlock = GetPooledPropertyBlock();
                renderer.GetPropertyBlock(currentPropertyBlock, i);

                _cache.Add((renderer, materials[i], i, currentPropertyBlock));

                var originalMaterial = materials[i];
                materials[i] = material;
                ReplaceTexture(renderer, originalMaterial, i, _propertyBlock, currentPropertyBlock);
            }

            renderer.sharedMaterials = materials;
        }

        private MaterialPropertyBlock GetPooledPropertyBlock()
        {
            if (_propertyBlockPoolIndex >= _propertyBlockPool.Count) _propertyBlockPool.Add(new MaterialPropertyBlock());

            var propertyBlock = _propertyBlockPool[_propertyBlockPoolIndex++];
            propertyBlock.Clear();

            return propertyBlock;
        }

        private static void ReplaceTexture(Renderer renderer, Material material, int materialIndex, MaterialPropertyBlock propertyBlock, MaterialPropertyBlock sourcePropertyBlock)
        {
            var texture = GetMainTexture(material);
            if (!texture) texture = sourcePropertyBlock?.GetTexture(MainTexturePropertyId);
            if (!texture) return;

            propertyBlock.Clear();
            propertyBlock.SetTexture(MainTexturePropertyId, texture);

            renderer.SetPropertyBlock(propertyBlock, materialIndex);
        }

        private static Texture GetMainTexture(Material material)
        {
            if (material.HasProperty(BaseMapPropertyId)) return material.GetTexture(BaseMapPropertyId);
            if (material.HasProperty(MainTexPropertyId)) return material.GetTexture(MainTexPropertyId);
            return material.HasProperty(MainTexturePropertyId) ? material.GetTexture(MainTexturePropertyId) : null;
        }
    }
}