A PropertyDrawer to assign one prefab per enum value

EnumToPrefab.cs

using System;

using UnityEngine;
using UnityEditor;

/*
* A custom PropertyDrawer to assign one prefab to each value in an enum.
* Makes for less error-prone editing, avoids a gazillion asserts at runtime,
* and even looks nice.
* Example usage in a MonoBehaviour:
* [EnumMapping(typeof(UIState))]
* public EnumMappingArray StateMappings;
*/

namespace EditorUtils
{
	// Unity won't let us define property drawers for arrays, so we have to use a workaround
	[Serializable]
	public struct EnumMappingArray
	{
		public GameObject[] data;
		
		// Prevents having to use .data when accessing values
		public GameObject this[int i]
		{
			get { return data[i]; }
		}
	}

	public class EnumMappingAttribute : PropertyAttribute
	{
		public Type enumType;

		public EnumMappingAttribute(Type enumType)
		{
			Debug.Assert(enumType.IsEnum);
			this.enumType = enumType;
		}
	}

	/*
	*	Example enum:
	*	public enum UIState
	*	{
	*		MainMenu,
	*		Hud,
	*		InGameMenu,
	*		Inventory,
	*		Count,
	*		None
	*	}
	*/

	[CustomPropertyDrawer(typeof(EnumMappingAttribute))]
	public class EnumToPrefabDrawer : PropertyDrawer
	{
		public override void OnGUI(Rect position, SerializedProperty property, GUIContent label)
		{
			Type enumType = ((EnumMappingAttribute)attribute).enumType;
			Debug.Assert(enumType.IsEnum);
			string[] enumNames = enumType.GetEnumNames();
			// Make sure the enum follows the format we expect
			Debug.Assert(enumNames[enumNames.Length - 2] == "Count" && enumNames[enumNames.Length - 1] ==  "None");

			Debug.Assert(property.type == typeof(EnumMappingArray).Name);
			SerializedProperty actualArray = property.FindPropertyRelative("data");
			Debug.Assert(actualArray != null && actualArray.isArray);
			if (actualArray.arraySize == 0)
			{
				// The array has not been populated yet
				int i = 0;
				foreach (string name in enumNames)
				{
					if (i >= enumNames.Length - 2)
						break;
					Debug.Assert(name != "Count" && name != "None");
					actualArray.InsertArrayElementAtIndex(i);
					SerializedProperty newProperty = actualArray.GetArrayElementAtIndex(i);
					i++;
				}
			}

			Rect newPosition = EditorGUI.PrefixLabel(position, GUIUtility.GetControlID(FocusType.Passive), label);
			EditorGUI.LabelField(newPosition, enumType.Name);
			EditorGUI.indentLevel++;
			for (int i = 0; i < actualArray.arraySize; ++i)
			{
				Rect row = new Rect(position.x, position.y + (i+1) * EditorGUIUtility.singleLineHeight + i * EditorGUIUtility.standardVerticalSpacing, position.width, EditorGUIUtility.singleLineHeight);
				row = EditorGUI.PrefixLabel(row, GUIUtility.GetControlID(FocusType.Passive), new GUIContent(enumType.GetEnumName(i)));
				EditorGUI.indentLevel--;
				EditorGUI.PropertyField(row, actualArray.GetArrayElementAtIndex(i), GUIContent.none);
				EditorGUI.indentLevel++;
			}
			EditorGUI.indentLevel--;
		}

		public override float GetPropertyHeight(SerializedProperty property, GUIContent label)
		{
			Type enumType = ((EnumMappingAttribute)attribute).enumType;
			Debug.Assert(enumType.IsEnum);
			// All relevant enums have Count and None fields that we ignore
			int arraySize = enumType.GetEnumNames().Length - 2;
			// Label + N lines
			return (arraySize + 1) * EditorGUIUtility.singleLineHeight + arraySize * EditorGUIUtility.standardVerticalSpacing;
		}
	}
}