IDMAnagement: Add owner ID pointer to embedded ones.

Add a dedicated `owner_id` pointer to ID types that can be embedded
(Collections and NodeTrees), and modify slightly come code to make
handling those more safe and consistent.

This implements first part of T69169.

Reviewed By: brecht

Differential Revision: https://developer.blender.org/D15838

8 files changed, 82 insertions, 55 deletions

diff --git a/source/blender/blenkernel/BKE_collection.h b/source/blender/blenkernel/BKE_collection.h
index feb3dc7de80..4a729311b6a 100644
--- a/source/blender/blenkernel/BKE_collection.h
+++ b/source/blender/blenkernel/BKE_collection.h
@@ -94,7 +94,7 @@ struct Collection *BKE_collection_duplicate(struct Main *bmain,
 /* Master Collection for Scene */
 
 #define BKE_SCENE_COLLECTION_NAME "Scene Collection"
-struct Collection *BKE_collection_master_add(void);
+struct Collection *BKE_collection_master_add(struct Scene *scene);
 
 /* Collection Objects */
 
@@ -296,7 +296,9 @@ void BKE_main_collections_parent_relations_rebuild(struct Main *bmain);
 /* .blend file I/O */
 
 void BKE_collection_blend_write_nolib(struct BlendWriter *writer, struct Collection *collection);
-void BKE_collection_blend_read_data(struct BlendDataReader *reader, struct Collection *collection);
+void BKE_collection_blend_read_data(struct BlendDataReader *reader,
+                                    struct Collection *collection,
+                                    struct ID *owner_id);
 void BKE_collection_blend_read_lib(struct BlendLibReader *reader, struct Collection *collection);
 void BKE_collection_blend_read_expand(struct BlendExpander *expander,
                                       struct Collection *collection);
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 8affbf0ca67..46303a4e19c 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -472,6 +472,11 @@ void ntreeSetTypes(const struct bContext *C, struct bNodeTree *ntree);
 
 struct bNodeTree *ntreeAddTree(struct Main *bmain, const char *name, const char *idname);
 
+struct bNodeTree *ntreeAddTreeEmbedded(struct Main *bmain,
+                                       struct ID *owner_id,
+                                       const char *name,
+                                       const char *idname);
+
 /* copy/free funcs, need to manage ID users */
 
 /**
@@ -540,7 +545,9 @@ void ntreeBlendWrite(struct BlendWriter *writer, struct bNodeTree *ntree);
 /**
  * \note `ntree` itself has been read!
  */
-void ntreeBlendReadData(struct BlendDataReader *reader, struct bNodeTree *ntree);
+void ntreeBlendReadData(struct BlendDataReader *reader,
+                        struct ID *owner_id,
+                        struct bNodeTree *ntree);
 void ntreeBlendReadLib(struct BlendLibReader *reader, struct bNodeTree *ntree);
 void ntreeBlendReadExpand(struct BlendExpander *expander, struct bNodeTree *ntree);
 
diff --git a/source/blender/blenkernel/intern/collection.c b/source/blender/blenkernel/intern/collection.c
index caf59a9b363..572a0caa3c7 100644
--- a/source/blender/blenkernel/intern/collection.c
+++ b/source/blender/blenkernel/intern/collection.c
@@ -143,6 +143,9 @@ static void collection_foreach_id(ID *id, LibraryForeachIDData *data)
 {
   Collection *collection = (Collection *)id;
 
+  BKE_LIB_FOREACHID_PROCESS_ID(
+      data, collection->owner_id, IDWALK_CB_LOOPBACK | IDWALK_CB_NEVER_SELF);
+
   LISTBASE_FOREACH (CollectionObject *, cob, &collection->gobject) {
     BKE_LIB_FOREACHID_PROCESS_IDSUPER(data, cob->ob, IDWALK_CB_USER);
   }
@@ -162,7 +165,7 @@ static void collection_foreach_id(ID *id, LibraryForeachIDData *data)
   }
 }
 
-static ID *collection_owner_get(Main *bmain, ID *id, ID *owner_id_hint)
+static ID *collection_owner_get(Main *bmain, ID *id, ID *UNUSED(owner_id_hint))
 {
   if ((id->flag & LIB_EMBEDDED_DATA) == 0) {
     return id;
@@ -171,20 +174,21 @@ static ID *collection_owner_get(Main *bmain, ID *id, ID *owner_id_hint)
 
   Collection *master_collection = (Collection *)id;
   BLI_assert((master_collection->flag & COLLECTION_IS_MASTER) != 0);
+  BLI_assert(master_collection->owner_id != NULL);
 
-  if (owner_id_hint != NULL && GS(owner_id_hint->name) == ID_SCE &&
-      ((Scene *)owner_id_hint)->master_collection == master_collection) {
-    return owner_id_hint;
-  }
-
+#ifndef NDEBUG
+  bool is_owner_found = false;
   LISTBASE_FOREACH (Scene *, scene, &bmain->scenes) {
     if (scene->master_collection == master_collection) {
-      return &scene->id;
+      BLI_assert(master_collection->owner_id == &scene->id);
+      BLI_assert(!is_owner_found);
+      is_owner_found = true;
     }
   }
+  BLI_assert(is_owner_found);
+#endif
 
-  BLI_assert_msg(0, "Embedded collection with no owner. Critical Main inconsistency.");
-  return NULL;
+  return master_collection->owner_id;
 }
 
 void BKE_collection_blend_write_nolib(BlendWriter *writer, Collection *collection)
@@ -233,8 +237,13 @@ void BKE_collection_compat_blend_read_data(BlendDataReader *reader, SceneCollect
 }
 #endif
 
-void BKE_collection_blend_read_data(BlendDataReader *reader, Collection *collection)
+void BKE_collection_blend_read_data(BlendDataReader *reader, Collection *collection, ID *owner_id)
 {
+  /* Special case for this pointer, do not rely on regular `lib_link` process here. Avoids needs
+   * for do_versioning, and ensures coherence of data in any case. */
+  BLI_assert((collection->id.flag & LIB_EMBEDDED_DATA) != 0 || owner_id == NULL);
+  collection->owner_id = owner_id;
+
   BLO_read_list(reader, &collection->gobject);
   BLO_read_list(reader, &collection->children);
 
@@ -265,7 +274,7 @@ void BKE_collection_blend_read_data(BlendDataReader *reader, Collection *collect
 static void collection_blend_read_data(BlendDataReader *reader, ID *id)
 {
   Collection *collection = (Collection *)id;
-  BKE_collection_blend_read_data(reader, collection);
+  BKE_collection_blend_read_data(reader, collection, NULL);
 }
 
 static void lib_link_collection_data(BlendLibReader *reader, Library *lib, Collection *collection)
@@ -849,14 +858,18 @@ Base *BKE_collection_or_layer_objects(const ViewLayer *view_layer, Collection *c
 /** \name Scene Master Collection
  * \{ */
 
-Collection *BKE_collection_master_add()
+Collection *BKE_collection_master_add(Scene *scene)
 {
+  BLI_assert(scene != NULL && scene->master_collection == NULL);
+
   /* Not an actual datablock, but owned by scene. */
   Collection *master_collection = BKE_libblock_alloc(
       NULL, ID_GR, BKE_SCENE_COLLECTION_NAME, LIB_ID_CREATE_NO_MAIN);
   master_collection->id.flag |= LIB_EMBEDDED_DATA;
+  master_collection->owner_id = &scene->id;
   master_collection->flag |= COLLECTION_IS_MASTER;
   master_collection->color_tag = COLLECTION_COLOR_NONE;
+
   return master_collection;
 }
 
diff --git a/source/blender/blenkernel/intern/linestyle.c b/source/blender/blenkernel/intern/linestyle.c
index 12a661d139b..776fe06edf7 100644
--- a/source/blender/blenkernel/intern/linestyle.c
+++ b/source/blender/blenkernel/intern/linestyle.c
@@ -2041,9 +2041,7 @@ void BKE_linestyle_default_shader(const bContext *C, FreestyleLineStyle *linesty
 
   BLI_assert(linestyle->nodetree == NULL);
 
-  ntree = ntreeAddTree(NULL, "stroke_shader", "ShaderNodeTree");
-
-  linestyle->nodetree = ntree;
+  ntree = ntreeAddTreeEmbedded(NULL, &linestyle->id, "stroke_shader", "ShaderNodeTree");
 
   uv_along_stroke = nodeAddStaticNode(C, ntree, SH_NODE_UVALONGSTROKE);
   uv_along_stroke->locx = 0.0f;
diff --git a/source/blender/blenkernel/intern/material.c b/source/blender/blenkernel/intern/material.c
index 248d292664a..442b31ff21e 100644
--- a/source/blender/blenkernel/intern/material.c
+++ b/source/blender/blenkernel/intern/material.c
@@ -1962,8 +1962,8 @@ static void material_default_surface_init(Material *ma)
 {
   strcpy(ma->id.name, "MADefault Surface");
 
-  bNodeTree *ntree = ntreeAddTree(NULL, "Shader Nodetree", ntreeType_Shader->idname);
-  ma->nodetree = ntree;
+  bNodeTree *ntree = ntreeAddTreeEmbedded(
+      NULL, &ma->id, "Shader Nodetree", ntreeType_Shader->idname);
   ma->use_nodes = true;
 
   bNode *principled = nodeAddStaticNode(NULL, ntree, SH_NODE_BSDF_PRINCIPLED);
@@ -1990,8 +1990,8 @@ static void material_default_volume_init(Material *ma)
 {
   strcpy(ma->id.name, "MADefault Volume");
 
-  bNodeTree *ntree = ntreeAddTree(NULL, "Shader Nodetree", ntreeType_Shader->idname);
-  ma->nodetree = ntree;
+  bNodeTree *ntree = ntreeAddTreeEmbedded(
+      NULL, &ma->id, "Shader Nodetree", ntreeType_Shader->idname);
   ma->use_nodes = true;
 
   bNode *principled = nodeAddStaticNode(NULL, ntree, SH_NODE_VOLUME_PRINCIPLED);
@@ -2015,8 +2015,8 @@ static void material_default_holdout_init(Material *ma)
 {
   strcpy(ma->id.name, "MADefault Holdout");
 
-  bNodeTree *ntree = ntreeAddTree(NULL, "Shader Nodetree", ntreeType_Shader->idname);
-  ma->nodetree = ntree;
+  bNodeTree *ntree = ntreeAddTreeEmbedded(
+      NULL, &ma->id, "Shader Nodetree", ntreeType_Shader->idname);
   ma->use_nodes = true;
 
   bNode *holdout = nodeAddStaticNode(NULL, ntree, SH_NODE_HOLDOUT);
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index e1eaed71f37..3cb2b80813f 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -325,6 +325,8 @@ static void node_foreach_id(ID *id, LibraryForeachIDData *data)
 {
   bNodeTree *ntree = (bNodeTree *)id;
 
+  BKE_LIB_FOREACHID_PROCESS_ID(data, ntree->owner_id, IDWALK_CB_LOOPBACK);
+
   BKE_LIB_FOREACHID_PROCESS_IDSUPER(data, ntree->gpd, IDWALK_CB_USER);
 
   LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
@@ -399,7 +401,7 @@ static void node_foreach_path(ID *id, BPathForeachPathData *bpath_data)
   }
 }
 
-static ID *node_owner_get(Main *bmain, ID *id, ID *owner_id_hint)
+static ID *node_owner_get(Main *UNUSED(bmain), ID *id, ID *UNUSED(owner_id_hint))
 {
   if ((id->flag & LIB_EMBEDDED_DATA) == 0) {
     return id;
@@ -408,30 +410,10 @@ static ID *node_owner_get(Main *bmain, ID *id, ID *owner_id_hint)
   // BLI_assert((id->tag & LIB_TAG_NO_MAIN) == 0);
 
   bNodeTree *ntree = reinterpret_cast<bNodeTree *>(id);
+  BLI_assert(ntree->owner_id != NULL);
+  BLI_assert(ntreeFromID(ntree->owner_id) == ntree);
 
-  if (owner_id_hint != nullptr && ntreeFromID(owner_id_hint) == ntree) {
-    return owner_id_hint;
-  }
-
-  ListBase *lists[] = {&bmain->materials,
-                       &bmain->lights,
-                       &bmain->worlds,
-                       &bmain->textures,
-                       &bmain->scenes,
-                       &bmain->linestyles,
-                       &bmain->simulations,
-                       nullptr};
-
-  for (int i = 0; lists[i] != nullptr; i++) {
-    LISTBASE_FOREACH (ID *, id_iter, lists[i]) {
-      if (ntreeFromID(id_iter) == ntree) {
-        return id_iter;
-      }
-    }
-  }
-
-  BLI_assert_msg(0, "Embedded node tree with no owner. Critical Main inconsistency.");
-  return nullptr;
+  return ntree->owner_id;
 }
 
 static void write_node_socket_default_value(BlendWriter *writer, bNodeSocket *sock)
@@ -667,8 +649,13 @@ static void direct_link_node_socket(BlendDataReader *reader, bNodeSocket *sock)
   sock->runtime = MEM_new<bNodeSocketRuntime>(__func__);
 }
 
-void ntreeBlendReadData(BlendDataReader *reader, bNodeTree *ntree)
+void ntreeBlendReadData(BlendDataReader *reader, ID *owner_id, bNodeTree *ntree)
 {
+  /* Special case for this pointer, do not rely on regular `lib_link` process here. Avoids needs
+   * for do_versioning, and ensures coherence of data in any case. */
+  BLI_assert((ntree->id.flag & LIB_EMBEDDED_DATA) != 0 || owner_id == nullptr);
+  ntree->owner_id = owner_id;
+
   /* NOTE: writing and reading goes in sync, for speed. */
   ntree->is_updating = false;
   ntree->typeinfo = nullptr;
@@ -830,7 +817,7 @@ void ntreeBlendReadData(BlendDataReader *reader, bNodeTree *ntree)
 static void ntree_blend_read_data(BlendDataReader *reader, ID *id)
 {
   bNodeTree *ntree = (bNodeTree *)id;
-  ntreeBlendReadData(reader, ntree);
+  ntreeBlendReadData(reader, nullptr, ntree);
 }
 
 static void lib_link_node_socket(BlendLibReader *reader, Library *lib, bNodeSocket *sock)
@@ -2575,12 +2562,12 @@ void nodePositionPropagate(bNode *node)
   }
 }
 
-bNodeTree *ntreeAddTree(Main *bmain, const char *name, const char *idname)
+static bNodeTree *ntreeAddTree_do(
+    Main *bmain, ID *owner_id, const bool is_embedded, const char *name, const char *idname)
 {
   /* trees are created as local trees for compositor, material or texture nodes,
    * node groups and other tree types are created as library data.
    */
-  const bool is_embedded = (bmain == nullptr);
   int flag = 0;
   if (is_embedded) {
     flag |= LIB_ID_CREATE_NO_MAIN;
@@ -2588,7 +2575,15 @@ bNodeTree *ntreeAddTree(Main *bmain, const char *name, const char *idname)
   bNodeTree *ntree = (bNodeTree *)BKE_libblock_alloc(bmain, ID_NT, name, flag);
   BKE_libblock_init_empty(&ntree->id);
   if (is_embedded) {
+    BLI_assert(owner_id != NULL);
     ntree->id.flag |= LIB_EMBEDDED_DATA;
+    ntree->owner_id = owner_id;
+    bNodeTree **ntree_owner_ptr = BKE_ntree_ptr_from_id(owner_id);
+    BLI_assert(ntree_owner_ptr != NULL);
+    *ntree_owner_ptr = ntree;
+  }
+  else {
+    BLI_assert(owner_id == NULL);
   }
 
   BLI_strncpy(ntree->idname, idname, sizeof(ntree->idname));
@@ -2597,6 +2592,19 @@ bNodeTree *ntreeAddTree(Main *bmain, const char *name, const char *idname)
   return ntree;
 }
 
+bNodeTree *ntreeAddTree(Main *bmain, const char *name, const char *idname)
+{
+  return ntreeAddTree_do(bmain, nullptr, false, name, idname);
+}
+
+bNodeTree *ntreeAddTreeEmbedded(Main *UNUSED(bmain),
+                                ID *owner_id,
+                                const char *name,
+                                const char *idname)
+{
+  return ntreeAddTree_do(nullptr, owner_id, true, name, idname);
+}
+
 bNodeTree *ntreeCopyTree_ex(const bNodeTree *ntree, Main *bmain, const bool do_id_user)
 {
   const int flag = do_id_user ? 0 : LIB_ID_CREATE_NO_USER_REFCOUNT | LIB_ID_CREATE_NO_MAIN;
diff --git a/source/blender/blenkernel/intern/scene.cc b/source/blender/blenkernel/intern/scene.cc
index 02f47a5ee35..71c77be5b46 100644
--- a/source/blender/blenkernel/intern/scene.cc
+++ b/source/blender/blenkernel/intern/scene.cc
@@ -216,7 +216,7 @@ static void scene_init_data(ID *id)
   }
 
   /* Master Collection */
-  scene->master_collection = BKE_collection_master_add();
+  scene->master_collection = BKE_collection_master_add(scene);
 
   BKE_view_layer_add(scene, "ViewLayer", nullptr, VIEWLAYER_ADD_NEW);
 }
diff --git a/source/blender/blenkernel/intern/simulation.cc b/source/blender/blenkernel/intern/simulation.cc
index 260d67de4d8..6afb698bcc0 100644
--- a/source/blender/blenkernel/intern/simulation.cc
+++ b/source/blender/blenkernel/intern/simulation.cc
@@ -51,8 +51,7 @@ static void simulation_init_data(ID *id)
 
   MEMCPY_STRUCT_AFTER(simulation, DNA_struct_default_get(Simulation), id);
 
-  bNodeTree *ntree = ntreeAddTree(nullptr, "Geometry Nodetree", ntreeType_Geometry->idname);
-  simulation->nodetree = ntree;
+  ntreeAddTreeEmbedded(nullptr, id, "Geometry Nodetree", ntreeType_Geometry->idname);
 }
 
 static void simulation_copy_data(Main *bmain, ID *id_dst, const ID *id_src, const int flag)