Fix some prefabs issues

Source/Editor/Editor.cs

@@ -254,7 +254,7 @@ namespace FlaxEditor
             if (loadingPreview != null)
             {
                 // Link it to the prefab preview to see it in the editor
-                loadingPreview.customControlLinked = control.Control;
+                loadingPreview.customControlLinked = control;
                 return loadingPreview;
             }
             return null;

Source/Editor/Viewport/Previews/PrefabPreview.cs

@@ -1,7 +1,6 @@
 // Copyright (c) 2012-2021 Wojciech Figat. All rights reserved.
 
 using FlaxEngine;
-using FlaxEngine.GUI;
 using Object = FlaxEngine.Object;
 
 namespace FlaxEditor.Viewport.Previews
@@ -19,7 +18,7 @@ namespace FlaxEditor.Viewport.Previews
 
         private Prefab _prefab;
         private Actor _instance;
-        internal Control customControlLinked;
+        internal UIControl customControlLinked;
 
         /// <summary>
         /// Gets or sets the prefab asset to preview.
@@ -29,16 +28,23 @@ namespace FlaxEditor.Viewport.Previews
             get => _prefab;
             set
             {
-                if (_prefab != value)
-                {
+                if (_prefab == value)
+                    return;
+
                 if (_instance)
                 {
-                        if (customControlLinked != null)
+                    // Unlink UI control
+                    if (customControlLinked)
                     {
-                            customControlLinked.Parent = null;
+                        if (customControlLinked.Control?.Parent == this)
+                            customControlLinked.Control.Parent = null;
                         customControlLinked = null;
                     }
+
+                    // Remove for preview
                     Task.RemoveCustomActor(_instance);
+
+                    // Delete
                     Object.Destroy(_instance);
                 }
 
@@ -46,7 +52,7 @@ namespace FlaxEditor.Viewport.Previews
 
                 if (_prefab)
                 {
-                        _prefab.WaitForLoaded(); // TODO: use lazy prefab spawning to reduce stalls
+                    _prefab.WaitForLoaded();
 
                     var prevPreview = LoadingPreview;
                     LoadingPreview = this;
@@ -60,11 +66,12 @@ namespace FlaxEditor.Viewport.Previews
                         _prefab = null;
                         throw new FlaxException("Failed to spawn a prefab for the preview.");
                     }
+
+                    // Add to preview
                     Task.AddCustomActor(_instance);
                 }
             }
         }
-        }
 
         /// <summary>
         /// Gets the instance of the prefab spawned for the preview.
@@ -79,11 +86,15 @@ namespace FlaxEditor.Viewport.Previews
 
                 if (_instance)
                 {
-                    if (customControlLinked != null)
+                    // Unlink UI control
+                    if (customControlLinked)
                     {
-                        customControlLinked.Parent = null;
+                        if (customControlLinked.Control?.Parent == this)
+                            customControlLinked.Control.Parent = null;
                         customControlLinked = null;
                     }
+
+                    // Remove for preview
                     Task.RemoveCustomActor(_instance);
                 }
 
@@ -91,6 +102,7 @@ namespace FlaxEditor.Viewport.Previews
 
                 if (_instance)
                 {
+                    // Add to preview
                     Task.AddCustomActor(_instance);
                 }
             }
@@ -108,7 +120,6 @@ namespace FlaxEditor.Viewport.Previews
         /// <inheritdoc />
         public override void OnDestroy()
         {
-            // Cleanup
             Prefab = null;
 
             base.OnDestroy();

Source/Editor/Windows/Assets/PrefabWindow.Hierarchy.cs

@@ -69,6 +69,7 @@ namespace FlaxEditor.Windows.Assets
             bool hasSthSelected = Selection.Count > 0;
             bool isSingleActorSelected = Selection.Count == 1 && Selection[0] is ActorNode;
             bool isRootSelected = isSingleActorSelected && Selection[0] == Graph.Main;
+            bool hasPrefabLink = isSingleActorSelected && (Selection[0] as ActorNode).HasPrefabLink;
 
             // Create popup
 
@@ -97,7 +98,7 @@ namespace FlaxEditor.Windows.Assets
             b.Enabled = hasSthSelected && !isRootSelected;
 
             b = contextMenu.AddButton("Set Root", SetRoot);
-            b.Enabled = isSingleActorSelected && !isRootSelected;
+            b.Enabled = isSingleActorSelected && !isRootSelected && hasPrefabLink;
 
             // Prefab options
 
@@ -108,8 +109,6 @@ namespace FlaxEditor.Windows.Assets
                         (Selection[0] as ActorNode).CanCreatePrefab &&
                         Editor.Windows.ContentWin.CurrentViewFolder.CanHaveAssets;
 
-            bool hasPrefabLink = isSingleActorSelected && (Selection[0] as ActorNode).HasPrefabLink;
-
             b = contextMenu.AddButton("Select Prefab", Editor.Prefabs.SelectPrefab);
             b.Enabled = hasPrefabLink;