Simulation: move responsibility of updating persistent data handles

This solver should not be responsible for updating the handles.

5 files changed, 101 insertions, 70 deletions

diff --git a/source/blender/simulation/intern/simulation_collect_influences.cc b/source/blender/simulation/intern/simulation_collect_influences.cc
index 764e587d157..bbc28077dff 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.cc
+++ b/source/blender/simulation/intern/simulation_collect_influences.cc
@@ -23,6 +23,8 @@
 
 #include "NOD_node_tree_multi_function.hh"
 
+#include "DEG_depsgraph_query.h"
+
 #include "BLI_rand.hh"
 
 namespace blender::sim {
@@ -434,8 +436,8 @@ static void prepare_particle_attribute_builders(nodes::MFNetworkTreeMap &network
   }
 }
 
-static void find_used_data_blocks(const nodes::DerivedNodeTree &tree,
-                                  SimulationInfluences &r_influences)
+static void find_used_persistent_data(const nodes::DerivedNodeTree &tree,
+                                      UsedPersistentData &r_used_persistent_data)
 {
   const bNodeSocketType *socktype = nodeSocketTypeFind("NodeSocketObject");
   BLI_assert(socktype != nullptr);
@@ -443,9 +445,9 @@ static void find_used_data_blocks(const nodes::DerivedNodeTree &tree,
   for (const nodes::DInputSocket *dsocket : tree.input_sockets()) {
     const bNodeSocket *bsocket = dsocket->bsocket();
     if (bsocket->typeinfo == socktype) {
-      Object *value = ((const bNodeSocketValueObject *)bsocket->default_value)->value;
-      if (value != nullptr) {
-        r_influences.used_data_blocks.add(&value->id);
+      Object *object = ((const bNodeSocketValueObject *)bsocket->default_value)->value;
+      if (object != nullptr) {
+        r_used_persistent_data.add(DEG_get_original_id(&object->id));
       }
     }
   }
@@ -454,7 +456,8 @@ static void find_used_data_blocks(const nodes::DerivedNodeTree &tree,
 void collect_simulation_influences(Simulation &simulation,
                                    ResourceCollector &resources,
                                    SimulationInfluences &r_influences,
-                                   RequiredStates &r_required_states)
+                                   RequiredStates &r_required_states,
+                                   UsedPersistentData &r_used_persistent_data)
 {
   nodes::NodeTreeRefMap tree_refs;
   const nodes::DerivedNodeTree tree{simulation.nodetree, tree_refs};
@@ -479,7 +482,7 @@ void collect_simulation_influences(Simulation &simulation,
     r_required_states.add(dnode_to_path(*dnode), SIM_TYPE_NAME_PARTICLE_SIMULATION);
   }
 
-  find_used_data_blocks(tree, r_influences);
+  find_used_persistent_data(tree, r_used_persistent_data);
 }
 
 }  // namespace blender::sim
diff --git a/source/blender/simulation/intern/simulation_collect_influences.hh b/source/blender/simulation/intern/simulation_collect_influences.hh
index 42cbea6977e..13744c3bc1a 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.hh
+++ b/source/blender/simulation/intern/simulation_collect_influences.hh
@@ -58,10 +58,29 @@ class RequiredStates {
   }
 };
 
+class UsedPersistentData {
+ private:
+  VectorSet<ID *> used_ids_;
+
+ public:
+  void add(ID *id)
+  {
+    BLI_assert(id != nullptr);
+    BLI_assert((id->tag & LIB_TAG_NO_MAIN) == 0);
+    used_ids_.add(id);
+  }
+
+  const VectorSet<ID *> &used_ids() const
+  {
+    return used_ids_;
+  }
+};
+
 void collect_simulation_influences(Simulation &simulation,
                                    ResourceCollector &resources,
                                    SimulationInfluences &r_influences,
-                                   RequiredStates &r_required_states);
+                                   RequiredStates &r_required_states,
+                                   UsedPersistentData &r_used_persistent_data);
 
 }  // namespace blender::sim
 
diff --git a/source/blender/simulation/intern/simulation_solver.cc b/source/blender/simulation/intern/simulation_solver.cc
index 1628d9cac15..9c70bdaa7b6 100644
--- a/source/blender/simulation/intern/simulation_solver.cc
+++ b/source/blender/simulation/intern/simulation_solver.cc
@@ -17,12 +17,13 @@
 #include "simulation_solver.hh"
 
 #include "BKE_customdata.h"
-#include "BKE_lib_id.h"
 #include "BKE_persistent_data_handle.hh"
 
 #include "BLI_rand.hh"
 #include "BLI_set.hh"
 
+#include "DEG_depsgraph_query.h"
+
 namespace blender::sim {
 
 ParticleForce::~ParticleForce()
@@ -246,57 +247,10 @@ BLI_NOINLINE static void remove_dead_and_add_new_particles(ParticleSimulationSta
   state.next_particle_id += allocator.total_allocated();
 }
 
-static void update_persistent_data_handles(Simulation &simulation,
-                                           const VectorSet<ID *> &used_data_blocks)
-{
-  Set<ID *> contained_ids;
-  Set<int> used_handles;
-
-  /* Remove handles that have been invalidated. */
-  LISTBASE_FOREACH_MUTABLE (
-      PersistentDataHandleItem *, handle_item, &simulation.persistent_data_handles) {
-    if (handle_item->id == nullptr) {
-      BLI_remlink(&simulation.persistent_data_handles, handle_item);
-      MEM_freeN(handle_item);
-      continue;
-    }
-    if (!used_data_blocks.contains(handle_item->id)) {
-      id_us_min(handle_item->id);
-      BLI_remlink(&simulation.persistent_data_handles, handle_item);
-      MEM_freeN(handle_item);
-      continue;
-    }
-    contained_ids.add_new(handle_item->id);
-    used_handles.add_new(handle_item->handle);
-  }
-
-  /* Add new handles that are not in the list yet. */
-  int next_handle = 0;
-  for (ID *id : used_data_blocks) {
-    if (contained_ids.contains(id)) {
-      continue;
-    }
-
-    /* Find the next available handle. */
-    while (used_handles.contains(next_handle)) {
-      next_handle++;
-    }
-    used_handles.add_new(next_handle);
-
-    PersistentDataHandleItem *handle_item = (PersistentDataHandleItem *)MEM_callocN(
-        sizeof(*handle_item), AT);
-    /* Cannot store const pointers in DNA. */
-    id_us_plus(id);
-    handle_item->id = id;
-    handle_item->handle = next_handle;
-
-    BLI_addtail(&simulation.persistent_data_handles, handle_item);
-  }
-}
-
 void initialize_simulation_states(Simulation &simulation,
                                   Depsgraph &UNUSED(depsgraph),
-                                  const SimulationInfluences &UNUSED(influences))
+                                  const SimulationInfluences &UNUSED(influences),
+                                  const bke::PersistentDataHandleMap &UNUSED(handle_map))
 {
   simulation.current_simulation_time = 0.0f;
 }
@@ -304,15 +258,9 @@ void initialize_simulation_states(Simulation &simulation,
 void solve_simulation_time_step(Simulation &simulation,
                                 Depsgraph &depsgraph,
                                 const SimulationInfluences &influences,
+                                const bke::PersistentDataHandleMap &handle_map,
                                 float time_step)
 {
-  update_persistent_data_handles(simulation, influences.used_data_blocks);
-
-  bke::PersistentDataHandleMap handle_map;
-  LISTBASE_FOREACH (PersistentDataHandleItem *, handle, &simulation.persistent_data_handles) {
-    handle_map.add(handle->handle, *handle->id);
-  }
-
   SimulationStateMap state_map;
   LISTBASE_FOREACH (SimulationState *, state, &simulation.states) {
     state_map.add(state);
diff --git a/source/blender/simulation/intern/simulation_solver.hh b/source/blender/simulation/intern/simulation_solver.hh
index 2dc0921ad9e..02c58da37f5 100644
--- a/source/blender/simulation/intern/simulation_solver.hh
+++ b/source/blender/simulation/intern/simulation_solver.hh
@@ -53,7 +53,6 @@ struct SimulationInfluences {
   Map<std::string, Vector<const ParticleForce *>> particle_forces;
   Map<std::string, fn::AttributesInfoBuilder *> particle_attributes_builder;
   Vector<const ParticleEmitter *> particle_emitters;
-  VectorSet<ID *> used_data_blocks;
 };
 
 class SimulationStateMap {
@@ -265,11 +264,13 @@ class ParticleForceContext {
 
 void initialize_simulation_states(Simulation &simulation,
                                   Depsgraph &depsgraph,
-                                  const SimulationInfluences &influences);
+                                  const SimulationInfluences &influences,
+                                  const bke::PersistentDataHandleMap &handle_map);
 
 void solve_simulation_time_step(Simulation &simulation,
                                 Depsgraph &depsgraph,
                                 const SimulationInfluences &influences,
+                                const bke::PersistentDataHandleMap &handle_map,
                                 float time_step);
 
 }  // namespace blender::sim
diff --git a/source/blender/simulation/intern/simulation_update.cc b/source/blender/simulation/intern/simulation_update.cc
index 7a748ee57d1..117a98d2285 100644
--- a/source/blender/simulation/intern/simulation_update.cc
+++ b/source/blender/simulation/intern/simulation_update.cc
@@ -17,6 +17,7 @@
 #include "SIM_simulation_update.hh"
 
 #include "BKE_customdata.h"
+#include "BKE_lib_id.h"
 #include "BKE_simulation.h"
 
 #include "DNA_scene_types.h"
@@ -29,6 +30,7 @@
 #include "BLI_listbase.h"
 #include "BLI_map.hh"
 #include "BLI_rand.h"
+#include "BLI_set.hh"
 #include "BLI_vector.hh"
 
 #include "particle_function.hh"
@@ -88,6 +90,54 @@ static void update_simulation_state_list(Simulation *simulation,
   add_missing_states(simulation, required_states);
 }
 
+static void update_persistent_data_handles(Simulation &simulation_orig,
+                                           const UsedPersistentData &used_persistent_data)
+{
+  Set<ID *> contained_ids;
+  Set<int> used_handles;
+
+  /* Remove handles that have been invalidated. */
+  LISTBASE_FOREACH_MUTABLE (
+      PersistentDataHandleItem *, handle_item, &simulation_orig.persistent_data_handles) {
+    if (handle_item->id == nullptr) {
+      BLI_remlink(&simulation_orig.persistent_data_handles, handle_item);
+      MEM_freeN(handle_item);
+      continue;
+    }
+    if (!used_persistent_data.used_ids().contains(handle_item->id)) {
+      id_us_min(handle_item->id);
+      BLI_remlink(&simulation_orig.persistent_data_handles, handle_item);
+      MEM_freeN(handle_item);
+      continue;
+    }
+    contained_ids.add_new(handle_item->id);
+    used_handles.add_new(handle_item->handle);
+  }
+
+  /* Add new handles that are not in the list yet. */
+  int next_handle = 0;
+  for (ID *id : used_persistent_data.used_ids()) {
+    if (contained_ids.contains(id)) {
+      continue;
+    }
+
+    /* Find the next available handle. */
+    while (used_handles.contains(next_handle)) {
+      next_handle++;
+    }
+    used_handles.add_new(next_handle);
+
+    PersistentDataHandleItem *handle_item = (PersistentDataHandleItem *)MEM_callocN(
+        sizeof(*handle_item), AT);
+    /* Cannot store const pointers in DNA. */
+    id_us_plus(id);
+    handle_item->id = id;
+    handle_item->handle = next_handle;
+
+    BLI_addtail(&simulation_orig.persistent_data_handles, handle_item);
+  }
+}
+
 void update_simulation_in_depsgraph(Depsgraph *depsgraph,
                                     Scene *scene_cow,
                                     Simulation *simulation_cow)
@@ -107,13 +157,23 @@ void update_simulation_in_depsgraph(Depsgraph *depsgraph,
   ResourceCollector resources;
   SimulationInfluences influences;
   RequiredStates required_states;
+  UsedPersistentData used_persistent_data;
+
+  collect_simulation_influences(
+      *simulation_cow, resources, influences, required_states, used_persistent_data);
+  update_persistent_data_handles(*simulation_orig, used_persistent_data);
 
-  collect_simulation_influences(*simulation_cow, resources, influences, required_states);
+  bke::PersistentDataHandleMap handle_map;
+  LISTBASE_FOREACH (
+      PersistentDataHandleItem *, handle_item, &simulation_orig->persistent_data_handles) {
+    ID *id_cow = DEG_get_evaluated_id(depsgraph, handle_item->id);
+    handle_map.add(handle_item->handle, *id_cow);
+  }
 
   if (current_frame == 1) {
     reinitialize_empty_simulation_states(simulation_orig, required_states);
 
-    initialize_simulation_states(*simulation_orig, *depsgraph, influences);
+    initialize_simulation_states(*simulation_orig, *depsgraph, influences, handle_map);
     simulation_orig->current_frame = 1;
 
     copy_states_to_cow(simulation_orig, simulation_cow);
@@ -122,7 +182,7 @@ void update_simulation_in_depsgraph(Depsgraph *depsgraph,
     update_simulation_state_list(simulation_orig, required_states);
 
     float time_step = 1.0f / 24.0f;
-    solve_simulation_time_step(*simulation_orig, *depsgraph, influences, time_step);
+    solve_simulation_time_step(*simulation_orig, *depsgraph, influences, handle_map, time_step);
     simulation_orig->current_frame = current_frame;
 
     copy_states_to_cow(simulation_orig, simulation_cow);