From 9c648c675f28e8d5f0bb9e195c0975267bdb436b Mon Sep 17 00:00:00 2001
From: mob-sakai <12690315+mob-sakai@users.noreply.github.com>
Date: Sat, 14 Dec 2024 21:18:20 +0900
Subject: [PATCH] fix: TextMeshPro objects appeared as black blocks when saving
 prefabs in prefab mode

close #285
---
 Packages/src/Runtime/UIEffectBase.cs | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/Packages/src/Runtime/UIEffectBase.cs b/Packages/src/Runtime/UIEffectBase.cs
index 608e677d..cc4c5419 100644
--- a/Packages/src/Runtime/UIEffectBase.cs
+++ b/Packages/src/Runtime/UIEffectBase.cs
@@ -1,4 +1,6 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using Coffee.UIEffectInternal;
 using UnityEngine;
@@ -9,6 +11,9 @@
 #if TMP_ENABLE
 using TMPro;
 #endif
+#if UNITY_EDITOR
+using UnityEditor.SceneManagement;
+#endif
 
 [assembly: InternalsVisibleTo("UIEffect")]
 [assembly: InternalsVisibleTo("Coffee.UIEffect.Editor")]
@@ -246,10 +251,11 @@ private static void OnPostprocessAllAssets(string[] _, string[] __, string[] ___
 
                 s_ShaderNameCache.Clear();
 #if UNITY_2021_3 || UNITY_2022_2_OR_NEWER
-                foreach (var effect in FindObjectsByType<UIEffectBase>(FindObjectsSortMode.None))
+                foreach (var effect in FindObjectsByType<UIEffectBase>(FindObjectsSortMode.None)
 #else
-                foreach (var effect in FindObjectsOfType<UIEffectBase>())
+                foreach (var effect in FindObjectsOfType<UIEffectBase>()
 #endif
+                             .Concat(GetAllComponentsInPrefabStage<UIEffectBase>()))
                 {
                     if (!effect.isActiveAndEnabled) continue;
                     if (!(effect.graphic is TextMeshProUGUI tmp) || !tmp.isActiveAndEnabled) continue;
@@ -259,6 +265,16 @@ private static void OnPostprocessAllAssets(string[] _, string[] __, string[] ___
                 EditorApplication.QueuePlayerLoopUpdate();
             }
         }
+
+        private static T[] GetAllComponentsInPrefabStage<T>() where T : Component
+        {
+            if (!PrefabStageUtility.GetCurrentPrefabStage()) return Array.Empty<T>();
+
+            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+            if (!prefabStage) return Array.Empty<T>();
+
+            return prefabStage.prefabContentsRoot.GetComponentsInChildren<T>(true);
+        }
 #endif
 
 #if UNITY_EDITOR