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

7 files changed, 152 insertions, 16 deletions

diff --git a/intern/cycles/blender/blender_object.cpp b/intern/cycles/blender/blender_object.cpp
index 54128cf82fc..619cdf98fa1 100644
--- a/intern/cycles/blender/blender_object.cpp
+++ b/intern/cycles/blender/blender_object.cpp
@@ -560,10 +560,12 @@ void BlenderSync::sync_objects(BL::Depsgraph &b_depsgraph,
   if (!cancel && !motion) {
     sync_background_light(b_v3d, use_portal);
 
-    /* handle removed data and modified pointers */
+    /* Handle removed data and modified pointers, as this may free memory, delete Nodes in the
+     * right order to ensure that dependant data is freed after their users. Objects should be
+     * freed before particle systems and geometries. */
     light_map.post_sync();
-    geometry_map.post_sync();
     object_map.post_sync();
+    geometry_map.post_sync();
     particle_system_map.post_sync();
   }
 
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
diff --git a/intern/cycles/render/background.cpp b/intern/cycles/render/background.cpp
index f0a779da012..b925e755434 100644
--- a/intern/cycles/render/background.cpp
+++ b/intern/cycles/render/background.cpp
@@ -59,6 +59,7 @@ Background::Background() : Node(get_node_type())
 
 Background::~Background()
 {
+  dereference_all_used_nodes();
 }
 
 void Background::device_update(Device *device, DeviceScene *dscene, Scene *scene)
diff --git a/intern/cycles/render/geometry.cpp b/intern/cycles/render/geometry.cpp
index bd2b8164cbb..aac5af06e9f 100644
--- a/intern/cycles/render/geometry.cpp
+++ b/intern/cycles/render/geometry.cpp
@@ -79,6 +79,7 @@ Geometry::Geometry(const NodeType *node_type, const Type type)
 
 Geometry::~Geometry()
 {
+  dereference_all_used_nodes();
   delete bvh;
 }
 
diff --git a/intern/cycles/render/light.cpp b/intern/cycles/render/light.cpp
index 858b177b69c..9b975daff0b 100644
--- a/intern/cycles/render/light.cpp
+++ b/intern/cycles/render/light.cpp
@@ -159,6 +159,7 @@ NODE_DEFINE(Light)
 
 Light::Light() : Node(get_node_type())
 {
+  dereference_all_used_nodes();
 }
 
 void Light::tag_update(Scene *scene)
diff --git a/intern/cycles/render/scene.cpp b/intern/cycles/render/scene.cpp
index 38e8d9145dc..dca545fa56c 100644
--- a/intern/cycles/render/scene.cpp
+++ b/intern/cycles/render/scene.cpp
@@ -143,21 +143,27 @@ void Scene::free_memory(bool final)
   delete bvh;
   bvh = NULL;
 
-  foreach (Shader *s, shaders)
-    delete s;
-  /* delete procedurals before other types as they may hold pointers to those types */
+  /* The order of deletion is important to make sure data is freed based on possible dependencies
+   * as the Nodes' reference counts are decremented in the destructors:
+   *
+   * - Procedurals can create and hold pointers to any other types.
+   * - Objects can hold pointers to Geometries and ParticleSystems
+   * - Lights and Geometries can hold pointers to Shaders.
+   *
+   * Similarly, we first delete all nodes and their associated device data, and then the managers
+   * and their associated device data.
+   */
   foreach (Procedural *p, procedurals)
     delete p;
-  foreach (Geometry *g, geometry)
-    delete g;
   foreach (Object *o, objects)
     delete o;
-  foreach (Light *l, lights)
-    delete l;
+  foreach (Geometry *g, geometry)
+    delete g;
   foreach (ParticleSystem *p, particle_systems)
     delete p;
+  foreach (Light *l, lights)
+    delete l;
 
-  shaders.clear();
   geometry.clear();
   objects.clear();
   lights.clear();
@@ -169,7 +175,25 @@ void Scene::free_memory(bool final)
     film->device_free(device, &dscene, this);
     background->device_free(device, &dscene);
     integrator->device_free(device, &dscene, true);
+  }
+
+  if (final) {
+    delete camera;
+    delete dicing_camera;
+    delete film;
+    delete background;
+    delete integrator;
+  }
+
+  /* Delete Shaders after every other nodes to ensure that we do not try to decrement the reference
+   * count on some dangling pointer. */
+  foreach (Shader *s, shaders)
+    delete s;
 
+  shaders.clear();
+
+  /* Now that all nodes have been deleted, we can safely delete managers and device data. */
+  if (device) {
     object_manager->device_free(device, &dscene, true);
     geometry_manager->device_free(device, &dscene, true);
     shader_manager->device_free(device, &dscene, this);
@@ -189,11 +213,6 @@ void Scene::free_memory(bool final)
 
   if (final) {
     delete lookup_tables;
-    delete camera;
-    delete dicing_camera;
-    delete film;
-    delete background;
-    delete integrator;
     delete object_manager;
     delete geometry_manager;
     delete shader_manager;