2 files changed, 113 insertions, 1 deletions

diff --git a/intern/cycles/graph/node.cpp b/intern/cycles/graph/node.cpp
index c926f6ab8ef..57f25283f85 100644
--- a/intern/cycles/graph/node.cpp
+++ b/intern/cycles/graph/node.cpp
@@ -367,9 +367,17 @@ void Node::copy_value(const SocketType &socket, const Node &other, const SocketT
       case SocketType::TRANSFORM_ARRAY:
         copy_array<Transform>(this, socket, &other, other_socket);
         break;
-      case SocketType::NODE_ARRAY:
+      case SocketType::NODE_ARRAY: {
         copy_array<void *>(this, socket, &other, other_socket);
+
+        array<Node *> &node_array = get_socket_value<array<Node *>>(this, socket);
+
+        for (Node *node : node_array) {
+          node->reference();
+        }
+
         break;
+      }
       default:
         assert(0);
         break;
@@ -379,6 +387,14 @@ void Node::copy_value(const SocketType &socket, const Node &other, const SocketT
     const void *src = ((char *)&other) + other_socket.struct_offset;
     void *dst = ((char *)this) + socket.struct_offset;
     memcpy(dst, src, socket.size());
+
+    if (socket.type == SocketType::NODE) {
+      Node *node = get_socket_value<Node *>(this, socket);
+
+      if (node) {
+        node->reference();
+      }
+    }
   }
 }
 
@@ -773,6 +789,26 @@ void Node::set_owner(const NodeOwner *owner_)
   owner = owner_;
 }
 
+void Node::dereference_all_used_nodes()
+{
+  foreach (const SocketType &socket, type->inputs) {
+    if (socket.type == SocketType::NODE) {
+      Node *node = get_socket_value<Node *>(this, socket);
+
+      if (node) {
+        node->dereference();
+      }
+    }
+    else if (socket.type == SocketType::NODE_ARRAY) {
+      const array<Node *> &nodes = get_socket_value<array<Node *>>(this, socket);
+
+      for (Node *node : nodes) {
+        node->dereference();
+      }
+    }
+  }
+}
+
 bool Node::socket_is_modified(const SocketType &input) const
 {
   return (socket_modified & input.modified_flag_bit) != 0;
@@ -803,6 +839,25 @@ template<typename T> void Node::set_if_different(const SocketType &input, T valu
   socket_modified |= input.modified_flag_bit;
 }
 
+void Node::set_if_different(const SocketType &input, Node *value)
+{
+  if (get_socket_value<Node *>(this, input) == value) {
+    return;
+  }
+
+  Node *old_node = get_socket_value<Node *>(this, input);
+  if (old_node) {
+    old_node->dereference();
+  }
+
+  if (value) {
+    value->reference();
+  }
+
+  get_socket_value<Node *>(this, input) = value;
+  socket_modified |= input.modified_flag_bit;
+}
+
 template<typename T> void Node::set_if_different(const SocketType &input, array<T> &value)
 {
   if (!socket_is_modified(input)) {
@@ -815,6 +870,27 @@ template<typename T> void Node::set_if_different(const SocketType &input, array<
   socket_modified |= input.modified_flag_bit;
 }
 
+void Node::set_if_different(const SocketType &input, array<Node *> &value)
+{
+  if (!socket_is_modified(input)) {
+    if (get_socket_value<array<Node *>>(this, input) == value) {
+      return;
+    }
+  }
+
+  array<Node *> &old_nodes = get_socket_value<array<Node *>>(this, input);
+  for (Node *old_node : old_nodes) {
+    old_node->dereference();
+  }
+
+  for (Node *new_node : value) {
+    new_node->reference();
+  }
+
+  get_socket_value<array<Node *>>(this, input).steal_data(value);
+  socket_modified |= input.modified_flag_bit;
+}
+
 void Node::print_modified_sockets() const
 {
   printf("Node : %s\n", name.c_str());
diff --git a/intern/cycles/graph/node.h b/intern/cycles/graph/node.h
index 2fc9a1e0281..aa365baeccd 100644
--- a/intern/cycles/graph/node.h
+++ b/intern/cycles/graph/node.h
@@ -177,8 +177,32 @@ struct Node {
   const NodeOwner *get_owner() const;
   void set_owner(const NodeOwner *owner_);
 
+  int reference_count() const
+  {
+    return ref_count;
+  }
+
+  void reference()
+  {
+    ref_count += 1;
+  }
+
+  void dereference()
+  {
+    ref_count -= 1;
+  }
+
+  /* Set the reference count to zero. This should only be called when we know for sure that the
+   * Node is not used by anyone else. For now, this is only the case when "deleting" shaders, as
+   * they are never actually deleted. */
+  void clear_reference_count()
+  {
+    ref_count = 0;
+  }
+
  protected:
   const NodeOwner *owner;
+  int ref_count{0};
 
   template<typename T> static T &get_socket_value(const Node *node, const SocketType &socket)
   {
@@ -189,7 +213,19 @@ struct Node {
 
   template<typename T> void set_if_different(const SocketType &input, T value);
 
+  /* Explicit overload for Node sockets so we can handle reference counting. The old Node is
+   * dereferenced, and the new one is referenced. */
+  void set_if_different(const SocketType &input, Node *value);
+
   template<typename T> void set_if_different(const SocketType &input, array<T> &value);
+
+  /* Explicit overload for Node sockets so we can handle reference counting. The old Nodes are
+   * dereferenced, and the new ones are referenced. */
+  void set_if_different(const SocketType &input, array<Node *> &value);
+
+  /* Call this function in derived classes' destructors to ensure that used Nodes are dereferenced
+   * properly. */
+  void dereference_all_used_nodes();
 };
 
 CCL_NAMESPACE_END