Fix #148032: vertex/weight paint undo broken

The issue was that the data-blocks of two different undo steps were detected to
be identical, even if the attributes changed. That's because even if the
implicitly-shared data was different, they were turned into the same pointer by
cadb3fe5c5 on write.

This patch makes it so that for undo steps, implicitly shared data does not use
the pointer stability feature (in a sense, implicit-sharing itself provides
pointer stability for undo steps already).

The main tricky aspect is that we need to know if a pointer is implicitly shared
in `writestruct_at_address_nr` and oftentimes that's called before the
corresponding shared data is actually written with `BLO_write_shared`. The
solution is to enforce that the blend-write code has to know what pointers are
implicitly-shared before they are written the first time. The simplest way to
ensure that is to call `BLO_write_shared` first. However, that's not always
possible, especially when the pointer is directly embedded in an ID. Therefore,
there is a new `BLO_write_shared_tag` function that can be used in such cases.

The undo performance for the file in #141262 is still fixed with this change.

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

source/blender/blenkernel/intern/attribute_storage.cc

@@ -634,19 +634,19 @@ void AttributeStorage::blend_write(BlendWriter &writer,
     switch (AttrStorageType(attr_dna.storage_type)) {
       case AttrStorageType::Single: {
         ::AttributeSingle *single_dna = static_cast<::AttributeSingle *>(attr_dna.data);
-        BLO_write_struct(&writer, AttributeSingle, single_dna);
         write_shared_array(
             writer, AttrType(attr_dna.data_type), single_dna->data, 1, single_dna->sharing_info);
+        BLO_write_struct(&writer, AttributeSingle, single_dna);
         break;
       }
       case AttrStorageType::Array: {
         ::AttributeArray *array_dna = static_cast<::AttributeArray *>(attr_dna.data);
-        BLO_write_struct(&writer, AttributeArray, array_dna);
         write_shared_array(writer,
                            AttrType(attr_dna.data_type),
                            array_dna->data,
                            array_dna->size,
                            array_dna->sharing_info);
+        BLO_write_struct(&writer, AttributeArray, array_dna);
         break;
       }
     }

source/blender/blenkernel/intern/curves.cc

@@ -122,6 +122,9 @@ static void curves_blend_write(BlendWriter *writer, ID *id, const void *id_addre
   blender::bke::CurvesGeometry::BlendWriteData write_data(scope);
   curves->geometry.wrap().blend_write_prepare(write_data);
 
+  BLO_write_shared_tag(writer, curves->geometry.curve_offsets);
+  BLO_write_shared_tag(writer, curves->geometry.custom_knots);
+
   /* Write LibData */
   BLO_write_id_struct(writer, Curves, id_address, &curves->id);
   BKE_id_blend_write(writer, &curves->id);

source/blender/blenkernel/intern/customdata.cc

@@ -5163,9 +5163,6 @@ void CustomData_blend_write(BlendWriter *writer,
     CustomData_external_write(data, id, cddata_mask, count, 0);
   }
 
-  BLO_write_struct_array_at_address(
-      writer, CustomDataLayer, data->totlayer, data->layers, layers_to_write.data());
-
   for (const CustomDataLayer &layer : layers_to_write) {
     const size_t size_in_bytes = CustomData_sizeof(eCustomDataType(layer.type)) * count;
     BLO_write_shared(writer, layer.data, size_in_bytes, layer.sharing_info, [&]() {
@@ -5173,6 +5170,9 @@ void CustomData_blend_write(BlendWriter *writer,
     });
   }
 
+  BLO_write_struct_array_at_address(
+      writer, CustomDataLayer, data->totlayer, data->layers, layers_to_write.data());
+
   if (data->external) {
     BLO_write_struct(writer, CustomDataExternal, data->external);
   }

source/blender/blenkernel/intern/mesh.cc

@@ -395,6 +395,8 @@ static void mesh_blend_write(BlendWriter *writer, ID *id, const void *id_address
   const blender::bke::MeshRuntime *mesh_runtime = mesh->runtime;
   mesh->runtime = nullptr;
 
+  BLO_write_shared_tag(writer, mesh->face_offset_indices);
+
   BLO_write_id_struct(writer, Mesh, id_address, &mesh->id);
   BKE_id_blend_write(writer, &mesh->id);
 

source/blender/blenkernel/intern/packedFile.cc

@@ -974,10 +974,10 @@ void BKE_packedfile_blend_write(BlendWriter *writer, const PackedFile *pf)
   if (pf == nullptr) {
     return;
   }
-  BLO_write_struct(writer, PackedFile, pf);
   BLO_write_shared(writer, pf->data, pf->size, pf->sharing_info, [&]() {
     BLO_write_raw(writer, pf->size, pf->data);
   });
+  BLO_write_struct(writer, PackedFile, pf);
 }
 
 void BKE_packedfile_blend_read(BlendDataReader *reader, PackedFile **pf_p, StringRefNull filepath)

source/blender/blenloader/BLO_read_write.hh

@@ -222,6 +222,10 @@ void BLO_write_string(BlendWriter *writer, const char *data_ptr);
  * user count of the sharing-info is increased making the data immutable. The provided callback
  * should serialize the potentially shared data. It is only called when necessary.
  *
+ * This should be called before the data is referenced in other written data (there is an assert
+ * that checks for this). If that's not possible, at least #BLO_write_shared_tag needs to be called
+ * before the pointer is first written.
+ *
  * \param approximate_size_in_bytes: Used to be able to approximate how large the undo step is in
  * total.
  * \param write_fn: Use the #BlendWrite to serialize the potentially shared data.
@@ -232,6 +236,11 @@ void BLO_write_shared(BlendWriter *writer,
                       const blender::ImplicitSharingInfo *sharing_info,
                       blender::FunctionRef<void()> write_fn);
 
+/**
+ * Needs to be called if the pointer is somewhere written before the call to #BLO_write_shared.
+ */
+void BLO_write_shared_tag(BlendWriter *writer, const void *data);
+
 /**
  * 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.

source/blender/blenloader/intern/writefile.cc

@@ -830,6 +830,9 @@ static void write_bhead(WriteData *wd, const BHead &bhead)
   mywrite(wd, &bh, sizeof(bh));
 }
 
+/** This bit is used to mark address ids that use implicit sharing during undo. */
+constexpr uint64_t implicit_sharing_address_id_flag = uint64_t(1) << 63;
+
 static uint64_t stable_id_from_hint(const uint64_t hint)
 {
   /* Add a stride. This is not strictly necessary but may help with debugging later on because it's
@@ -839,6 +842,8 @@ static uint64_t stable_id_from_hint(const uint64_t hint)
     /* Null values are reserved for nullptr. */
     stable_id = (1 << 4);
   }
+  /* Remove the first bit as it reserved for pointers for implicit sharing.*/
+  stable_id &= ~implicit_sharing_address_id_flag;
   return stable_id;
 }
 
@@ -857,6 +862,27 @@ static uint64_t get_next_stable_address_id(WriteData &wd)
   return stable_id;
 }
 
+/**
+ * When writing an undo step, implicitly shared pointers do not use stable-pointers because that
+ * would lead to incorrect detection if a data-block has been changed between undo steps. That's
+ * because different shared data could be mapped to the same stable pointer, leading to
+ * #is_memchunk_identical to being true even if the referenced data is actually different.
+ *
+ * Another way to look at it is that implicit-sharing is a system for stable pointers (at runtime)
+ * itself. So it does not need an additional layer of stable pointers on top.
+ */
+static uint64_t get_address_id_for_implicit_sharing_data(const void *data)
+{
+  BLI_assert(data != nullptr);
+  uint64_t address_id = uint64_t(data);
+  /* Adding this bit so that it never overlap with an id generated by #stable_id_from_hint.
+   * Assuming that the given pointer is an actual pointer, it will stay unique when the
+   * #implicit_sharing_address_id_flag bit is set. That's because the upper bits of the pointer
+   * are effectively unused nowadays. */
+  address_id |= implicit_sharing_address_id_flag;
+  return address_id;
+}
+
 static uint64_t get_address_id_int(WriteData &wd, const void *address)
 {
   if (address == nullptr) {
@@ -2248,6 +2274,21 @@ void BLO_write_string(BlendWriter *writer, const char *data_ptr)
   }
 }
 
+void BLO_write_shared_tag(BlendWriter *writer, const void *data)
+{
+  if (!data) {
+    return;
+  }
+  if (!BLO_write_is_undo(writer)) {
+    return;
+  }
+  const uint64_t address_id = get_address_id_for_implicit_sharing_data(data);
+  /* Check that the pointer has not been written before it was tagged as being shared. */
+  BLI_assert(writer->wd->stable_address_ids.pointer_map.lookup_default(data, address_id) ==
+             address_id);
+  writer->wd->stable_address_ids.pointer_map.add(data, address_id);
+}
+
 void BLO_write_shared(BlendWriter *writer,
                       const void *data,
                       const size_t approximate_size_in_bytes,
@@ -2257,6 +2298,9 @@ void BLO_write_shared(BlendWriter *writer,
   if (data == nullptr) {
     return;
   }
+  if (sharing_info) {
+    BLO_write_shared_tag(writer, data);
+  }
   const uint64_t address_id = get_address_id_int(*writer->wd, data);
   if (BLO_write_is_undo(writer)) {
     MemFile &memfile = *writer->wd->mem.written_memfile;

source/blender/modifiers/intern/MOD_correctivesmooth.cc

@@ -792,8 +792,6 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
     }
   }
 
-  BLO_write_struct_at_address(writer, CorrectiveSmoothModifierData, md, &csmd);
-
   if (csmd.bind_coords != nullptr) {
     BLO_write_shared(writer,
                      csmd.bind_coords,
@@ -804,6 +802,8 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
                            writer, csmd.bind_coords_num, (const float *)csmd.bind_coords);
                      });
   }
+
+  BLO_write_struct_at_address(writer, CorrectiveSmoothModifierData, md, &csmd);
 }
 
 static void blend_read(BlendDataReader *reader, ModifierData *md)

source/blender/modifiers/intern/MOD_laplaciandeform.cc

@@ -832,14 +832,14 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
     }
   }
 
-  BLO_write_struct_at_address(writer, LaplacianDeformModifierData, md, &lmd);
-
   if (lmd.vertexco != nullptr) {
     BLO_write_shared(
         writer, lmd.vertexco, sizeof(float[3]) * lmd.verts_num, lmd.vertexco_sharing_info, [&]() {
           BLO_write_float3_array(writer, lmd.verts_num, lmd.vertexco);
         });
   }
+
+  BLO_write_struct_at_address(writer, LaplacianDeformModifierData, md, &lmd);
 }
 
 static void blend_read(BlendDataReader *reader, ModifierData *md)

source/blender/modifiers/intern/MOD_meshdeform.cc

@@ -563,8 +563,6 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
 
   const int size = mmd.dyngridsize;
 
-  BLO_write_struct_at_address(writer, MeshDeformModifierData, md, &mmd);
-
   BLO_write_shared(writer,
                    mmd.bindinfluences,
                    sizeof(MDefInfluence) * mmd.influences_num,
@@ -609,6 +607,8 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
                    sizeof(MDefInfluence) * mmd.verts_num,
                    mmd.dynverts_sharing_info,
                    [&]() { BLO_write_int32_array(writer, mmd.verts_num, mmd.dynverts); });
+
+  BLO_write_struct_at_address(writer, MeshDeformModifierData, md, &mmd);
 }
 
 static void blend_read(BlendDataReader *reader, ModifierData *md)

source/blender/modifiers/intern/MOD_surfacedeform.cc

@@ -1652,8 +1652,6 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
     }
   }
 
-  BLO_write_struct_at_address(writer, SurfaceDeformModifierData, md, &smd);
-
   if (smd.verts != nullptr) {
     BLO_write_shared(
         writer, smd.verts, sizeof(SDefVert) * smd.bind_verts_num, smd.verts_sharing_info, [&]() {
@@ -1684,6 +1682,8 @@ static void blend_write(BlendWriter *writer, const ID *id_owner, const ModifierD
           }
         });
   }
+
+  BLO_write_struct_at_address(writer, SurfaceDeformModifierData, md, &smd);
 }
 
 static void blend_read(BlendDataReader *reader, ModifierData *md)