riversong-code-showcase/Source/Riversong/Game/World/Terrain/TerrainShaderParametersSystem.cs

using System;
using System.Threading.Tasks;
using Cysharp.Threading.Tasks;
using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using Object = UnityEngine.Object;

namespace DanieleMarotta.RiversongCodeShowcase
{
    [RequiresWorldReadyForUpdate]
    public class TerrainShaderParametersSystem : GameSystem, IInitializable, IDisposable, IUpdatable, IOnWorldGenerationCompletedCallback
    {
        private static readonly int BlockedTilesMaskID = Shader.PropertyToID("_Blocked_Tiles_Mask");

        private static readonly int CliffsMaskID = Shader.PropertyToID("_Cliffs_Mask");

        private static readonly int WindMapID = Shader.PropertyToID("_Wind_Map");

        private static readonly int WindMapScaleID = Shader.PropertyToID("_Wind_Map_Scale");

        private static readonly int WindMapSpeedID = Shader.PropertyToID("_Wind_Map_Speed");

        private static readonly int WindBendID = Shader.PropertyToID("_Wind_Bend");

        [InjectService]
        private GameConfig _config;

        [InjectService]
        private EditingState _editingState;

        [InjectService]
        private World _world;

        [InjectService]
        private ISignalBus _signalBus;

        private Texture2D _blockedTilesMask;

        private Texture2D _windMap;

        private Texture2D _cliffsMask;

        public TerrainShaderParametersSystem(IServiceLocator serviceLocator) : base(serviceLocator)
        {
        }

        public UniTask InitializeAsync()
        {
            _signalBus.Subscribe<WorldGenerationCompletedSignal>(OnWorldGenerationCompleted);

            return UniTask.CompletedTask;
        }

        public void Dispose()
        {
            _signalBus.Unsubscribe<WorldGenerationCompletedSignal>(OnWorldGenerationCompleted);

            if (_blockedTilesMask)
            {
                Object.Destroy(_blockedTilesMask);
                _blockedTilesMask = null;
            }

            if (_cliffsMask)
            {
                Object.Destroy(_cliffsMask);
                _cliffsMask = null;
            }
        }

        private void OnWorldGenerationCompleted(WorldGenerationCompletedSignal signal)
        {
            signal.Callbacks.Add(this);
        }

        public async UniTask OnWorldGenerationCompletedAsync(World world)
        {
            InitializedBlockedTilesMask();
            await InitializeCliffsMaskAsync();
            await InitializeWindAsync();
            InitializeDebugGUI();
        }

        private void InitializedBlockedTilesMask()
        {
            _blockedTilesMask = new Texture2D(_world.Size.x, _world.Size.y, TextureFormat.RGBA32, false);
            Shader.SetGlobalTexture(BlockedTilesMaskID, _blockedTilesMask);

            var black = new NativeArray<Color32>(_world.Size.x * _world.Size.y, Allocator.Temp);
            _blockedTilesMask.SetPixelData(black, 0);
        }

        private async UniTask InitializeCliffsMaskAsync()
        {
            var size = _world.Size;

            _cliffsMask = new Texture2D(size.x + 1, size.y + 1, TextureFormat.RGBA32, false);
            Shader.SetGlobalTexture(CliffsMaskID, _cliffsMask);

            var pixelData = new NativeArray<Color32>((size.x + 1) * (size.y + 1), Allocator.Persistent);
            var heightmap = _world.Heightmap.GetNativeArray();

            Parallel.ForEach(
                TileRange.From(0, size),
                tile =>
                {
                    var color = default(Color32);

                    var tileIndex = math.mad(tile.y, size.x, tile.x);
                    var tr = tile.x < size.x && tile.y < size.y ? heightmap[tileIndex] : 0;
                    var tl = tile.x > 0 && tile.y < size.y ? heightmap[tileIndex - 1] : 0;
                    var br = tile.x < size.x && tile.y > 0 ? heightmap[tileIndex - size.x] : 0;
                    var bl = tile is { x: > 0, y: > 0 } ? heightmap[tileIndex - size.x - 1] : 0;

                    if (tr != tl || br != bl) color.r = 255;
                    if (tr != br || tl != bl) color.g = 255;

                    var pixelIndex = math.mad(tile.y, size.x + 1, tile.x);
                    // ReSharper disable once AccessToDisposedClosure
                    pixelData[pixelIndex] = color;
                });

            await UniTask.NextFrame();

            _cliffsMask.SetPixelData(pixelData, 0);
            _cliffsMask.Apply();

            pixelData.Dispose();
        }

        private async UniTask InitializeWindAsync()
        {
            _windMap = await _config.Terrain.Wind.Map.LoadAssetAsync<Texture2D>();

            var config = _config.Terrain.Wind;

            Shader.SetGlobalTexture(WindMapID, _windMap);
            Shader.SetGlobalFloat(WindMapScaleID, config.MapScale);
            Shader.SetGlobalFloat(WindMapSpeedID, config.Speed);
            Shader.SetGlobalVector(WindBendID, config.BendDirection);
        }

        private void InitializeDebugGUI()
        {
#if DEBUG
            var gui = TerrainShaderDebugGUI.Create();
            gui.BlockedTilesMaskTexture = _blockedTilesMask;
            gui.CliffsMaskTexture = _cliffsMask;
#endif
        }

        public void Update()
        {
            UpdateBlockedTiles();
        }

        private void UpdateBlockedTiles()
        {
            var pixelData = _blockedTilesMask.GetPixelData<Color32>(0);

            new BlockMapJob
                {
                    BlockMap = _world.BlockMap.GetNativeArray(),
                    Fertility = _world.Fertility.GetNativeArray(),
                    BlockTexture = pixelData,
                    GrowthQ8_8 = (int)math.round(_config.Terrain.GrassGrowthRate * Time.deltaTime * byte.MaxValue * 256)
                }.Schedule(pixelData.Length, _blockedTilesMask.width)
                .Complete();

            if (_editingState.ActiveTool != null)
                foreach (var (tile, _) in _editingState.ActiveTool.AffectedTiles)
                {
                    if (_world.BlockMap.IsBlocked(tile, BlockReason.InvalidElevation)) continue;

                    var i = tile.x + tile.y * _blockedTilesMask.width;

                    var color = pixelData[i];
                    color.g = byte.MaxValue;
                    pixelData[i] = color;
                }

            _blockedTilesMask.Apply();
        }

        [BurstCompile]
        private struct BlockMapJob : IJobParallelFor
        {
            [ReadOnly]
            public NativeArray<BlockReason> BlockMap;

            [ReadOnly]
            public NativeArray<FertilityMapValue> Fertility;

            public NativeArray<Color32> BlockTexture;

            public int GrowthQ8_8;

            [BurstCompile]
            public void Execute(int index)
            {
                var color = BlockTexture[index];

                var (currentFertility, maxFertility) = Fertility[index];
                var clear = (currentFertility > 0 && currentFertility < maxFertility) || (BlockMap[index] & BlockReason.ClearGrassMask) != 0;

                if (clear)
                {
                    color.r = 255;
                    color.b = 0;
                }
                else
                {
                    var q = (color.r << 8) | color.b;
                    q = math.max(0, q - GrowthQ8_8);
                    color.r = (byte)(q >> 8);
                    color.b = (byte)(q & 0xFF);
                }

                color.g = 0;

                BlockTexture[index] = color;
            }
        }
    }
}