libquery: add optional handling of 'UI' ID pointers.

Handling those through different ways /might/ be needed sometimes, but
in most case this is just a nest of issues, since you can easily forget
to take them into account.

Note that this should be a 'non-functional' change, as this new behavior
is not used anywhere yet.

5 files changed, 183 insertions, 8 deletions

diff --git a/source/blender/blenkernel/BKE_lib_query.h b/source/blender/blenkernel/BKE_lib_query.h
index d47ba2497d4..35c42e62810 100644
--- a/source/blender/blenkernel/BKE_lib_query.h
+++ b/source/blender/blenkernel/BKE_lib_query.h
@@ -106,7 +106,8 @@ typedef int (*LibraryIDLinkCallback)(LibraryIDLinkCallbackData *cb_data);
 enum {
   IDWALK_NOP = 0,
   IDWALK_READONLY = (1 << 0),
-  IDWALK_RECURSE = (1 << 1), /* Also implies IDWALK_READONLY. */
+  IDWALK_RECURSE = (1 << 1),    /* Also implies IDWALK_READONLY. */
+  IDWALK_INCLUDE_UI = (1 << 2), /* Include UI pointers (from WM and screens editors). */
 
   IDWALK_NO_INDIRECT_PROXY_DATA_USAGE = (1 << 8), /* Ugly special case :(((( */
 };
diff --git a/source/blender/blenkernel/BKE_main.h b/source/blender/blenkernel/BKE_main.h
index 66cbe6d838f..8b46939c6d1 100644
--- a/source/blender/blenkernel/BKE_main.h
+++ b/source/blender/blenkernel/BKE_main.h
@@ -73,10 +73,17 @@ typedef struct MainIDRelations {
   struct GHash *id_user_to_used;
   struct GHash *id_used_to_user;
 
+  short flag;
+
   /* Private... */
   struct BLI_mempool *entry_pool;
 } MainIDRelations;
 
+enum {
+  /* Those bmain relations include pointers/usages from editors. */
+  MAINIDRELATIONS_INCLUDE_UI = 1 << 0,
+};
+
 typedef struct Main {
   struct Main *next, *prev;
   char name[1024];                   /* 1024 = FILE_MAX */
@@ -149,7 +156,7 @@ 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_create(struct Main *bmain, const short flag);
 void BKE_main_relations_free(struct Main *bmain);
 
 struct GSet *BKE_main_gset_create(struct Main *bmain, struct GSet *gset);
diff --git a/source/blender/blenkernel/intern/lib_id.c b/source/blender/blenkernel/intern/lib_id.c
index 1cb5905cc94..051b43a177e 100644
--- a/source/blender/blenkernel/intern/lib_id.c
+++ b/source/blender/blenkernel/intern/lib_id.c
@@ -2022,7 +2022,7 @@ void BKE_main_id_refcount_recompute(struct Main *bmain, const bool do_linked_onl
                                 id,
                                 id_refcount_recompute_callback,
                                 POINTER_FROM_INT((int)do_linked_only),
-                                IDWALK_READONLY);
+                                IDWALK_READONLY | IDWALK_INCLUDE_UI);
   }
   FOREACH_MAIN_ID_END;
 }
@@ -2128,7 +2128,7 @@ void BKE_library_make_local(Main *bmain,
   TIMEIT_START(make_local);
 #endif
 
-  BKE_main_relations_create(bmain);
+  BKE_main_relations_create(bmain, 0);
 
 #ifdef DEBUG_TIME
   printf("Pre-compute current ID relations: Done.\n");
diff --git a/source/blender/blenkernel/intern/lib_query.c b/source/blender/blenkernel/intern/lib_query.c
index 0f782717f79..e1868d88351 100644
--- a/source/blender/blenkernel/intern/lib_query.c
+++ b/source/blender/blenkernel/intern/lib_query.c
@@ -41,11 +41,13 @@
 #include "DNA_mask_types.h"
 #include "DNA_node_types.h"
 #include "DNA_object_force_types.h"
+#include "DNA_outliner_types.h"
 #include "DNA_lightprobe_types.h"
 #include "DNA_rigidbody_types.h"
 #include "DNA_scene_types.h"
 #include "DNA_sequence_types.h"
 #include "DNA_screen_types.h"
+#include "DNA_space_types.h"
 #include "DNA_speaker_types.h"
 #include "DNA_sound_types.h"
 #include "DNA_text_types.h"
@@ -375,6 +377,148 @@ static void library_foreach_collection(LibraryForeachIDData *data, Collection *c
   FOREACH_FINALIZE_VOID;
 }
 
+static void library_foreach_dopesheet(LibraryForeachIDData *data, bDopeSheet *ads)
+{
+  if (ads != NULL) {
+    FOREACH_CALLBACK_INVOKE_ID(data, ads->source, IDWALK_CB_NOP);
+    FOREACH_CALLBACK_INVOKE(data, ads->filter_grp, IDWALK_CB_NOP);
+  }
+
+  FOREACH_FINALIZE_VOID;
+}
+
+static void library_foreach_screen_area(LibraryForeachIDData *data, ScrArea *area)
+{
+  FOREACH_CALLBACK_INVOKE(data, area->full, IDWALK_CB_NOP);
+
+  for (SpaceLink *sl = area->spacedata.first; sl; sl = sl->next) {
+    switch (sl->spacetype) {
+      case SPACE_VIEW3D: {
+        View3D *v3d = (View3D *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, v3d->camera, IDWALK_CB_NOP);
+        FOREACH_CALLBACK_INVOKE(data, v3d->ob_centre, IDWALK_CB_NOP);
+
+        if (v3d->localvd) {
+          FOREACH_CALLBACK_INVOKE(data, v3d->localvd->camera, IDWALK_CB_NOP);
+        }
+        break;
+      }
+      case SPACE_GRAPH: {
+        SpaceGraph *sipo = (SpaceGraph *)sl;
+
+        library_foreach_dopesheet(data, sipo->ads);
+        break;
+      }
+      case SPACE_PROPERTIES: {
+        SpaceProperties *sbuts = (SpaceProperties *)sl;
+
+        FOREACH_CALLBACK_INVOKE_ID(data, sbuts->pinid, IDWALK_CB_NOP);
+        break;
+      }
+      case SPACE_FILE:
+        break;
+      case SPACE_ACTION: {
+        SpaceAction *saction = (SpaceAction *)sl;
+
+        library_foreach_dopesheet(data, &saction->ads);
+        FOREACH_CALLBACK_INVOKE(data, saction->action, IDWALK_CB_NOP);
+        break;
+      }
+      case SPACE_IMAGE: {
+        SpaceImage *sima = (SpaceImage *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, sima->image, IDWALK_CB_USER_ONE);
+        FOREACH_CALLBACK_INVOKE(data, sima->mask_info.mask, IDWALK_CB_USER_ONE);
+        FOREACH_CALLBACK_INVOKE(data, sima->gpd, IDWALK_CB_USER);
+        break;
+      }
+      case SPACE_SEQ: {
+        SpaceSeq *sseq = (SpaceSeq *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, sseq->gpd, IDWALK_CB_USER);
+        break;
+      }
+      case SPACE_NLA: {
+        SpaceNla *snla = (SpaceNla *)sl;
+
+        library_foreach_dopesheet(data, snla->ads);
+        break;
+      }
+      case SPACE_TEXT: {
+        SpaceText *st = (SpaceText *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, st->text, IDWALK_CB_NOP);
+        break;
+      }
+      case SPACE_SCRIPT: {
+        SpaceScript *scpt = (SpaceScript *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, scpt->script, IDWALK_CB_NOP);
+        break;
+      }
+      case SPACE_OUTLINER: {
+        SpaceOutliner *so = (SpaceOutliner *)sl;
+
+        FOREACH_CALLBACK_INVOKE_ID(data, so->search_tse.id, IDWALK_CB_NOP);
+
+        if (so->treestore != NULL) {
+          TreeStoreElem *tselem;
+          BLI_mempool_iter iter;
+
+          BLI_mempool_iternew(so->treestore, &iter);
+          while ((tselem = BLI_mempool_iterstep(&iter))) {
+            FOREACH_CALLBACK_INVOKE_ID(data, tselem->id, IDWALK_CB_NOP);
+          }
+        }
+        break;
+      }
+      case SPACE_NODE: {
+        SpaceNode *snode = (SpaceNode *)sl;
+        bNodeTreePath *path;
+
+        const bool is_private_nodetree = snode->id != NULL &&
+                                         ntreeFromID(snode->id) == snode->nodetree;
+
+        FOREACH_CALLBACK_INVOKE_ID(data, snode->id, IDWALK_CB_NOP);
+        FOREACH_CALLBACK_INVOKE_ID(data, snode->from, IDWALK_CB_NOP);
+
+        FOREACH_CALLBACK_INVOKE(
+            data, snode->nodetree, is_private_nodetree ? IDWALK_CB_PRIVATE : IDWALK_CB_USER);
+
+        for (path = snode->treepath.first; path; path = path->next) {
+          if (path == snode->treepath.first) {
+            /* first nodetree in path is same as snode->nodetree */
+            FOREACH_CALLBACK_INVOKE(
+                data, path->nodetree, is_private_nodetree ? IDWALK_CB_PRIVATE : IDWALK_CB_NOP);
+          }
+          else {
+            FOREACH_CALLBACK_INVOKE(data, path->nodetree, IDWALK_CB_USER);
+          }
+
+          if (path->nodetree == NULL) {
+            break;
+          }
+        }
+
+        FOREACH_CALLBACK_INVOKE(data, snode->edittree, IDWALK_CB_NOP);
+        break;
+      }
+      case SPACE_CLIP: {
+        SpaceClip *sclip = (SpaceClip *)sl;
+
+        FOREACH_CALLBACK_INVOKE(data, sclip->clip, IDWALK_CB_USER_ONE);
+        FOREACH_CALLBACK_INVOKE(data, sclip->mask_info.mask, IDWALK_CB_USER_ONE);
+        break;
+      }
+      default:
+        break;
+    }
+  }
+
+  FOREACH_FINALIZE_VOID;
+}
+
 static void library_foreach_ID_as_subdata_link(ID **id_pp,
                                                LibraryIDLinkCallback callback,
                                                void *user_data,
@@ -456,7 +600,9 @@ static void library_foreach_ID_link(Main *bmain,
       data.cb_flag_clear = inherit_data->cb_flag_clear;
     }
 
-    if (bmain != NULL && bmain->relations != NULL && (flag & IDWALK_READONLY)) {
+    if (bmain != NULL && bmain->relations != NULL && (flag & IDWALK_READONLY) &&
+        (((bmain->relations->flag & MAINIDRELATIONS_INCLUDE_UI) == 0) ==
+         ((data.flag & IDWALK_INCLUDE_UI) == 0))) {
       /* Note that this is minor optimization, even in worst cases (like id being an object with
        * lots of drivers and constraints and modifiers, or material etc. with huge node tree),
        * but we might as well use it (Main->relations is always assumed valid,
@@ -1020,6 +1166,7 @@ static void library_foreach_ID_link(Main *bmain,
         }
         break;
       }
+
       case ID_AC: {
         bAction *act = (bAction *)id;
 
@@ -1042,6 +1189,11 @@ static void library_foreach_ID_link(Main *bmain,
             /* allow callback to set a different workspace */
             BKE_workspace_active_set(win->workspace_hook, (WorkSpace *)workspace);
           }
+          if (data.flag & IDWALK_INCLUDE_UI) {
+            for (ScrArea *area = win->global_areas.areabase.first; area; area = area->next) {
+              library_foreach_screen_area(&data, area);
+            }
+          }
         }
         break;
       }
@@ -1062,6 +1214,7 @@ static void library_foreach_ID_link(Main *bmain,
         }
         break;
       }
+
       case ID_GD: {
         bGPdata *gpencil = (bGPdata *)id;
         /* materials */
@@ -1077,8 +1230,18 @@ static void library_foreach_ID_link(Main *bmain,
         break;
       }
 
+      case ID_SCR: {
+        if (data.flag & IDWALK_INCLUDE_UI) {
+          bScreen *screen = (bScreen *)id;
+
+          for (ScrArea *area = screen->areabase.first; area; area = area->next) {
+            library_foreach_screen_area(&data, area);
+          }
+        }
+        break;
+      }
+
       /* Nothing needed for those... */
-      case ID_SCR:
       case ID_IM:
       case ID_VF:
       case ID_TXT:
diff --git a/source/blender/blenkernel/intern/main.c b/source/blender/blenkernel/intern/main.c
index 75bf74f63bf..659c3944edb 100644
--- a/source/blender/blenkernel/intern/main.c
+++ b/source/blender/blenkernel/intern/main.c
@@ -246,7 +246,7 @@ static int main_relations_create_idlink_cb(LibraryIDLinkCallbackData *cb_data)
 }
 
 /** Generate the mappings between used IDs and their users, and vice-versa. */
-void BKE_main_relations_create(Main *bmain)
+void BKE_main_relations_create(Main *bmain, const short flag)
 {
   if (bmain->relations != NULL) {
     BKE_main_relations_free(bmain);
@@ -262,10 +262,14 @@ void BKE_main_relations_create(Main *bmain)
 
   ID *id;
   FOREACH_MAIN_ID_BEGIN (bmain, id) {
+    const int idwalk_flag = IDWALK_READONLY |
+                            ((flag & MAINIDRELATIONS_INCLUDE_UI) != 0 ? IDWALK_INCLUDE_UI : 0);
     BKE_library_foreach_ID_link(
-        NULL, id, main_relations_create_idlink_cb, bmain->relations, IDWALK_READONLY);
+        NULL, id, main_relations_create_idlink_cb, bmain->relations, idwalk_flag);
   }
   FOREACH_MAIN_ID_END;
+
+  bmain->relations->flag = flag;
 }
 
 void BKE_main_relations_free(Main *bmain)