From d425fae2a76077fcb12b917c619bcec71ef47b77 Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 15 Dec 2020 11:59:52 +0100
Subject: [PATCH] Add Branch On Enum node

---
 Source/Editor/Surface/Archetypes/Flow.cs      | 126 ++++++++++++++++++
 Source/Engine/Content/Assets/VisualScript.cpp |  22 +++
 2 files changed, 148 insertions(+)

diff --git a/Source/Editor/Surface/Archetypes/Flow.cs b/Source/Editor/Surface/Archetypes/Flow.cs
index a214319ed..b357fc8ff 100644
--- a/Source/Editor/Surface/Archetypes/Flow.cs
+++ b/Source/Editor/Surface/Archetypes/Flow.cs
@@ -1,5 +1,9 @@
 // Copyright (c) 2012-2020 Wojciech Figat. All rights reserved.
 
+using System.Collections.Generic;
+using FlaxEditor.GUI;
+using FlaxEditor.Scripting;
+using FlaxEditor.Surface.Elements;
 using FlaxEngine;
 using FlaxEngine.GUI;
 
@@ -87,6 +91,112 @@ namespace FlaxEditor.Surface.Archetypes
             }
         }
 
+
+        private class BranchOnEnumNode : SurfaceNode
+        {
+            public BranchOnEnumNode(uint id, VisjectSurfaceContext context, NodeArchetype nodeArch, GroupArchetype groupArch)
+            : base(id, context, nodeArch, groupArch)
+            {
+            }
+
+            public override void OnLoaded()
+            {
+                base.OnLoaded();
+
+                // Restore saved input boxes layout
+                if (Values[0] is byte[] data)
+                {
+                    for (int i = 0; i < data.Length / 4; i++)
+                        AddBox(true, i + 2, i, string.Empty, new ScriptType(typeof(void)), true);
+                }
+            }
+
+            public override void OnSurfaceLoaded()
+            {
+                base.OnSurfaceLoaded();
+
+                UpdateBoxes();
+                GetBox(1).CurrentTypeChanged += box => UpdateBoxes();
+            }
+
+            public override void ConnectionTick(Box box)
+            {
+                base.ConnectionTick(box);
+
+                UpdateBoxes();
+            }
+
+            private unsafe void UpdateBoxes()
+            {
+                var valueBox = (InputBox)GetBox(1);
+                const int firstBox = 2;
+                const int maxBoxes = 40;
+                bool isInvalid = false;
+                var data = Utils.GetEmptyArray<byte>();
+
+                if (valueBox.HasAnyConnection)
+                {
+                    var valueType = valueBox.CurrentType;
+                    if (valueType.IsEnum)
+                    {
+                        // Get enum entries
+                        var entries = new List<EnumComboBox.Entry>();
+                        EnumComboBox.BuildEntriesDefault(valueType.Type, entries);
+
+                        // Setup switch value inputs
+                        int id = firstBox;
+                        for (; id < entries.Count + firstBox; id++)
+                        {
+                            var e = entries[id - firstBox];
+                            var box = AddBox(true, id, id - firstBox, e.Name, new ScriptType(typeof(void)), true);
+                            if (!string.IsNullOrEmpty(e.Tooltip))
+                                box.TooltipText = e.Tooltip;
+                        }
+                        for (; id < maxBoxes; id++)
+                        {
+                            var box = GetBox(id);
+                            if (box == null)
+                                break;
+                            RemoveElement(box);
+                        }
+
+                        // Build data about enum entries values for the runtime
+                        if (Values[0] is byte[] dataPrev && dataPrev.Length == entries.Count * 4)
+                            data = dataPrev;
+                        else
+                            data = new byte[entries.Count * 4];
+                        fixed (byte* dataPtr = data)
+                        {
+                            int* dataValues = (int*)dataPtr;
+                            for (int i = 0; i < entries.Count; i++)
+                                dataValues[i] = entries[i].Value;
+                        }
+                    }
+                    else
+                    {
+                        // Not an enum
+                        isInvalid = true;
+                    }
+                }
+
+                if (isInvalid)
+                {
+                    // Remove input values boxes if switch value is unused
+                    for (int id = firstBox; id < maxBoxes; id++)
+                    {
+                        var box = GetBox(id);
+                        if (box == null)
+                            break;
+                        RemoveElement(box);
+                    }
+                }
+
+                Values[0] = data;
+
+                ResizeAuto();
+            }
+        }
+
         /// <summary>
         /// The nodes for that group.
         /// </summary>
@@ -176,6 +286,22 @@ namespace FlaxEditor.Surface.Archetypes
                     NodeElementArchetype.Factory.Output(16, string.Empty, typeof(void), 17, true),
                 }
             },
+            new NodeArchetype
+            {
+                TypeID = 5,
+                Title = "Branch On Enum",
+                Create = (id, context, arch, groupArch) => new BranchOnEnumNode(id, context, arch, groupArch),
+                Description = "Performs the flow logic branch based on the enum value",
+                Flags = NodeFlags.VisualScriptGraph,
+                Size = new Vector2(160, 60),
+                DefaultValues = new object[] { Utils.GetEmptyArray<byte>() },
+                ConnectionsHints = ConnectionsHint.Enum,
+                Elements = new[]
+                {
+                    NodeElementArchetype.Factory.Input(0, string.Empty, false, typeof(void), 0),
+                    NodeElementArchetype.Factory.Input(1, "Value", true, null, 1),
+                }
+            },
         };
     }
 }
diff --git a/Source/Engine/Content/Assets/VisualScript.cpp b/Source/Engine/Content/Assets/VisualScript.cpp
index 2cab30892..bb27a233e 100644
--- a/Source/Engine/Content/Assets/VisualScript.cpp
+++ b/Source/Engine/Content/Assets/VisualScript.cpp
@@ -1129,6 +1129,28 @@ void VisualScriptExecutor::ProcessGroupFlow(Box* boxBase, Node* node, Value& val
                 eatBox(node, boxBase->FirstConnection());
         }
         break;
+    }
+        // Branch On Enum
+    case 5:
+    {
+        const Value v = tryGetValue(node->GetBox(1), Value::Null);
+        if (v.Type.Type == VariantType::Enum && node->Values.Count() == 1 && node->Values[0].Type.Type == VariantType::Blob)
+        {
+            int32* dataValues = (int32*)node->Values[0].AsBlob.Data;
+            int32 dataValuesCount = node->Values[0].AsBlob.Length / 4;
+            int32 vAsInt = (int32)v;
+            for (int32 i = 0; i < dataValuesCount; i++)
+            {
+                if (dataValues[i] == vAsInt)
+                {
+                    boxBase = node->GetBox(i + 2);
+                    if (boxBase->HasConnection())
+                        eatBox(node, boxBase->FirstConnection());
+                    break;
+                }
+            }
+        }
+        break;
     }
     }
 }