AtlasPrefabReference.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.U2D;

[CreateAssetMenu(menuName = "ScriptableObjects/AtlasPrefabReference")]
public class AtlasPrefabReference : ScriptableObject
{
    [SerializeField] public SpriteAtlas Atlas;
    [SerializeField] public List<GameObject> Prefabs;
}

AtlasPrefabReferenceUpdater.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEngine;
using UnityEngine.UI;

[InitializeOnLoad]
public class AtlasPrefabReferenceUpdater : UnityEditor.AssetModificationProcessor
{
    static AtlasPrefabReferenceUpdater()
    {
        UpdateAtlasReferencesDelayed(GetAtlasPrefabReferences());
        PrefabUtility.prefabInstanceUpdated += PrefabInstanceUpdated;
    }

    private static void PrefabInstanceUpdated(GameObject instance)
    {
        var prefab = PrefabUtility.GetPrefabParent(instance) as GameObject;
        if (!prefab) { return; }

        var atlases = GetAtlasPrefabReferences().Where(r => r.Prefabs.Contains(prefab));
        UpdateAtlasReferencesDelayed(atlases);
    }

    private static AssetDeleteResult OnWillDeleteAsset(string assetPath, RemoveAssetOptions options)
    {
        var deletedSprite = AssetDatabase.LoadAssetAtPath<Sprite>(assetPath);
        if (!deletedSprite) { return AssetDeleteResult.DidNotDelete; }

        var atlases = GetAtlasPrefabReferences().Where(r => r.Atlas.ContainsSprite(deletedSprite));
        UpdateAtlasReferencesDelayed(atlases);

        return AssetDeleteResult.DidNotDelete;
    }

    private static string[] OnWillSaveAssets(string[] paths)
    {
        var atlases = paths.Select(AssetDatabase.LoadAssetAtPath<AtlasPrefabReference>);
        UpdateAtlasReferencesDelayed(atlases);
        return paths;
    }

    private static IEnumerable<T> FindAssets<T>() where T : Object
    {
        var type = typeof(T).FullName;
        foreach (var assetGuid in AssetDatabase.FindAssets("t:" + type))
        {
            var obj = AssetDatabase.LoadAssetAtPath<T>(AssetDatabase.GUIDToAssetPath(assetGuid));
            if (obj != null)
            {
                yield return obj;
            }
        }
    }

    private static IEnumerable<AtlasPrefabReference> GetAtlasPrefabReferences()
    {
        foreach (var reference in FindAssets<AtlasPrefabReference>())
        {
            if (reference.Atlas == null)
            {
                Debug.LogWarning("No sprite atlas referenced: " + reference.name);
                continue;
            }
            if (reference.Prefabs.Any(o => o == null))
            {
                Debug.LogWarning("One or more prefab references are null: " + reference.name);
                continue;
            }
            yield return reference;
        }
    }

    public static void UpdateAtlasReferencesDelayed(IEnumerable<AtlasPrefabReference> references)
    {
        foreach (var reference in references)
        {
            if (!reference) { continue; }
            var localReference = reference;
            EditorApplication.delayCall += () => UpdateAtlasReferences(localReference);
        }
    }

    private static IEnumerable<Sprite> GetDistinctSprites(IEnumerable<GameObject> prefabs)
    {
        var images = prefabs.SelectMany(p => p.GetComponentsInChildren<Image>());
        return images.Select(i => i.sprite).Where(i => i != null).Distinct();
    }

    public static void UpdateAtlasReferences(AtlasPrefabReference reference)
    {
        var sprites = GetDistinctSprites(reference.Prefabs);
        reference.Atlas.SetSprites(sprites.ToList());
    }
}

SerializedPropertyExtensions.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.
using System.Collections.Generic;
using UnityEditor;
using UnityEngine;

public static class SerializedPropertyExtensions
{
    public static void SetObjects<T>(this SerializedProperty property, IList<T> objects) where T : Object
    {
        property.ClearArray();
        for (var i = 0; i < objects.Count; i++)
        {
            property.InsertArrayElementAtIndex(i);
            property.GetArrayElementAtIndex(i).objectReferenceValue = objects[i];
        }
    }
}

SpriteAtlasExtensions.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.
using System.Collections.Generic;
using UnityEditor;
using UnityEngine;
using UnityEngine.U2D;

public static class SpriteAtlasExtensions
{
    public const string SpritePackables = "m_EditorData.packables";

    public static void SetSprites(this SpriteAtlas spriteAtlas, IList<Sprite> sprites)
    {
        var serializedObject = new SerializedObject(spriteAtlas);
        var packables = serializedObject.FindProperty(SpritePackables);
        packables.SetObjects(sprites);
        serializedObject.ApplyModifiedProperties();
    }

    public static bool ContainsSprite(this SpriteAtlas spriteAtlas, Sprite sprite)
    {
        var serializedObject = new SerializedObject(spriteAtlas);
        var packables = serializedObject.FindProperty(SpritePackables);
        for (var i = 0; i < packables.arraySize; i++)
        {
            var containedSprite = packables.GetArrayElementAtIndex(i).objectReferenceValue as Sprite;
            if (sprite == containedSprite)
            {
                return true;
            }
        }
        return false;
    }
}