Cycles: add reference counting to Nodes

This adds a reference count to Nodes which is incremented or decremented
whenever they are added to or removed from a socket, which will help us
track used Nodes throughout the scene graph generically without having to
add an explicit count or flag on specific Node types. This is especially
useful to track Nodes defined through Procedurals out of Cycles' control.

This also modifies the order in which nodes are deleted to ensure that
upon deletion, a Node does not attempt to decrement the reference
count of another Node which was already freed or deleted.

This is not currently used, but will be in the next commit.

Reviewed By: brecht

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

2 files changed, 113 insertions, 1 deletions

diff --git a/intern/cycles/graph/node.cpp b/intern/cycles/graph/node.cpp
index c926f6ab8ef..c5b4bb471bb 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