Geometry Nodes: support baking data block references

With this patch, materials are kept intact in simulation zones and bake nodes
without any additional user action.

This implements the design proposed in #108410 to support referencing
data-blocks (only materials for now) in the baked data. The task also describes
why this is not a trivial issue. A previous attempt was implemented in #109703
but it didn't work well-enough.

The solution is to have an explicit `name (+ library name) -> data-block`
mapping that is stored in the modifier for each bake node and simulation zone.
The `library name` is necessary for it to be unique within a .blend file. Note
that this refers to the name of the `Library` data-block and not a file path.
The baked data only contains the names of the used data-blocks. When the baked
data is loaded, the correct material data-block is looked up from the mapping.

### Automatic Mapping Generation

The most tricky aspect of this approach is to make it feel mostly automatic.
From the user point-of-view, it should just work. Therefore, we don't want the
user to have to create the mapping manually in the majority of cases. Creating
the mapping automatically is difficult because the data-blocks that should
become part of the mapping are only known during depsgraph evaluation. So we
somehow have to gather the missing data blocks during evaluation and then write
the new mappings back to the original data.

While writing back to original data is something we do in some cases already,
the situation here is different, because we are actually creating new relations
between data-blocks. This also means that we'll have to do user-counting. Since
user counts in data-blocks are *not* atomic, we can't do that from multiple
threads at the same time. Also, under some circumstances, it may be necessary to
trigger depsgraph evaluation again after the write-back because it actually
affects the result.

To solve this, a small new API is added in `DEG_depsgraph_writeback_sync.hh`. It
allows gathering tasks which write back to original data in a synchronous way
which may also require a reevaluation.

### Accessing the Mapping

A new `BakeDataBlockMap` is passed to geometry nodes evaluation by the modifier.
This map allows getting the `ID` pointer that should be used for a specific
data-block name that is stored in baked data. It's also used to gather all the
missing data mappings during evaluation.

### Weak ID References

The baked/cached geometries may have references to other data-blocks (currently
only materials, but in the future also e.g. instanced objects/collections).
However, the pointers of these data-blocks are not stable over time. That is
especially true when storing/loading the data from disk, but also just when
playing back the animation. Therefore, the used data-blocks have to referenced
in a different way at run-time.

This is solved by adding `std::unique_ptr<bake::BakeMaterialsList>` to the
run-time data of various geometry data-blocks. If the data-block is cached over
a longer period of time (such that material pointers can't be used directly), it
stores the material name (+ library name) used by each material slot. When the
geometry is used again, the material pointers are restored using these weak name
references and the `BakeDataBlockMap`.

### Manual Mapping Management

There is a new `Data-Blocks` panel in the bake settings in the node editor
sidebar that allows inspecting and modifying the data-blocks that are used when
baking. The user can change what data-block a specific name is mapped to.

Pull Request: https://projects.blender.org/blender/blender/pulls/117043

source/blender/blenkernel/BKE_bake_data_block_id.hh

@@ -0,0 +1,56 @@
+/* SPDX-FileCopyrightText: 2024 Blender Authors
+ *
+ * SPDX-License-Identifier: GPL-2.0-or-later */
+
+#pragma once
+
+/** \file
+ * \ingroup bke
+ */
+
+#include "BLI_string_ref.hh"
+#include "BLI_struct_equality_utils.hh"
+#include "BLI_vector.hh"
+
+#include "DNA_ID_enums.h"
+
+struct ID;
+struct NodesModifierDataBlock;
+
+namespace blender::bke::bake {
+
+/**
+ * Unique weak reference to a data block within a #Main. It's used when caching/baking data-block
+ * references. Data-block pointers can't be used directly, because they are not stable over time
+ * and between Blender sessions.
+ */
+struct BakeDataBlockID {
+  ID_Type type;
+  /**
+   * Name of the data-block, without the type prefix.
+   */
+  std::string id_name;
+  /**
+   * Name of the library data-block that the data-block is in. This refers to `Library.id.name` and
+   * not the file path. The type prefix of the name is omitted. If this is empty, the data-block is
+   * expected to be local and not linked.
+   */
+  std::string lib_name;
+
+  BakeDataBlockID(ID_Type type, std::string id_name, std::string lib_name);
+  BakeDataBlockID(const ID &id);
+  BakeDataBlockID(const NodesModifierDataBlock &data_block);
+
+  uint64_t hash() const;
+
+  friend std::ostream &operator<<(std::ostream &stream, const BakeDataBlockID &id);
+
+  BLI_STRUCT_EQUALITY_OPERATORS_3(BakeDataBlockID, type, id_name, lib_name)
+};
+
+/**
+ * A list of weak data-block references for material slots.
+ */
+struct BakeMaterialsList : public Vector<std::optional<BakeDataBlockID>> {};
+
+}  // namespace blender::bke::bake

source/blender/blenkernel/BKE_bake_data_block_map.hh

@@ -0,0 +1,43 @@
+/* SPDX-FileCopyrightText: 2024 Blender Authors
+ *
+ * SPDX-License-Identifier: GPL-2.0-or-later */
+
+#pragma once
+
+/** \file
+ * \ingroup bke
+ */
+
+#include <optional>
+
+#include "BLI_string_ref.hh"
+#include "BLI_struct_equality_utils.hh"
+
+#include "BKE_bake_data_block_id.hh"
+
+#include "DNA_ID_enums.h"
+
+namespace blender::bke::bake {
+
+/**
+ * Maps #BakeDataBlockID to the corresponding data-blocks. This is used during depsgraph evaluation
+ * to remap weak data-block references stored in baked data to the actual data-blocks at run-time.
+ *
+ * Also it keeps track of missing data-blocks, so that they can be added later.
+ */
+struct BakeDataBlockMap {
+ public:
+  /**
+   * Tries to retrieve the data block for the given key. If it's not explicitly mapped, it might be
+   * added to the mapping. If it's still not found, null is returned.
+   */
+  virtual ID *lookup_or_remember_missing(const BakeDataBlockID &key) = 0;
+
+  /**
+   * Tries to add the data block to the map. This may not succeed in all cases, e.g. if the
+   * implementation does not allow inserting new mapping items.
+   */
+  virtual void try_add(ID &id) = 0;
+};
+
+}  // namespace blender::bke::bake

source/blender/blenkernel/BKE_bake_items.hh

@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include "BKE_bake_data_block_map.hh"
 #include "BKE_geometry_set.hh"
 
 namespace blender::bke::bake {
@@ -42,13 +43,18 @@ class GeometryBakeItem : public BakeItem {
   GeometryBakeItem(GeometrySet geometry);
 
   /**
-   * Removes parts of the geometry that can't be stored in the simulation state:
-   * - Anonymous attributes can't be stored because it is not known which of them will or will not
-   * be used in the future.
-   * - Materials can't be stored directly, because they are linked ID data blocks that can't be
-   *   restored from baked data currently.
+   * Removes parts of the geometry that can't be baked/cached (anonymous attributes) and replaces
+   * data-block pointers with #BakeDataBlockID.
    */
-  static void cleanup_geometry(GeometrySet &geometry);
+  static void prepare_geometry_for_bake(GeometrySet &geometry, BakeDataBlockMap *data_block_map);
+
+  /**
+   * The baked data does not have raw pointers to referenced data-blocks because those would become
+   * dangling quickly. Instead it has weak name-based references (#BakeDataBlockID). This function
+   * attempts to restore the actual data block pointers based on the weak references using the
+   * given mapping.
+   */
+  static void try_restore_data_blocks(GeometrySet &geometry, BakeDataBlockMap *data_block_map);
 };
 
 /**

source/blender/blenkernel/BKE_bake_items_socket.hh

@@ -35,8 +35,8 @@ struct BakeSocketConfig {
  * Create new bake items from the socket values. The socket values are not destructed, but they may
  * be in a moved-from state afterwards.
  */
-Array<std::unique_ptr<BakeItem>> move_socket_values_to_bake_items(Span<void *> socket_values,
-                                                                  const BakeSocketConfig &config);
+Array<std::unique_ptr<BakeItem>> move_socket_values_to_bake_items(
+    Span<void *> socket_values, const BakeSocketConfig &config, BakeDataBlockMap *data_block_map);
 
 /**
  * Create socket values from bake items.
@@ -52,6 +52,7 @@ Array<std::unique_ptr<BakeItem>> move_socket_values_to_bake_items(Span<void *> s
 void move_bake_items_to_socket_values(
     Span<BakeItem *> bake_items,
     const BakeSocketConfig &config,
+    BakeDataBlockMap *data_block_map,
     FunctionRef<std::shared_ptr<AnonymousAttributeFieldInput>(int socket_index, const CPPType &)>
         make_attribute_field,
     Span<void *> r_socket_values);
@@ -63,6 +64,7 @@ void move_bake_items_to_socket_values(
 void copy_bake_items_to_socket_values(
     Span<const BakeItem *> bake_items,
     const BakeSocketConfig &config,
+    BakeDataBlockMap *data_block_map,
     FunctionRef<std::shared_ptr<AnonymousAttributeFieldInput>(int, const CPPType &)>
         make_attribute_field,
     Span<void *> r_socket_values);

source/blender/blenkernel/BKE_curves.hh

@@ -33,6 +33,9 @@ class AttributeAccessor;
 class MutableAttributeAccessor;
 enum class AttrDomain : int8_t;
 }  // namespace blender::bke
+namespace blender::bke::bake {
+struct BakeMaterialsList;
+}
 
 namespace blender::bke {
 
@@ -111,6 +114,9 @@ class CurvesGeometryRuntime {
 
   /** Normal direction vectors for each evaluated point. */
   mutable SharedCache<Vector<float3>> evaluated_normal_cache;
+
+  /** Stores weak references to material data blocks. */
+  std::unique_ptr<bake::BakeMaterialsList> bake_materials;
 };
 
 /**

source/blender/blenkernel/BKE_lib_id.hh

@@ -204,6 +204,10 @@ void BKE_libblock_ensure_unique_name(Main *bmain, ID *id) ATTR_NONNULL();
 ID *BKE_libblock_find_name(Main *bmain, short type, const char *name) ATTR_WARN_UNUSED_RESULT
     ATTR_NONNULL();
 ID *BKE_libblock_find_session_uid(Main *bmain, short type, uint32_t session_uid);
+ID *BKE_libblock_find_name_and_library(Main *bmain,
+                                       short type,
+                                       const char *name,
+                                       const char *lib_name);
 /**
  * Duplicate (a.k.a. deep copy) common processing options.
  * See also eDupli_ID_Flags for options controlling what kind of IDs to duplicate.

source/blender/blenkernel/BKE_mesh_types.hh

@@ -29,6 +29,9 @@ struct SubsurfRuntimeData;
 namespace blender::bke {
 struct EditMeshData;
 }
+namespace blender::bke::bake {
+struct BakeMaterialsList;
+}
 
 /** #MeshRuntime.wrapper_type */
 enum eMeshWrapperType {
@@ -207,6 +210,9 @@ struct MeshRuntime {
    */
   BitVector<> subsurf_optimal_display_edges;
 
+  /** Stores weak references to material data blocks. */
+  std::unique_ptr<bake::BakeMaterialsList> bake_materials;
+
   MeshRuntime();
   ~MeshRuntime();
 };

source/blender/blenkernel/BKE_pointcloud.hh

@@ -22,6 +22,9 @@ struct Main;
 struct Object;
 struct PointCloud;
 struct Scene;
+namespace blender::bke::bake {
+struct BakeMaterialsList;
+}
 
 /* PointCloud datablock */
 extern const char *POINTCLOUD_ATTR_POSITION;
@@ -37,6 +40,9 @@ struct PointCloudRuntime {
    */
   mutable SharedCache<Bounds<float3>> bounds_cache;
 
+  /** Stores weak references to material data blocks. */
+  std::unique_ptr<bake::BakeMaterialsList> bake_materials;
+
   MEM_CXX_CLASS_ALLOC_FUNCS("PointCloudRuntime");
 };
 

source/blender/blenkernel/CMakeLists.txt

@@ -70,6 +70,7 @@ set(SRC
   intern/attribute_math.cc
   intern/autoexec.cc
   intern/bake_geometry_nodes_modifier.cc
+  intern/bake_data_block_map.cc
   intern/bake_items.cc
   intern/bake_items_paths.cc
   intern/bake_items_serialize.cc
@@ -336,6 +337,7 @@ set(SRC
   BKE_attribute_math.hh
   BKE_autoexec.hh
   BKE_bake_geometry_nodes_modifier.hh
+  BKE_bake_data_block_map.hh
   BKE_bake_items.hh
   BKE_bake_items_paths.hh
   BKE_bake_items_serialize.hh

source/blender/blenkernel/intern/bake_data_block_map.cc

@@ -0,0 +1,47 @@
+/* SPDX-FileCopyrightText: 2005 Blender Authors
+ *
+ * SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include <ostream>
+
+#include "BLI_hash.hh"
+
+#include "BKE_bake_data_block_map.hh"
+
+#include "DNA_ID.h"
+#include "DNA_modifier_types.h"
+
+namespace blender::bke::bake {
+
+BakeDataBlockID::BakeDataBlockID(ID_Type type, std::string id_name, std::string lib_name)
+    : type(type), id_name(std::move(id_name)), lib_name(std::move(lib_name))
+{
+}
+
+BakeDataBlockID::BakeDataBlockID(const ID &id)
+{
+  this->type = GS(id.name);
+  this->id_name = id.name + 2;
+  if (id.lib) {
+    this->lib_name = id.lib->id.name + 2;
+  }
+}
+
+BakeDataBlockID::BakeDataBlockID(const NodesModifierDataBlock &data_block)
+    : BakeDataBlockID(ID_Type(data_block.id_type),
+                      StringRef(data_block.id_name),
+                      StringRef(data_block.lib_name))
+{
+}
+
+std::ostream &operator<<(std::ostream &stream, const BakeDataBlockID &id)
+{
+  return stream << "(" << id.id_name << ", Lib: " << id.lib_name << ")";
+}
+
+uint64_t BakeDataBlockID::hash() const
+{
+  return get_default_hash(this->type, this->id_name, this->lib_name);
+}
+
+}  // namespace blender::bke::bake

source/blender/blenkernel/intern/bake_items.cc

@@ -27,27 +27,49 @@ using DictionaryValuePtr = std::shared_ptr<DictionaryValue>;
 
 GeometryBakeItem::GeometryBakeItem(GeometrySet geometry) : geometry(std::move(geometry)) {}
 
-static void remove_materials(Material ***materials, short *materials_num)
+static std::unique_ptr<BakeMaterialsList> materials_to_weak_references(
+    Material ***materials, short *materials_num, BakeDataBlockMap *data_block_map)
 {
+  if (*materials_num == 0) {
+    return {};
+  }
+  auto materials_list = std::make_unique<BakeMaterialsList>();
+  materials_list->resize(*materials_num);
+  for (const int i : materials_list->index_range()) {
+    Material *material = (*materials)[i];
+    if (material) {
+      (*materials_list)[i] = BakeDataBlockID(material->id);
+      if (data_block_map) {
+        data_block_map->try_add(material->id);
+      }
+    }
+  }
+
   MEM_SAFE_FREE(*materials);
   *materials_num = 0;
+
+  return materials_list;
 }
 
-void GeometryBakeItem::cleanup_geometry(GeometrySet &main_geometry)
+void GeometryBakeItem::prepare_geometry_for_bake(GeometrySet &main_geometry,
+                                                 BakeDataBlockMap *data_block_map)
 {
   main_geometry.ensure_owns_all_data();
   main_geometry.modify_geometry_sets([&](GeometrySet &geometry) {
     if (Mesh *mesh = geometry.get_mesh_for_write()) {
       mesh->attributes_for_write().remove_anonymous();
-      remove_materials(&mesh->mat, &mesh->totcol);
+      mesh->runtime->bake_materials = materials_to_weak_references(
+          &mesh->mat, &mesh->totcol, data_block_map);
     }
     if (Curves *curves = geometry.get_curves_for_write()) {
       curves->geometry.wrap().attributes_for_write().remove_anonymous();
-      remove_materials(&curves->mat, &curves->totcol);
+      curves->geometry.runtime->bake_materials = materials_to_weak_references(
+          &curves->mat, &curves->totcol, data_block_map);
     }
     if (PointCloud *pointcloud = geometry.get_pointcloud_for_write()) {
       pointcloud->attributes_for_write().remove_anonymous();
-      remove_materials(&pointcloud->mat, &pointcloud->totcol);
+      pointcloud->runtime->bake_materials = materials_to_weak_references(
+          &pointcloud->mat, &pointcloud->totcol, data_block_map);
     }
     if (bke::Instances *instances = geometry.get_instances_for_write()) {
       instances->attributes_for_write().remove_anonymous();
@@ -59,6 +81,53 @@ void GeometryBakeItem::cleanup_geometry(GeometrySet &main_geometry)
   });
 }
 
+static void restore_materials(Material ***materials,
+                              short *materials_num,
+                              std::unique_ptr<BakeMaterialsList> materials_list,
+                              BakeDataBlockMap *data_block_map)
+{
+  if (!materials_list) {
+    return;
+  }
+  BLI_assert(*materials == nullptr);
+  *materials_num = materials_list->size();
+  *materials = MEM_cnew_array<Material *>(materials_list->size(), __func__);
+  if (!data_block_map) {
+    return;
+  }
+
+  for (const int i : materials_list->index_range()) {
+    const std::optional<BakeDataBlockID> &data_block_id = (*materials_list)[i];
+    if (data_block_id) {
+      (*materials)[i] = reinterpret_cast<Material *>(
+          data_block_map->lookup_or_remember_missing(*data_block_id));
+    }
+  }
+}
+
+void GeometryBakeItem::try_restore_data_blocks(GeometrySet &main_geometry,
+                                               BakeDataBlockMap *data_block_map)
+{
+  main_geometry.modify_geometry_sets([&](GeometrySet &geometry) {
+    if (Mesh *mesh = geometry.get_mesh_for_write()) {
+      restore_materials(
+          &mesh->mat, &mesh->totcol, std::move(mesh->runtime->bake_materials), data_block_map);
+    }
+    if (Curves *curves = geometry.get_curves_for_write()) {
+      restore_materials(&curves->mat,
+                        &curves->totcol,
+                        std::move(curves->geometry.runtime->bake_materials),
+                        data_block_map);
+    }
+    if (PointCloud *pointcloud = geometry.get_pointcloud_for_write()) {
+      restore_materials(&pointcloud->mat,
+                        &pointcloud->totcol,
+                        std::move(pointcloud->runtime->bake_materials),
+                        data_block_map);
+    }
+  });
+}
+
 PrimitiveBakeItem::PrimitiveBakeItem(const CPPType &type, const void *value) : type_(type)
 {
   value_ = MEM_mallocN_aligned(type.size(), type.alignment(), __func__);

source/blender/blenkernel/intern/bake_items_serialize.cc

@@ -365,6 +365,32 @@ template<typename T>
   return *r_data != nullptr;
 }
 
+[[nodiscard]] static bool load_materials(const io::serialize::ArrayValue &io_materials,
+                                         std::unique_ptr<BakeMaterialsList> &materials)
+{
+  if (io_materials.elements().is_empty()) {
+    return true;
+  }
+  materials = std::make_unique<BakeMaterialsList>();
+  for (const auto &io_material_value : io_materials.elements()) {
+    if (io_material_value->type() == io::serialize::eValueType::Null) {
+      materials->append(std::nullopt);
+      continue;
+    }
+    const auto *io_material = io_material_value->as_dictionary_value();
+    if (!io_material) {
+      return false;
+    }
+    std::optional<std::string> id_name = io_material->lookup_str("name");
+    if (!id_name) {
+      return false;
+    }
+    std::string lib_name = io_material->lookup_str("lib_name").value_or("");
+    materials->append(BakeDataBlockID(ID_MA, std::move(*id_name), std::move(lib_name)));
+  }
+  return true;
+}
+
 [[nodiscard]] static bool load_attributes(const io::serialize::ArrayValue &io_attributes,
                                           MutableAttributeAccessor &attributes,
                                           const BlobReader &blob_reader,
@@ -450,6 +476,12 @@ static PointCloud *try_load_pointcloud(const DictionaryValue &io_geometry,
   if (!load_attributes(*io_attributes, attributes, blob_reader, blob_sharing)) {
     return cancel();
   }
+
+  if (const io::serialize::ArrayValue *io_materials = io_pointcloud->lookup_array("materials")) {
+    if (!load_materials(*io_materials, pointcloud->runtime->bake_materials)) {
+      return cancel();
+    }
+  }
   return pointcloud;
 }
 
@@ -499,6 +531,12 @@ static Curves *try_load_curves(const DictionaryValue &io_geometry,
     return cancel();
   }
 
+  if (const io::serialize::ArrayValue *io_materials = io_curves->lookup_array("materials")) {
+    if (!load_materials(*io_materials, curves.runtime->bake_materials)) {
+      return cancel();
+    }
+  }
+
   curves.update_curve_types();
 
   return curves_id;
@@ -554,6 +592,12 @@ static Mesh *try_load_mesh(const DictionaryValue &io_geometry,
     return cancel();
   }
 
+  if (const io::serialize::ArrayValue *io_materials = io_mesh->lookup_array("materials")) {
+    if (!load_materials(*io_materials, mesh->runtime->bake_materials)) {
+      return cancel();
+    }
+  }
+
   return mesh;
 }
 
@@ -630,20 +674,23 @@ static GeometrySet load_geometry(const DictionaryValue &io_geometry,
   return geometry;
 }
 
-static std::shared_ptr<io::serialize::ArrayValue> serialize_material_slots(
-    const Span<const Material *> material_slots)
+static std::shared_ptr<io::serialize::ArrayValue> serialize_materials(
+    const std::unique_ptr<BakeMaterialsList> &materials)
 {
   auto io_materials = std::make_shared<io::serialize::ArrayValue>();
-  for (const Material *material : material_slots) {
-    if (material == nullptr) {
-      io_materials->append_null();
+  if (!materials) {
+    return io_materials;
+  }
+  for (const std::optional<BakeDataBlockID> &material : *materials) {
+    if (material) {
+      auto io_material = io_materials->append_dict();
+      io_material->append_str("name", material->id_name);
+      if (!material->lib_name.empty()) {
+        io_material->append_str("lib_name", material->lib_name);
+      }
     }
     else {
-      auto io_material = io_materials->append_dict();
-      io_material->append_str("name", material->id.name + 2);
-      if (material->id.lib != nullptr) {
-        io_material->append_str("lib_name", material->id.lib->id.name + 2);
-      }
+      io_materials->append_null();
     }
   }
   return io_materials;
@@ -707,7 +754,7 @@ static std::shared_ptr<DictionaryValue> serialize_geometry_set(const GeometrySet
                                                      mesh.runtime->face_offsets_sharing_info));
     }
 
-    auto io_materials = serialize_material_slots({mesh.mat, mesh.totcol});
+    auto io_materials = serialize_materials(mesh.runtime->bake_materials);
     io_mesh->append("materials", io_materials);
 
     auto io_attributes = serialize_attributes(mesh.attributes(), blob_writer, blob_sharing, {});
@@ -719,7 +766,7 @@ static std::shared_ptr<DictionaryValue> serialize_geometry_set(const GeometrySet
 
     io_pointcloud->append_int("num_points", pointcloud.totpoint);
 
-    auto io_materials = serialize_material_slots({pointcloud.mat, pointcloud.totcol});
+    auto io_materials = serialize_materials(pointcloud.runtime->bake_materials);
     io_pointcloud->append("materials", io_materials);
 
     auto io_attributes = serialize_attributes(
@@ -744,7 +791,7 @@ static std::shared_ptr<DictionaryValue> serialize_geometry_set(const GeometrySet
                                          curves.runtime->curve_offsets_sharing_info));
     }
 
-    auto io_materials = serialize_material_slots({curves_id.mat, curves_id.totcol});
+    auto io_materials = serialize_materials(curves.runtime->bake_materials);
     io_curves->append("materials", io_materials);
 
     auto io_attributes = serialize_attributes(curves.attributes(), blob_writer, blob_sharing, {});

source/blender/blenkernel/intern/bake_items_socket.cc

@@ -3,16 +3,15 @@
  * SPDX-License-Identifier: GPL-2.0-or-later */
 
 #include "BKE_bake_items_socket.hh"
-
 #include "BKE_geometry_fields.hh"
 #include "BKE_node.hh"
-
 #include "BKE_node_socket_value.hh"
 
 namespace blender::bke::bake {
 
 Array<std::unique_ptr<BakeItem>> move_socket_values_to_bake_items(const Span<void *> socket_values,
-                                                                  const BakeSocketConfig &config)
+                                                                  const BakeSocketConfig &config,
+                                                                  BakeDataBlockMap *data_block_map)
 {
   BLI_assert(socket_values.size() == config.types.size());
   BLI_assert(socket_values.size() == config.geometries_by_attribute.size());
@@ -99,7 +98,7 @@ Array<std::unique_ptr<BakeItem>> move_socket_values_to_bake_items(const Span<voi
       continue;
     }
     GeometrySet &geometry = static_cast<GeometryBakeItem *>(bake_items[i].get())->geometry;
-    GeometryBakeItem::cleanup_geometry(geometry);
+    GeometryBakeItem::prepare_geometry_for_bake(geometry, data_block_map);
   }
 
   return bake_items;
@@ -192,6 +191,14 @@ static void rename_attributes(const Span<GeometrySet *> geometries,
   }
 }
 
+static void restore_data_blocks(const Span<GeometrySet *> geometries,
+                                BakeDataBlockMap *data_block_map)
+{
+  for (GeometrySet *main_geometry : geometries) {
+    GeometryBakeItem::try_restore_data_blocks(*main_geometry, data_block_map);
+  }
+}
+
 static void default_initialize_socket_value(const eNodeSocketDatatype socket_type, void *r_value)
 {
   const char *socket_idname = nodeStaticSocketType(socket_type, 0);
@@ -203,6 +210,7 @@ static void default_initialize_socket_value(const eNodeSocketDatatype socket_typ
 void move_bake_items_to_socket_values(
     const Span<BakeItem *> bake_items,
     const BakeSocketConfig &config,
+    BakeDataBlockMap *data_block_map,
     FunctionRef<std::shared_ptr<AnonymousAttributeFieldInput>(int, const CPPType &)>
         make_attribute_field,
     const Span<void *> r_socket_values)
@@ -237,11 +245,13 @@ void move_bake_items_to_socket_values(
   }
 
   rename_attributes(geometries, attribute_map);
+  restore_data_blocks(geometries, data_block_map);
 }
 
 void copy_bake_items_to_socket_values(
     const Span<const BakeItem *> bake_items,
     const BakeSocketConfig &config,
+    BakeDataBlockMap *data_block_map,
     FunctionRef<std::shared_ptr<AnonymousAttributeFieldInput>(int, const CPPType &)>
         make_attribute_field,
     const Span<void *> r_socket_values)
@@ -273,6 +283,7 @@ void copy_bake_items_to_socket_values(
   }
 
   rename_attributes(geometries, attribute_map);
+  restore_data_blocks(geometries, data_block_map);
 }
 
 }  // namespace blender::bke::bake

source/blender/blenkernel/intern/curves_geometry.cc

@@ -26,6 +26,7 @@
 
 #include "BKE_attribute.hh"
 #include "BKE_attribute_math.hh"
+#include "BKE_bake_data_block_id.hh"
 #include "BKE_curves.hh"
 #include "BKE_curves_utils.hh"
 #include "BKE_customdata.hh"
@@ -119,6 +120,11 @@ static void copy_curves_geometry(CurvesGeometry &dst, const CurvesGeometry &src)
   dst.runtime->evaluated_length_cache = src.runtime->evaluated_length_cache;
   dst.runtime->evaluated_tangent_cache = src.runtime->evaluated_tangent_cache;
   dst.runtime->evaluated_normal_cache = src.runtime->evaluated_normal_cache;
+
+  if (src.runtime->bake_materials) {
+    dst.runtime->bake_materials = std::make_unique<bake::BakeMaterialsList>(
+        *src.runtime->bake_materials);
+  }
 }
 
 CurvesGeometry::CurvesGeometry(const CurvesGeometry &other) : CurvesGeometry()

source/blender/blenkernel/intern/lib_id.cc

@@ -1489,6 +1489,31 @@ ID *BKE_libblock_find_session_uid(Main *bmain, const short type, const uint32_t
   return nullptr;
 }
 
+ID *BKE_libblock_find_name_and_library(Main *bmain,
+                                       const short type,
+                                       const char *name,
+                                       const char *lib_name)
+{
+  ListBase *lb = which_libbase(bmain, type);
+  BLI_assert(lb != nullptr);
+  LISTBASE_FOREACH (ID *, id, lb) {
+    if (!STREQ(id->name + 2, name)) {
+      continue;
+    }
+    if (lib_name == nullptr || lib_name[0] == '\0') {
+      if (id->lib == nullptr) {
+        return id;
+      }
+      return nullptr;
+    }
+    if (!STREQ(id->lib->id.name + 2, lib_name)) {
+      continue;
+    }
+    return id;
+  }
+  return nullptr;
+}
+
 void id_sort_by_name(ListBase *lb, ID *id, ID *id_sorting_hint)
 {
 #define ID_SORT_STEP_SIZE 512

source/blender/blenkernel/intern/mesh.cc

@@ -44,6 +44,7 @@
 
 #include "BKE_anim_data.h"
 #include "BKE_attribute.hh"
+#include "BKE_bake_data_block_id.hh"
 #include "BKE_bpath.h"
 #include "BKE_deform.hh"
 #include "BKE_editmesh.hh"
@@ -146,6 +147,10 @@ static void mesh_copy_data(Main *bmain, ID *id_dst, const ID *id_src, const int
   mesh_dst->runtime->vert_to_face_map_cache = mesh_src->runtime->vert_to_face_map_cache;
   mesh_dst->runtime->vert_to_corner_map_cache = mesh_src->runtime->vert_to_corner_map_cache;
   mesh_dst->runtime->corner_to_face_map_cache = mesh_src->runtime->corner_to_face_map_cache;
+  if (mesh_src->runtime->bake_materials) {
+    mesh_dst->runtime->bake_materials = std::make_unique<blender::bke::bake::BakeMaterialsList>(
+        *mesh_src->runtime->bake_materials);
+  }
 
   /* Only do tessface if we have no faces. */
   const bool do_tessface = ((mesh_src->totface_legacy != 0) && (mesh_src->faces_num == 0));

source/blender/blenkernel/intern/mesh_runtime.cc

@@ -17,6 +17,7 @@
 #include "BLI_task.hh"
 #include "BLI_timeit.hh"
 
+#include "BKE_bake_data_block_id.hh"
 #include "BKE_bvhutils.hh"
 #include "BKE_customdata.hh"
 #include "BKE_editmesh_cache.hh"

source/blender/blenkernel/intern/pointcloud.cc

@@ -25,6 +25,7 @@
 #include "BLI_vector.hh"
 
 #include "BKE_anim_data.h"
+#include "BKE_bake_data_block_id.hh"
 #include "BKE_customdata.hh"
 #include "BKE_geometry_set.hh"
 #include "BKE_global.h"
@@ -86,6 +87,11 @@ static void pointcloud_copy_data(Main * /*bmain*/,
 
   pointcloud_dst->runtime = new blender::bke::PointCloudRuntime();
   pointcloud_dst->runtime->bounds_cache = pointcloud_src->runtime->bounds_cache;
+  if (pointcloud_src->runtime->bake_materials) {
+    pointcloud_dst->runtime->bake_materials =
+        std::make_unique<blender::bke::bake::BakeMaterialsList>(
+            *pointcloud_src->runtime->bake_materials);
+  }
 
   pointcloud_dst->batch_cache = nullptr;
 }

source/blender/blenkernel/intern/scene.cc

@@ -24,6 +24,7 @@
 #include "DNA_mask_types.h"
 #include "DNA_material_types.h"
 #include "DNA_mesh_types.h"
+#include "DNA_modifier_types.h"
 #include "DNA_node_types.h"
 #include "DNA_object_types.h"
 #include "DNA_rigidbody_types.h"
@@ -79,6 +80,7 @@
 #include "BKE_linestyle.h"
 #include "BKE_main.hh"
 #include "BKE_mask.h"
+#include "BKE_modifier.hh"
 #include "BKE_node.hh"
 #include "BKE_node_runtime.hh"
 #include "BKE_object.hh"
@@ -97,6 +99,7 @@
 #include "DEG_depsgraph_build.hh"
 #include "DEG_depsgraph_debug.hh"
 #include "DEG_depsgraph_query.hh"
+#include "DEG_depsgraph_writeback_sync.hh"
 
 #include "RE_engine.h"
 
@@ -2567,7 +2570,7 @@ static void scene_graph_update_tagged(Depsgraph *depsgraph, Main *bmain, bool on
     prepare_mesh_for_viewport_render(bmain, scene, view_layer);
     /* Update all objects: drivers, matrices, etc. flags set
      * by depsgraph or manual, no layer check here, gets correct flushed. */
-    DEG_evaluate_on_refresh(depsgraph);
+    DEG_evaluate_on_refresh(depsgraph, DEG_EVALUATE_SYNC_WRITEBACK_YES);
     /* Update sound system. */
     BKE_scene_update_sound(depsgraph, bmain);
     /* Notify python about depsgraph update. */
@@ -2647,10 +2650,10 @@ void BKE_scene_graph_update_for_newframe_ex(Depsgraph *depsgraph, const bool cle
      * lose any possible unkeyed changes made by the handler. */
     if (pass == 0) {
       const float frame = BKE_scene_frame_get(scene);
-      DEG_evaluate_on_framechange(depsgraph, frame);
+      DEG_evaluate_on_framechange(depsgraph, frame, DEG_EVALUATE_SYNC_WRITEBACK_YES);
     }
     else {
-      DEG_evaluate_on_refresh(depsgraph);
+      DEG_evaluate_on_refresh(depsgraph, DEG_EVALUATE_SYNC_WRITEBACK_YES);
     }
     /* Update sound system animation. */
     BKE_scene_update_sound(depsgraph, bmain);