Support Menu sockets in the for-each-element zone

Adds support for the "Menu" socket type in for-each-element zones. This
only includes field inputs and their matching per-element values, but
not outputting attributes of type Menu (Menu attributes are not
generally supported at this point).

A dedicated enum propagation function is added for the zone input node.
This isn't technically necessary: the first 2 inputs and outputs should
be ignored but are not menus anyway. However, this is clearer and
provides a place for future changes.

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

source/blender/blenkernel/intern/node_socket_value.cc

@@ -63,6 +63,45 @@ template<typename T> static std::optional<eNodeSocketDatatype> static_type_to_so
   return std::nullopt;
 }
 
+/**
+ * Check if a socket type stores the static C++ type.
+ */
+template<typename T>
+static bool static_type_is_base_socket_type(const eNodeSocketDatatype socket_type)
+{
+  switch (socket_type) {
+    case SOCK_INT:
+      return std::is_same_v<T, int>;
+    case SOCK_FLOAT:
+      return std::is_same_v<T, float>;
+    case SOCK_BOOLEAN:
+      return std::is_same_v<T, bool>;
+    case SOCK_VECTOR:
+      return std::is_same_v<T, float3>;
+    case SOCK_RGBA:
+      return std::is_same_v<T, ColorGeometry4f>;
+    case SOCK_ROTATION:
+      return std::is_same_v<T, math::Quaternion>;
+    case SOCK_MATRIX:
+      return std::is_same_v<T, float4x4>;
+    case SOCK_STRING:
+      return std::is_same_v<T, std::string>;
+    case SOCK_MENU:
+      return std::is_same_v<T, int>;
+    case SOCK_CUSTOM:
+    case SOCK_SHADER:
+    case SOCK_OBJECT:
+    case SOCK_IMAGE:
+    case SOCK_GEOMETRY:
+    case SOCK_COLLECTION:
+    case SOCK_TEXTURE:
+    case SOCK_MATERIAL:
+      return false;
+  }
+  BLI_assert_unreachable();
+  return false;
+}
+
 template<typename T> T SocketValueVariant::extract()
 {
   if constexpr (std::is_same_v<T, fn::GField>) {
@@ -86,7 +125,7 @@ template<typename T> T SocketValueVariant::extract()
     }
   }
   else if constexpr (fn::is_field_v<T>) {
-    BLI_assert(socket_type_ == static_type_to_socket_type<typename T::base_type>());
+    BLI_assert(static_type_is_base_socket_type<typename T::base_type>(socket_type_));
     return T(this->extract<fn::GField>());
   }
 #ifdef WITH_OPENVDB
@@ -109,12 +148,12 @@ template<typename T> T SocketValueVariant::extract()
     }
   }
   else if constexpr (is_VolumeGrid_v<T>) {
-    BLI_assert(socket_type_ == static_type_to_socket_type<typename T::base_type>());
+    BLI_assert(static_type_is_base_socket_type<typename T::base_type>(socket_type_));
     return this->extract<GVolumeGrid>().typed<typename T::base_type>();
   }
 #endif
   else {
-    BLI_assert(socket_type_ == static_type_to_socket_type<T>());
+    BLI_assert(static_type_is_base_socket_type<T>(socket_type_));
     if (kind_ == Kind::Single) {
       return std::move(value_.get<T>());
     }
@@ -302,6 +341,8 @@ void *SocketValueVariant::allocate_single(const eNodeSocketDatatype socket_type)
       return value_.allocate<ColorGeometry4f>();
     case SOCK_STRING:
       return value_.allocate<std::string>();
+    case SOCK_MENU:
+      return value_.allocate<int>();
     default: {
       BLI_assert_unreachable();
       return nullptr;

source/blender/blenkernel/intern/node_tree_update.cc

@@ -954,6 +954,23 @@ class NodeTreeMainUpdater {
           }
         }
       }
+      else if (node->type == GEO_NODE_FOREACH_GEOMETRY_ELEMENT_INPUT) {
+        /* Propagate menu from element inputs to field inputs. */
+        BLI_assert(node->input_sockets().size() == node->output_sockets().size());
+        /* Inputs Geometry, Selection and outputs Index, Element are ignored. */
+        const IndexRange sockets = node->input_sockets().index_range().drop_front(2);
+        for (const int socket_i : sockets) {
+          bNodeSocket *input = node->input_sockets()[socket_i];
+          bNodeSocket *output = node->output_sockets()[socket_i];
+          if (input->is_available() && input->type == SOCK_MENU && output->is_available() &&
+              output->type == SOCK_MENU)
+          {
+            this->update_socket_enum_definition(
+                *input->default_value_typed<bNodeSocketValueMenu>(),
+                *output->default_value_typed<bNodeSocketValueMenu>());
+          }
+        }
+      }
       else {
         /* Propagate over internal relations. */
         /* XXX Placeholder implementation just propagates all outputs

source/blender/nodes/geometry/include/NOD_geo_foreach_geometry_element.hh

@@ -60,7 +60,8 @@ struct ForeachGeometryElementInputItemsAccessor {
                 SOCK_BOOLEAN,
                 SOCK_ROTATION,
                 SOCK_MATRIX,
-                SOCK_INT);
+                SOCK_INT,
+                SOCK_MENU);
   }
 
   static void init_with_socket_type_and_name(bNode &node,