Add optional, free-after-use usages mapping of IDs to Main.

The new MainIDRelations stores two mappings, one from ID users to ID
used, the other vice-versa.

That data is assumed to be short-living runtime, code creating it is
responsible to clear it asap. It will be much useful in places where we
handle relations between IDs for a lot of them at once.

Note: This commit is not fully functional, that is, the infamous, ugly,
PoS non-ID nodetrees will not be handled correctly when building relations.
Fix needed here is a bit noisy, so will be done in next own commit.

3 files changed, 103 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_library.h b/source/blender/blenkernel/BKE_library.h
index 2d9c35f7fd0..34035673f88 100644
--- a/source/blender/blenkernel/BKE_library.h
+++ b/source/blender/blenkernel/BKE_library.h
@@ -105,6 +105,9 @@ void BKE_main_free(struct Main *mainvar);
 void BKE_main_lock(struct Main *bmain);
 void BKE_main_unlock(struct Main *bmain);
 
+void BKE_main_relations_create(struct Main *bmain);
+void BKE_main_relations_free(struct Main *bmain);
+
 struct BlendThumbnail *BKE_main_thumbnail_from_imbuf(struct Main *bmain, struct ImBuf *img);
 struct ImBuf *BKE_main_thumbnail_to_imbuf(struct Main *bmain, struct BlendThumbnail *data);
 void BKE_main_thumbnail_create(struct Main *bmain);
diff --git a/source/blender/blenkernel/BKE_main.h b/source/blender/blenkernel/BKE_main.h
index a4f5c425282..387045878f3 100644
--- a/source/blender/blenkernel/BKE_main.h
+++ b/source/blender/blenkernel/BKE_main.h
@@ -51,6 +51,8 @@ extern "C" {
 struct EvaluationContext;
 struct Library;
 struct MainLock;
+struct GHash;
+struct BLI_mempool;
 
 /* Blender thumbnail, as written on file (width, height, and data as char RGBA). */
 /* We pack pixel data after that struct. */
@@ -59,6 +61,22 @@ typedef struct BlendThumbnail {
 	char rect[0];
 } BlendThumbnail;
 
+/* Structs caching relations between data-blocks in a given Main. */
+typedef struct MainIDRelationsEntry {
+	struct MainIDRelationsEntry *next;
+	/* WARNING! for user_to_used, that pointer is really an ID** one, but for used_to_user, it’s only an ID* one! */
+	struct ID **id_pointer;
+	int usage_flag;  /* Using IDWALK_ enums, in BKE_library_query.h */
+} MainIDRelationsEntry;
+
+typedef struct MainIDRelations {
+	struct GHash *id_user_to_used;
+	struct GHash *id_used_to_user;
+
+	/* Private... */
+	struct BLI_mempool *entry_pool;
+} MainIDRelations;
+
 typedef struct Main {
 	struct Main *next, *prev;
 	char name[1024]; /* 1024 = FILE_MAX */
@@ -111,6 +129,11 @@ typedef struct Main {
 	/* Evaluation context used by viewport */
 	struct EvaluationContext *eval_ctx;
 
+	/* Must be generated, used and freed by same code - never assume this is valid data unless you know
+	 * when, who and how it was created.
+	 * Used by code doing a lot of remapping etc. at once to speed things up. */
+	struct MainIDRelations *relations;
+
 	struct MainLock *lock;
 } Main;
 
diff --git a/source/blender/blenkernel/intern/library.c b/source/blender/blenkernel/intern/library.c
index 77013a55d18..a06386f2f6a 100644
--- a/source/blender/blenkernel/intern/library.c
+++ b/source/blender/blenkernel/intern/library.c
@@ -76,6 +76,7 @@
 #include "BLI_ghash.h"
 #include "BLI_linklist.h"
 #include "BLI_memarena.h"
+#include "BLI_mempool.h"
 #include "BLI_string_utils.h"
 
 #include "BLI_threads.h"
@@ -1252,6 +1253,10 @@ void BKE_main_free(Main *mainvar)
 		}
 	}
 
+	if (mainvar->relations) {
+		BKE_main_relations_free(mainvar);
+	}
+
 	BLI_spin_end((SpinLock *)mainvar->lock);
 	MEM_freeN(mainvar->lock);
 	DEG_evaluation_context_free(mainvar->eval_ctx);
@@ -1268,6 +1273,78 @@ void BKE_main_unlock(struct Main *bmain)
 	BLI_spin_unlock((SpinLock *) bmain->lock);
 }
 
+
+static int main_relations_create_cb(void *user_data, ID *id_self, ID **id_pointer, int cd_flag)
+{
+	MainIDRelations *rel = user_data;
+
+	if (*id_pointer) {
+		MainIDRelationsEntry *entry, **entry_p;
+
+		entry = BLI_mempool_alloc(rel->entry_pool);
+		if (BLI_ghash_ensure_p(rel->id_user_to_used, id_self, (void ***)&entry_p)) {
+			entry->next = *entry_p;
+		}
+		else {
+			entry->next = NULL;
+		}
+		entry->id_pointer = id_pointer;
+		entry->usage_flag = cd_flag;
+		*entry_p = entry;
+
+		entry = BLI_mempool_alloc(rel->entry_pool);
+		if (BLI_ghash_ensure_p(rel->id_used_to_user, *id_pointer, (void ***)&entry_p)) {
+			entry->next = *entry_p;
+		}
+		else {
+			entry->next = NULL;
+		}
+		entry->id_pointer = (ID **)id_self;
+		entry->usage_flag = cd_flag;
+		*entry_p = entry;
+	}
+
+	return IDWALK_RET_NOP;
+}
+
+/** Generate the mappings between used IDs and their users, and vice-versa. */
+void BKE_main_relations_create(Main *bmain)
+{
+	ListBase *lbarray[MAX_LIBARRAY];
+	ID *id;
+	int a;
+
+	if (bmain->relations != NULL) {
+		BKE_main_relations_free(bmain);
+	}
+
+	bmain->relations = MEM_mallocN(sizeof(*bmain->relations), __func__);
+	bmain->relations->id_used_to_user = BLI_ghash_new(BLI_ghashutil_ptrhash, BLI_ghashutil_ptrcmp, __func__);
+	bmain->relations->id_user_to_used = BLI_ghash_new(BLI_ghashutil_ptrhash, BLI_ghashutil_ptrcmp, __func__);
+	bmain->relations->entry_pool = BLI_mempool_create(sizeof(MainIDRelationsEntry), 128, 128, BLI_MEMPOOL_NOP);
+
+	for (a = set_listbasepointers(bmain, lbarray); a--; ) {
+		for (id = lbarray[a]->first; id; id = id->next) {
+			BKE_library_foreach_ID_link(id, main_relations_create_cb, bmain->relations, IDWALK_READONLY);
+		}
+	}
+}
+
+void BKE_main_relations_free(Main *bmain)
+{
+	if (bmain->relations) {
+		if (bmain->relations->id_used_to_user) {
+			BLI_ghash_free(bmain->relations->id_used_to_user, NULL, NULL);
+		}
+		if (bmain->relations->id_user_to_used) {
+			BLI_ghash_free(bmain->relations->id_user_to_used, NULL, NULL);
+		}
+		BLI_mempool_destroy(bmain->relations->entry_pool);
+		MEM_freeN(bmain->relations);
+		bmain->relations = NULL;
+	}
+}
+
 /**
  * Generates a raw .blend file thumbnail data from given image.
  *