Unity Shader - How to efficiently recolor specific coordinates?

Question

First, please allow me to explain what I've got and then I'll go over what I'm trying to figure out next.

What I've got

I've got a textured custom mesh with some edges that exactly align with integer world coordinates in Unity. To the mesh I've added my own crude yet effective custom surface shader that looks like this:

    Shader "Custom/GridHighlightShader"
{
    Properties
    {
        [HideInInspector]_SelectionColor("SelectionColor", Color) = (0.1,0.1,0.1,1)
        [HideInInspector]_MovementColor("MovementColor", Color) = (0,0.205,1,1)
        [HideInInspector]_AttackColor("AttackColor", Color) = (1,0,0,1)
        [HideInInspector]_GlowInterval("_GlowInterval", float) = 1
        _MainTex("Albedo (RGB)", 2D) = "white" {}
        _Glossiness("Smoothness", Range(0,1)) = 0.5
        _Metallic("Metallic", Range(0,1)) = 0.0
    }
        SubShader
        {
            Tags { "RenderType" = "Opaque" }
            LOD 200

            CGPROGRAM
            // Physically based Standard lighting model, and enable shadows on all light types
            #pragma surface surf Standard fullforwardshadows

            // Use shader model 3.0 target, to get nicer looking lighting
            #pragma target 3.0

            struct Input
            {
                float2 uv_MainTex;
                float3 worldNormal;
                float3 worldPos;
            };

            sampler2D _MainTex;
            half _Glossiness;
            half _Metallic;
            fixed4 _SelectionColor;
            fixed4 _MovementColor;
            fixed4 _AttackColor;
            half _GlowInterval;
            half _ColorizationArrayLength = 0;
            float4 _ColorizationArray[600];
            half _isPixelInColorizationArray = 0;

            // Add instancing support for this shader. You need to check 'Enable Instancing' on materials that use the shader.
            // See https://docs.unity3d.com/Manual/GPUInstancing.html for more information about instancing.
            // #pragma instancing_options assumeuniformscaling
            UNITY_INSTANCING_BUFFER_START(Props)
                // put more per-instance properties here
            UNITY_INSTANCING_BUFFER_END(Props)

                void surf(Input IN, inout SurfaceOutputStandard o)
                {
                fixed4 c = tex2D(_MainTex, IN.uv_MainTex);

                // Update only the normals facing up and down
                if (abs(IN.worldNormal.x) <= 0.5 && (abs(IN.worldNormal.z) <= 0.5))
                {
                    // If no colors were passed in, reset all of the colors
                    if (_ColorizationArray[0].w == 0)
                    {
                        _isPixelInColorizationArray = 0;
                    }
                    else
                    {
                        for (int i = 0; i < _ColorizationArrayLength; i++)
                        {
                            if (abs(IN.worldPos.x) >= _ColorizationArray[i].x && abs(IN.worldPos.x) < _ColorizationArray[i].x + 1
                                && abs(IN.worldPos.z) >= _ColorizationArray[i].z && abs(IN.worldPos.z) < _ColorizationArray[i].z + 1
                                )
                            {
                                _isPixelInColorizationArray = _ColorizationArray[i].w;
                            }
                        }
                    }

                    if (_isPixelInColorizationArray > 0)
                    {
                        if (_isPixelInColorizationArray == 1)
                        {
                            c = tex2D(_MainTex, IN.uv_MainTex) + (_SelectionColor * _GlowInterval) - 1;
                        }
                        else if (_isPixelInColorizationArray == 2)
                        {
                            c = tex2D(_MainTex, IN.uv_MainTex) + (_MovementColor * _GlowInterval);
                        }
                        else if (_isPixelInColorizationArray == 3)
                        {
                            c = tex2D(_MainTex, IN.uv_MainTex) + (_AttackColor * _GlowInterval);
                        }
                    }
                }
                o.Albedo = c.rgb;
                o.Metallic = _Metallic;
                o.Smoothness = _Glossiness;
                o.Alpha = c.a;

            }
            ENDCG
        }
            FallBack "Diffuse"
}

Into the shader I feed a float that simply oscilates between 2 and 3 over time using some maths, this is done from a simple update function in Unity:

private void Update() 
{
    var t = (2 + ((Mathf.Sin(Time.time))));
    meshRenderer.material.SetFloat("_GlowInterval", t);
}

I also feed the shader an array of Vector4 called _ColorizationArray that stores 0 to 600 coordinates, each representing a tile to be colored at runtime. These tiles may or may not be highlighted depending on their selectionMode value at runtime. Here's the method I'm using to do that:

    public void SetColorizationCollectionForShader()
    {
        var coloredTilesArray = Battlemap.Instance.tiles.Where(x => x.selectionMode != TileSelectionMode.None).ToArray();

        // https://docs.unity3d.com/ScriptReference/Material.SetVectorArray.html                
        // Set the tile count in the shader's own integer variable
        meshRenderer.material.SetInt("_ColorizationArrayLength", coloredTilesArray.Length);

        // Loop through the tiles to be colored only and grab their world coordinates
        for(int i = 0; i < coloredTilesArray.Length; i++)
        {
            // Also grab the selection mode as the w value of a float4
            colorizationArray[i] = new Vector4(coloredTilesArray[i].x - Battlemap.HALF_TILE_SIZE, coloredTilesArray[i].y, coloredTilesArray[i].z - Battlemap.HALF_TILE_SIZE, (float)coloredTilesArray[i].selectionMode);            
        }

        // Feed the overwritten array into the shader
        meshRenderer.material.SetVectorArray("_ColorizationArray", colorizationArray);
    }

And the result is this blue glowy set of tiles that are set and changed dynamically at runtime:

My goal with all of this is to highlight squares (or tiles, if you will) on a mesh as part of a grid-based tactical game where units can move to any tile within the highlighted area. After each unit moves it can then undergo an attack where tiles are highlighted red, and then the next unit takes it's turn and so on. Since I expect AI, and movement calculations, and particle effects to take up the majority of processing time I need to highlight the tiles both dynamically and very efficiently at runtime.

What I'd like to do next

Whew, ok. Now, if you know anything about shaders (which I certainly don't, I only started looking at cg code yesterday) you're probably thinking "Oh dear god what an inefficient mess. What are you doing?! If statements?! In a shader?" And I wouldn't blame you.

What I'd really like to do is pretty much the same thing, only much more efficiently. Using specific tile indices I'd like to tell the shader "color the surface blue inside of these tiles, and these tiles only" and do it in a way that is efficient for both the GPU and the CPU.

How can I achieve this? I'm already computing tile world coordinates in C# code and providing the coordinates to the shader, but beyond that I'm at a loss. I realize I should maybe switch to a vertex/frag shader, but I'd also like to avoid losing any of the default dynamic lighting on the mesh if possible.

Also, is there a type of variable that would allow the shader to color the mesh blue using the local mesh coordinates rather than world coordinates? Would be nice to be able to move the mesh around without having to worry about shader code.

Edit: In the 2 weeks since posting this question I've edited the shader by passing in an array of Vector4s and a half to represent how much of the array to actually process, _ColorizationArrayLength, it works well, but is hardly more efficient - this is producing GPU spikes that take about 17ms to process on a fairly modern graphics card. I've updated the shader code above as well as portions of the original question.

Answer 1

Since your colorizing only cares about 2d position in a grid of equally sized squares that are all aligned to the same grid, we can pass in a 2d texture whose coloring says what the ground should be colored like.

In your shader, add a 2D _ColorizeMap, and a Vector _WorldSpaceRange. The map will be used to pass in which sections of the world should be colorized, and the range will tell the shader how to convert between world space and UV (texture) space. Since the game grid is aligned to the world x/y axes, we can just linearly scale the coordinates from world space to UV space.

Then, when the normal is facing upwards (which you can check if the normal's y is high enough), get an inverse lerp of the world position, and sample from _ColorizeMap to get how/whether it should be colored.

Shader "Custom/GridHighlightShader"
{
    Properties
    {
        [HideInInspector]_GlowInterval("_GlowInterval", float) = 1
        _MainTex("Albedo (RGB)", 2D) = "white" {}
        _Glossiness("Smoothness", Range(0,1)) = 0.5
        _Metallic("Metallic", Range(0,1)) = 0.0
        [HideInInspector]_ColorizeMap("Colorize Map", 2D) = "black" {} 
        _WorldSpaceRange("World Space Range", Vector) = (0,0,100,100)
    }
        SubShader
        {
            Tags { "RenderType" = "Opaque" }
            LOD 200

            CGPROGRAM
            // Physically based Standard lighting model, 
            // and enable shadows on all light types
            #pragma surface surf Standard fullforwardshadows

            // Use shader model 3.0 target, to get nicer looking lighting
            #pragma target 3.0

            struct Input
            {
                float2 uv_MainTex;
                float3 worldNormal;
                float3 worldPos;
            };

            sampler2D _MainTex;
            half _Glossiness;
            half _Metallic;
            half _GlowInterval;

            sampler2D _ColorizeMap;
            fixed4 _WorldSpaceRange;


            // Add instancing support for this shader. 
            // You need to check 'Enable Instancing' on materials that use the shader.
            // See https://docs.unity3d.com/Manual/GPUInstancing.html 
            // for more information about instancing.
            // #pragma instancing_options assumeuniformscaling
            UNITY_INSTANCING_BUFFER_START(Props)
                // put more per-instance properties here
            UNITY_INSTANCING_BUFFER_END(Props)

            void surf(Input IN, inout SurfaceOutputStandard o)
            {
                fixed4 c = tex2D(_MainTex, IN.uv_MainTex);

                // Update only the normals facing up and down
                if (abs(IN.worldNormal.y) >= 0.866)) // abs(y) >= sin(60 degrees)
                {
                    fixed4 colorizedMapUV = (IN.worldPos.xz-_WorldSpaceRange.xy) 
                            / (_WorldSpaceRange.zw-_WorldSpaceRange.xy);

                    half4 colorType = tex2D(_ColorizeMap, colorizedMapUV);

                    c = c + (colorType * _GlowInterval); 
                }
                o.Albedo = c.rgb;
                o.Metallic = _Metallic;
                o.Smoothness = _Glossiness;
                o.Alpha = c.a;

            }
            ENDCG
        }
    FallBack "Diffuse"
}

And remove the branching:

Shader "Custom/GridHighlightShader"
{
    Properties
    {
        [HideInInspector]_GlowInterval("_GlowInterval", float) = 1
        _MainTex("Albedo (RGB)", 2D) = "white" {}
        _Glossiness("Smoothness", Range(0,1)) = 0.5
        _Metallic("Metallic", Range(0,1)) = 0.0
        [HideInInspector]_ColorizeMap("Colorize Map", 2D) = "black" {}
        _WorldSpaceRange("World Space Range", Vector) = (0,0,100,100)
    }
        SubShader
        {
            Tags { "RenderType" = "Opaque" }
            LOD 200
            CGPROGRAM
            // Physically based Standard lighting model, 
            // and enable shadows on all light types
            #pragma surface surf Standard fullforwardshadows

            // Use shader model 3.0 target, to get nicer looking lighting
            #pragma target 3.0

            struct Input
            {
                float2 uv_MainTex;
                float3 worldNormal;
                float3 worldPos;
            };

            sampler2D _MainTex;
            half _Glossiness;
            half _Metallic;
            half _GlowInterval;

            sampler2D _ColorizeMap;
            fixed4 _WorldSpaceRange;

            // Add instancing support for this shader.
            // You need to check 'Enable Instancing' on materials that use the shader.
            // See https://docs.unity3d.com/Manual/GPUInstancing.html 
            // for more information about instancing.
            // #pragma instancing_options assumeuniformscaling
            UNITY_INSTANCING_BUFFER_START(Props)
                // put more per-instance properties here
            UNITY_INSTANCING_BUFFER_END(Props)

            void surf(Input IN, inout SurfaceOutputStandard o)
            {

                half4 c = tex2D(_MainTex, IN.uv_MainTex);
                float2 colorizedMapUV = (IN.worldPos.xz - _WorldSpaceRange.xy)
                        / (_WorldSpaceRange.zw - _WorldSpaceRange.xy);
                half4 colorType = tex2D(_ColorizeMap, colorizedMapUV);

                // abs(y) >= sin(60 degrees) = 0.866
                c = c + step(0.866, abs(IN.worldNormal.y)) * colorType * _GlowInterval;

                o.Albedo = c.rgb;
                o.Metallic = _Metallic;
                o.Smoothness = _Glossiness;
                o.Alpha = c.a;

            }
            ENDCG
        }
            FallBack "Diffuse"
}

And then in your C# code, create a texture without filtering. Start the texture all black, then add colors to the texture depending on how highlighting should be done. Also, tell the shader the range in world space (minX,minZ,maxX,maxZ) that the color map represents:

    public void SetColorizationCollectionForShader()
{   
    Color[] selectionColors = new Color[4] { Color.clear, new Color(0.5f, 0.5f, 0.5f, 0.5f), Color.blue, Color.red };
    float leftMostTileX = 0f + Battlemap.HALF_TILE_SIZE;
    float backMostTileZ = 0f + Battlemap.HALF_TILE_SIZE;

    float rightMostTileX = leftMostTileX + (Battlemap.Instance.GridMaxX - 1)
            * Battlemap.TILE_SIZE;
    float forwardMostTileZ = backMostTileZ + (Battlemap.Instance.GridMaxZ - 1)
            * Battlemap.TILE_SIZE;

    Texture2D colorTex = new Texture2D(Battlemap.Instance.GridMaxX, Battlemap.Instance.GridMaxZ);
    colorTex.filterMode = FilterMode.Point;

    Vector4 worldRange = new Vector4(
            leftMostTileX - Battlemap.HALF_TILE_SIZE,
            backMostTileZ - Battlemap.HALF_TILE_SIZE,
            rightMostTileX + Battlemap.HALF_TILE_SIZE,
            forwardMostTileZ + Battlemap.HALF_TILE_SIZE);

    meshRenderer.material.SetVector("_WorldSpaceRange", worldRange);        

    // Loop through the tiles to be colored only and grab their world coordinates
    for (int i = 0; i < Battlemap.Instance.tiles.Length; i++)
    {
        // determine pixel index from position
        float xT = Mathf.InverseLerp(leftMostTileX, rightMostTileX,
                Battlemap.Instance.tiles[i].x);
        int texXPos = Mathf.RoundToInt(Mathf.Lerp(0f, Battlemap.Instance.GridMaxX - 1.0f, xT));

        float yT = Mathf.InverseLerp(backMostTileZ, forwardMostTileZ,
                Battlemap.Instance.tiles[i].z);
        int texYPos = Mathf.RoundToInt(Mathf.Lerp(0f, Battlemap.Instance.GridMaxZ - 1.0f, yT));

        colorTex.SetPixel(texXPos, texYPos, selectionColors[(int)Battlemap.Instance.tiles[i].selectionMode]);
    }
    colorTex.Apply();

    // Feed the color map into the shader
    meshRenderer.material.SetTexture("_ColorizeMap", colorTex);
}

There might be some wonkiness at the borders of the tiles, and there might be some alignment issues between texture space/world space but this should get you started.