From c00b32364ab97a6c8434d7fe02d17c0691e8eb7c Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 2 Feb 2021 12:51:36 +0100
Subject: [PATCH] Add `SceneGraphNode.OnContextMenu` for customizations

---
 Source/Editor/SceneGraph/SceneGraphNode.cs    |  7 +++++++
 .../Windows/SceneTreeWindow.ContextMenu.cs    | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/Source/Editor/SceneGraph/SceneGraphNode.cs b/Source/Editor/SceneGraph/SceneGraphNode.cs
index 7ddf2d876..cc47c20f2 100644
--- a/Source/Editor/SceneGraph/SceneGraphNode.cs
+++ b/Source/Editor/SceneGraph/SceneGraphNode.cs
@@ -317,6 +317,13 @@ namespace FlaxEditor.SceneGraph
         {
         }
 
+        /// <summary>
+        /// Called when scene tree window wants to show the context menu. Allows to add custom options.
+        /// </summary>
+        public virtual void OnContextMenu(FlaxEditor.GUI.ContextMenu.ContextMenu contextMenu)
+        {
+        }
+
         /// <summary>
         /// The scene graph node state container. Used for Editor undo actions (eg. restoring deleted node).
         /// </summary>
diff --git a/Source/Editor/Windows/SceneTreeWindow.ContextMenu.cs b/Source/Editor/Windows/SceneTreeWindow.ContextMenu.cs
index b6756b438..f407017df 100644
--- a/Source/Editor/Windows/SceneTreeWindow.ContextMenu.cs
+++ b/Source/Editor/Windows/SceneTreeWindow.ContextMenu.cs
@@ -1,6 +1,7 @@
 // Copyright (c) 2012-2021 Wojciech Figat. All rights reserved.
 
 using System;
+using System.Collections.Generic;
 using FlaxEditor.GUI.ContextMenu;
 using FlaxEditor.SceneGraph;
 using FlaxEngine;
@@ -113,6 +114,24 @@ namespace FlaxEditor.Windows
 
             // Custom options
 
+            bool showCustomNodeOptions = Editor.SceneEditing.Selection.Count == 1;
+            if (!showCustomNodeOptions && Editor.SceneEditing.Selection.Count != 0)
+            {
+                showCustomNodeOptions = true;
+                for (int i = 1; i < Editor.SceneEditing.Selection.Count; i++)
+                {
+                    if (Editor.SceneEditing.Selection[0].GetType() != Editor.SceneEditing.Selection[i].GetType())
+                    {
+                        showCustomNodeOptions = false;
+                        break;
+                    }
+                }
+            }
+            if (showCustomNodeOptions)
+            {
+                Editor.SceneEditing.Selection[0].OnContextMenu(contextMenu);
+            }
+
             ContextMenuShow?.Invoke(contextMenu);
 
             return contextMenu;