Skip to content

Instantly share code, notes, and snippets.

@2DArray
Last active October 12, 2020 13:56
Show Gist options
  • Save 2DArray/1567899b3a75b7beefc17bbf2a1332f5 to your computer and use it in GitHub Desktop.
Save 2DArray/1567899b3a75b7beefc17bbf2a1332f5 to your computer and use it in GitHub Desktop.
[CopyToGPU] Attribute for structs and consts
// 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment