Refactor Editor gizmo modes ownership to support using them in prefab window

Source/Editor/Gizmo/GizmosCollection.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using FlaxEditor.Viewport.Modes;
 using FlaxEngine;
 
 namespace FlaxEditor.Gizmo
@@ -13,7 +14,10 @@ namespace FlaxEditor.Gizmo
     [HideInEditor]
     public class GizmosCollection : List<GizmoBase>
     {
+        private IGizmoOwner _owner;
         private GizmoBase _active;
+        private EditorGizmoMode _activeMode;
+        private readonly List<EditorGizmoMode> _modes = new List<EditorGizmoMode>();
 
         /// <summary>
         /// Occurs when active gizmo tool gets changed.
@@ -31,7 +35,7 @@ namespace FlaxEditor.Gizmo
                 if (_active == value)
                     return;
                 if (value != null && !Contains(value))
-                    throw new InvalidOperationException("Invalid Gizmo.");
+                    throw new ArgumentException("Not added.");
 
                 _active?.OnDeactivated();
                 _active = value;
@@ -40,6 +44,42 @@ namespace FlaxEditor.Gizmo
             }
         }
 
+        /// <summary>
+        /// Gets the active gizmo mode.
+        /// </summary>
+        public EditorGizmoMode ActiveMode
+        {
+            get => _activeMode;
+            set
+            {
+                if (_activeMode == value)
+                    return;
+                if (value != null)
+                {
+                    if (!_modes.Contains(value))
+                        throw new ArgumentException("Not added.");
+                    if (value.Owner != _owner)
+                        throw new InvalidOperationException();
+                }
+
+                _activeMode?.OnDeactivated();
+                Active = null;
+                _activeMode = value;
+                _activeMode?.OnActivated();
+                ActiveModeChanged?.Invoke(value);
+            }
+        }
+
+        /// <summary>
+        /// Occurs when active mode gets changed.
+        /// </summary>
+        public event Action<EditorGizmoMode> ActiveModeChanged;
+
+        public GizmosCollection(IGizmoOwner owner)
+        {
+            _owner = owner;
+        }
+
         /// <summary>
         /// Removes the specified item.
         /// </summary>
@@ -57,7 +97,65 @@ namespace FlaxEditor.Gizmo
         public new void Clear()
         {
             Active = null;
+            ActiveMode = null;
+            foreach (var mode in _modes)
+                mode.Dispose();
+            _modes.Clear();
+
             base.Clear();
         }
+
+        /// <summary>
+        /// Adds the mode to the viewport.
+        /// </summary>
+        /// <param name="mode">The mode.</param>
+        public void AddMode(EditorGizmoMode mode)
+        {
+            if (mode == null)
+                throw new ArgumentNullException(nameof(mode));
+            if (_modes.Contains(mode))
+                throw new ArgumentException("Already added.");
+            if (mode.Owner != null)
+                throw new ArgumentException("Already added to other viewport.");
+
+            _modes.Add(mode);
+            mode.Init(_owner);
+        }
+
+        /// <summary>
+        /// Removes the mode from the viewport.
+        /// </summary>
+        /// <param name="mode">The mode.</param>
+        public void RemoveMode(EditorGizmoMode mode)
+        {
+            if (mode == null)
+                throw new ArgumentNullException(nameof(mode));
+            if (!_modes.Contains(mode))
+                throw new ArgumentException("Not added.");
+            if (mode.Owner != _owner)
+                throw new ArgumentException("Not added to this viewport.");
+
+            if (_activeMode == mode)
+                ActiveMode = null;
+            _modes.Remove(mode);
+        }
+
+        /// <summary>
+        /// Sets the active mode.
+        /// </summary>
+        /// <typeparam name="T">The mode type.</typeparam>
+        /// <returns>The activated mode.</returns>
+        public T SetActiveMode<T>() where T : EditorGizmoMode
+        {
+            for (int i = 0; i < _modes.Count; i++)
+            {
+                if (_modes[i] is T mode)
+                {
+                    ActiveMode = mode;
+                    return mode;
+                }
+            }
+            throw new ArgumentException("Not added mode to activate.");
+        }
     }
 }

Source/Editor/Tools/Foliage/EditFoliageGizmoMode.cs

@@ -1,6 +1,7 @@
 // Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
 
 using System;
+using FlaxEditor.Gizmo;
 using FlaxEditor.SceneGraph.Actors;
 using FlaxEditor.Viewport;
 using FlaxEditor.Viewport.Modes;
@@ -60,11 +61,11 @@ namespace FlaxEditor.Tools.Foliage
         public event Action SelectedInstanceIndexChanged;
 
         /// <inheritdoc />
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner owner)
         {
-            base.Init(viewport);
+            base.Init(owner);
 
-            Gizmo = new EditFoliageGizmo(viewport, this);
+            Gizmo = new EditFoliageGizmo(owner, this);
             SelectionOutline = FlaxEngine.Object.New<EditFoliageSelectionOutline>();
             SelectionOutline.GizmoMode = this;
         }
@@ -82,15 +83,15 @@ namespace FlaxEditor.Tools.Foliage
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
-            Viewport.OverrideSelectionOutline(SelectionOutline);
+            ((MainEditorGizmoViewport)Owner).OverrideSelectionOutline(SelectionOutline);
             SelectedInstanceIndex = -1;
         }
 
         /// <inheritdoc />
         public override void OnDeactivated()
         {
-            Viewport.OverrideSelectionOutline(null);
+            ((MainEditorGizmoViewport)Owner).OverrideSelectionOutline(null);
 
             base.OnDeactivated();
         }

Source/Editor/Tools/Foliage/FoliageTab.cs

@@ -242,13 +242,13 @@ namespace FlaxEditor.Tools.Foliage
             switch (_modes.SelectedTabIndex)
             {
             case 0:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<NoGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<NoGizmoMode>();
                 break;
             case 1:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<PaintFoliageGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<PaintFoliageGizmoMode>();
                 break;
             case 2:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<EditFoliageGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<EditFoliageGizmoMode>();
                 break;
             default: throw new IndexOutOfRangeException("Invalid foliage tab mode.");
             }

Source/Editor/Tools/Foliage/PaintFoliageGizmoMode.cs

@@ -1,5 +1,6 @@
 // Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
 
+using FlaxEditor.Gizmo;
 using FlaxEditor.SceneGraph.Actors;
 using FlaxEditor.Viewport;
 using FlaxEditor.Viewport.Modes;
@@ -69,11 +70,11 @@ namespace FlaxEditor.Tools.Foliage
         }
 
         /// <inheritdoc />
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner owner)
         {
-            base.Init(viewport);
+            base.Init(owner);
 
-            Gizmo = new PaintFoliageGizmo(viewport, this);
+            Gizmo = new PaintFoliageGizmo(owner, this);
         }
 
         /// <inheritdoc />
@@ -81,7 +82,7 @@ namespace FlaxEditor.Tools.Foliage
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
             ClearCursor();
         }
 

Source/Editor/Tools/Terrain/CarveTab.cs

@@ -191,13 +191,13 @@ namespace FlaxEditor.Tools.Terrain
             switch (_modes.SelectedTabIndex)
             {
             case 0:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<SculptTerrainGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<SculptTerrainGizmoMode>();
                 break;
             case 1:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<PaintTerrainGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<PaintTerrainGizmoMode>();
                 break;
             case 2:
-                Editor.Windows.EditWin.Viewport.SetActiveMode<EditTerrainGizmoMode>();
+                Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<EditTerrainGizmoMode>();
                 break;
             default: throw new IndexOutOfRangeException("Invalid carve tab mode.");
             }

Source/Editor/Tools/Terrain/EditTerrainGizmoMode.cs

@@ -1,6 +1,7 @@
 // Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
 
 using System;
+using FlaxEditor.Gizmo;
 using FlaxEditor.Viewport;
 using FlaxEditor.Viewport.Modes;
 using FlaxEngine;
@@ -84,12 +85,12 @@ namespace FlaxEditor.Tools.Terrain
         public event Action SelectedChunkCoordChanged;
 
         /// <inheritdoc />
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner owner)
         {
-            base.Init(viewport);
+            base.Init(owner);
 
             EditMode = Modes.Edit;
-            Gizmo = new EditTerrainGizmo(viewport, this);
+            Gizmo = new EditTerrainGizmo(owner, this);
         }
 
         /// <inheritdoc />
@@ -97,7 +98,7 @@ namespace FlaxEditor.Tools.Terrain
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
         }
 
         /// <summary>

Source/Editor/Tools/Terrain/PaintTerrainGizmoMode.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
+using FlaxEditor.Gizmo;
 using FlaxEditor.SceneGraph.Actors;
 using FlaxEditor.Tools.Terrain.Brushes;
 using FlaxEditor.Tools.Terrain.Paint;
@@ -251,11 +252,11 @@ namespace FlaxEditor.Tools.Terrain
         internal EditTerrainMapAction CurrentEditUndoAction => _activeAction;
 
         /// <inheritdoc />
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner owner)
         {
-            base.Init(viewport);
+            base.Init(owner);
 
-            Gizmo = new PaintTerrainGizmo(viewport, this);
+            Gizmo = new PaintTerrainGizmo(owner, this);
             Gizmo.PaintStarted += OnPaintStarted;
             Gizmo.PaintEnded += OnPaintEnded;
         }
@@ -265,7 +266,7 @@ namespace FlaxEditor.Tools.Terrain
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
             ClearCursor();
         }
 

Source/Editor/Tools/Terrain/SculptTerrainGizmoMode.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
+using FlaxEditor.Gizmo;
 using FlaxEditor.SceneGraph.Actors;
 using FlaxEditor.Tools.Terrain.Brushes;
 using FlaxEditor.Tools.Terrain.Sculpt;
@@ -270,11 +271,11 @@ namespace FlaxEditor.Tools.Terrain
         internal EditTerrainMapAction CurrentEditUndoAction => _activeAction;
 
         /// <inheritdoc />
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner owner)
         {
-            base.Init(viewport);
+            base.Init(owner);
 
-            Gizmo = new SculptTerrainGizmo(viewport, this);
+            Gizmo = new SculptTerrainGizmo(owner, this);
             Gizmo.PaintStarted += OnPaintStarted;
             Gizmo.PaintEnded += OnPaintEnded;
         }
@@ -284,7 +285,7 @@ namespace FlaxEditor.Tools.Terrain
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
             ClearCursor();
         }
 

Source/Editor/Tools/VertexPainting.cs

@@ -244,9 +244,9 @@ namespace FlaxEditor.Tools
                 {
                     Tab = this,
                 };
-                _editor.Windows.EditWin.Viewport.AddMode(_gizmoMode);
+                _editor.Windows.EditWin.Viewport.Gizmos.AddMode(_gizmoMode);
             }
-            _editor.Windows.EditWin.Viewport.SetActiveMode(_gizmoMode);
+            _editor.Windows.EditWin.Viewport.Gizmos.ActiveMode = _gizmoMode;
         }
 
         /// <inheritdoc />
@@ -317,7 +317,7 @@ namespace FlaxEditor.Tools
         public Color PaintColor = Color.White;
         public VertexColorsMask PaintMask = VertexColorsMask.RGB;
 
-        public override void Init(MainEditorGizmoViewport viewport)
+        public override void Init(IGizmoOwner viewport)
         {
             base.Init(viewport);
 
@@ -328,7 +328,7 @@ namespace FlaxEditor.Tools
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Gizmo;
+            Owner.Gizmos.Active = Gizmo;
         }
     }
 

Source/Editor/Viewport/EditorGizmoViewport.cs

@@ -27,6 +27,7 @@ namespace FlaxEditor.Viewport
         {
             Undo = undo;
             SceneGraphRoot = sceneGraphRoot;
+            Gizmos = new GizmosCollection(this);
 
             SetUpdate(ref _update, OnUpdate);
         }
@@ -40,7 +41,7 @@ namespace FlaxEditor.Viewport
         }
 
         /// <inheritdoc />
-        public GizmosCollection Gizmos { get; } = new GizmosCollection();
+        public GizmosCollection Gizmos { get; }
 
         /// <inheritdoc />
         public SceneRenderTask RenderTask => Task;
@@ -96,5 +97,16 @@ namespace FlaxEditor.Viewport
 
             root.UpdateCallbacksToRemove.Add(_update);
         }
+
+        /// <inheritdoc />
+        public override void OnDestroy()
+        {
+            if (IsDisposing)
+                return;
+
+            Gizmos.Clear();
+
+            base.OnDestroy();
+        }
     }
 }

Source/Editor/Viewport/MainEditorGizmoViewport.Modes.cs

@@ -1,158 +0,0 @@
-// Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
-
-using System;
-using System.Collections.Generic;
-using FlaxEditor.Viewport.Modes;
-
-namespace FlaxEditor.Viewport
-{
-    public partial class MainEditorGizmoViewport
-    {
-        private EditorGizmoMode _activeMode;
-        private readonly List<EditorGizmoMode> _modes = new List<EditorGizmoMode>();
-
-        /// <summary>
-        /// Gets the active gizmo mode.
-        /// </summary>
-        public EditorGizmoMode ActiveMode => _activeMode;
-
-        /// <summary>
-        /// Occurs when active mode gets changed.
-        /// </summary>
-        public event Action<EditorGizmoMode> ActiveModeChanged;
-
-        /// <summary>
-        /// The sculpt terrain gizmo.
-        /// </summary>
-        public Tools.Terrain.SculptTerrainGizmoMode SculptTerrainGizmo;
-
-        /// <summary>
-        /// The paint terrain gizmo.
-        /// </summary>
-        public Tools.Terrain.PaintTerrainGizmoMode PaintTerrainGizmo;
-
-        /// <summary>
-        /// The edit terrain gizmo.
-        /// </summary>
-        public Tools.Terrain.EditTerrainGizmoMode EditTerrainGizmo;
-
-        /// <summary>
-        /// The paint foliage gizmo.
-        /// </summary>
-        public Tools.Foliage.PaintFoliageGizmoMode PaintFoliageGizmo;
-
-        /// <summary>
-        /// The edit foliage gizmo.
-        /// </summary>
-        public Tools.Foliage.EditFoliageGizmoMode EditFoliageGizmo;
-
-        private void InitModes()
-        {
-            // Add default modes used by the editor
-            _modes.Add(new TransformGizmoMode());
-            _modes.Add(new NoGizmoMode());
-            _modes.Add(SculptTerrainGizmo = new Tools.Terrain.SculptTerrainGizmoMode());
-            _modes.Add(PaintTerrainGizmo = new Tools.Terrain.PaintTerrainGizmoMode());
-            _modes.Add(EditTerrainGizmo = new Tools.Terrain.EditTerrainGizmoMode());
-            _modes.Add(PaintFoliageGizmo = new Tools.Foliage.PaintFoliageGizmoMode());
-            _modes.Add(EditFoliageGizmo = new Tools.Foliage.EditFoliageGizmoMode());
-            for (int i = 0; i < _modes.Count; i++)
-            {
-                _modes[i].Init(this);
-            }
-
-            // Activate transform mode first
-            _activeMode = _modes[0];
-        }
-
-        private void DisposeModes()
-        {
-            // Cleanup
-            _activeMode = null;
-            for (int i = 0; i < _modes.Count; i++)
-                _modes[i].Dispose();
-            _modes.Clear();
-        }
-
-        /// <summary>
-        /// Adds the mode to the viewport.
-        /// </summary>
-        /// <param name="mode">The mode.</param>
-        public void AddMode(EditorGizmoMode mode)
-        {
-            if (mode == null)
-                throw new ArgumentNullException(nameof(mode));
-            if (_modes.Contains(mode))
-                throw new ArgumentException("Already added.");
-            if (mode.Viewport != null)
-                throw new ArgumentException("Already added to other viewport.");
-
-            _modes.Add(mode);
-            mode.Init(this);
-        }
-
-        /// <summary>
-        /// Removes the mode from the viewport.
-        /// </summary>
-        /// <param name="mode">The mode.</param>
-        public void RemoveMode(EditorGizmoMode mode)
-        {
-            if (mode == null)
-                throw new ArgumentNullException(nameof(mode));
-            if (!_modes.Contains(mode))
-                throw new ArgumentException("Not added.");
-            if (mode.Viewport != this)
-                throw new ArgumentException("Not added to this viewport.");
-
-            if (_activeMode == mode)
-                SetActiveMode(null);
-
-            _modes.Remove(mode);
-        }
-
-        /// <summary>
-        /// Sets the active mode.
-        /// </summary>
-        /// <param name="mode">The mode.</param>
-        public void SetActiveMode(EditorGizmoMode mode)
-        {
-            if (mode == _activeMode)
-                return;
-            if (mode != null)
-            {
-                if (!_modes.Contains(mode))
-                    throw new ArgumentException("Not added.");
-                if (mode.Viewport != this)
-                    throw new ArgumentException("Not added to this viewport.");
-            }
-
-            _activeMode?.OnDeactivated();
-
-            Gizmos.Active = null;
-            _activeMode = mode;
-
-            _activeMode?.OnActivated();
-
-            ActiveModeChanged?.Invoke(mode);
-        }
-
-        /// <summary>
-        /// Sets the active mode.
-        /// </summary>
-        /// <typeparam name="T">The mode type.</typeparam>
-        /// <returns>The activated mode.</returns>
-        public T SetActiveMode<T>() where T : EditorGizmoMode
-        {
-            for (int i = 0; i < _modes.Count; i++)
-            {
-                if (_modes[i] is T mode)
-                {
-                    SetActiveMode(mode);
-                    return mode;
-                }
-            }
-
-            throw new ArgumentException("Not added mode to activate.");
-        }
-    }
-}

Source/Editor/Viewport/MainEditorGizmoViewport.cs

@@ -10,6 +10,7 @@ using FlaxEditor.SceneGraph;
 using FlaxEditor.SceneGraph.Actors;
 using FlaxEditor.Scripting;
 using FlaxEditor.Viewport.Cameras;
+using FlaxEditor.Viewport.Modes;
 using FlaxEditor.Viewport.Widgets;
 using FlaxEditor.Windows;
 using FlaxEngine;
@@ -22,7 +23,7 @@ namespace FlaxEditor.Viewport
     /// Main editor gizmo viewport used by the <see cref="EditGameWindow"/>.
     /// </summary>
     /// <seealso cref="FlaxEditor.Viewport.EditorGizmoViewport" />
-    public partial class MainEditorGizmoViewport : EditorGizmoViewport, IEditorPrimitivesOwner, IGizmoOwner
+    public class MainEditorGizmoViewport : EditorGizmoViewport, IEditorPrimitivesOwner
     {
         private readonly Editor _editor;
 
@@ -183,6 +184,31 @@ namespace FlaxEditor.Viewport
             set => _showNavigationButton.Checked = value;
         }
 
+        /// <summary>
+        /// The sculpt terrain gizmo.
+        /// </summary>
+        public Tools.Terrain.SculptTerrainGizmoMode SculptTerrainGizmo;
+
+        /// <summary>
+        /// The paint terrain gizmo.
+        /// </summary>
+        public Tools.Terrain.PaintTerrainGizmoMode PaintTerrainGizmo;
+
+        /// <summary>
+        /// The edit terrain gizmo.
+        /// </summary>
+        public Tools.Terrain.EditTerrainGizmoMode EditTerrainGizmo;
+
+        /// <summary>
+        /// The paint foliage gizmo.
+        /// </summary>
+        public Tools.Foliage.PaintFoliageGizmoMode PaintFoliageGizmo;
+
+        /// <summary>
+        /// The edit foliage gizmo.
+        /// </summary>
+        public Tools.Foliage.EditFoliageGizmoMode EditFoliageGizmo;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="MainEditorGizmoViewport"/> class.
         /// </summary>
@@ -380,7 +406,20 @@ namespace FlaxEditor.Viewport
             DragHandlers.Add(_dragActorType);
             DragHandlers.Add(_dragAssets);
 
-            InitModes();
+            // Init gizmo modes
+            {
+                // Add default modes used by the editor
+                Gizmos.AddMode(new TransformGizmoMode());
+                Gizmos.AddMode(new NoGizmoMode());
+                Gizmos.AddMode(SculptTerrainGizmo = new Tools.Terrain.SculptTerrainGizmoMode());
+                Gizmos.AddMode(PaintTerrainGizmo = new Tools.Terrain.PaintTerrainGizmoMode());
+                Gizmos.AddMode(EditTerrainGizmo = new Tools.Terrain.EditTerrainGizmoMode());
+                Gizmos.AddMode(PaintFoliageGizmo = new Tools.Foliage.PaintFoliageGizmoMode());
+                Gizmos.AddMode(EditFoliageGizmo = new Tools.Foliage.EditFoliageGizmoMode());
+
+                // Activate transform mode first
+                Gizmos.SetActiveMode<TransformGizmoMode>();
+            }
 
             // Setup input actions
             InputActions.Add(options => options.TranslateMode, () => TransformGizmo.ActiveMode = TransformGizmoBase.Mode.Translate);
@@ -1117,7 +1156,6 @@ namespace FlaxEditor.Viewport
             if (IsDisposing)
                 return;
 
-            DisposeModes();
             _debugDrawData.Dispose();
             if (_task != null)
             {

Source/Editor/Viewport/Modes/EditorGizmoMode.cs

@@ -1,6 +1,7 @@
 // Copyright (c) 2012-2023 Wojciech Figat. All rights reserved.
 
 using System;
+using FlaxEditor.Gizmo;
 using FlaxEngine;
 
 namespace FlaxEditor.Viewport.Modes
@@ -14,12 +15,12 @@ namespace FlaxEditor.Viewport.Modes
     [HideInEditor]
     public abstract class EditorGizmoMode
     {
-        private MainEditorGizmoViewport _viewport;
+        private IGizmoOwner _owner;
 
         /// <summary>
-        /// Gets the viewport.
+        /// Gets the gizmos owner viewport.
         /// </summary>
-        public MainEditorGizmoViewport Viewport => _viewport;
+        public IGizmoOwner Owner => _owner;
 
         /// <summary>
         /// Occurs when mode gets activated.
@@ -34,10 +35,10 @@ namespace FlaxEditor.Viewport.Modes
         /// <summary>
         /// Initializes the specified mode and links it to the viewport.
         /// </summary>
-        /// <param name="viewport">The viewport.</param>
+        /// <param name="owner">The gizmos owner.</param>
-        public virtual void Init(MainEditorGizmoViewport viewport)
+        public virtual void Init(IGizmoOwner owner)
         {
-            _viewport = viewport;
+            _owner = owner;
         }
 
         /// <summary>
@@ -45,7 +46,7 @@ namespace FlaxEditor.Viewport.Modes
         /// </summary>
         public virtual void Dispose()
         {
-            _viewport = null;
+            _owner = null;
         }
 
         /// <summary>

Source/Editor/Viewport/Modes/NoGizmoMode.cs

@@ -13,7 +13,7 @@ namespace FlaxEditor.Viewport.Modes
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = null;
+            Owner.Gizmos.Active = null;
         }
     }
 }

Source/Editor/Viewport/Modes/TransformGizmoMode.cs

@@ -13,7 +13,7 @@ namespace FlaxEditor.Viewport.Modes
         {
             base.OnActivated();
 
-            Viewport.Gizmos.Active = Viewport.TransformGizmo;
+            Owner.Gizmos.Active = ((MainEditorGizmoViewport)Owner).TransformGizmo;
         }
     }
 }

Source/Editor/Viewport/PrefabWindowViewport.cs

@@ -84,6 +84,7 @@ namespace FlaxEditor.Viewport
             _dragAssets = new DragAssets(ValidateDragItem);
             ShowDebugDraw = true;
             ShowEditorPrimitives = true;
+            Gizmos = new GizmosCollection(this);
 
             // Prepare rendering task
             Task.ActorsSource = ActorsSources.CustomActors;
@@ -305,7 +306,7 @@ namespace FlaxEditor.Viewport
         }
 
         /// <inheritdoc />
-        public GizmosCollection Gizmos { get; } = new GizmosCollection();
+        public GizmosCollection Gizmos { get; }
 
         /// <inheritdoc />
         public SceneRenderTask RenderTask => Task;

Source/Editor/Windows/ToolboxWindow.cs

@@ -116,7 +116,7 @@ namespace FlaxEditor.Windows
         : base(string.Empty, icon)
         {
             Editor = editor;
-            Selected += tab => Editor.Windows.EditWin.Viewport.SetActiveMode<TransformGizmoMode>();
+            Selected += tab => Editor.Windows.EditWin.Viewport.Gizmos.SetActiveMode<TransformGizmoMode>();
             ScriptsBuilder.ScriptsReload += OnScriptsReload;
             ScriptsBuilder.ScriptsReloadEnd += OnScriptsReloadEnd;