ThomagicRenderer/Runtime/Scripts/TreeRenderer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

namespace Assets.ThoMagic.Renderer
{
    public class TreeRenderer : TreeStreamer
    {
        private Terrain terrain;
        private TerrainData terrainData;
        private List<int> treePrototypesRegistered;
        private Dictionary<int, TreeInfo> treeToInstances = new Dictionary<int, TreeInfo>();
        private HashSet<int> touchedTrees = new HashSet<int>();
        private bool isDirty;
        private float maxTreeDistance = 0;

        class TreeInfo
        {
            internal int instanceId;
            internal int objectId;
        }

        public TreeRenderer(Terrain terrain)
            : base(terrain)
        {
            this.terrain = terrain;
            terrainData = terrain.terrainData;
        }

        public void LateUpdate()
        {
            if (isDirty)
            {
                isDirty = false;
                ApplyChanges();
            }

            if (Application.isEditor && !Application.isPlaying)
                RebuildChangedPrototypes();
        }

        int GetTreeTerrainInstanceId(TreeInstance tree)
        {
            return HashCode.Combine(tree.position.x, tree.position.y, tree.position.z, tree.rotation, tree.widthScale, tree.heightScale, tree.prototypeIndex);
        }

        public void ApplyChanges()
        {
            Vector3 size = terrainData.size;
            Vector3 offset = terrain.transform.position;

            var treeInstances = terrainData.treeInstances;
            touchedTrees.Clear();
            foreach (var treeInstance in treeInstances)
            {
                if (treeInstance.widthScale == 0)
                    continue;

                var treeTerrainInstanceId = GetTreeTerrainInstanceId(treeInstance);
                if (!treeToInstances.ContainsKey(treeTerrainInstanceId))
                {
                    var treeId = treePrototypesRegistered[treeInstance.prototypeIndex];
                    var prefab = RendererPool.GetObject(treeId);
                    var quaternion = Quaternion.Euler(0.0f, 57.2957801818848f * treeInstance.rotation, 0.0f);
                    var instanceId = AddInstance(treeId, new Vector3(treeInstance.position.x * size.x,
                        treeInstance.position.y * size.y,
                        treeInstance.position.z * size.z) + offset, quaternion, 
                        treeInstance.widthScale * prefab.transform.localScale.x, 
                        treeInstance.heightScale * prefab.transform.localScale.y
                        );
                    treeToInstances[treeTerrainInstanceId] = new TreeInfo { instanceId = instanceId, objectId = treeId };
                    touchedTrees.Add(treeTerrainInstanceId);
                }
                else
                {
                    touchedTrees.Add(treeTerrainInstanceId);
                }
            }

            foreach(var treeTerrainInstanceId in treeToInstances.Keys.ToList())
            {
                if(!touchedTrees.Contains(treeTerrainInstanceId))
                {                    
                    RemoveInstance(treeToInstances[treeTerrainInstanceId].objectId, treeToInstances[treeTerrainInstanceId].instanceId);
                    treeToInstances.Remove(treeTerrainInstanceId);
                }
            }

            Build(maxTreeDistance);
        }

        public void Load()
        {
            maxTreeDistance = 0;

            Vector3 size = terrainData.size;
            Vector3 offset = terrain.transform.position;

            Recycle();

            int ownerHash = GetHashCode();

            if (treePrototypesRegistered == null)
                treePrototypesRegistered = new List<int>();

            var treePrototypes = terrainData.treePrototypes;
            if(treePrototypes.Length > 0)
            {
                foreach(var treePrototype in treePrototypes)
                {
                    if(treePrototype.prefab != null)
                    {
                        var settings = GetSettingsOrDefault(treePrototype.prefab);
                        treePrototypesRegistered.Add(RendererPool.RegisterObject(treePrototype.prefab, settings, this, ownerHash));

                        if (settings.Settings.RenderDistance == 0)
                            maxTreeDistance = float.MaxValue;

                        if(settings.Settings.RenderDistance > maxTreeDistance)
                            maxTreeDistance = settings.Settings.RenderDistance;
                    }
                }

                var treeInstances = terrainData.treeInstances;
                foreach (var treeInstance in treeInstances)
                {
                    var treeTerrainInstanceId = GetTreeTerrainInstanceId(treeInstance);                    
                    var treeId = treePrototypesRegistered[treeInstance.prototypeIndex];
                    var prefab = RendererPool.GetObject(treeId);
                    var quaternion = Quaternion.Euler(0.0f, 57.2957801818848f * treeInstance.rotation, 0.0f);
                    var instanceId = AddInstance(treeId, new Vector3(treeInstance.position.x * size.x,
                        treeInstance.position.y * size.y,
                        treeInstance.position.z * size.z) + offset, quaternion,
                        treeInstance.widthScale * prefab.transform.localScale.x,
                        treeInstance.heightScale * prefab.transform.localScale.y);
                    treeToInstances[treeTerrainInstanceId] = new TreeInfo { instanceId = instanceId, objectId = treeId };
                }
            }

            Build(maxTreeDistance);
        }

        private void Recycle()
        {
            if (treePrototypesRegistered == null)
                return;

            treeToInstances.Clear();

            Clear();

            var ownerHash = GetHashCode();

            for (int index = 0; index < treePrototypesRegistered.Count; ++index)
            {
                var objectId = treePrototypesRegistered[index];
                RendererPool.RemoveObject(objectId, ownerHash);
            }
            
            treePrototypesRegistered.Clear();
        }

        private void RebuildChangedPrototypes()
        {
            TreePrototype[] treePrototypes = terrainData.treePrototypes;
            if (treePrototypes.Length != treePrototypesRegistered.Count)
            {
                Load();
            }
            else
            {
                int ownerHash = GetHashCode();
                maxTreeDistance = 0;
                for (int index = 0; index < treePrototypesRegistered.Count; ++index)
                {
                    int treeId = treePrototypesRegistered[index];

                    GameObject prefab = treePrototypes[index].prefab;
                    GameObject gameObject = RendererPool.GetObject(treeId);
                    if (prefab != gameObject)
                    {
                        RendererPool.RemoveObject(treeId, ownerHash);
                        if (prefab != null)
                        {
                            var settings = GetSettingsOrDefault(prefab);

                            if (settings.Settings.RenderDistance == 0)
                                maxTreeDistance = float.MaxValue;

                            if (settings.Settings.RenderDistance > maxTreeDistance)
                                maxTreeDistance = settings.Settings.RenderDistance;

                            treePrototypesRegistered[index] = RendererPool.RegisterObject(prefab, settings, this, ownerHash);
                        }
                    }
                    else if (RendererPool.ContentHashChanged(treeId))
                    {
                        RendererPool.RemoveObject(treeId, ownerHash);
                        if (prefab != null)
                        {
                            var settings = GetSettingsOrDefault(prefab);

                            if (settings.Settings.RenderDistance == 0)
                                maxTreeDistance = float.MaxValue;

                            if (settings.Settings.RenderDistance > maxTreeDistance)
                                maxTreeDistance = settings.Settings.RenderDistance;

                            treePrototypesRegistered[index] = RendererPool.RegisterObject(prefab, settings, this, ownerHash);
                        }
                    }
                    else if (prefab != null)
                    {
                        var settings = GetSettingsOrDefault(prefab);

                        if (settings.Settings.RenderDistance == 0)
                            maxTreeDistance = float.MaxValue;

                        if (settings.Settings.RenderDistance > maxTreeDistance)
                            maxTreeDistance = settings.Settings.RenderDistance;

                        RendererPool.SetObjectSettings(prefab, settings);
                    }
                }
            }
        }

        public void Destroy()
        {
            Recycle();           
        }

        public void OnTerrainChanged(TerrainChangedFlags flags)
        {
            if(flags.HasFlag(TerrainChangedFlags.TreeInstances))
                isDirty = true; 
            else if (flags.HasFlag(TerrainChangedFlags.DelayedHeightmapUpdate))
                isDirty = true;
            else if (flags.HasFlag(TerrainChangedFlags.Heightmap))
                isDirty = true;
            else if (flags.HasFlag(TerrainChangedFlags.HeightmapResolution))
                isDirty = true;
            else if (flags.HasFlag(TerrainChangedFlags.Holes))
                isDirty = true;
            else if (flags.HasFlag(TerrainChangedFlags.DelayedHolesUpdate))
                isDirty = true;
            else if (flags.HasFlag(TerrainChangedFlags.FlushEverythingImmediately))
                isDirty = true;
        }

        private IInstanceRenderSettings GetSettingsOrDefault(GameObject gameObject)
        {
            IInstanceRenderSettings instanceRenderSettings;
            return gameObject.TryGetComponent(out instanceRenderSettings) ? instanceRenderSettings : new DefaultRenderSettings();
        }

        private class DefaultRenderSettings : IInstanceRenderSettings
        {
            public InstanceRenderSettings Settings => new InstanceRenderSettings()
            {
                Render = true,
                DensityInDistance = 1f,
                DensityInDistanceFalloff = Vector2.zero,
                RenderDistance = 0.0f,
                ShadowDistance = 0.0f,
                Shadows = true,
                Supported = true
            };
        }
    }
}