Shader Nodes: support Menu Switch node

This implements the Menu Switch node in shader nodes. It's the same node that is
used in Geometry Nodes and the Compositor.

The Menu Switch node is purely handled during preprocessing and thus builds on
top of #141936. Hence, it's input has to be a single value, just like the
iteration count for repeat zones. This limitation can be lifted in the future,
but currently there is no way to produce a non-single menu value in shader
nodes. This will become possible if other Switch nodes are added though.

Pull Request: https://projects.blender.org/blender/blender/pulls/146896

scripts/startup/bl_ui/node_add_menu_shader.py

@@ -404,6 +404,8 @@ class NODE_MT_shader_node_utilities_base(node_add_menu.NodeMenu):
         self.node_operator(layout, "NodeEvaluateClosure")
         self.node_operator(layout, "NodeCombineBundle")
         self.node_operator(layout, "NodeSeparateBundle")
+        layout.separator()
+        self.node_operator(layout, "GeometryNodeMenuSwitch")
 
 
 class NODE_MT_shader_node_all_base(node_add_menu.NodeMenu):

source/blender/blenkernel/intern/node_tree_update.cc

@@ -525,7 +525,7 @@ class NodeTreeMainUpdater {
     this->make_node_previews_dirty(ntree);
 
     this->propagate_runtime_flags(ntree);
-    if (ELEM(ntree.type, NTREE_GEOMETRY, NTREE_COMPOSIT)) {
+    if (ELEM(ntree.type, NTREE_GEOMETRY, NTREE_COMPOSIT, NTREE_SHADER)) {
       if (this->propagate_enum_definitions(ntree)) {
         result.interface_changed = true;
       }

source/blender/nodes/geometry/nodes/node_geo_menu_switch.cc

@@ -3,7 +3,9 @@
  * SPDX-License-Identifier: GPL-2.0-or-later */
 
 #include "BLI_array_utils.hh"
+
 #include "node_geometry_util.hh"
+#include "shader/node_shader_util.hh"
 
 #include "DNA_node_types.h"
 
@@ -553,7 +555,7 @@ static void register_node()
 {
   static blender::bke::bNodeType ntype;
 
-  geo_cmp_node_type_base(&ntype, "GeometryNodeMenuSwitch", GEO_NODE_MENU_SWITCH);
+  common_node_type_base(&ntype, "GeometryNodeMenuSwitch", GEO_NODE_MENU_SWITCH);
   ntype.ui_name = "Menu Switch";
   ntype.ui_description = "Select from multiple inputs by name";
   ntype.enum_name_legacy = "MENU_SWITCH";

source/blender/nodes/intern/shader_nodes_inline.cc

@@ -14,6 +14,7 @@
 #include "BLI_math_vector.h"
 #include "BLI_stack.hh"
 
+#include "NOD_menu_value.hh"
 #include "NOD_multi_function.hh"
 #include "NOD_node_declaration.hh"
 #include "NOD_node_in_compute_context.hh"
@@ -33,7 +34,7 @@ struct NodeAndSocket {
 };
 
 struct PrimitiveSocketValue {
-  std::variant<int, float, bool, ColorGeometry4f, float3> value;
+  std::variant<int, float, bool, ColorGeometry4f, float3, MenuValue> value;
 
   const void *buffer() const
   {
@@ -63,6 +64,9 @@ struct PrimitiveSocketValue {
     if (type.is<float3>()) {
       return {*static_cast<const float3 *>(value.get())};
     }
+    if (type.is<MenuValue>()) {
+      return {*static_cast<const MenuValue *>(value.get())};
+    }
     BLI_assert_unreachable();
     return {};
   }
@@ -127,6 +131,9 @@ struct SocketValue {
         case SOCK_RGBA:
           return PrimitiveSocketValue{
               ColorGeometry4f(socket.default_value_typed<bNodeSocketValueRGBA>()->value)};
+        case SOCK_MENU:
+          return PrimitiveSocketValue{
+              MenuValue(socket.default_value_typed<bNodeSocketValueMenu>()->value)};
         default:
           return std::nullopt;
       }
@@ -438,6 +445,10 @@ class ShaderNodesInliner {
       this->handle_output_socket__separate_bundle(socket);
       return;
     }
+    if (node->is_type("GeometryNodeMenuSwitch")) {
+      this->handle_output_socket__menu_switch(socket);
+      return;
+    }
     this->handle_output_socket__eval(socket);
   }
 
@@ -852,6 +863,53 @@ class ShaderNodesInliner {
     this->store_socket_value_fallback(socket);
   }
 
+  void handle_output_socket__menu_switch(const SocketInContext &socket)
+  {
+    const NodeInContext node = socket.owner_node();
+    const auto &storage = *static_cast<const NodeMenuSwitch *>(node->storage);
+
+    const SocketInContext menu_input = node.input_socket(0);
+    const SocketValue *menu_socket_value = value_by_socket_.lookup_ptr(menu_input);
+    if (!menu_socket_value) {
+      /* The menu value is not known yet, so schedule it for now. */
+      this->schedule_socket(menu_input);
+      return;
+    }
+
+    const std::optional<PrimitiveSocketValue> menu_value_opt = menu_socket_value->to_primitive(
+        *menu_input->typeinfo);
+    if (!menu_value_opt) {
+      /* This limitation may be lifted in the future. Menu Switch nodes could be supported natively
+       * by render engines or we convert them to a bunch of mix nodes. */
+      this->store_socket_value_fallback(socket);
+      params_.r_error_messages.append({node.node, TIP_("Menu value has to be a constant value")});
+      return;
+    }
+    const MenuValue menu_value = std::get<MenuValue>(menu_value_opt->value);
+    /* Find the selected item index. */
+    std::optional<int> selected_index;
+    for (const int item_i : IndexRange(storage.enum_definition.items_num)) {
+      const NodeEnumItem &item = storage.enum_definition.items_array[item_i];
+      if (MenuValue(item.identifier) == menu_value) {
+        selected_index = item_i;
+        break;
+      }
+    }
+    if (!selected_index.has_value()) {
+      /* The input value does not exist in the menu. */
+      this->store_socket_value_fallback(socket);
+      return;
+    }
+    if (socket->index() == 0) {
+      /* Handle forwarding the selected value. */
+      this->forward_value_or_schedule(socket, node.input_socket(*selected_index + 1));
+      return;
+    }
+    /* Set the value of the mask output. */
+    const bool is_selected = selected_index == socket->index() - 1;
+    this->store_socket_value(socket, {PrimitiveSocketValue{is_selected}});
+  }
+
   /**
    * Evaluate a node to compute the value of the given output socket. This may also compute all the
    * other outputs of the node.

source/blender/nodes/shader/node_shader_tree.cc

@@ -158,7 +158,9 @@ static bool shader_validate_link(eNodeSocketDatatype from, eNodeSocketDatatype t
   if (from == SOCK_SHADER) {
     return to == SOCK_SHADER;
   }
-  if (ELEM(to, SOCK_BUNDLE, SOCK_CLOSURE) || ELEM(from, SOCK_BUNDLE, SOCK_CLOSURE)) {
+  if (ELEM(to, SOCK_BUNDLE, SOCK_CLOSURE, SOCK_MENU) ||
+      ELEM(from, SOCK_BUNDLE, SOCK_CLOSURE, SOCK_MENU))
+  {
     return from == to;
   }
   return true;
@@ -175,7 +177,8 @@ static bool shader_node_tree_socket_type_valid(blender::bke::bNodeTreeType * /*n
                                                                         SOCK_RGBA,
                                                                         SOCK_SHADER,
                                                                         SOCK_BUNDLE,
-                                                                        SOCK_CLOSURE);
+                                                                        SOCK_CLOSURE,
+                                                                        SOCK_MENU);
 }
 
 blender::bke::bNodeTreeType *ntreeType_Shader;