UI: Use socket type info color to draw links

Currently, colored links overlay only supports standard sockets defined
by Blender. Some add-ons like Animation Nodes defines custom sockets for
everything and hence doesn't get colored sockets.

This patch uses the draw color from the socket type info to draw links
in order to support custom sockets.

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

Reviewed By: Hans Goudey

3 files changed, 32 insertions, 21 deletions

diff --git a/source/blender/editors/space_node/drawnode.cc b/source/blender/editors/space_node/drawnode.cc
index 24f5decacdf..83b7c119b62 100644
--- a/source/blender/editors/space_node/drawnode.cc
+++ b/source/blender/editors/space_node/drawnode.cc
@@ -4268,7 +4268,8 @@ static void nodelink_batch_add_link(const SpaceNode *snode,
 }
 
 /* don't do shadows if th_col3 is -1. */
-void node_draw_link_bezier(const View2D *v2d,
+void node_draw_link_bezier(const bContext *C,
+                           const View2D *v2d,
                            const SpaceNode *snode,
                            const bNodeLink *link,
                            int th_col1,
@@ -4310,18 +4311,21 @@ void node_draw_link_bezier(const View2D *v2d,
         snode->overlay.flag & SN_OVERLAY_SHOW_WIRE_COLORS &&
         ((link->fromsock == nullptr || link->fromsock->typeinfo->type >= 0) &&
          (link->tosock == nullptr || link->tosock->typeinfo->type >= 0))) {
+      PointerRNA from_node_ptr, to_node_ptr;
+      RNA_pointer_create((ID *)snode->edittree, &RNA_Node, link->fromnode, &from_node_ptr);
+      RNA_pointer_create((ID *)snode->edittree, &RNA_Node, link->tonode, &to_node_ptr);
       if (link->fromsock) {
-        copy_v4_v4(colors[1], std_node_socket_colors[link->fromsock->typeinfo->type]);
+        node_socket_color_get(C, snode->edittree, &from_node_ptr, link->fromsock, colors[1]);
       }
       else {
-        copy_v4_v4(colors[1], std_node_socket_colors[link->tosock->typeinfo->type]);
+        node_socket_color_get(C, snode->edittree, &to_node_ptr, link->tosock, colors[1]);
       }
 
       if (link->tosock) {
-        copy_v4_v4(colors[2], std_node_socket_colors[link->tosock->typeinfo->type]);
+        node_socket_color_get(C, snode->edittree, &to_node_ptr, link->tosock, colors[2]);
       }
       else {
-        copy_v4_v4(colors[2], std_node_socket_colors[link->fromsock->typeinfo->type]);
+        node_socket_color_get(C, snode->edittree, &from_node_ptr, link->fromsock, colors[2]);
       }
     }
     else {
@@ -4392,7 +4396,7 @@ void node_draw_link_bezier(const View2D *v2d,
 }
 
 /* NOTE: this is used for fake links in groups too. */
-void node_draw_link(View2D *v2d, SpaceNode *snode, bNodeLink *link)
+void node_draw_link(const bContext *C, View2D *v2d, SpaceNode *snode, bNodeLink *link)
 {
   int th_col1 = TH_WIRE_INNER, th_col2 = TH_WIRE_INNER, th_col3 = TH_WIRE;
 
@@ -4436,7 +4440,7 @@ void node_draw_link(View2D *v2d, SpaceNode *snode, bNodeLink *link)
     }
   }
 
-  node_draw_link_bezier(v2d, snode, link, th_col1, th_col2, th_col3);
+  node_draw_link_bezier(C, v2d, snode, link, th_col1, th_col2, th_col3);
 }
 
 void ED_node_draw_snap(View2D *v2d, const float cent[2], float size, NodeBorder border, uint pos)
diff --git a/source/blender/editors/space_node/node_draw.cc b/source/blender/editors/space_node/node_draw.cc
index a6496294f96..f76c29ef0d9 100644
--- a/source/blender/editors/space_node/node_draw.cc
+++ b/source/blender/editors/space_node/node_draw.cc
@@ -725,12 +725,15 @@ int node_get_colorid(bNode *node)
   }
 }
 
-static void node_draw_mute_line(const View2D *v2d, const SpaceNode *snode, const bNode *node)
+static void node_draw_mute_line(const bContext *C,
+                                const View2D *v2d,
+                                const SpaceNode *snode,
+                                const bNode *node)
 {
   GPU_blend(GPU_BLEND_ALPHA);
 
   LISTBASE_FOREACH (const bNodeLink *, link, &node->internal_links) {
-    node_draw_link_bezier(v2d, snode, link, TH_WIRE_INNER, TH_WIRE_INNER, TH_WIRE);
+    node_draw_link_bezier(C, v2d, snode, link, TH_WIRE_INNER, TH_WIRE_INNER, TH_WIRE);
   }
 
   GPU_blend(GPU_BLEND_NONE);
@@ -825,13 +828,13 @@ static void node_socket_outline_color_get(const bool selected,
 /* Usual convention here would be node_socket_get_color(), but that's already used (for setting a
  * color property socket). */
 void node_socket_color_get(
-    bContext *C, bNodeTree *ntree, PointerRNA *node_ptr, bNodeSocket *sock, float r_color[4])
+    const bContext *C, bNodeTree *ntree, PointerRNA *node_ptr, bNodeSocket *sock, float r_color[4])
 {
   PointerRNA ptr;
   BLI_assert(RNA_struct_is_a(node_ptr->type, &RNA_Node));
   RNA_pointer_create((ID *)ntree, &RNA_NodeSocket, sock, &ptr);
 
-  sock->typeinfo->draw_color(C, &ptr, node_ptr, r_color);
+  sock->typeinfo->draw_color((bContext *)C, &ptr, node_ptr, r_color);
 }
 
 struct SocketTooltipData {
@@ -1049,7 +1052,7 @@ static void node_socket_draw_nested(const bContext *C,
   float color[4];
   float outline_color[4];
 
-  node_socket_color_get((bContext *)C, ntree, node_ptr, sock, color);
+  node_socket_color_get(C, ntree, node_ptr, sock, color);
   node_socket_outline_color_get(selected, sock->type, outline_color);
 
   node_socket_draw(sock,
@@ -1464,7 +1467,7 @@ void node_draw_sockets(const View2D *v2d,
 
     float color[4];
     float outline_color[4];
-    node_socket_color_get((bContext *)C, ntree, &node_ptr, socket, color);
+    node_socket_color_get(C, ntree, &node_ptr, socket, color);
     node_socket_outline_color_get(selected, socket->type, outline_color);
 
     node_socket_draw_multi_input(color, outline_color, width, height, socket->locx, socket->locy);
@@ -1762,7 +1765,7 @@ static void node_draw_basis(const bContext *C,
 
   /* Wire across the node when muted/disabled. */
   if (node->flag & NODE_MUTED) {
-    node_draw_mute_line(v2d, snode, node);
+    node_draw_mute_line(C, v2d, snode, node);
   }
 
   /* Body. */
@@ -1891,7 +1894,7 @@ static void node_draw_hidden(const bContext *C,
 
   /* Wire across the node when muted/disabled. */
   if (node->flag & NODE_MUTED) {
-    node_draw_mute_line(v2d, snode, node);
+    node_draw_mute_line(C, v2d, snode, node);
   }
 
   /* Body. */
@@ -2202,7 +2205,7 @@ void node_draw_nodetree(const bContext *C,
 
   LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
     if (!nodeLinkIsHidden(link)) {
-      node_draw_link(&region->v2d, snode, link);
+      node_draw_link(C, &region->v2d, snode, link);
     }
   }
   nodelink_batch_end(snode);
@@ -2387,7 +2390,7 @@ void node_draw_space(const bContext *C, ARegion *region)
     GPU_line_smooth(true);
     LISTBASE_FOREACH (bNodeLinkDrag *, nldrag, &snode->runtime->linkdrag) {
       LISTBASE_FOREACH (LinkData *, linkdata, &nldrag->links) {
-        node_draw_link(v2d, snode, (bNodeLink *)linkdata->data);
+        node_draw_link(C, v2d, snode, (bNodeLink *)linkdata->data);
       }
     }
     GPU_line_smooth(false);
diff --git a/source/blender/editors/space_node/node_intern.h b/source/blender/editors/space_node/node_intern.h
index c0d50e753ff..383fe5afdf9 100644
--- a/source/blender/editors/space_node/node_intern.h
+++ b/source/blender/editors/space_node/node_intern.h
@@ -121,7 +121,7 @@ void node_draw_sockets(const struct View2D *v2d,
 void node_update_default(const struct bContext *C, struct bNodeTree *ntree, struct bNode *node);
 int node_select_area_default(struct bNode *node, int x, int y);
 int node_tweak_area_default(struct bNode *node, int x, int y);
-void node_socket_color_get(struct bContext *C,
+void node_socket_color_get(const struct bContext *C,
                            struct bNodeTree *ntree,
                            struct PointerRNA *node_ptr,
                            struct bNodeSocket *sock,
@@ -186,8 +186,12 @@ void NODE_OT_backimage_sample(struct wmOperatorType *ot);
 void nodelink_batch_start(struct SpaceNode *snode);
 void nodelink_batch_end(struct SpaceNode *snode);
 
-void node_draw_link(struct View2D *v2d, struct SpaceNode *snode, struct bNodeLink *link);
-void node_draw_link_bezier(const struct View2D *v2d,
+void node_draw_link(const struct bContext *C,
+                    struct View2D *v2d,
+                    struct SpaceNode *snode,
+                    struct bNodeLink *link);
+void node_draw_link_bezier(const struct bContext *C,
+                           const struct View2D *v2d,
                            const struct SpaceNode *snode,
                            const struct bNodeLink *link,
                            int th_col1,
@@ -348,4 +352,4 @@ extern const char *node_context_dir[];
 namespace blender::ed::space_node {
 Vector<ui::ContextPathItem> context_path_for_space_node(const bContext &C);
 }
-#endif
\ No newline at end of file
+#endif