rj-action-library/Runtime/Shading/Generators/Generic/ShaderVariant.cs

using Godot;
using System.Reflection;
using System.Collections.Generic;

namespace Rokojori
{ 
  public class ShaderVariant
  { 
    public string variantName = "";
    public List<ShaderCode> shaderCode = [];

    public string GetCode( ShaderGenerationContext context )
    {
      
      ResolveVariableOccurances( context );

      var cleanedUpCode = RemoveDuplicateIncludes( shaderCode, context );
      
      
      var hasScreen = false;
      var hasDepth = false;
      var hasNormalRoughness = false;

      cleanedUpCode = cleanedUpCode.Filter(
        ( c )=>
        {
          if ( c is BuiltInTextureShaderCode bsc )
          {
            if ( BuiltInTextureShaderCode.TextureType.Screen == bsc.textureType )
            {
              hasScreen = true;
            }

            if ( BuiltInTextureShaderCode.TextureType.Depth == bsc.textureType )
            {
              hasDepth = true;
            }

            if ( BuiltInTextureShaderCode.TextureType.NormalRoughness == bsc.textureType )
            {
              hasNormalRoughness = true;
            }

            return false;
          }

          return true;
        }
      );

      var insertionIndex = cleanedUpCode.FindIndex( 
        ( c )=>
        {
          return c.GetCode( context ).Contains( "group_uniforms" );
        }
      );



      if ( hasNormalRoughness )
      {
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
        cleanedUpCode.Insert( insertionIndex, BuiltInTextureShaderCode.NormalRoughness.GetBuiltInCode() );
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
      }

      if ( hasDepth )
      {
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
        cleanedUpCode.Insert( insertionIndex, BuiltInTextureShaderCode.Depth.GetBuiltInCode() );
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
      }

      if ( hasScreen )
      {
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
        cleanedUpCode.Insert( insertionIndex, BuiltInTextureShaderCode.Screen.GetBuiltInCode() );
        cleanedUpCode.Insert( insertionIndex, ShaderGenerationModule.ToCode( "\n" )[ 0 ]  );
      }

  

      var blocks = GetCodeBlocks( cleanedUpCode );

      var sortedCode = new List<ShaderCode>();

      blocks.ForEach(
        ( b )=>
        {
          if ( ! b[ 0 ].sortableCode )
          {
            sortedCode.AddRange( b );
            return;
          }      

          sortedCode.AddRange( SortBlock( b ) );
        }
      );      

      return sortedCode.Map( s => s.GetCode( context ) ).Join( "" );
    }

    void ResolveVariableOccurances( ShaderGenerationContext context )
    {
      shaderCode.ForEach( s => s.ResolveVariableOccurances( context ) );
    }

    
    List<List<ShaderCode>> GetCodeBlocks( List<ShaderCode> shaderCode )
    {
      var blocks = new List<List<ShaderCode>>();
      List<ShaderCode> currentBlock = null;

      for ( int i = 0; i < shaderCode.Count; i++ )
      {
        if ( ! shaderCode[ i ].sortableCode || 
            currentBlock != null && currentBlock[ 0 ].phase != shaderCode[ i ].phase )
        {
          if ( currentBlock != null && currentBlock.Count > 0 )
          {
            blocks.Add( currentBlock );
          }          

          currentBlock = null;

          if ( ! shaderCode[ i ].sortableCode )
          {            
            blocks.Add( [ shaderCode[ i ] ] );

            continue;
          }
          
        }

        if ( currentBlock == null )
        {
          currentBlock = new List<ShaderCode>();
        }

        currentBlock.Add( shaderCode[ i ] );
      }

      if ( currentBlock != null )
      {
        blocks.Add( currentBlock );
      }

      return blocks;
    } 

    List<ShaderCode> RemoveDuplicateIncludes( List<ShaderCode> shaderCode, ShaderGenerationContext context )
    {
      var cleanedUp = new List<ShaderCode>();

      var includes = new HashSet<string>();


      var includeStart = "#include \"res://addons/rokojori_action_library/Runtime/Shading/Library/";

      shaderCode.ForEach(
        ( sc )=>
        {
          var code = sc.GetCode( context );

          if ( ! code.Contains( includeStart ) )
          {
            cleanedUp.Add( sc );
            return;
          }


          var include = code.GetInnerMatch( includeStart, ".gdshaderinc" );

          if ( include != null && includes.Contains( include ) )
          {
            return;
          }

          cleanedUp.Add( sc );

          if ( include != null )
          {
            includes.Add( include );
          }

          
        }
      );

      return cleanedUp;
    }

    List<ShaderCode> SortBlock( List<ShaderCode> blockShaderCode )
    {
      var g = new Graph<ShaderCode>( blockShaderCode );

      var reads = new MapList<string,ShaderCode>();
      var assignments = new MapList<string,ShaderCode>();
      var modifications = new MapList<string,ShaderCode>();

      blockShaderCode.ForEach(
        ( sc )=>
        {
          sc.variableOccurances.ForEach(
            ( occ )=>
            {
              if ( ShaderCodeVariableOccuranceType.Read == occ.type )
              {
                reads.Add( occ.variable.variableName, sc );
              }

              if ( ShaderCodeVariableOccuranceType.Assignment == occ.type )
              {
                assignments.Add( occ.variable.variableName, sc );
              }

              if ( ShaderCodeVariableOccuranceType.Modification == occ.type )
              {
                modifications.Add( occ.variable.variableName, sc );
              }
            }
          );
        }
      );

      blockShaderCode.ForEach(
        ( s ) =>
        {
          s.variableOccurances.ForEach(
            ( occ )=>
            {
              if ( ShaderCodeVariableOccuranceType.Assignment == occ.type )
              {
                return;
              }

              if ( ! assignments.ContainsKey( occ.variable.variableName ) )
              {
                return;
              }

              var assigningCodes = assignments[ occ.variable.variableName ];

              assigningCodes.ForEach( sc => g.Connect( sc, s ) );

              RJLog.Log( "Set dependency:", s, occ.variable.variableName );
            }
          );
        }
      );

      return g.ComputeOrder();
    } 
 
    public ShaderVariant Extend( ShaderCode shaderCode )
    {
      var variant = new ShaderVariant();
      variant.variantName = variantName + shaderCode.variantName;

      variant.shaderCode = this.shaderCode.Clone();
      variant.shaderCode.Add( shaderCode );

      return variant;
    }

    public ShaderVariant Extend( ShaderVariant shaderVariant )
    {
      var extendedVariant = new ShaderVariant();
      extendedVariant.variantName = variantName + shaderVariant.variantName;

      extendedVariant.shaderCode = shaderCode.Clone();
      extendedVariant.shaderCode.AddRange( shaderVariant.shaderCode );

      return extendedVariant;
    }

    public static List<ShaderVariant> CombineVariants( List<ShaderVariant> before, List<ShaderVariant> addition )
    {
      if ( before == null || before.Count == 0 )
      {
        return addition;
      }

      if ( addition == null || addition.Count == 0 )
      {
        return before;
      }

      return Lists.Combine( before, addition, ( a, b ) => a.Extend( b ) );
    }
  }

}