source/blender/blenkernel/intern/customdata.cc

@@ -5100,6 +5100,13 @@ void CustomData_blend_write(BlendWriter *writer,
       writer, CustomDataLayer, data->totlayer, data->layers, layers_to_write.data());
 
   for (const CustomDataLayer &layer : layers_to_write) {
+    if (BLO_write_is_undo(writer) && layer.cow != nullptr) {
+      BLO_write_cow(
+          writer, layer.cow, layer.data, [type = eCustomDataType(layer.type), count](void *data) {
+            free_layer_data(type, data, count);
+          });
+      continue;
+    }
     switch (layer.type) {
       case CD_MDEFORMVERT:
         BKE_defvert_blend_write(writer, count, static_cast<const MDeformVert *>(layer.data));
@@ -5216,6 +5223,11 @@ void CustomData_blend_read(BlendDataReader *reader, CustomData *data, const int
     }
 
     if (CustomData_verify_versions(data, i)) {
+      if (BLO_read_is_cow_data(reader, layer->data)) {
+        BLI_assert(layer->cow != nullptr);
+        BLI_cow_user_add(layer->cow);
+        continue;
+      }
       BLO_read_data_address(reader, &layer->data);
       layer->cow = BLI_cow_new(1);
       if (CustomData_layer_ensure_data_exists(layer, count)) {

source/blender/blenloader/BLO_read_write.h

@@ -31,6 +31,10 @@
 
 #include "DNA_windowmanager_types.h" /* for eReportType */
 
+#ifdef __cplusplus
+#  include <functional>
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -42,6 +46,7 @@ typedef struct BlendWriter BlendWriter;
 
 struct BlendFileReadReport;
 struct Main;
+struct bCopyOnWrite;
 
 /* -------------------------------------------------------------------- */
 /** \name Blend Write API
@@ -175,6 +180,18 @@ void BLO_write_string(BlendWriter *writer, const char *data_ptr);
 
 /* Misc. */
 
+#ifdef __cplusplus
+/**
+ * Give (shared) ownership of the data to the undo system so that the data does not have to be
+ * copied. A free-function has to be provided, which is called when the undo step is freed and no
+ * one else references the data anymore.
+ */
+void BLO_write_cow(BlendWriter *writer,
+                   const bCopyOnWrite *cow,
+                   const void *cow_data,
+                   std::function<void(void *data)> free_data_fn);
+#endif
+
 /**
  * Sometimes different data is written depending on whether the file is saved to disk or used for
  * undo. This function returns true when the current file-writing is done for undo.
@@ -236,6 +253,13 @@ void BLO_read_float3_array(BlendDataReader *reader, int array_size, float **ptr_
 void BLO_read_double_array(BlendDataReader *reader, int array_size, double **ptr_p);
 void BLO_read_pointer_array(BlendDataReader *reader, void **ptr_p);
 
+/**
+ * True when the pointer refers to valid data and can still be used. In order to take
+ * (shared) ownership of the data, the user count of the corresponding #bCopyOnWrite should be
+ * increased.
+ */
+bool BLO_read_is_cow_data(BlendDataReader *reader, const void *cow_data);
+
 /* Misc. */
 
 int BLO_read_fileversion_get(BlendDataReader *reader);

source/blender/blenloader/BLO_undofile.h

@@ -29,9 +29,24 @@ typedef struct {
   uint id_session_uuid;
 } MemFileChunk;
 
+#ifdef __cplusplus
+#  include <functional>
+
+#  include "BLI_copy_on_write.h"
+#  include "BLI_map.hh"
+struct MemFileCowStorage {
+  blender::Map<const void *, std::pair<const bCopyOnWrite *, std::function<void(void *)>>> map;
+
+  ~MemFileCowStorage();
+};
+#else
+typedef struct MemFileCowStorage MemFileCowStorage;
+#endif
+
 typedef struct MemFile {
   ListBase chunks;
   size_t size;
+  MemFileCowStorage *cow_storage;
 } MemFile;
 
 typedef struct MemFileWriteData {

source/blender/blenloader/intern/readfile.cc

@@ -5124,4 +5124,20 @@ void BLO_expand_id(BlendExpander *expander, ID *id)
   expand_doit(expander->fd, expander->main, id);
 }
 
+bool BLO_read_is_cow_data(BlendDataReader *reader, const void *cow_data)
+{
+  if (!BLO_read_data_is_undo(reader)) {
+    return false;
+  }
+  if (!(reader->fd->flags & FD_FLAGS_IS_MEMFILE)) {
+    return false;
+  }
+  UndoReader *undo_reader = reinterpret_cast<UndoReader *>(reader->fd->file);
+  MemFile &memfile = *undo_reader->memfile;
+  if (memfile.cow_storage == nullptr) {
+    return false;
+  }
+  return memfile.cow_storage->map.contains(cow_data);
+}
+
 /** \} */

source/blender/blenloader/intern/undofile.cc

@@ -48,9 +48,23 @@ void BLO_memfile_free(MemFile *memfile)
     }
     MEM_freeN(chunk);
   }
+  if (memfile->cow_storage) {
+    MEM_delete(memfile->cow_storage);
+    memfile->cow_storage = nullptr;
+  }
   memfile->size = 0;
 }
 
+MemFileCowStorage::~MemFileCowStorage()
+{
+  for (auto item : this->map.items()) {
+    if (BLI_cow_user_remove(item.value.first)) {
+      item.value.second(const_cast<void *>(item.key));
+      BLI_cow_free(item.value.first);
+    }
+  }
+}
+
 void BLO_memfile_merge(MemFile *first, MemFile *second)
 {
   /* We use this mapping to store the memory buffers from second memfile chunks which are not owned

source/blender/blenloader/intern/writefile.cc

@@ -85,6 +85,7 @@
 
 #include "BLI_bitmap.h"
 #include "BLI_blenlib.h"
+#include "BLI_copy_on_write.h"
 #include "BLI_endian_defines.h"
 #include "BLI_endian_switch.h"
 #include "BLI_link_utils.h"
@@ -1686,4 +1687,19 @@ bool BLO_write_is_undo(BlendWriter *writer)
   return writer->wd->use_memfile;
 }
 
+void BLO_write_cow(BlendWriter *writer,
+                   const bCopyOnWrite *cow,
+                   const void *cow_data,
+                   std::function<void(void *data)> free_data_fn)
+{
+  BLI_assert(writer->wd->use_memfile);
+  MemFile &memfile = *writer->wd->mem.written_memfile;
+  if (memfile.cow_storage == nullptr) {
+    memfile.cow_storage = MEM_new<MemFileCowStorage>(__func__);
+  }
+  if (memfile.cow_storage->map.add(cow_data, {cow, free_data_fn})) {
+    BLI_cow_user_add(cow);
+  }
+}
+
 /** \} */

source/blender/editors/sculpt_paint/sculpt_undo.cc

@@ -1856,7 +1856,8 @@ static void sculpt_undo_set_active_layer(struct bContext *C, SculptAttrRef *attr
     CustomData *cdata = attr->domain == ATTR_DOMAIN_POINT ? &me->vdata : &me->ldata;
     int totelem = attr->domain == ATTR_DOMAIN_POINT ? me->totvert : me->totloop;
 
-    CustomData_add_layer_named(cdata, attr->type, CD_SET_DEFAULT, totelem, attr->name);
+    CustomData_add_layer_named(
+        cdata, eCustomDataType(attr->type), CD_SET_DEFAULT, totelem, attr->name);
     layer = BKE_id_attribute_find(&me->id, attr->name, attr->type, attr->domain);
   }