Geometry Nodes: new Switch node

This is a first iteration of a switch node. It can only switch between
two inputs values based on a boolean. A more sophisticated switch
node that has an integer selector will probably come later.

Currently, the geometry nodes evaluator does not support lazy evaluation
of individual inputs. Therefore, all inputs will be computed currently.
An improvement to the evaluator will be worked on separately.

Ref: T85374

Differential Revision: https://developer.blender.org/D10460

9 files changed, 224 insertions, 0 deletions

diff --git a/release/scripts/startup/nodeitems_builtins.py b/release/scripts/startup/nodeitems_builtins.py
index ab4d9353e1b..ad2de08bcf2 100644
--- a/release/scripts/startup/nodeitems_builtins.py
+++ b/release/scripts/startup/nodeitems_builtins.py
@@ -550,6 +550,7 @@ geometry_node_categories = [
         NodeItem("ShaderNodeMath"),
         NodeItem("FunctionNodeBooleanMath"),
         NodeItem("FunctionNodeFloatCompare"),
+        NodeItem("GeometryNodeSwitch"),
     ]),
     GeometryNodeCategory("GEO_VECTOR", "Vector", items=[
         NodeItem("ShaderNodeSeparateXYZ"),
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index d6c4ad037e2..2a33c4819e3 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1413,6 +1413,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_ATTRIBUTE_MAP_RANGE 1040
 #define GEO_NODE_ATTRIBUTE_CLAMP 1041
 #define GEO_NODE_BOUNDING_BOX 1042
+#define GEO_NODE_SWITCH 1043
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 02195e0d60f..473be34d69a 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -4968,6 +4968,7 @@ static void registerGeometryNodes()
   register_node_type_geo_sample_texture();
   register_node_type_geo_subdivide();
   register_node_type_geo_subdivision_surface();
+  register_node_type_geo_switch();
   register_node_type_geo_transform();
   register_node_type_geo_triangulate();
   register_node_type_geo_volume_to_mesh();
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index 282d71f6a87..f57515c0e93 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -1308,6 +1308,11 @@ typedef struct NodeGeometryMeshLine {
   uint8_t count_mode;
 } NodeGeometryMeshLine;
 
+typedef struct NodeSwitch {
+  /* NodeSwitch. */
+  uint8_t input_type;
+} NodeSwitch;
+
 /* script node mode */
 #define NODE_SCRIPT_INTERNAL 0
 #define NODE_SCRIPT_EXTERNAL 1
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index b0254ce2ef3..1da711f260c 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -65,6 +65,20 @@
 const EnumPropertyItem rna_enum_node_socket_in_out_items[] = {
     {SOCK_IN, "IN", 0, "Input", ""}, {SOCK_OUT, "OUT", 0, "Output", ""}, {0, NULL, 0, NULL, NULL}};
 
+static const EnumPropertyItem node_socket_data_type_items[] = {
+    {SOCK_FLOAT, "FLOAT", 0, "Float", ""},
+    {SOCK_INT, "INT", 0, "Int", ""},
+    {SOCK_BOOLEAN, "BOOLEAN", 0, "Boolean", ""},
+    {SOCK_VECTOR, "VECTOR", 0, "Vector", ""},
+    {SOCK_STRING, "STRING", 0, "String", ""},
+    {SOCK_RGBA, "RGBA", 0, "Color", ""},
+    {SOCK_OBJECT, "OBJECT", 0, "Object", ""},
+    {SOCK_IMAGE, "IMAGE", 0, "Image", ""},
+    {SOCK_GEOMETRY, "GEOMETRY", 0, "Geometry", ""},
+    {SOCK_COLLECTION, "COLLECTION", 0, "Collection", ""},
+    {0, NULL, 0, NULL, NULL},
+};
+
 #ifndef RNA_RUNTIME
 static const EnumPropertyItem rna_enum_node_socket_display_shape_items[] = {
     {SOCK_DISPLAY_SHAPE_CIRCLE, "CIRCLE", 0, "Circle", ""},
@@ -1913,6 +1927,29 @@ static const EnumPropertyItem *itemf_function_check(
   return item_array;
 }
 
+static bool switch_type_supported(const EnumPropertyItem *item)
+{
+  return ELEM(item->value,
+              SOCK_FLOAT,
+              SOCK_INT,
+              SOCK_BOOLEAN,
+              SOCK_VECTOR,
+              SOCK_STRING,
+              SOCK_RGBA,
+              SOCK_GEOMETRY,
+              SOCK_OBJECT,
+              SOCK_COLLECTION);
+}
+
+static const EnumPropertyItem *rna_GeometryNodeSwitch_type_itemf(bContext *UNUSED(C),
+                                                                 PointerRNA *UNUSED(ptr),
+                                                                 PropertyRNA *UNUSED(prop),
+                                                                 bool *r_free)
+{
+  *r_free = true;
+  return itemf_function_check(node_socket_data_type_items, switch_type_supported);
+}
+
 static bool attribute_clamp_type_supported(const EnumPropertyItem *item)
 {
   return ELEM(item->value, CD_PROP_FLOAT, CD_PROP_FLOAT3, CD_PROP_INT32, CD_PROP_COLOR);
@@ -9630,6 +9667,19 @@ static void def_geo_mesh_line(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
 }
 
+static void def_geo_switch(StructRNA *srna)
+{
+  PropertyRNA *prop;
+
+  RNA_def_struct_sdna_from(srna, "NodeSwitch", "storage");
+  prop = RNA_def_property(srna, "input_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "input_type");
+  RNA_def_property_enum_items(prop, node_socket_data_type_items);
+  RNA_def_property_enum_funcs(prop, NULL, NULL, "rna_GeometryNodeSwitch_type_itemf");
+  RNA_def_property_ui_text(prop, "Input Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+}
+
 /* -------------------------------------------------------------------------- */
 
 static void rna_def_shader_node(BlenderRNA *brna)
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 8168f0dba73..7a966b660b4 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -180,6 +180,7 @@ set(SRC
   geometry/nodes/node_geo_points_to_volume.cc
   geometry/nodes/node_geo_subdivide.cc
   geometry/nodes/node_geo_subdivision_surface.cc
+  geometry/nodes/node_geo_switch.cc
   geometry/nodes/node_geo_transform.cc
   geometry/nodes/node_geo_triangulate.cc
   geometry/nodes/node_geo_volume_to_mesh.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index b84c80e916c..74d77787eea 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -69,6 +69,7 @@ void register_node_type_geo_points_to_volume(void);
 void register_node_type_geo_sample_texture(void);
 void register_node_type_geo_subdivide(void);
 void register_node_type_geo_subdivision_surface(void);
+void register_node_type_geo_switch(void);
 void register_node_type_geo_transform(void);
 void register_node_type_geo_triangulate(void);
 void register_node_type_geo_volume_to_mesh(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index bf6bf34325a..12e14768c35 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -310,6 +310,7 @@ DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_GRID, 0, "MESH_PRIMITIVE_GRID", Me
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_MAP_RANGE, def_geo_attribute_map_range, "ATTRIBUTE_MAP_RANGE", AttributeMapRange, "Attribute Map Range", "")
 DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_CLAMP, def_geo_attribute_clamp, "ATTRIBUTE_CLAMP", AttributeClamp, "Attribute Clamp", "")
 DefNode(GeometryNode, GEO_NODE_BOUNDING_BOX, 0, "BOUNDING_BOX", BoundBox, "Bounding Box", "")
+DefNode(GeometryNode, GEO_NODE_SWITCH, def_geo_switch, "SWITCH", Switch, "Switch", "")
 
 /* undefine macros */
 #undef DefNode
diff --git a/source/blender/nodes/geometry/nodes/node_geo_switch.cc b/source/blender/nodes/geometry/nodes/node_geo_switch.cc
new file mode 100644
index 00000000000..6736e963184
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_switch.cc
@@ -0,0 +1,163 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "node_geometry_util.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+static bNodeSocketTemplate geo_node_switch_in[] = {
+    {SOCK_BOOLEAN, N_("Switch")},
+
+    {SOCK_FLOAT, N_("A"), 0.0, 0.0, 0.0, 0.0, -FLT_MAX, FLT_MAX},
+    {SOCK_FLOAT, N_("B"), 0.0, 0.0, 0.0, 0.0, -FLT_MAX, FLT_MAX},
+    {SOCK_INT, N_("A"), 0, 0, 0, 0, -100000, 100000},
+    {SOCK_INT, N_("B"), 0, 0, 0, 0, -100000, 100000},
+    {SOCK_BOOLEAN, N_("A")},
+    {SOCK_BOOLEAN, N_("B")},
+    {SOCK_VECTOR, N_("A"), 0.0, 0.0, 0.0, 0.0, -FLT_MAX, FLT_MAX},
+    {SOCK_VECTOR, N_("B"), 0.0, 0.0, 0.0, 0.0, -FLT_MAX, FLT_MAX},
+    {SOCK_RGBA, N_("A"), 0.8, 0.8, 0.8, 1.0},
+    {SOCK_RGBA, N_("B"), 0.8, 0.8, 0.8, 1.0},
+    {SOCK_STRING, N_("A")},
+    {SOCK_STRING, N_("B")},
+    {SOCK_GEOMETRY, N_("A")},
+    {SOCK_GEOMETRY, N_("B")},
+    {SOCK_OBJECT, N_("A")},
+    {SOCK_OBJECT, N_("B")},
+    {SOCK_COLLECTION, N_("A")},
+    {SOCK_COLLECTION, N_("B")},
+    {-1, ""},
+};
+
+static bNodeSocketTemplate geo_node_switch_out[] = {
+    {SOCK_FLOAT, N_("Output")},
+    {SOCK_INT, N_("Output")},
+    {SOCK_BOOLEAN, N_("Output")},
+    {SOCK_VECTOR, N_("Output")},
+    {SOCK_RGBA, N_("Output")},
+    {SOCK_STRING, N_("Output")},
+    {SOCK_GEOMETRY, N_("Output")},
+    {SOCK_OBJECT, N_("Output")},
+    {SOCK_COLLECTION, N_("Output")},
+    {-1, ""},
+};
+
+static void geo_node_switch_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "input_type", 0, "", ICON_NONE);
+}
+
+static void geo_node_switch_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  NodeSwitch *data = (NodeSwitch *)MEM_callocN(sizeof(NodeSwitch), __func__);
+  data->input_type = SOCK_FLOAT;
+  node->storage = data;
+}
+
+namespace blender::nodes {
+
+template<typename T>
+void output_input(GeoNodeExecParams &params,
+                  const bool input,
+                  const StringRef input_suffix,
+                  const StringRef output_identifier)
+{
+  if (input) {
+    params.set_output(output_identifier, params.extract_input<T>("B" + input_suffix));
+  }
+  else {
+    params.set_output(output_identifier, params.extract_input<T>("A" + input_suffix));
+  }
+}
+
+static void geo_node_switch_update(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  NodeSwitch *node_storage = (NodeSwitch *)node->storage;
+  int index = 0;
+  LISTBASE_FOREACH (bNodeSocket *, socket, &node->inputs) {
+    nodeSetSocketAvailability(
+        socket, index == 0 || socket->type == (eNodeSocketDatatype)node_storage->input_type);
+    index++;
+  }
+  LISTBASE_FOREACH (bNodeSocket *, socket, &node->outputs) {
+    nodeSetSocketAvailability(socket,
+                              socket->type == (eNodeSocketDatatype)node_storage->input_type);
+  }
+}
+
+static void geo_node_switch_exec(GeoNodeExecParams params)
+{
+  const NodeSwitch &storage = *(const NodeSwitch *)params.node().storage;
+  const bool input = params.extract_input<bool>("Switch");
+  switch ((eNodeSocketDatatype)storage.input_type) {
+    case SOCK_FLOAT: {
+      output_input<float>(params, input, "", "Output");
+      break;
+    }
+    case SOCK_INT: {
+      output_input<int>(params, input, "_001", "Output_001");
+      break;
+    }
+    case SOCK_BOOLEAN: {
+      output_input<bool>(params, input, "_002", "Output_002");
+      break;
+    }
+    case SOCK_VECTOR: {
+      output_input<float3>(params, input, "_003", "Output_003");
+      break;
+    }
+    case SOCK_RGBA: {
+      output_input<Color4f>(params, input, "_004", "Output_004");
+      break;
+    }
+    case SOCK_STRING: {
+      output_input<std::string>(params, input, "_005", "Output_005");
+      break;
+    }
+    case SOCK_GEOMETRY: {
+      output_input<GeometrySet>(params, input, "_006", "Output_006");
+      break;
+    }
+    case SOCK_OBJECT: {
+      output_input<bke::PersistentObjectHandle>(params, input, "_007", "Output_007");
+      break;
+    }
+    case SOCK_COLLECTION: {
+      output_input<bke::PersistentCollectionHandle>(params, input, "_008", "Output_008");
+      break;
+    }
+    default:
+      BLI_assert_unreachable();
+      break;
+  }
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_switch()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_SWITCH, "Switch", NODE_CLASS_GEOMETRY, 0);
+  node_type_socket_templates(&ntype, geo_node_switch_in, geo_node_switch_out);
+  node_type_init(&ntype, geo_node_switch_init);
+  node_type_update(&ntype, blender::nodes::geo_node_switch_update);
+  node_type_storage(&ntype, "NodeSwitch", node_free_standard_storage, node_copy_standard_storage);
+  ntype.geometry_node_execute = blender::nodes::geo_node_switch_exec;
+  ntype.draw_buttons = geo_node_switch_layout;
+  nodeRegisterType(&ntype);
+}