RebasePrefab - Unity editor tool to convert a prefab into a variant of a new empty prefab

RebasePrefab.cs

/*
 * @brief Editor menu item for turning a prefab into a variant of a new empty prefab.
 *
 * @usage Access from the MenuItem defined in the code. As right-click or from the menu bar.
 * 
 * @author Aurelio Provedo
 * Contact: [email protected]
 *
 */

using System;
using System.IO;
using System.Text.RegularExpressions;
using UnityEngine;
using UnityEditor;

namespace kraj0tEditor
{
    public static class RebasePrefab
    {
        private const string BASE_FILENAME_SUFFIX = " Base";
            
        private const string TEMPLATE_FILENAME_ROOT = "RebasePrefab_Template_";

        private const string TEMPLATE_ORIGINAL_ROOT_GAMEOBJECT_NAME = "ORIGINAL_ROOT_GAMEOBJECT_NAME";
        private const string TEMPLATE_ORIGINAL_ROOT_GAMEOBJECT_FILEID = "ORIGINAL_ROOT_GAMEOBJECT_FILEID";
        private const string TEMPLATE_ORIGINAL_TRANSFORM_FILEID = "ORIGINAL_TRANSFORM_FILEID";
        
        private const string TEMPLATE_BASE_GUID = "BASE_GUID";
        private const string TEMPLATE_BASE_ROOT_GAMEOBJECT_FILEID = "BASE_ROOT_GAMEOBJECT_FILEID";
        private const string TEMPLATE_BASE_TRANSFORM_FILEID = "BASE_TRANSFORM_FILEID";
        
        private const string TEMPLATE_PREFAB_INSTANCE_FILEID = "PREFAB_INSTANCE_FILEID";
        
        [MenuItem("kraj0t/Rebase Selected Prefab (Create Empty Base)", true, -600)]
        [MenuItem("Assets/kraj0t/Rebase Prefab (Create Empty Base)", true, -600)]
        private static bool RebaseSelectedPrefabValidation()
        {
            return IsSinglePrefabSelected();
        }

        public static bool IsSinglePrefabSelected()
        {
            if (Selection.objects.Length != 1 || Selection.assetGUIDs.Length != 1)
                return false;
            
            var go = Selection.activeObject as GameObject;
            if (go == null)
                return false;
              
            var type = PrefabUtility.GetPrefabAssetType(go);
            return ((type == PrefabAssetType.Regular || type == PrefabAssetType.Variant) &&
                    (PrefabUtility.IsPartOfRegularPrefab(go) || PrefabUtility.IsPartOfVariantPrefab(go)) &&
                    PrefabUtility.IsPartOfPrefabAsset(go) && PrefabUtility.IsPartOfPrefabThatCanBeAppliedTo(go) &&
                    EditorUtility.IsPersistent(go));
        }

        [MenuItem("kraj0t/Rebase Selected Prefab (Create Empty Base)", false, -600)]
        [MenuItem("Assets/kraj0t/Rebase Prefab (Create Empty Base)", false, -600)]
        private static void RebasePrefabSelection()
        {
            // TODO: after success, explain to the user that they must now apply the overrides to the base as they see fit.

            // No need to check for valid input here. The validate method will do that.

            var go = Selection.activeObject as GameObject;

            RebasePrefabToNewPrefab(go, out var newVariantPrefab, out var newBasePrefab);

            if (newBasePrefab == null)
            {
                EditorGUIFunctions.DisplayMultiNotification(null, "Error creating a base prefab for " + go.name, 
                    LogType.Error, go, "Error rebasing prefab", true);
            }
            else
            {
                // TODO: fix selection not being handled correctly.
                EditorGUIFunctions.DisplayMultiNotification(new []{newVariantPrefab},go.name + " is now a variant of " + newBasePrefab.name, 
                    LogType.Log, newVariantPrefab, "Prefab rebased", true);
            }
        }
        
        /// <summary>
        /// Use this to make a prefab become a variant of a new empty prefab.
        /// In other words, a new prefab will become the base of the provided prefab.
        /// References across the whole project will be preserved. 
        /// Important: this function will fail if a prefab instance is passed.
        /// </summary>
        /// <param name="originalPrefab">The prefab asset that will become a variant. This reference will be null after this function returns. See newVariantPrefab for the new reference.</param>
        /// <param name="newVariantPrefab">The original prefab that will now be a variant of the new base.</param>
        /// <param name="newBasePrefab">The newly-created empty prefab that is the base of the original one.</param>
        public static void RebasePrefabToNewPrefab(GameObject originalPrefab, out GameObject newVariantPrefab, out GameObject newBasePrefab)
        {
            var origFilePath = PrefabUtility.GetPrefabAssetPathOfNearestInstanceRoot(originalPrefab);

            var origGO = originalPrefab;
            var origTransform = origGO.transform;
            
            // Before making any modification, we will instantiate a copy of the original prefab in the scene.
            // This will allow us to access the original values later.
            var tempOrigGOCopy = GameObject.Instantiate(originalPrefab);
            
            // Get ids of the root GameObject and its Transform in the original prefab file.
            string origGuidUnused;
            long origGOFileID;
            long origTransformFileID;
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(origGO, out origGuidUnused, out origGOFileID);
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(origTransform, out origGuidUnused, out origTransformFileID);
            
            var baseFilePath = origFilePath.Insert(origFilePath.Length - 7, BASE_FILENAME_SUFFIX);
            
            // Create the empty prefab that will become the base of the original one.
            // A variant prefab cannot override the type of its base's Transform. This is a Unity limitation.
            // Example: if the base has a Transform, the variant cannot change it to a RectTransform.
            // Therefore, the type of Transform must be the same in the base.
            var origTransformType = origTransform.GetType(); 
            var baseGOTempInstance = new GameObject(originalPrefab.name + " Base");
            if (baseGOTempInstance.transform.GetType() != origTransformType)
            {
                baseGOTempInstance.AddComponent(origTransformType);
            }
            var baseGO = PrefabUtility.SaveAsPrefabAsset(baseGOTempInstance, baseFilePath); 
            var baseTransform = baseGO.transform;
            GameObject.DestroyImmediate(baseGOTempInstance, false);
            
            // Get ids of the root GameObject and its Transform in the newly created empty prefab file.
            string baseGuid;
            long baseGOFileID;
            long baseTransformFileID;
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(baseGO, out baseGuid, out baseGOFileID);
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(baseTransform, out baseGuid, out baseTransformFileID);
            
            // TODO: optimize the whole read-write operation. StreamReader/Writer should be faster.
            
            var origYAMLText = File.ReadAllText(origFilePath);
            var templateYAMLPath = FindRebasePrefabYAMLTemplateByTransformType(origTransformType);
            var templateYAMLText = File.ReadAllText(templateYAMLPath);
            
            // We need to assign a value for the PrefabInstance fileID.
            // However, we must ensure that the fileID was not being used already in the original prefab.
            // That would be a huge coincidence, but not an impossibility.
            var prefabInstanceFileID = YAMLFunctions.GenerateUniqueFileIDForAssetYAML(origYAMLText, 1)[0];
                
            // Remove the original YAML blocks for the root GameObject and its Transform.
            // Both of them will now be defined in the base prefab, and the variant will only reference them.
            // That is what the "stripped" blocks in the template are for.
            var rebasedYAMLText = origYAMLText + templateYAMLText;
            rebasedYAMLText = RemoveYAMLObjectDefinition(rebasedYAMLText, origGOFileID);
            rebasedYAMLText = RemoveYAMLObjectDefinition(rebasedYAMLText, origTransformFileID);
                
            // Replace the fileIDs from the template with the actual ones from the prefabs.
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_BASE_GUID, baseGuid);
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_ORIGINAL_ROOT_GAMEOBJECT_NAME, origGO.name);
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_ORIGINAL_ROOT_GAMEOBJECT_FILEID, origGOFileID.ToString());
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_ORIGINAL_TRANSFORM_FILEID, origTransformFileID.ToString());
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_BASE_ROOT_GAMEOBJECT_FILEID, baseGOFileID.ToString());
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_BASE_TRANSFORM_FILEID, baseTransformFileID.ToString());
            rebasedYAMLText = rebasedYAMLText.Replace(TEMPLATE_PREFAB_INSTANCE_FILEID, prefabInstanceFileID.ToString());
        
            // Save the prefab's YAML content and tell Unity to refresh it.
            AssetDatabase.StartAssetEditing();
            File.WriteAllText(origFilePath, rebasedYAMLText);
            AssetDatabase.StopAssetEditing();
            AssetDatabase.ImportAsset(baseFilePath);
            AssetDatabase.ImportAsset(origFilePath);
            
            // Assign the return values.
            newBasePrefab = AssetDatabase.LoadAssetAtPath<GameObject>(baseFilePath);
            newVariantPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(origFilePath);

            // Copy the original values of the transform from the temporary instance we created at the start of this call.
            // The YAML templates already contain the default overrides for the transform properties. It is Unity's
            // normal behaviour to add all default overrides to a newly created variant, so we do the same by
            // modifying the default overrides with the actual original values here.
            ComponentFunctions.CopyComponentUsingReflection(newVariantPrefab.transform, tempOrigGOCopy.transform,
                new[] {"parent", "parentInternal", "name"}, new[] {"m_Father", "m_Name"});

            // Destroy the temporary instance.
            GameObject.DestroyImmediate(tempOrigGOCopy);

            // Apply the changes to the root's transform values.
            newVariantPrefab = PrefabUtility.SavePrefabAsset(newVariantPrefab);

            // After saving the YAML text directly, Unity modifies the fileIDs of the root gameobject and its transform.
            // This was found empirically. An educated guess would say that Unity detects that the prefab has turned into a variant, and
            // therefore it needs new fileIDs for its stripped objects.
            // So, here we need force our desired fileIDs a second time.
            var secondPassYAMLText = File.ReadAllText(origFilePath);
            long unityGeneratedGOFileID;
            long unityGeneratedTransformFileID;
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(newVariantPrefab, out origGuidUnused, out unityGeneratedGOFileID);
            AssetDatabase.TryGetGUIDAndLocalFileIdentifier(newVariantPrefab.transform, out origGuidUnused, out unityGeneratedTransformFileID);
            secondPassYAMLText = secondPassYAMLText.Replace(unityGeneratedGOFileID.ToString(), origGOFileID.ToString());
            secondPassYAMLText = secondPassYAMLText.Replace(unityGeneratedTransformFileID.ToString(), origTransformFileID.ToString());
            AssetDatabase.StartAssetEditing();
            File.WriteAllText(origFilePath, secondPassYAMLText);
            AssetDatabase.StopAssetEditing();            
            AssetDatabase.ImportAsset(origFilePath);
            newVariantPrefab = AssetDatabase.LoadAssetAtPath<GameObject>(origFilePath);
        }

        private static string FindRebasePrefabYAMLTemplateByTransformType(Type origTransformType)
        {
            var templateFullFilename = TEMPLATE_FILENAME_ROOT + origTransformType.Name;
            var foundTextAssets = AssetDatabase.FindAssets("t:TextAsset " + templateFullFilename);
            if (foundTextAssets.Length == 0)
            {
                throw new FileLoadException(
                    "No template file was found with the YAML text of rebased prefabs! Please reinstall the RebasePrefab assets.");
            }
            else if (foundTextAssets.Length > 1)
            {
                throw new FileLoadException(
                    "More than one template file was was found with the YAML text of rebased prefabs! The type of the searched Transform was \"" +
                    origTransformType.Name + "\". Please reinstall the RebasePrefab assets.");
            }

            return AssetDatabase.GUIDToAssetPath(foundTextAssets[0]);
        }

        private static string RemoveYAMLObjectDefinition(string YAMLText, long fileID)
        {
            // TODO: give this a second thought. Is this the best way to find the template files? At least, use a better regex that does everything in one go.

            var objectBlockStartRegexPattern = "--- !u!(.*)" + fileID;
            var startMatch = Regex.Match(YAMLText, objectBlockStartRegexPattern);
            var fromStart = YAMLText.Substring(startMatch.Index);

            var nextObjectRegexPattern = "--- !u!(.*) &(?!" + fileID + ")";
            var endMatch = Regex.Match(fromStart, nextObjectRegexPattern);

            var fromEnd = fromStart.Substring(endMatch.Index);
            var untilStart = YAMLText.Substring(0, startMatch.Index);

            return untilStart + fromEnd;
        }
    }
}

RebasePrefab_Template_RectTransform.txt

--- !u!1001 &PREFAB_INSTANCE_FILEID
PrefabInstance:
  m_ObjectHideFlags: 0
  serializedVersion: 2
  m_Modification:
    m_TransformParent: {fileID: 0}
    m_Modifications:
    - target: {fileID: BASE_ROOT_GAMEOBJECT_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_Name
      value: ORIGINAL_ROOT_GAMEOBJECT_NAME
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.z
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.z
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.w
      value: 1
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_RootOrder
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.z
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchoredPosition.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchoredPosition.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_SizeDelta.x
      value: 100
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_SizeDelta.y
      value: 100
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchorMin.x
      value: 0.5
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchorMin.y
      value: 0.5
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchorMax.x
      value: 0.5
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_AnchorMax.y
      value: 0.5
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_Pivot.x
      value: 0.5
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_Pivot.y
      value: 0.5
      objectReference: {fileID: 0}
    m_RemovedComponents: []
  m_SourcePrefab: {fileID: 100100000, guid: BASE_GUID, type: 3}
--- !u!1 &ORIGINAL_ROOT_GAMEOBJECT_FILEID stripped
GameObject:
  m_CorrespondingSourceObject: {fileID: BASE_ROOT_GAMEOBJECT_FILEID, guid: BASE_GUID,
    type: 3}
  m_PrefabInstance: {fileID: PREFAB_INSTANCE_FILEID}
  m_PrefabAsset: {fileID: 0}
--- !u!224 &ORIGINAL_TRANSFORM_FILEID stripped
RectTransform:
  m_CorrespondingSourceObject: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
    type: 3}
  m_PrefabInstance: {fileID: PREFAB_INSTANCE_FILEID}
  m_PrefabAsset: {fileID: 0}

RebasePrefab_Template_Transform.txt

--- !u!1001 &PREFAB_INSTANCE_FILEID
PrefabInstance:
  m_ObjectHideFlags: 0
  serializedVersion: 2
  m_Modification:
    m_TransformParent: {fileID: 0}
    m_Modifications:
    - target: {fileID: BASE_ROOT_GAMEOBJECT_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_Name
      value: ORIGINAL_ROOT_GAMEOBJECT_NAME
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalPosition.z
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.z
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalRotation.w
      value: 1
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_RootOrder
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.x
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.y
      value: 0
      objectReference: {fileID: 0}
    - target: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
        type: 3}
      propertyPath: m_LocalEulerAnglesHint.z
      value: 0
      objectReference: {fileID: 0}
    m_RemovedComponents: []
  m_SourcePrefab: {fileID: 100100000, guid: BASE_GUID, type: 3}
--- !u!1 &ORIGINAL_ROOT_GAMEOBJECT_FILEID stripped
GameObject:
  m_CorrespondingSourceObject: {fileID: BASE_ROOT_GAMEOBJECT_FILEID, guid: BASE_GUID,
    type: 3}
  m_PrefabInstance: {fileID: PREFAB_INSTANCE_FILEID}
  m_PrefabAsset: {fileID: 0}
--- !u!4 &ORIGINAL_TRANSFORM_FILEID stripped
Transform:
  m_CorrespondingSourceObject: {fileID: BASE_TRANSFORM_FILEID, guid: BASE_GUID,
    type: 3}
  m_PrefabInstance: {fileID: PREFAB_INSTANCE_FILEID}
  m_PrefabAsset: {fileID: 0}