[CopyToGPU] Attribute for structs and consts

AutoCopyToGPU.cs

// AutoCopyToGPU, written by Eli Piilonen in 2020, MIT license

// Whenever your C# scripts compile, structs/consts/enums that you
// mark with [CopyToGPU] will be copied to a shader-include file.

// DO NOT put this file in an Editor folder! It will break.

///// C# Usage:
//
// [CopyToGPU]
// struct MySharedGPUStruct {
//     int myIntField;
//     uint anotherOne;
//     Vector2 oldschoolFloat2;
//     float2 fancyNewFloat2;
//     Color thisIsAFloat4InAShader;
//     SomeEnum enumField;
// }
//
// [CopyToGPU]
// const int mySharedConstValue = 64;
//
// [CopyToGPU]
// const float thisOneIsAFloat = 123.45f;
//
//
///// About Enums
//
//  If your copied structs include enum fields,
//  their enum types will be copied to the GPU
//  as sets of const values - this only works if
//  your enum's underlying value type is int (the
//  default for enums, if unspecified) or uint.
//  
//  If your enum-type is Difficulty and its
//  values are Easy and Hard, then
//  they will be copied to your shader as:
//
//  static const int Difficulty_Easy = 0;
//  static const int Difficulty_Hard = 1;
// 
// // legal:
// enum Difficulty {
//     Easy,
//     Hard
// }
//
// // also legal:
// enum Difficulty : uint {
//     Easy = 1,
//     Hard = 10,
// }
//
///// Shader Usage (HLSL or Cg):
//
// #include "Assets/Shader Includes/GPUStructs.hlsl"
// 
// //   Loading this .hlsl file from a Cg shader is legal, because
// //   in this case the languages have identical syntax.
//
// //   (you can change the path by editing EXPORT_PATH, below)
//


using System;
#if UNITY_EDITOR
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Text;
using UnityEngine;
using UnityEditor;
#endif

public class CopyToGPUAttribute : Attribute { }

#if UNITY_EDITOR
public class AutoCopyToGPU {
    const string EXPORT_PATH = "Assets/Shader Includes/GPUStructs.hlsl";
    
    static HashSet<string> constNames = new HashSet<string>();
    static Dictionary<string, EnumInfo> enumLookup = new Dictionary<string, EnumInfo>();
    
    static string[] illegalFieldNames = {
        "matrix"
    };
    
    struct EnumInfo {
        public Type enumType;
        public string[] valueNames;
        public int[] intValues;
        public uint[] uintValues;
        
        public EnumInfo(Type type) {
            enumType = type;
            valueNames = Enum.GetNames(type);
            Type valueType = Enum.GetUnderlyingType(type);
            uintValues = null;
            intValues = null;
            if (valueType == typeof(int)) {
                intValues = (int[])Enum.GetValues(type);
            } else if (valueType == typeof(uint)) {
                uintValues = (uint[])Enum.GetValues(type);
            }
        }
    }
    
    [UnityEditor.Callbacks.DidReloadScripts]
    static void OnRecompile() {
        StringBuilder shaderText = new StringBuilder();
        constNames.Clear();
        
        Assembly[] assemblies = AppDomain.CurrentDomain.GetAssemblies();
        List<FieldInfo> constants = new List<FieldInfo>();
        for (int i=0; i<assemblies.Length; i++) {
            Type[] types = assemblies[i].GetTypes();
            for (int j=0; j<types.Length; j++) {
                Type type = types[j];
                
                if (type.IsEnum) {
                    if (Attribute.IsDefined(type, typeof(CopyToGPUAttribute))) {
                        RegisterEnum(type);
                    }
                } else if (type.IsValueType) {
                    if (Attribute.IsDefined(type, typeof(CopyToGPUAttribute))) {
                        WriteStruct(type, shaderText);
                    }
                }
                
                FieldInfo[] staticFields = type.GetFields(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static);
                for (int k=0; k<staticFields.Length; k++) {
                    FieldInfo field = staticFields[k];
                    if (field.IsLiteral && field.IsInitOnly == false) {
                        if (Attribute.IsDefined(field, typeof(CopyToGPUAttribute))) {
                            constants.Add(field);
                            if (constNames.Contains(field.Name) == false) {
                                constNames.Add(field.Name);
                            } else {
                                Debug.LogError($"GPU const name {field.Name} is used more than once!");
                            }
                        }
                    }
                }
            }
        }
        
        if (enumLookup.Count > 0) {
            StringBuilder enumsText = new StringBuilder();
            var enumerator = enumLookup.GetEnumerator();
            while (enumerator.MoveNext()) {
                WriteEnum(enumerator.Current.Value, enumsText);
                enumsText.Append('\n');
            }
            
            enumsText.Append(shaderText);
            shaderText = enumsText;
        }
        
        if (constants.Count > 0) {
            StringBuilder constantsText = new StringBuilder();
            for (int i=0; i<constants.Count; i++) {
                WriteConstant(constants[i], constantsText);
            }
            
            constantsText.Append('\n');
            constantsText.Append(shaderText);
            shaderText = constantsText;
        }
        
        string folderName = Path.GetDirectoryName(EXPORT_PATH);
        if (folderName != null && Directory.Exists(folderName) == false) {
            Directory.CreateDirectory(folderName);
        }
        
        string finalOutput = "// Generated by AutoCopyToGPU.cs\n\n" + shaderText;
        
        File.WriteAllText(EXPORT_PATH, finalOutput);
        AssetDatabase.ImportAsset(EXPORT_PATH);
    }
    
    static void WriteStruct(Type type, StringBuilder text) {
        //Debug.Log($"writing {type.Name}");
        
        // Maybe you want to filter out some of your struct's fields here?
        
        text.Append($"struct {type.Name} {{\n");
            FieldInfo[] fields = type.GetFields();
            for (int i=0; i<fields.Length; i++) {
                FieldInfo field = fields[i];
                if (Array.IndexOf(illegalFieldNames, field.Name) != -1) {
                    Debug.LogError($"AutoCopyToGPU:  the field name {field.Name} is illegal in shader code! (in struct {type.Name})");
                }
                string gpuType = GetGPUTypeName(field);
                text.Append($"\t{gpuType} {field.Name};\n");
            }
        text.Append("};\n\n");
    }
    
    static void WriteEnum(EnumInfo enumInfo, StringBuilder text) {
        Type valueType = Enum.GetUnderlyingType(enumInfo.enumType);
        for (int i=0; i<enumInfo.valueNames.Length; i++) {
            string value;
            string typeName;
            if (valueType == typeof(int)) {
                value = enumInfo.intValues[i].ToString();
                typeName = "int";
            } else {
                value = enumInfo.uintValues[i].ToString();
                typeName = "uint";
            }
            text.Append($"static const {typeName} {enumInfo.enumType.Name}_{enumInfo.valueNames[i]} = {value};\n");
        }
    }
    
    static void WriteConstant(FieldInfo fieldInfo, StringBuilder text) {
        //Debug.Log($"writing const {fieldInfo.Name}");
        
        string typeName = GetGPUTypeName(fieldInfo);
        object value = fieldInfo.GetRawConstantValue();
        text.Append($"static const {typeName} {fieldInfo.Name} = {value};\n");
    }
    
    static string GetGPUTypeName(FieldInfo field) {
        
        // You can add support for more field-types here
        // (as long as the GPU allows them!)
        
        Type type = field.FieldType;
        if (type.Namespace == "Unity.Mathematics") {
            return type.Name;
        } else if (type == typeof(int)) {
            return "int";
        } else if (type == typeof(float)) {
            return "float";
        } else if (type == typeof(uint)) {
            return "uint";
        } else if (type == typeof(Vector2)) {
            return "float2";
        } else if (type == typeof(Vector3)) {
            return "float3";
        } else if (type == typeof(Vector4)) {
            return "float4";
        } else if (type == typeof(Color)) {
            return "float4";
        } else if (type == typeof(Vector2Int)) {
            return "int2";
        } else if (type == typeof(Vector3Int)) {
            return "int3";
        } else if (type == typeof(Matrix4x4)) {
            return "float4x4";
        } else if (type.IsEnum) {
            RegisterEnum(type);
            Type valueType = Enum.GetUnderlyingType(type);
            if (valueType == typeof(int)) {
                return "int";
            } else if (valueType == typeof(uint)) {
                return "uint";
            } else {
                Debug.LogError($"AutoCopyToGPU does not know how to copy an enum which uses {valueType.Name} as an underlying value type (field: {field.DeclaringType}.{field.Name}");
                return "!UnknownEnumValueType!";
            }
        } else {
            Debug.LogError($"AutoCopyToGPU does not know how to copy a {type} (field: {field.DeclaringType.Name}.{field.Name})");
            return "!UnknownType!";
        }
    }
    
    static void RegisterEnum(Type enumType) {
        string enumTypeName = enumType.Name;
        if (enumLookup.TryGetValue(enumTypeName, out EnumInfo enumInfo) == false) {
            enumLookup.Add(enumTypeName, new EnumInfo(enumType));
        } else {
            if (enumInfo.enumType != enumType) {
                Debug.LogError($"More than one enum-type is being copied to the GPU with the same name: {enumTypeName}");
            }
        }
        
    }
}
#endif