Nodes: Add Compare node operations to link drag search menu

Exposes compare operations via rna emums.
This uses the rna enum to build the search list using
named operations linked to socket A.
This also weights the Math Node comparison operations lower
for geometry node trees.

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

2 files changed, 36 insertions, 17 deletions

diff --git a/source/blender/nodes/function/nodes/node_fn_compare.cc b/source/blender/nodes/function/nodes/node_fn_compare.cc
index e1c15ac2be1..0607acc95f2 100644
--- a/source/blender/nodes/function/nodes/node_fn_compare.cc
+++ b/source/blender/nodes/function/nodes/node_fn_compare.cc
@@ -130,21 +130,33 @@ static void node_compare_gather_link_searches(GatherLinkSearchOpParams &params)
 
   const eNodeSocketDatatype type = static_cast<eNodeSocketDatatype>(params.other_socket().type);
 
-  if (ELEM(type, SOCK_FLOAT, SOCK_BOOLEAN, SOCK_RGBA, SOCK_VECTOR, SOCK_INT)) {
-    params.add_item(IFACE_("A"), SocketSearchOp{"A", type, NODE_COMPARE_GREATER_THAN});
-    params.add_item(IFACE_("B"), SocketSearchOp{"B", type, NODE_COMPARE_GREATER_THAN});
-    params.add_item(
-        IFACE_("C"),
-        SocketSearchOp{
-            "C", SOCK_VECTOR, NODE_COMPARE_GREATER_THAN, NODE_COMPARE_MODE_DOT_PRODUCT});
-    params.add_item(
-        IFACE_("Angle"),
-        SocketSearchOp{
-            "Angle", SOCK_VECTOR, NODE_COMPARE_GREATER_THAN, NODE_COMPARE_MODE_DIRECTION});
-  }
-  else if (type == SOCK_STRING) {
-    params.add_item(IFACE_("A"), SocketSearchOp{"A", type, NODE_COMPARE_EQUAL});
-    params.add_item(IFACE_("B"), SocketSearchOp{"B", type, NODE_COMPARE_EQUAL});
+  if (ELEM(type, SOCK_BOOLEAN, SOCK_FLOAT, SOCK_RGBA, SOCK_VECTOR, SOCK_INT, SOCK_STRING)) {
+    const eNodeSocketDatatype mode_type = (type == SOCK_BOOLEAN) ? SOCK_FLOAT : type;
+    const bool string_type = (type == SOCK_STRING);
+    /* Add socket A compare operations. */
+    for (const EnumPropertyItem *item = rna_enum_node_compare_operation_items;
+         item->identifier != nullptr;
+         item++) {
+      if (item->name != nullptr && item->identifier[0] != '\0') {
+        if (!string_type &&
+            ELEM(item->value, NODE_COMPARE_COLOR_BRIGHTER, NODE_COMPARE_COLOR_DARKER)) {
+          params.add_item(IFACE_(item->name),
+                          SocketSearchOp{"A", SOCK_RGBA, (NodeCompareOperation)item->value});
+        }
+        else if ((!string_type) ||
+                 (string_type && ELEM(item->value, NODE_COMPARE_EQUAL, NODE_COMPARE_NOT_EQUAL))) {
+          params.add_item(IFACE_(item->name),
+                          SocketSearchOp{"A", mode_type, (NodeCompareOperation)item->value});
+        }
+      }
+    }
+    /* Add Angle socket. */
+    if (!string_type) {
+      params.add_item(
+          IFACE_("Angle"),
+          SocketSearchOp{
+              "Angle", SOCK_VECTOR, NODE_COMPARE_GREATER_THAN, NODE_COMPARE_MODE_DIRECTION});
+    }
   }
 }
 
diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index cff3d5f347a..5cbecfcccc1 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -71,14 +71,21 @@ static void sh_node_math_gather_link_searches(GatherLinkSearchOpParams &params)
   if (params.node_tree().typeinfo->validate_link(
           static_cast<eNodeSocketDatatype>(params.other_socket().type), SOCK_FLOAT)) {
 
+    const bool is_geometry_node_tree = params.node_tree().type == NTREE_GEOMETRY;
     const int weight = ELEM(params.other_socket().type, SOCK_FLOAT, SOCK_BOOLEAN, SOCK_INT) ? 0 :
                                                                                               -1;
 
     for (const EnumPropertyItem *item = rna_enum_node_math_items; item->identifier != nullptr;
          item++) {
       if (item->name != nullptr && item->identifier[0] != '\0') {
-        params.add_item(
-            IFACE_(item->name), SocketSearchOp{"Value", (NodeMathOperation)item->value}, weight);
+        const int gn_weight =
+            (is_geometry_node_tree &&
+             ELEM(item->value, NODE_MATH_COMPARE, NODE_MATH_GREATER_THAN, NODE_MATH_LESS_THAN)) ?
+                -1 :
+                weight;
+        params.add_item(IFACE_(item->name),
+                        SocketSearchOp{"Value", (NodeMathOperation)item->value},
+                        gn_weight);
       }
     }
   }