Tip of the Week #4: Working with Bitflag Enums in Unity Editor

Enums are one of C#’s basic data types and represent a set of named integer constants, e.g. attack types or player states.

public enum PlayerStates
{
    None,
    IsDead,
    IsJumping,
    IsGood,
    HasCompletedTheGame,
    HasAllAchievments,
}


The unity inspector supports enums out of the box, so the previous enum would look like this. We can select exactly one of all possible options. But what if we want to select more than one option. A player can certainly be jumping and have all achievements. To do that we can use a custom property drawer and bitwise operations.

To do that we can use a custom property drawer and bitwise operations. First of all, we have to change our enum so that each value is a Flag that corresponds to one unique bit. This allows us to use this enum like a 32-bit big bitmask (the underlying type of an enum is int32 by default).

[Flags]
public enum PlayerStates
{
    None = 0,
    IsDead = 1 << 0,
    IsJumpting = 1 << 1,
    IsGood = 1 << 2,
    HasCompletedTheGame = 1 << 3,
    HasAllAchievments = 1 << 4,
}

None is a special case because it represents the case when all other flags are unset so it corresponds to zero. We can also mark our enum with the Flag attribute, which changes the ToString of this enum to correctly display a value with multiple active Flags.

If we want to visualize a field with a custom drawer we have to create a custom Attribute with which to tag the fields we want to use that drawer on. To do that we simply inherit from Unity PropteryAttribute.

using UnityEngine;

public class EnumFlagsAttribute : PropertyAttribute
{
}

Now we tag all field that should use our custom drawer with that attribute (we can omit the Attribute part of the name). It stacks with other attributes, but we cannot add it multiple times.

[EnumFlags]
[SerializeField]
private PlayerStates playerStates;

[EnumFlagsAttribute]
public DamageTypes DamageType;

Nothing has changed yet because we haven’t yet created a custom Property Drawer that corresponds to our attribute so the default enum drawer is used. To this custom drawer, we have to create a class inside an Editor folder that inherits from  ProptertyDrawer and is a CustomPorpertyDrawer of the type EnumFlagsAttribute.

using UnityEditor;
using UnityEngine;
[CustomPropertyDrawer(typeof(EnumFlagsAttribute))]
public class EnumFlagsAttributeDrawer : PropertyDrawer
{
}

The two important methods we need to override are
float GetPropertyHeight(SerializedProperty property, GUIContent label))
and
void OnGUI(Rect position, SerializedProperty property, GUIContent label).

I delegate these methods to an internal ReorderableList that also handles adding and removing actions. If you are interested in how ReorderableLists work, check out this blog post by Valentin Simonov.

The full EnumFlagAttributeDrawer in action looks like this. Please note that the PlayerStates enum has a None value, while the DamageTypes enum has no None value and therefore can be in a state that corresponds to no value.

Full source code:

using System;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEditorInternal;
using UnityEngine;

[CustomPropertyDrawer(typeof(EnumFlagsAttribute))]
public class EnumFlagsAttributeDrawer : PropertyDrawer
{
    private string[] enumNames;
    private readonly Dictionary<string, int> enumNameToValue = new Dictionary<string, int>();
    private readonly Dictionary<string, string> enumNameToDisplayName = new Dictionary<string, string>();
    private readonly Dictionary<string, string> enumNameToTooltip = new Dictionary<string, string>();
    private readonly List<string> activeEnumNames = new List<string>();

    private SerializedProperty serializedProperty;
    private ReorderableList reorderableList;

    private bool firstTime = true;

    private Type EnumType
    {
        get { return fieldInfo.FieldType; }
    }

    public override float GetPropertyHeight(SerializedProperty property, GUIContent label)
    {
        serializedProperty = property;
        SetupIfFirstTime();
        return reorderableList.GetHeight();
    }

    private void SetupIfFirstTime()
    {
        if (!firstTime)
        {
            return;
        }

        enumNames = serializedProperty.enumNames;

        CacheEnumMetadata();
        ParseActiveEnumNames();

        reorderableList = GenerateReorderableList();
        firstTime = false;
    }

    private void CacheEnumMetadata()
    {
        for (var index = 0; index < enumNames.Length; index++)
        {
            enumNameToDisplayName[enumNames[index]] = serializedProperty.enumDisplayNames[index];
        }

        foreach (string enumName in enumNames)
        {
            enumNameToTooltip[enumName] = EnumType.Name + "." + enumName;
        }

        foreach (string name in enumNames)
        {
            enumNameToValue.Add(name, (int)Enum.Parse(EnumType, name));
        }
    }

    public override void OnGUI(Rect position, SerializedProperty property, GUIContent label)
    {
        EditorGUI.BeginDisabledGroup(serializedProperty.hasMultipleDifferentValues);
        reorderableList.DoList(position);
        EditorGUI.EndDisabledGroup();
    }

    private ReorderableList GenerateReorderableList()
    {
        return new ReorderableList(activeEnumNames, typeof(string), false, true, true, true)
        {

            drawHeaderCallback = rect =>
            {
                EditorGUI.LabelField(rect, new GUIContent(serializedProperty.displayName, "EnumType: " + EnumType.Name));
            },
            drawElementCallback = (rect, index, isActive, isFocused) =>
            {
                rect.y += 2;
                EditorGUI.LabelField(
                    new Rect(rect.x, rect.y, rect.width, EditorGUIUtility.singleLineHeight),
                    new GUIContent(enumNameToDisplayName[activeEnumNames[index]], enumNameToTooltip[activeEnumNames[index]]),
                    EditorStyles.label);

            },
            onAddDropdownCallback = (Rect buttonRect, ReorderableList l) =>
            {
                var menu = new GenericMenu();
                foreach (string enumName in enumNames)
                {
                    if (activeEnumNames.Contains(enumName) == false)
                    {
                        menu.AddItem(new GUIContent(enumNameToDisplayName[enumName]),
                            false, data =>
                            {
                                if (enumNameToValue[(string)data] == 0)
                                {
                                    activeEnumNames.Clear();
                                }
                                activeEnumNames.Add((string)data);
                                SaveActiveValues();
                                ParseActiveEnumNames();
                            },
                            enumName);
                    }
                }
                menu.ShowAsContext();
            },
            onRemoveCallback = l =>
            {
                ReorderableList.defaultBehaviours.DoRemoveButton(l);
                SaveActiveValues();
                ParseActiveEnumNames();
            }
        };
    }


    private void ParseActiveEnumNames()
    {
        activeEnumNames.Clear();
        foreach (string enumValue in enumNames)
        {
            if (IsFlagSet(enumValue))
            {
                activeEnumNames.Add(enumValue);
            }
        }
    }

    private bool IsFlagSet(string enumValue)
    {
        if (enumNameToValue[enumValue] == 0)
        {
            return serializedProperty.intValue == 0;
        }
        return (serializedProperty.intValue & enumNameToValue[enumValue]) == enumNameToValue[enumValue];
    }

    private void SaveActiveValues()
    {
        serializedProperty.intValue = ConvertActiveNamesToInt();
        serializedProperty.serializedObject.ApplyModifiedProperties();
    }

    private int ConvertActiveNamesToInt()
    {
        return activeEnumNames.Aggregate(0, (current, activeEnumName) => current | enumNameToValue[activeEnumName]);
    }

}