'users of ID' py API.

This mainly adds bpy.data.user_map() method, which goes over the whole Main database
to build a mapping (dict) {ID: {users_of_that_ID}}.

Very handy to check and debug ID usages, but could also be really valuable for py addons
creating temporary scenes, or some exporters, etc.

Note: current code in master's libquery misses some IDs (and reports some it should not,
like nodetrees), this is fixed in id-remap but still needs serious review before going to master.
This basically means that current bpy.data.user_map() **will not** report a complete and exhaustive
state of dependencies between IDs. Should work OK in most cases though.

Original work/idea comes from id-remap branch, was heavily reworked by @campbellbarton
and myself for master.

Reviewers: campbellbarton, sergey

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

10 files changed, 434 insertions, 0 deletions

diff --git a/release/scripts/modules/bpy_types.py b/release/scripts/modules/bpy_types.py
index 1c08fc5e51f..830db976822 100644
--- a/release/scripts/modules/bpy_types.py
+++ b/release/scripts/modules/bpy_types.py
@@ -26,6 +26,7 @@ StructMetaPropGroup = bpy_types.bpy_struct_meta_idprop
 # StructRNA = bpy_types.Struct
 
 bpy_types.BlendDataLibraries.load = _bpy._library_load
+bpy_types.BlendData.user_map = _bpy._rna_id_collection_user_map
 
 
 class Context(StructRNA):
diff --git a/source/blender/blenkernel/BKE_library_query.h b/source/blender/blenkernel/BKE_library_query.h
index 194b6f0e6b3..f07644dc833 100644
--- a/source/blender/blenkernel/BKE_library_query.h
+++ b/source/blender/blenkernel/BKE_library_query.h
@@ -65,4 +65,6 @@ enum {
 void BKE_library_foreach_ID_link(struct ID *id, LibraryIDLinkCallback callback, void *user_data, int flag);
 void BKE_library_update_ID_link_user(struct ID *id_dst, struct ID *id_src, const int cd_flag);
 
+int BKE_library_ID_use_ID(struct ID *id_user, struct ID *id_used);
+
 #endif  /* __BKE_LIBRARY_QUERY_H__ */
diff --git a/source/blender/blenkernel/intern/library_query.c b/source/blender/blenkernel/intern/library_query.c
index d54f382f4c0..7936aab6700 100644
--- a/source/blender/blenkernel/intern/library_query.c
+++ b/source/blender/blenkernel/intern/library_query.c
@@ -665,3 +665,49 @@ void BKE_library_update_ID_link_user(ID *id_dst, ID *id_src, const int cd_flag)
 		id_us_ensure_real(id_dst);
 	}
 }
+
+/* ***** ID users iterator. ***** */
+typedef struct IDUsersIter {
+	ID *id;
+
+	ListBase *lb_array[MAX_LIBARRAY];
+	int lb_idx;
+
+	ID *curr_id;
+	int count;  /* Set by callback. */
+} IDUsersIter;
+
+static bool foreach_libblock_id_users_callback(void *user_data, ID **id_p, int UNUSED(cb_flag))
+{
+	IDUsersIter *iter = user_data;
+
+	if (*id_p && (*id_p == iter->id)) {
+		iter->count++;
+	}
+
+	return true;
+}
+
+/**
+ * Return the number of times given \a id_user uses/references \a id_used.
+ *
+ * \note This only checks for pointer references of an ID, shallow usages (like e.g. by RNA paths, as done
+ *       for FCurves) are not detected at all.
+ *
+ * \param id_user the ID which is supposed to use (reference) \a id_used.
+ * \param id_used the ID which is supposed to be used (referenced) by \a id_user.
+ * \return the number of direct usages/references of \a id_used by \a id_user.
+ */
+int BKE_library_ID_use_ID(ID *id_user, ID *id_used)
+{
+	IDUsersIter iter;
+
+	/* We do not care about iter.lb_array/lb_idx here... */
+	iter.id = id_used;
+	iter.curr_id = id_user;
+	iter.count = 0;
+
+	BKE_library_foreach_ID_link(iter.curr_id, foreach_libblock_id_users_callback, (void *)&iter, IDWALK_NOP);
+
+	return iter.count;
+}
diff --git a/source/blender/makesrna/intern/rna_ID.c b/source/blender/makesrna/intern/rna_ID.c
index 3a163300264..63ea836e3de 100644
--- a/source/blender/makesrna/intern/rna_ID.c
+++ b/source/blender/makesrna/intern/rna_ID.c
@@ -90,6 +90,7 @@ EnumPropertyItem rna_enum_id_type_items[] = {
 #include "BKE_font.h"
 #include "BKE_idprop.h"
 #include "BKE_library.h"
+#include "BKE_library_query.h"
 #include "BKE_animsys.h"
 #include "BKE_material.h"
 #include "BKE_depsgraph.h"
@@ -976,6 +977,14 @@ static void rna_def_ID(BlenderRNA *brna)
 	RNA_def_function_ui_description(func, "Clear the user count of a data-block so its not saved, "
 	                                "on reload the data will be removed");
 
+	func = RNA_def_function(srna, "user_of_id", "BKE_library_ID_use_ID");
+	RNA_def_function_ui_description(func, "Count the number of times that ID uses/references given one");
+	parm = RNA_def_pointer(func, "id", "ID", "", "ID to count usages");
+	RNA_def_property_flag(parm, PROP_NEVER_NULL);
+	parm = RNA_def_int(func, "count", 0, 0, INT_MAX,
+	                   "", "Number of usages/references of given id by current datablock", 0, INT_MAX);
+	RNA_def_function_return(func, parm);
+
 	func = RNA_def_function(srna, "animation_data_create", "rna_ID_animation_data_create");
 	RNA_def_function_flag(func, FUNC_USE_MAIN);
 	RNA_def_function_ui_description(func, "Create animation data to this ID, note that not all ID types support this");
diff --git a/source/blender/python/intern/CMakeLists.txt b/source/blender/python/intern/CMakeLists.txt
index f04bca75a8c..cbfbe0a8768 100644
--- a/source/blender/python/intern/CMakeLists.txt
+++ b/source/blender/python/intern/CMakeLists.txt
@@ -69,6 +69,7 @@ set(SRC
 	bpy_rna_anim.c
 	bpy_rna_array.c
 	bpy_rna_callback.c
+	bpy_rna_id_collection.c
 	bpy_traceback.c
 	bpy_util.c
 	bpy_utils_previews.c
@@ -95,6 +96,7 @@ set(SRC
 	bpy_rna.h
 	bpy_rna_anim.h
 	bpy_rna_callback.h
+	bpy_rna_id_collection.h
 	bpy_traceback.h
 	bpy_util.h
 	bpy_utils_previews.h
diff --git a/source/blender/python/intern/bpy.c b/source/blender/python/intern/bpy.c
index bdae2a7cb17..930950a6600 100644
--- a/source/blender/python/intern/bpy.c
+++ b/source/blender/python/intern/bpy.c
@@ -45,6 +45,7 @@
 #include "bpy_util.h"
 #include "bpy_rna.h"
 #include "bpy_app.h"
+#include "bpy_rna_id_collection.h"
 #include "bpy_props.h"
 #include "bpy_library.h"
 #include "bpy_operator.h"
@@ -322,6 +323,7 @@ void BPy_init_modules(void)
 
 	/* needs to be first so bpy_types can run */
 	BPY_library_module(mod);
+	BPY_rna_id_collection_module(mod);
 
 	bpy_import_test("bpy_types");
 	PyModule_AddObject(mod, "data", BPY_rna_module()); /* imports bpy_types by running this */
diff --git a/source/blender/python/intern/bpy_rna.c b/source/blender/python/intern/bpy_rna.c
index 45a57352376..d338b499114 100644
--- a/source/blender/python/intern/bpy_rna.c
+++ b/source/blender/python/intern/bpy_rna.c
@@ -37,6 +37,7 @@
 
 #include "RNA_types.h"
 
+#include "BLI_bitmap.h"
 #include "BLI_dynstr.h"
 #include "BLI_string.h"
 #include "BLI_listbase.h"
@@ -1176,6 +1177,69 @@ static int pyrna_string_to_enum(PyObject *item, PointerRNA *ptr, PropertyRNA *pr
 	return 0;
 }
 
+/**
+ * Takes a set of strings and map it to and array of booleans.
+ *
+ * Useful when the values aren't flags.
+ *
+ * \param type_convert_sign: Maps signed to unsuigned range,
+ * needed when we want to use the full range of a signed short/char.
+ */
+BLI_bitmap *pyrna_set_to_enum_bitmap(
+        EnumPropertyItem *items, PyObject *value,
+        int type_size, bool type_convert_sign,
+        int bitmap_size,
+        const char *error_prefix)
+{
+	/* set looping */
+	Py_ssize_t pos = 0;
+	Py_ssize_t hash = 0;
+	PyObject *key;
+
+	BLI_bitmap *bitmap = BLI_BITMAP_NEW(bitmap_size, __func__);
+
+	while (_PySet_NextEntry(value, &pos, &key, &hash)) {
+		const char *param = _PyUnicode_AsString(key);
+		if (param == NULL) {
+			PyErr_Format(PyExc_TypeError,
+			             "%.200s expected a string, not %.200s",
+			             error_prefix, Py_TYPE(key)->tp_name);
+			goto error;
+		}
+
+		int ret;
+		if (pyrna_enum_value_from_id(items, param, &ret, error_prefix) == -1) {
+			goto error;
+		}
+
+		int index = ret;
+
+		if (type_convert_sign) {
+			if (type_size == 2) {
+				union { signed short as_signed; unsigned short as_unsigned; } ret_convert;
+				ret_convert.as_signed = (signed short)ret;
+				index = (int)ret_convert.as_unsigned;
+			}
+			else if (type_size == 1) {
+				union { signed char as_signed; unsigned char as_unsigned; } ret_convert;
+				ret_convert.as_signed = (signed char)ret;
+				index = (int)ret_convert.as_unsigned;
+			}
+			else {
+				BLI_assert(0);
+			}
+		}
+		BLI_assert(index < bitmap_size);
+		BLI_BITMAP_ENABLE(bitmap, index);
+	}
+
+	return bitmap;
+
+error:
+	MEM_freeN(bitmap);
+	return NULL;
+}
+
 /* 'value' _must_ be a set type, error check before calling */
 int pyrna_set_to_enum_bitfield(EnumPropertyItem *items, PyObject *value, int *r_value, const char *error_prefix)
 {
diff --git a/source/blender/python/intern/bpy_rna.h b/source/blender/python/intern/bpy_rna.h
index f546c2955e5..c5d4a346f56 100644
--- a/source/blender/python/intern/bpy_rna.h
+++ b/source/blender/python/intern/bpy_rna.h
@@ -184,6 +184,11 @@ bool      pyrna_id_FromPyObject(PyObject *obj, struct ID **id);
 int pyrna_pydict_to_props(PointerRNA *ptr, PyObject *kw, int all_args, const char *error_prefix);
 PyObject *pyrna_prop_to_py(PointerRNA *ptr, PropertyRNA *prop);
 
+unsigned int *pyrna_set_to_enum_bitmap(
+        struct EnumPropertyItem *items, PyObject *value,
+        int type_size, bool type_convert_sign,
+        int bitmap_size,
+        const char *error_prefix);
 PyObject *pyrna_enum_bitfield_to_py(struct EnumPropertyItem *items, int value);
 int pyrna_set_to_enum_bitfield(EnumPropertyItem *items, PyObject *value, int *r_value, const char *error_prefix);
 
diff --git a/source/blender/python/intern/bpy_rna_id_collection.c b/source/blender/python/intern/bpy_rna_id_collection.c
new file mode 100644
index 00000000000..a80190d8162
--- /dev/null
+++ b/source/blender/python/intern/bpy_rna_id_collection.c
@@ -0,0 +1,271 @@
+/*
+ * ***** BEGIN GPL LICENSE BLOCK *****
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ *
+ * Contributor(s): Bastien Montagne
+ *
+ * ***** END GPL LICENSE BLOCK *****
+ */
+
+/** \file blender/python/intern/bpy_rna_id_collection.c
+ *  \ingroup pythonintern
+ *
+ * This file adds some helpers related to ID/Main handling, that cannot fit well in RNA itself.
+ */
+
+#include <Python.h>
+#include <stddef.h>
+
+#include "MEM_guardedalloc.h"
+
+#include "BLI_utildefines.h"
+#include "BLI_bitmap.h"
+
+#include "BKE_global.h"
+#include "BKE_main.h"
+#include "BKE_library.h"
+#include "BKE_library_query.h"
+
+#include "DNA_ID.h"
+
+#include "bpy_util.h"
+#include "bpy_rna_id_collection.h"
+
+#include "../generic/py_capi_utils.h"
+#include "../generic/python_utildefines.h"
+
+#include "RNA_access.h"
+#include "RNA_types.h"
+#include "RNA_enum_types.h"
+
+#include "bpy_rna.h"
+
+typedef struct IDUserMapData {
+	/* place-holder key only used for lookups to avoid creating new data only for lookups
+	 * (never return its contents) */
+	PyObject *py_id_key_lookup_only;
+
+	/* we loop over data-blocks that this ID points to (do build a reverse lookup table) */
+	PyObject *py_id_curr;
+	ID *id_curr;
+
+	/* filter the values we add into the set */
+	BLI_bitmap *types_bitmap;
+
+	PyObject *user_map;     /* set to fill in as we iterate */
+	bool is_subset;         /* true when we're only mapping a subset of all the ID's (subset arg is passed) */
+} IDUserMapData;
+
+
+static int id_code_as_index(const short idcode)
+{
+	return (int)*((unsigned short *)&idcode);
+}
+
+static bool id_check_type(const ID *id, const BLI_bitmap *types_bitmap)
+{
+	return BLI_BITMAP_TEST_BOOL(types_bitmap, id_code_as_index(GS(id->name)));
+}
+
+static bool foreach_libblock_id_user_map_callback(void *user_data, ID **id_p, int UNUSED(cb_flag))
+{
+	IDUserMapData *data = user_data;
+
+	if (*id_p) {
+
+		if (data->types_bitmap) {
+			if (!id_check_type(*id_p, data->types_bitmap)) {
+				return true;
+			}
+		}
+
+		/* pyrna_struct_hash() uses ptr.data only,
+		 * but pyrna_struct_richcmp() uses also ptr.type,
+		 * so we need to create a valid PointerRNA here...
+		 */
+		PyObject *key = data->py_id_key_lookup_only;
+		RNA_id_pointer_create(*id_p, &((BPy_StructRNA *)key)->ptr);
+
+		PyObject *set;
+		if ((set = PyDict_GetItem(data->user_map, key)) == NULL) {
+
+			/* limit to key's added already */
+			if (data->is_subset) {
+				return true;
+			}
+
+			/* Cannot use our placeholder key here! */
+			key = pyrna_id_CreatePyObject(*id_p);
+			set = PySet_New(NULL);
+			PyDict_SetItem(data->user_map, key, set);
+			Py_DECREF(set);
+			Py_DECREF(key);
+		}
+
+		if (data->py_id_curr == NULL) {
+			data->py_id_curr = pyrna_id_CreatePyObject(data->id_curr);
+		}
+
+		PySet_Add(set, data->py_id_curr);
+	}
+
+	return true;
+}
+
+PyDoc_STRVAR(bpy_user_map_doc,
+".. method:: user_map([subset=(id1, id2, ...)], key_types={..}, value_types={..})\n"
+"\n"
+"   Returns a mapping of all ID datablocks in current ``bpy.data`` to a set of all datablocks using them.\n"
+"\n"
+"   For list of valid set members for key_types & value_types, see: :class:`bpy.types.KeyingSetPath.id_type`.\n"
+"\n"
+"   :arg subset: When passed, only these data-blocks and their users will be included as keys/values in the map.\n"
+"   :type subset: sequence\n"
+"   :arg key_types: Filter the keys mapped by ID types.\n"
+"   :type key_types: set of strings\n"
+"   :arg value_types: Filter the values in the set by ID types.\n"
+"   :type value_types: set of strings\n"
+"   :return: dictionary of :class:`bpy.types.ID` instances, with sets of ID's as their values.\n"
+"   :rtype: dict\n"
+);
+static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *kwds)
+{
+#if 0  /* If someone knows how to get a proper 'self' in that case... */
+	BPy_StructRNA *pyrna = (BPy_StructRNA *)self;
+	Main *bmain = pyrna->ptr.data;
+#else
+	Main *bmain = G.main;  /* XXX Ugly, but should work! */
+#endif
+
+	static const char *kwlist[] = {"subset", "key_types", "value_types", NULL};
+	PyObject *subset = NULL;
+
+	PyObject *key_types = NULL;
+	PyObject *val_types = NULL;
+	BLI_bitmap *key_types_bitmap = NULL;
+	BLI_bitmap *val_types_bitmap = NULL;
+
+	PyObject *ret = NULL;
+
+
+	if (!PyArg_ParseTupleAndKeywords(
+	        args, kwds, "|O$O!O!:user_map", (char **)kwlist,
+	        &subset,
+	        &PySet_Type, &key_types,
+	        &PySet_Type, &val_types))
+	{
+		return NULL;
+	}
+
+	if (key_types) {
+		key_types_bitmap = pyrna_set_to_enum_bitmap(
+		        rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types");
+		if (key_types_bitmap == NULL) {
+			goto error;
+		}
+	}
+
+	if (val_types) {
+		val_types_bitmap = pyrna_set_to_enum_bitmap(
+		        rna_enum_id_type_items, val_types, sizeof(short), true, USHRT_MAX, "value types");
+		if (val_types_bitmap == NULL) {
+			goto error;
+		}
+	}
+
+	IDUserMapData data_cb = {NULL};
+
+	if (subset) {
+		PyObject *subset_fast = PySequence_Fast(subset, "user_map");
+		if (subset_fast == NULL) {
+			goto error;
+		}
+
+		PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast);
+		Py_ssize_t subset_len = PySequence_Fast_GET_SIZE(subset_fast);
+
+		data_cb.user_map = PyDict_New();
+		data_cb.is_subset = true;
+		for (; subset_len; subset_array++, subset_len--) {
+			PyObject *set = PySet_New(NULL);
+			PyDict_SetItem(data_cb.user_map, *subset_array, set);
+			Py_DECREF(set);
+		}
+		Py_DECREF(subset_fast);
+	}
+	else {
+		data_cb.user_map = PyDict_New();
+	}
+
+	data_cb.types_bitmap = key_types_bitmap;
+
+	ListBase *lb_array[MAX_LIBARRAY];
+	int lb_index;
+	lb_index = set_listbasepointers(bmain, lb_array);
+
+	while (lb_index--) {
+
+		if (val_types_bitmap && lb_array[lb_index]->first) {
+			if (!id_check_type(lb_array[lb_index]->first, val_types_bitmap)) {
+				continue;
+			}
+		}
+
+		for (ID *id = lb_array[lb_index]->first; id; id = id->next) {
+			/* One-time init, ID is just used as placeholder here, we abuse this in iterator callback
+			 * to avoid having to rebuild a complete bpyrna object each time for the key searching
+			 * (where only ID pointer value is used). */
+			if (data_cb.py_id_key_lookup_only == NULL) {
+				data_cb.py_id_key_lookup_only = pyrna_id_CreatePyObject(id);
+			}
+
+			data_cb.id_curr = id;
+			BKE_library_foreach_ID_link(id, foreach_libblock_id_user_map_callback, &data_cb, IDWALK_NOP);
+			if (data_cb.py_id_curr) {
+				Py_DECREF(data_cb.py_id_curr);
+				data_cb.py_id_curr = NULL;
+			}
+		}
+	}
+
+	ret = data_cb.user_map;
+
+
+error:
+
+	Py_XDECREF(data_cb.py_id_key_lookup_only);
+
+	if (key_types_bitmap) {
+		MEM_freeN(key_types_bitmap);
+	}
+
+	if (val_types_bitmap) {
+		MEM_freeN(val_types_bitmap);
+	}
+
+	return ret;
+
+}
+
+int BPY_rna_id_collection_module(PyObject *mod_par)
+{
+	static PyMethodDef user_map = {
+	    "user_map", (PyCFunction)bpy_user_map, METH_VARARGS | METH_KEYWORDS, bpy_user_map_doc};
+
+	PyModule_AddObject(mod_par, "_rna_id_collection_user_map", PyCFunction_New(&user_map, NULL));
+
+	return 0;
+}
diff --git a/source/blender/python/intern/bpy_rna_id_collection.h b/source/blender/python/intern/bpy_rna_id_collection.h
new file mode 100644
index 00000000000..c293289c821
--- /dev/null
+++ b/source/blender/python/intern/bpy_rna_id_collection.h
@@ -0,0 +1,32 @@
+/*
+ * ***** BEGIN GPL LICENSE BLOCK *****
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ *
+ * Contributor(s): Bastien Montagne
+ *
+ * ***** END GPL LICENSE BLOCK *****
+ */
+
+/** \file blender/python/intern/bpy_rna_id_collection.h
+ *  \ingroup pythonintern
+ */
+
+#ifndef __BPY_RNA_ID_COLLECTION_H__
+#define __BPY_RNA_ID_COLLECTION_H__
+
+int BPY_rna_id_collection_module(PyObject *);
+
+#endif  /* __BPY_RNA_ID_COLLECTION_H__ */