From 03544ed54f8dc5fb28d2c6b655a4153286c0137f Mon Sep 17 00:00:00 2001
From: Jacques Lucke <jacques@blender.org>
Date: Mon, 14 Jun 2021 10:04:32 +0200
Subject: Fix T88807: crash when there are multiple links between the same
 sockets

This commit does two things:

* Disallows creating more than one link from one socket to a multi socket input.
* Properly count links if there happen to be more than one link between the same sockets.

The new link counting should also be more efficient asymptotically.

Differential Revision: https://developer.blender.org/D11570
---
 source/blender/editors/space_node/node_draw.cc | 38 ++++++++++++++------------
 source/blender/makesrna/intern/rna_nodetree.c  |  7 +++++
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/source/blender/editors/space_node/node_draw.cc b/source/blender/editors/space_node/node_draw.cc
index fd9c0f42f2d..1dc8e1412af 100644
--- a/source/blender/editors/space_node/node_draw.cc
+++ b/source/blender/editors/space_node/node_draw.cc
@@ -82,6 +82,7 @@
 #  include "COM_compositor.h"
 #endif
 
+using blender::Map;
 using blender::Set;
 using blender::Span;
 using blender::Vector;
@@ -1767,26 +1768,27 @@ static void node_update(const bContext *C, bNodeTree *ntree, bNode *node)
 
 static void count_mutli_input_socket_links(bNodeTree *ntree, SpaceNode *snode)
 {
+  Map<bNodeSocket *, int> counts;
+  LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
+    if (link->tosock->flag & SOCK_MULTI_INPUT) {
+      int &count = counts.lookup_or_add(link->tosock, 0);
+      count++;
+    }
+  }
+  /* Count temporary links going into this socket. */
+  LISTBASE_FOREACH (bNodeLinkDrag *, nldrag, &snode->runtime->linkdrag) {
+    LISTBASE_FOREACH (LinkData *, linkdata, &nldrag->links) {
+      bNodeLink *link = (bNodeLink *)linkdata->data;
+      if (link->tosock && (link->tosock->flag & SOCK_MULTI_INPUT)) {
+        int &count = counts.lookup_or_add(link->tosock, 0);
+        count++;
+      }
+    }
+  }
   LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
-    LISTBASE_FOREACH (struct bNodeSocket *, socket, &node->inputs) {
+    LISTBASE_FOREACH (bNodeSocket *, socket, &node->inputs) {
       if (socket->flag & SOCK_MULTI_INPUT) {
-        Set<bNodeSocket *> visited_from_sockets;
-        socket->total_inputs = 0;
-        LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
-          if (link->tosock == socket) {
-            visited_from_sockets.add(link->fromsock);
-          }
-        }
-        /* Count temporary links going into this socket. */
-        LISTBASE_FOREACH (bNodeLinkDrag *, nldrag, &snode->runtime->linkdrag) {
-          LISTBASE_FOREACH (LinkData *, linkdata, &nldrag->links) {
-            bNodeLink *link = (bNodeLink *)linkdata->data;
-            if (link->tosock == socket) {
-              visited_from_sockets.add(link->fromsock);
-            }
-          }
-        }
-        socket->total_inputs = visited_from_sockets.size();
+        socket->total_inputs = counts.lookup_default(socket, 0);
       }
     }
   }
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index b625f0fc204..2a1ea3d6716 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -1283,6 +1283,13 @@ static bNodeLink *rna_NodeTree_link_new(bNodeTree *ntree,
     if (nodeCountSocketLinks(ntree, tosock) + 1 > nodeSocketLinkLimit(tosock)) {
       nodeRemSocketLinks(ntree, tosock);
     }
+    if (tosock->flag & SOCK_MULTI_INPUT) {
+      LISTBASE_FOREACH_MUTABLE (bNodeLink *, link, &ntree->links) {
+        if (link->fromsock == fromsock && link->tosock == tosock) {
+          nodeRemLink(ntree, link);
+        }
+      }
+    }
   }
 
   ret = nodeAddLink(ntree, fromnode, fromsock, tonode, tosock);
-- 
cgit v1.2.3