Geometry Nodes: Generalized Compare Node

Replace compare floats node with a generalized compare node. The node
allows for the comparison of float, int, string, color, and vector.

The datatypes support the following operators:
Float, Int: <, >, <=, >=, ==, !=
String: ==, !=
Color: ==, !=, lighter, darker
    (using rgb_to_grayscale value as the brightness value)

Vector Supports 5 comparison modes for: ==, !=, <, >, <=, >=
Average: The average of the components of the vectors are compared.
Dot Product: The dot product of the vectors are compared.
Direction: The angle between the vectors is compared to an angle
Element-wise: The individual components of the vectors are compared.
Length: The lengths of the vectors are compared.

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

2 files changed, 178 insertions, 13 deletions

diff --git a/source/blender/makesrna/RNA_enum_items.h b/source/blender/makesrna/RNA_enum_items.h
index 78dd977af20..05f4816fcad 100644
--- a/source/blender/makesrna/RNA_enum_items.h
+++ b/source/blender/makesrna/RNA_enum_items.h
@@ -172,6 +172,7 @@ DEF_ENUM(rna_enum_mapping_type_items)
 DEF_ENUM(rna_enum_node_vec_math_items)
 DEF_ENUM(rna_enum_node_boolean_math_items)
 DEF_ENUM(rna_enum_node_float_compare_items)
+DEF_ENUM(rna_enum_node_compare_operation_items)
 DEF_ENUM(rna_enum_node_filter_items)
 DEF_ENUM(rna_enum_node_float_to_int_items)
 DEF_ENUM(rna_enum_node_map_range_items)
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 2db517d6b58..6469607f292 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -305,36 +305,68 @@ const EnumPropertyItem rna_enum_node_boolean_math_items[] = {
 };
 
 const EnumPropertyItem rna_enum_node_float_compare_items[] = {
-    {NODE_FLOAT_COMPARE_LESS_THAN,
+    {NODE_COMPARE_LESS_THAN,
      "LESS_THAN",
      0,
      "Less Than",
      "True when the first input is smaller than second input"},
-    {NODE_FLOAT_COMPARE_LESS_EQUAL,
+    {NODE_COMPARE_LESS_EQUAL,
      "LESS_EQUAL",
      0,
      "Less Than or Equal",
      "True when the first input is smaller than the second input or equal"},
-    {NODE_FLOAT_COMPARE_GREATER_THAN,
+    {NODE_COMPARE_GREATER_THAN,
      "GREATER_THAN",
      0,
      "Greater Than",
      "True when the first input is greater than the second input"},
-    {NODE_FLOAT_COMPARE_GREATER_EQUAL,
+    {NODE_COMPARE_GREATER_EQUAL,
      "GREATER_EQUAL",
      0,
      "Greater Than or Equal",
      "True when the first input is greater than the second input or equal"},
-    {NODE_FLOAT_COMPARE_EQUAL,
-     "EQUAL",
+    {NODE_COMPARE_EQUAL, "EQUAL", 0, "Equal", "True when both inputs are approximately equal"},
+    {NODE_COMPARE_NOT_EQUAL,
+     "NOT_EQUAL",
+     0,
+     "Not Equal",
+     "True when both inputs are not approximately equal"},
+    {0, NULL, 0, NULL, NULL},
+};
+
+const EnumPropertyItem rna_enum_node_compare_operation_items[] = {
+    {NODE_COMPARE_LESS_THAN,
+     "LESS_THAN",
+     0,
+     "Less Than",
+     "True when the first input is smaller than second input"},
+    {NODE_COMPARE_LESS_EQUAL,
+     "LESS_EQUAL",
      0,
-     "Equal",
-     "True when both inputs are approximately equal"},
-    {NODE_FLOAT_COMPARE_NOT_EQUAL,
+     "Less Than or Equal",
+     "True when the first input is smaller than the second input or equal"},
+    {NODE_COMPARE_GREATER_THAN,
+     "GREATER_THAN",
+     0,
+     "Greater Than",
+     "True when the first input is greater than the second input"},
+    {NODE_COMPARE_GREATER_EQUAL,
+     "GREATER_EQUAL",
+     0,
+     "Greater Than or Equal",
+     "True when the first input is greater than the second input or equal"},
+    {NODE_COMPARE_EQUAL, "EQUAL", 0, "Equal", "True when both inputs are approximately equal"},
+    {NODE_COMPARE_NOT_EQUAL,
      "NOT_EQUAL",
      0,
      "Not Equal",
      "True when both inputs are not approximately equal"},
+    {NODE_COMPARE_COLOR_BRIGHTER,
+     "BRIGHTER",
+     0,
+     "Brighter",
+     "True when the first input is brighter"},
+    {NODE_COMPARE_COLOR_DARKER, "DARKER", 0, "Darker", "True when the first input is darker"},
     {0, NULL, 0, NULL, NULL},
 };
 
@@ -2070,10 +2102,76 @@ static const EnumPropertyItem *rna_GeometryNodeSwitch_type_itemf(bContext *UNUSE
   return itemf_function_check(node_socket_data_type_items, switch_type_supported);
 }
 
+static bool compare_type_supported(const EnumPropertyItem *item)
+{
+  return ELEM(item->value, SOCK_FLOAT, SOCK_INT, SOCK_VECTOR, SOCK_STRING, SOCK_RGBA);
+}
+
+static bool compare_main_operation_supported(const EnumPropertyItem *item)
+{
+  return !ELEM(item->value, NODE_COMPARE_COLOR_BRIGHTER, NODE_COMPARE_COLOR_DARKER);
+}
+
+static bool compare_rgba_operation_supported(const EnumPropertyItem *item)
+{
+  return ELEM(item->value,
+              NODE_COMPARE_EQUAL,
+              NODE_COMPARE_NOT_EQUAL,
+              NODE_COMPARE_COLOR_BRIGHTER,
+              NODE_COMPARE_COLOR_DARKER);
+}
+
+static bool compare_string_operation_supported(const EnumPropertyItem *item)
+{
+  return ELEM(item->value, NODE_COMPARE_EQUAL, NODE_COMPARE_NOT_EQUAL);
+}
+
+static bool compare_other_operation_supported(const EnumPropertyItem *UNUSED(item))
+{
+  return false;
+}
+
+static const EnumPropertyItem *rna_FunctionNodeCompare_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, compare_type_supported);
+}
+
+static const EnumPropertyItem *rna_FunctionNodeCompare_operation_itemf(bContext *UNUSED(C),
+                                                                       PointerRNA *ptr,
+                                                                       PropertyRNA *UNUSED(prop),
+                                                                       bool *r_free)
+{
+  *r_free = true;
+  bNode *node = ptr->data;
+  NodeFunctionCompare *data = (NodeFunctionCompare *)node->storage;
+
+  if (ELEM(data->data_type, SOCK_FLOAT, SOCK_INT, SOCK_VECTOR)) {
+    return itemf_function_check(rna_enum_node_compare_operation_items,
+                                compare_main_operation_supported);
+  }
+  else if (data->data_type == SOCK_STRING) {
+    return itemf_function_check(rna_enum_node_compare_operation_items,
+                                compare_string_operation_supported);
+  }
+  else if (data->data_type == SOCK_RGBA) {
+    return itemf_function_check(rna_enum_node_compare_operation_items,
+                                compare_rgba_operation_supported);
+  }
+  else {
+    return itemf_function_check(rna_enum_node_compare_operation_items,
+                                compare_other_operation_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);
 }
+
 static const EnumPropertyItem *rna_GeometryNodeAttributeClamp_type_itemf(bContext *UNUSED(C),
                                                                          PointerRNA *UNUSED(ptr),
                                                                          PropertyRNA *UNUSED(prop),
@@ -2150,6 +2248,30 @@ static void rna_GeometryNodeAttributeRandomize_data_type_update(Main *bmain,
   rna_Node_socket_update(bmain, scene, ptr);
 }
 
+static void rna_GeometryNodeCompare_data_type_update(Main *bmain, Scene *scene, PointerRNA *ptr)
+{
+  bNode *node = ptr->data;
+  NodeFunctionCompare *node_storage = (NodeFunctionCompare *)node->storage;
+
+  if (node_storage->data_type == SOCK_RGBA && !ELEM(node_storage->operation,
+                                                    NODE_COMPARE_EQUAL,
+                                                    NODE_COMPARE_NOT_EQUAL,
+                                                    NODE_COMPARE_COLOR_BRIGHTER,
+                                                    NODE_COMPARE_COLOR_DARKER)) {
+    node_storage->operation = NODE_COMPARE_EQUAL;
+  }
+  else if (node_storage->data_type == SOCK_STRING &&
+           !ELEM(node_storage->operation, NODE_COMPARE_EQUAL, NODE_COMPARE_NOT_EQUAL)) {
+    node_storage->operation = NODE_COMPARE_EQUAL;
+  }
+  else if (node_storage->data_type != SOCK_RGBA &&
+           ELEM(node_storage->operation, NODE_COMPARE_COLOR_BRIGHTER, NODE_COMPARE_COLOR_DARKER)) {
+    node_storage->operation = NODE_COMPARE_EQUAL;
+  }
+
+  rna_Node_socket_update(bmain, scene, ptr);
+}
+
 static bool attribute_convert_type_supported(const EnumPropertyItem *item)
 {
   return ELEM(item->value,
@@ -4865,15 +4987,57 @@ static void def_boolean_math(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
 }
 
-static void def_float_compare(StructRNA *srna)
+static void def_compare(StructRNA *srna)
 {
+
+  static const EnumPropertyItem mode_items[] = {
+      {NODE_COMPARE_MODE_ELEMENT,
+       "ELEMENT",
+       0,
+       "Element-Wise",
+       "Compare each element of the input vectors"},
+      {NODE_COMPARE_MODE_LENGTH, "LENGTH", 0, "Length", "Compare the length of the input vectors"},
+      {NODE_COMPARE_MODE_AVERAGE,
+       "AVERAGE",
+       0,
+       "Average",
+       "Compare the average of the input vectors elements"},
+      {NODE_COMPARE_MODE_DOT_PRODUCT,
+       "DOT_PRODUCT",
+       0,
+       "Dot Product",
+       "Compare the dot products of the input vectors"},
+      {NODE_COMPARE_MODE_DIRECTION,
+       "DIRECTION",
+       0,
+       "Direction",
+       "Compare the direction of the input vectors"},
+      {0, NULL, 0, NULL, NULL},
+  };
+
   PropertyRNA *prop;
 
+  RNA_def_struct_sdna_from(srna, "NodeFunctionCompare", "storage");
+
   prop = RNA_def_property(srna, "operation", PROP_ENUM, PROP_NONE);
-  RNA_def_property_enum_sdna(prop, NULL, "custom1");
-  RNA_def_property_enum_items(prop, rna_enum_node_float_compare_items);
+  RNA_def_property_enum_funcs(prop, NULL, NULL, "rna_FunctionNodeCompare_operation_itemf");
+  RNA_def_property_enum_items(prop, rna_enum_node_compare_operation_items);
+  RNA_def_property_enum_default(prop, NODE_COMPARE_EQUAL);
   RNA_def_property_ui_text(prop, "Operation", "");
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+
+  prop = RNA_def_property(srna, "data_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_funcs(prop, NULL, NULL, "rna_FunctionNodeCompare_type_itemf");
+  RNA_def_property_enum_items(prop, node_socket_data_type_items);
+  RNA_def_property_enum_default(prop, SOCK_FLOAT);
+  RNA_def_property_ui_text(prop, "Input Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_GeometryNodeCompare_data_type_update");
+
+  prop = RNA_def_property(srna, "mode", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_items(prop, mode_items);
+  RNA_def_property_enum_default(prop, NODE_COMPARE_MODE_ELEMENT);
+  RNA_def_property_ui_text(prop, "Mode", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
 }
 
 static void def_float_to_int(StructRNA *srna)
@@ -9589,7 +9753,7 @@ static void def_geo_attribute_attribute_compare(StructRNA *srna)
 
   prop = RNA_def_property(srna, "operation", PROP_ENUM, PROP_NONE);
   RNA_def_property_enum_items(prop, rna_enum_node_float_compare_items);
-  RNA_def_property_enum_default(prop, NODE_FLOAT_COMPARE_GREATER_THAN);
+  RNA_def_property_enum_default(prop, NODE_COMPARE_GREATER_THAN);
   RNA_def_property_ui_text(prop, "Operation", "");
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");