6 changed files with 191 additions and 128 deletions
--- a/source/blender/blenkernel/BKE_compute_cache.hh
+++ b/source/blender/blenkernel/BKE_compute_cache.hh
@ -9,108 +9,185 @@

 #include "BKE_geometry_set.hh"

-namespace blender::bke {
+namespace blender::bke::sim {

-struct GeometryCacheValue {
+/* TODO: Take into account subframes. But do we make caches for subframes? Probably not. */
+struct TimePoint {
  int frame;
  float time;
-  GeometrySet geometry_set;
+
+  uint64_t hash()
+  {
+    return get_default_hash(this->frame);
+  }
+
+  friend bool operator==(const TimePoint &a, const TimePoint &b)
+  {
+    return a.frame == b.frame;
+  }
 };

 /* TODO: Clear cache when editing nodes? Only sometimes, when persistent caching is turned off. */
-struct SimulationCache {
-  Vector<GeometryCacheValue> geometry_per_frame;
+class SimulationCache {

-  const GeometryCacheValue *value_at_or_before_time(const int frame) const
-  {
-    const int index = this->index_at_or_before_time(frame);
-    if (index >= geometry_per_frame.size()) {
-      return nullptr;
+  struct CacheValues {
+    /* TODO: This will need to be a generic value at some point. */
+    /* Map from simulation time index (see #SimulationCache::times) to value. */
+    Vector<GeometrySet> persistent_cache;
+
+    std::optional<GeometrySet> non_persistent_value;
+
+    GeometrySet lookup_index(const int index) const
+    {
+      if (!persistent_cache.index_range().contains(index)) {
+        return persistent_cache[index];
+      }
+      return {};
    }
-    return &geometry_per_frame[index];
+  };
+
+  /* Map from cache data identifier (socket name) to values at stored times. */
+  Map<std::string, CacheValues> caches_;
+
+  /* Ordered list of cached simulation frames. */
+  Vector<TimePoint> persistent_times_;
+
+  std::optional<TimePoint> start_time_;
+
+  std::optional<TimePoint> last_run_time_;
+
+ public:
+  std::optional<TimePoint> start_time() const
+  {
+    return start_time_;
+  }
+  std::optional<TimePoint> last_run_time() const
+  {
+    return last_run_time_;
  }

-  GeometryCacheValue *value_at_or_before_time(const int frame)
+  std::optional<GeometrySet> value_at_or_before_time(StringRef data_name, TimePoint time)
  {
-    const int index = this->index_at_or_before_time(frame);
-    if (index >= geometry_per_frame.size()) {
-      return nullptr;
+    const CacheValues *values = caches_.lookup_ptr(data_name);
+    if (!values) {
+      return std::nullopt;
    }
-    return &geometry_per_frame[index];
-  }
-
-  const GeometryCacheValue *value_before_time(const int frame) const
-  {
-    const int index = this->index_before_time(frame);
-    if (index >= geometry_per_frame.size()) {
-      return nullptr;
-    }
-    return &geometry_per_frame[index];
-  }
-
-  GeometryCacheValue *value_at_time(const int frame)
-  {
-    for (const int i : geometry_per_frame.index_range()) {
-      if (geometry_per_frame[i].frame == frame) {
-        return &geometry_per_frame[i];
+    if (last_run_time_->time < time.time) {
+      if (values->non_persistent_value) {
+        return std::move(values->non_persistent_value);
      }
    }
-    return nullptr;
+    const int index = this->index_at_or_before_time(time);
+    if (!values->persistent_cache.index_range().contains(index)) {
+      return std::nullopt;
+    }
+    return values->persistent_cache[index];
  }

-  GeometryCacheValue &value_at_time_ensure(const int frame)
+  std::optional<GeometrySet> value_before_time(StringRef data_name, TimePoint time)
  {
-    for (const int i : geometry_per_frame.index_range()) {
-      if (geometry_per_frame[i].frame == frame) {
-        return geometry_per_frame[i];
+    const CacheValues *values = caches_.lookup_ptr(data_name);
+    if (!values) {
+      return std::nullopt;
+    }
+    if (last_run_time_->time < time.time) {
+      if (values->non_persistent_value) {
+        return std::move(values->non_persistent_value);
      }
    }
-    const int index = this->index_before_time(frame);
-    GeometryCacheValue value{};
-    geometry_per_frame.insert(index, value);
-    return geometry_per_frame[index];
+    const int index = this->index_before_time(time);
+    if (!values->persistent_cache.index_range().contains(index)) {
+      return std::nullopt;
+    }
+    return values->persistent_cache[index];
  }

-  void insert(GeometrySet &geometry_set, const int frame, const float time)
+  std::optional<GeometrySet> value_at_time(StringRef data_name, TimePoint time)
  {
-    BLI_assert(!this->value_at_time(frame));
-    GeometryCacheValue value{};
-    value.frame = frame;
-    value.time = time;
-    value.geometry_set = geometry_set;
-    const int index = this->index_before_time(frame);
-    geometry_per_frame.insert(index, value);
+    const CacheValues *values = caches_.lookup_ptr(data_name);
+    if (!values) {
+      return std::nullopt;
+    }
+    const std::optional<int> index = this->index_at_time(time);
+    if (!index) {
+      return std::nullopt;
+    }
+    return values->persistent_cache[*index];
+  }
+
+  void store_temporary(const StringRef data_name, const TimePoint time, GeometrySet value)
+  {
+    last_run_time_.emplace(time);
+    if (!start_time_) {
+      start_time_.emplace(time);
+    }
+    CacheValues &values = caches_.lookup_or_add_default_as(data_name);
+    values.non_persistent_value.emplace(std::move(value));
+  }
+
+  void store_persistent(const StringRef data_name, const TimePoint time, GeometrySet value)
+  {
+    last_run_time_.emplace(time);
+    if (!start_time_) {
+      start_time_.emplace(time);
+    }
+    const int index = this->index_before_time(time);
+    persistent_times_.resize(index);
+    persistent_times_[index] = time;
+    CacheValues &values = caches_.lookup_or_add_default_as(data_name);
+    values.persistent_cache.resize(index);
+    values.persistent_cache[index] = std::move(value);
+  }
+
+  void clear()
+  {
+    persistent_times_.clear();
+    caches_.clear();
+  }
+
+  bool is_empty() const
+  {
+    return persistent_times_.is_empty();
  }

 private:
-  int index_at_or_before_time(const int frame) const
+  int index_at_or_before_time(const TimePoint time) const
  {
-    if (geometry_per_frame.is_empty()) {
+    if (persistent_times_.is_empty()) {
      return 0;
    }
    int insert_index = 0;
-    for (const int i : geometry_per_frame.index_range()) {
-      if (geometry_per_frame[i].frame <= frame) {
+    for (const int i : persistent_times_.index_range()) {
+      if (persistent_times_[i].frame <= time.frame) {
        break;
      }
      insert_index++;
    }
    return insert_index;
  }
-  int index_before_time(const int frame) const
+  int index_before_time(const TimePoint time) const
  {
-    if (geometry_per_frame.is_empty()) {
+    if (persistent_times_.is_empty()) {
      return 0;
    }
    int insert_index = 0;
-    for (const int i : geometry_per_frame.index_range()) {
-      if (geometry_per_frame[i].frame < frame) {
+    for (const int i : persistent_times_.index_range()) {
+      if (persistent_times_[i].frame < time.frame) {
        break;
      }
      insert_index++;
    }
    return insert_index;
  }
+  std::optional<int> index_at_time(const TimePoint time) const
+  {
+    for (const int i : persistent_times_.index_range()) {
+      if (persistent_times_[i].frame == time.frame) {
+        return i;
+      }
+    }
+    return std::nullopt;
+  }
 };

 struct ComputeCaches {
@ -125,6 +202,11 @@ struct ComputeCaches {
    cache_per_context = other.cache_per_context;
  }

+  SimulationCache *lookup_context(const ComputeContextHash &context_hash)
+  {
+    std::scoped_lock lock{mutex};
+    return cache_per_context.lookup_ptr(context_hash);
+  }
  const SimulationCache *lookup_context(const ComputeContextHash &context_hash) const
  {
    std::scoped_lock lock{mutex};
@ -143,4 +225,4 @@ struct ComputeCaches {
  }
 };

-}  // namespace blender::bke
+}  // namespace blender::bke::sim
--- a/source/blender/makesdna/DNA_modifier_types.h
+++ b/source/blender/makesdna/DNA_modifier_types.h
@ -11,10 +11,10 @@
 #include "DNA_session_uuid_types.h"

 #ifdef __cplusplus
-namespace blender::bke {
+namespace blender::bke::sim {
 struct ComputeCaches;
 }
-using ComputeCachesHandle = blender::bke::ComputeCaches;
+using ComputeCachesHandle = blender::bke::sim::ComputeCaches;
 #else
 typedef struct ComputeCachesHandle ComputeCachesHandle;
 #endif
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@ -1095,7 +1095,7 @@ static GeometrySet compute_geometry(
    const blender::nodes::GeometryNodesLazyFunctionGraphInfo &lf_graph_info,
    const bNode &output_node,
    GeometrySet input_geometry_set,
-    blender::bke::ComputeCaches &compute_caches,
+    blender::bke::sim::ComputeCaches &compute_caches,
    NodesModifierData *nmd,
    const ModifierEvalContext *ctx)
 {
@ -1295,7 +1295,7 @@ static void modifyGeometry(ModifierData *md,
  NodesModifierData *orig_nmd = reinterpret_cast<NodesModifierData *>(
      BKE_modifier_get_original(ctx->object, md));
  if (!orig_nmd->simulation_caches) {
-    orig_nmd->simulation_caches = new blender::bke::ComputeCaches();
+    orig_nmd->simulation_caches = new blender::bke::sim::ComputeCaches();
  }

  geometry_set = compute_geometry(tree,
@ -1861,9 +1861,9 @@ static void copyData(const ModifierData *md, ModifierData *target, const int fla

  tnmd->runtime_eval_log = nullptr;
  if (nmd->simulation_caches) {
-    const blender::bke::ComputeCaches &src_caches = *static_cast<blender::bke::ComputeCaches *>(
-        nmd->simulation_caches);
-    tnmd->simulation_caches = new blender::bke::ComputeCaches(src_caches);
+    const blender::bke::sim::ComputeCaches &src_caches =
+        *static_cast<blender::bke::sim::ComputeCaches *>(nmd->simulation_caches);
+    tnmd->simulation_caches = new blender::bke::sim::ComputeCaches(src_caches);
  }

  if (nmd->settings.properties != nullptr) {
@ -1879,7 +1879,7 @@ static void freeData(ModifierData *md)
    nmd->settings.properties = nullptr;
  }

-  delete static_cast<blender::bke::ComputeCaches *>(nmd->simulation_caches);
+  delete static_cast<blender::bke::sim::ComputeCaches *>(nmd->simulation_caches);

  clear_runtime_data(nmd);
 }
--- a/source/blender/nodes/NOD_geometry_nodes_lazy_function.hh
+++ b/source/blender/nodes/NOD_geometry_nodes_lazy_function.hh
@ -47,7 +47,7 @@ struct GeoNodesModifierData {
  /** Optional logger. */
  geo_eval_log::GeoModifierLog *eval_log = nullptr;

-  bke::ComputeCaches *cache_per_frame;
+  bke::sim::ComputeCaches *cache_per_frame;

  /**
   * Some nodes should be executed even when their output is not used (e.g. active viewer nodes and
--- a/source/blender/nodes/geometry/nodes/node_geo_simulation_input.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_simulation_input.cc
@ -72,27 +72,31 @@ static void node_geo_exec(GeoNodeExecParams params)

  const Scene *scene = DEG_get_input_scene(params.depsgraph());
  const float scene_ctime = BKE_scene_ctime_get(scene);
-  const int scene_frame = int(scene_ctime);
+  const bke::sim::TimePoint time{int(scene_ctime), scene_ctime};

  const GeoNodesLFUserData &lf_data = *params.user_data();
-  bke::ComputeCaches &all_caches = *lf_data.modifier_data->cache_per_frame;
+  bke::sim::ComputeCaches &all_caches = *lf_data.modifier_data->cache_per_frame;

  const bke::NodeGroupComputeContext cache_context(lf_data.compute_context, sim_output_node_id);
-  const bke::SimulationCache *cache = all_caches.lookup_context(cache_context.hash());
+  bke::sim::SimulationCache *cache = all_caches.lookup_context(cache_context.hash());
  if (!cache) {
    params.set_output("Geometry", params.extract_input<GeometrySet>("Geometry"));
    return;
  }

-  if (const bke::GeometryCacheValue *data = cache->value_before_time(scene_frame)) {
+  const float elapsed_time = cache->is_empty() ? 0.0f : scene_ctime - cache->start_time()->time;
+  const float delta_time = cache->is_empty() ? 0.0f : scene_ctime - cache->last_run_time()->time;
+
+  if (params.lazy_output_is_required("Delta Time")) {
+    params.set_output("Delta Time", delta_time);
+  }
+  if (params.lazy_output_is_required("Elapsed Time")) {
+    params.set_output("Elapsed Time", elapsed_time);
+  }
+
+  if (std::optional<GeometrySet> cached_value = cache->value_before_time("Geometry", time)) {
    if (params.lazy_output_is_required("Geometry")) {
-      params.set_output("Geometry", std::move(data->geometry_set));
-    }
-    if (params.lazy_output_is_required("Delta Time")) {
-      params.set_output("Delta Time", scene_ctime - data->time);
-    }
-    if (params.lazy_output_is_required("Elapsed Time")) {
-      params.set_output("Elapsed Time", scene_ctime - cache->geometry_per_frame.first().time);
+      params.set_output("Geometry", std::move(*cached_value));
    }
    return;
  }
@ -102,17 +106,6 @@ static void node_geo_exec(GeoNodeExecParams params)
  }

  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
-  if (params.lazy_output_is_required("Delta Time")) {
-    params.set_output("Delta Time", -1.0f); /* TODO: How to get this?*/
-  }
-  if (params.lazy_output_is_required("Elapsed Time")) {
-    if (cache->geometry_per_frame.is_empty()) {
-      params.set_output("Elapsed Time", 0.0f);
-    }
-    else {
-      params.set_output("Elapsed Time", scene_ctime - cache->geometry_per_frame.first().time);
-    }
-  }
  params.set_output("Geometry", std::move(geometry_set));
 }

--- a/source/blender/nodes/geometry/nodes/node_geo_simulation_output.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_simulation_output.cc
@ -47,47 +47,39 @@ static void node_geo_exec(GeoNodeExecParams params)
  const NodeGeometrySimulationOutput &storage = node_storage(node);
  const Scene *scene = DEG_get_input_scene(params.depsgraph());
  const float scene_ctime = BKE_scene_ctime_get(scene);
-  const int scene_frame = int(scene_ctime);
+  const bke::sim::TimePoint time{int(scene_ctime), scene_ctime};

  const GeoNodesLFUserData &lf_data = *params.user_data();
-  bke::ComputeCaches &all_caches = *lf_data.modifier_data->cache_per_frame;
+  bke::sim::ComputeCaches &all_caches = *lf_data.modifier_data->cache_per_frame;

  const bke::NodeGroupComputeContext cache_context(lf_data.compute_context, node.identifier);
-  bke::SimulationCache &cache = all_caches.ensure_for_context(cache_context.hash());
-
-  if (cache.geometry_per_frame.is_empty()) {
-    if (params.lazy_output_is_required("Started")) {
-      params.set_output("Started", false);
-    }
-  }
-  else {
-    if (params.lazy_output_is_required("Elapsed Time")) {
-      params.set_output("Elapsed Time", scene_ctime - cache.geometry_per_frame.first().time);
-    }
-    if (params.lazy_output_is_required("Started")) {
-      params.set_output("Started", true);
-    }
-  }
+  bke::sim::SimulationCache &cache = all_caches.ensure_for_context(cache_context.hash());

+  const float elapsed_time = cache.is_empty() ? 0.0f : scene_ctime - cache.start_time()->time;
  const bool run = params.get_input<bool>("Run");
-  if (run) {
-    if (params.lazy_output_is_required("Ended")) {
-      params.set_output("Ended", false);
-    }
+  const bool started = !cache.is_empty();
+  const bool ended = !cache.is_empty() && !run;
+
+  if (params.lazy_output_is_required("Elapsed Time")) {
+    params.set_output("Elapsed Time", elapsed_time);
  }
-  else {
-    if (params.lazy_output_is_required("Ended")) {
-      params.set_output("Ended", true);
-    }
-    if (const bke::GeometryCacheValue *data = cache.value_at_or_before_time(scene_frame)) {
-      params.set_output("Geometry", data->geometry_set);
+  if (params.lazy_output_is_required("Started")) {
+    params.set_output("Started", started);
+  }
+  if (params.lazy_output_is_required("Ended")) {
+    params.set_output("Ended", ended);
+  }
+
+  if (!run) {
+    if (std::optional<GeometrySet> value = cache.value_at_or_before_time("Geometry", time)) {
+      params.set_output("Geometry", std::move(*value));
      params.set_input_unused("Geometry");
      return;
    }
  }

-  if (const bke::GeometryCacheValue *data = cache.value_at_time(scene_frame)) {
-    params.set_output("Geometry", data->geometry_set);
+  if (std::optional<GeometrySet> value = cache.value_at_time("Geometry", time)) {
+    params.set_output("Geometry", std::move(*value));
    params.set_input_unused("Geometry");
    return;
  }
@ -98,19 +90,15 @@ static void node_geo_exec(GeoNodeExecParams params)

  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
  geometry_set.ensure_owns_direct_data();
-  /* TODO: The "Use cache" input should probably become a "Persistent Cache" option. */
-  if (storage.use_persistent_cache || cache.geometry_per_frame.is_empty()) {
+  if (storage.use_persistent_cache) {
    /* If using the cache or there is no cached data yet, write the input in a new cache value. */
-    cache.insert(geometry_set, scene_frame, scene_ctime);
+    cache.store_persistent("Geometry", time, geometry_set);
  }
  else {
-    /* If we aren't using the cache, overrite the cache to only store the last frame. */
-    /* TODO: This breaks the elapsed time. */
-    /* TODO: Should we allow turning on and off caching procedurally? In that case we wouldn't be
-     * able to clear the whole cache, we would have to check if persistent caching was enabled for
-     * the *last* frame and then decide to replace it or not. */
-    cache.geometry_per_frame.clear();
-    cache.insert(geometry_set, scene_frame, scene_ctime);
+    /* TODO: Maybe don't clear the whole cache here. */
+    /* TODO: Move the geometry set here if the output isn't needed. */
+    cache.clear();
+    cache.store_temporary("Geometry", time, geometry_set);
  }

  params.set_output("Geometry", std::move(geometry_set));