Nodes: Add node group assets in add menu

This patch builds on the work from bdb57541475f to add node group
assets directly in the node editor add menu. Assets are added after
separators to distinguish them, but otherwise they look like any other
node. The catalog trees from all configured libraries are used to build
the menu hierarchy. Only catalogs with matching asset types are used
though.

There are a few limitations of this initial version. For now this only
supports geometry nodes. Support for other built-in node systems just
requires some refactoring of the corresponding add menu though. Lazy
loading will be added in a followup commit. For now there is a label
the first time the menu is opened.

Like the search menu integration, re-saving asset library files in 3.4
is required, if it hasn't been done already.

Implementation wise, there is a some ugly code here. A lot of that is
because the asset system isn't complete. The RNA API doesn't work well
yet, and the system isn't built to interact with multiple libraries at
once. It's also ugly because of the way we combine automatic menu
generation with builtin menus. As noted in a code comment, these two
systems could be merged completely so that the menus for builtin nodes
are also generated in the same way.

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

1 files changed, 297 insertions, 0 deletions

diff --git a/source/blender/editors/space_node/add_menu_assets.cc b/source/blender/editors/space_node/add_menu_assets.cc
new file mode 100644
index 00000000000..4e305ac8af4
--- /dev/null
+++ b/source/blender/editors/space_node/add_menu_assets.cc
@@ -0,0 +1,297 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "BLI_multi_value_map.hh"
+
+#include "DNA_screen_types.h"
+#include "DNA_space_types.h"
+
+#include "BKE_asset.h"
+#include "BKE_asset_catalog.hh"
+#include "BKE_asset_library.hh"
+#include "BKE_idprop.h"
+#include "BKE_screen.h"
+
+#include "BLT_translation.h"
+
+#include "RNA_access.h"
+#include "RNA_prototypes.h"
+
+#include "ED_asset.h"
+
+#include "node_intern.hh"
+
+namespace blender::ed::space_node {
+
+static bool node_add_menu_poll(const bContext *C, MenuType * /*mt*/)
+{
+  return CTX_wm_space_node(C);
+}
+
+struct LibraryAsset {
+  AssetLibraryReference library_ref;
+  AssetHandle handle;
+};
+
+struct LibraryCatalog {
+  bke::AssetLibrary *library;
+  const bke::AssetCatalog *catalog;
+};
+
+struct AssetItemTree {
+  bke::AssetCatalogTree catalogs;
+  MultiValueMap<bke::AssetCatalogPath, LibraryAsset> assets_per_path;
+  Map<const bke::AssetCatalogTreeItem *, bke::AssetCatalogPath> full_catalog_per_tree_item;
+};
+
+static bool all_loading_finished()
+{
+  for (const AssetLibraryReference &library : bke::all_valid_asset_library_refs()) {
+    if (!ED_assetlist_is_loaded(&library)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+static AssetItemTree build_catalog_tree(const bContext &C, const bNodeTree *node_tree)
+{
+  if (!node_tree) {
+    return {};
+  }
+  const Main &bmain = *CTX_data_main(&C);
+  const Vector<AssetLibraryReference> all_libraries = bke::all_valid_asset_library_refs();
+
+  /* Merge catalogs from all libraries to deduplicate menu items. Also store the catalog and
+   * library for each asset ID in order to use them later when retrieving assets and removing
+   * empty catalogs.  */
+  Map<bke::CatalogID, LibraryCatalog> id_to_catalog_map;
+  bke::AssetCatalogTree catalogs_from_all_libraries;
+  for (const AssetLibraryReference &library_ref : all_libraries) {
+    if (bke::AssetLibrary *library = BKE_asset_library_load(&bmain, library_ref)) {
+      if (bke::AssetCatalogTree *tree = library->catalog_service->get_catalog_tree()) {
+        tree->foreach_item([&](bke::AssetCatalogTreeItem &item) {
+          const bke::CatalogID &id = item.get_catalog_id();
+          bke::AssetCatalog *catalog = library->catalog_service->find_catalog(id);
+          catalogs_from_all_libraries.insert_item(*catalog);
+          id_to_catalog_map.add(item.get_catalog_id(), LibraryCatalog{library, catalog});
+        });
+      }
+    }
+  }
+
+  /* Find all the matching node group assets for every catalog path. */
+  MultiValueMap<bke::AssetCatalogPath, LibraryAsset> assets_per_path;
+  for (const AssetLibraryReference &library_ref : all_libraries) {
+    AssetFilterSettings type_filter{};
+    type_filter.id_types = FILTER_ID_NT;
+
+    ED_assetlist_storage_fetch(&library_ref, &C);
+    ED_assetlist_ensure_previews_job(&library_ref, &C);
+    ED_assetlist_iterate(library_ref, [&](AssetHandle asset) {
+      if (!ED_asset_filter_matches_asset(&type_filter, &asset)) {
+        return true;
+      }
+      const AssetMetaData &meta_data = *ED_asset_handle_get_metadata(&asset);
+      const IDProperty *tree_type = BKE_asset_metadata_idprop_find(&meta_data, "type");
+      if (tree_type == nullptr || IDP_Int(tree_type) != node_tree->type) {
+        return true;
+      }
+      if (BLI_uuid_is_nil(meta_data.catalog_id)) {
+        return true;
+      }
+      const LibraryCatalog &library_catalog = id_to_catalog_map.lookup(meta_data.catalog_id);
+      assets_per_path.add(library_catalog.catalog->path, LibraryAsset{library_ref, asset});
+      return true;
+    });
+  }
+
+  /* Build the final tree without any of the catalogs that don't have proper node group assets. */
+  bke::AssetCatalogTree catalogs_with_node_assets;
+  catalogs_from_all_libraries.foreach_item([&](bke::AssetCatalogTreeItem &item) {
+    if (!assets_per_path.lookup(item.catalog_path()).is_empty()) {
+      const bke::CatalogID &id = item.get_catalog_id();
+      const LibraryCatalog &library_catalog = id_to_catalog_map.lookup(id);
+      bke::AssetCatalog *catalog = library_catalog.library->catalog_service->find_catalog(id);
+      catalogs_with_node_assets.insert_item(*catalog);
+    }
+  });
+
+  /* Build another map storing full asset paths for each tree item, in order to have stable
+   * pointers to asset catalog paths to use for context pointers. This is necessary because
+   * #bke::AssetCatalogTreeItem doesn't store its full path directly. */
+  Map<const bke::AssetCatalogTreeItem *, bke::AssetCatalogPath> full_catalog_per_tree_item;
+  catalogs_with_node_assets.foreach_item([&](bke::AssetCatalogTreeItem &item) {
+    full_catalog_per_tree_item.add_new(&item, item.catalog_path());
+  });
+
+  return {std::move(catalogs_with_node_assets),
+          std::move(assets_per_path),
+          std::move(full_catalog_per_tree_item)};
+}
+
+static void node_add_catalog_assets_draw(const bContext *C, Menu *menu)
+{
+  bScreen &screen = *CTX_wm_screen(C);
+  const SpaceNode &snode = *CTX_wm_space_node(C);
+  if (!snode.runtime->assets_for_menu) {
+    BLI_assert_unreachable();
+    return;
+  }
+  AssetItemTree &tree = *snode.runtime->assets_for_menu;
+  const bNodeTree *edit_tree = snode.edittree;
+  if (!edit_tree) {
+    return;
+  }
+
+  const PointerRNA menu_path_ptr = CTX_data_pointer_get(C, "asset_catalog_path");
+  if (RNA_pointer_is_null(&menu_path_ptr)) {
+    return;
+  }
+  const bke::AssetCatalogPath &menu_path = *static_cast<const bke::AssetCatalogPath *>(
+      menu_path_ptr.data);
+
+  const Span<LibraryAsset> asset_items = tree.assets_per_path.lookup(menu_path);
+  bke::AssetCatalogTreeItem *catalog_item = tree.catalogs.find_item(menu_path);
+  BLI_assert(catalog_item != nullptr);
+
+  if (asset_items.is_empty() && !catalog_item->has_children()) {
+    return;
+  }
+
+  uiLayout *layout = menu->layout;
+  uiItemS(layout);
+
+  for (const LibraryAsset &item : asset_items) {
+    uiLayout *col = uiLayoutColumn(layout, false);
+    PointerRNA file{
+        &screen.id, &RNA_FileSelectEntry, const_cast<FileDirEntry *>(item.handle.file_data)};
+    uiLayoutSetContextPointer(col, "active_file", &file);
+
+    PointerRNA library_ptr{&screen.id,
+                           &RNA_AssetLibraryReference,
+                           const_cast<AssetLibraryReference *>(&item.library_ref)};
+    uiLayoutSetContextPointer(col, "asset_library_ref", &library_ptr);
+
+    uiItemO(col, ED_asset_handle_get_name(&item.handle), ICON_NONE, "NODE_OT_add_group_asset");
+  }
+
+  catalog_item->foreach_child([&](bke::AssetCatalogTreeItem &child_item) {
+    const bke::AssetCatalogPath &path = tree.full_catalog_per_tree_item.lookup(&child_item);
+    PointerRNA path_ptr{
+        &screen.id, &RNA_AssetCatalogPath, const_cast<bke::AssetCatalogPath *>(&path)};
+    uiLayout *col = uiLayoutColumn(layout, false);
+    uiLayoutSetContextPointer(col, "asset_catalog_path", &path_ptr);
+    uiItemM(col, "NODE_MT_node_add_catalog_assets", path.name().c_str(), ICON_NONE);
+  });
+}
+
+static void add_root_catalogs_draw(const bContext *C, Menu *menu)
+{
+  bScreen &screen = *CTX_wm_screen(C);
+  SpaceNode &snode = *CTX_wm_space_node(C);
+  const bNodeTree *edit_tree = snode.edittree;
+  uiLayout *layout = menu->layout;
+
+  snode.runtime->assets_for_menu = std::make_shared<AssetItemTree>(
+      build_catalog_tree(*C, edit_tree));
+
+  const bool loading_finished = all_loading_finished();
+
+  AssetItemTree &tree = *snode.runtime->assets_for_menu;
+  if (tree.catalogs.is_empty() && loading_finished) {
+    return;
+  }
+
+  uiItemS(layout);
+
+  if (!loading_finished) {
+    uiItemL(layout, IFACE_("Loading Asset Libraries"), ICON_INFO);
+  }
+
+  /* Avoid adding a separate root catalog when the assets have already been added to one of the
+   * builtin menus.
+   * TODO: The need to define the builtin menu labels here is completely non-ideal. We don't have
+   * any UI introspection that can do this though. This can be solved in the near future by
+   * removing the need to define the add menu completely, instead using a per-node-type path which
+   * can be merged with catalog tree.
+   */
+  static Set<std::string> all_builtin_menus = []() {
+    Set<std::string> menus;
+    menus.add_new("Attribute");
+    menus.add_new("Color");
+    menus.add_new("Curve");
+    menus.add_new("Curve Primitives");
+    menus.add_new("Curve Topology");
+    menus.add_new("Geometry");
+    menus.add_new("Input");
+    menus.add_new("Instances");
+    menus.add_new("Material");
+    menus.add_new("Mesh");
+    menus.add_new("Mesh Primitives");
+    menus.add_new("Mesh Topology");
+    menus.add_new("Output");
+    menus.add_new("Point");
+    menus.add_new("Text");
+    menus.add_new("Texture");
+    menus.add_new("Utilities");
+    menus.add_new("UV");
+    menus.add_new("Vector");
+    menus.add_new("Volume");
+    menus.add_new("Group");
+    menus.add_new("Layout");
+    return menus;
+  }();
+
+  tree.catalogs.foreach_root_item([&](bke::AssetCatalogTreeItem &item) {
+    if (all_builtin_menus.contains(item.get_name())) {
+      return;
+    }
+    const bke::AssetCatalogPath &path = tree.full_catalog_per_tree_item.lookup(&item);
+    PointerRNA path_ptr{
+        &screen.id, &RNA_AssetCatalogPath, const_cast<bke::AssetCatalogPath *>(&path)};
+    uiLayout *col = uiLayoutColumn(layout, false);
+    uiLayoutSetContextPointer(col, "asset_catalog_path", &path_ptr);
+    uiItemM(col, "NODE_MT_node_add_catalog_assets", path.name().c_str(), ICON_NONE);
+  });
+}
+
+MenuType add_catalog_assets_menu_type()
+{
+  MenuType type{};
+  BLI_strncpy(type.idname, "NODE_MT_node_add_catalog_assets", sizeof(type.idname));
+  type.poll = node_add_menu_poll;
+  type.draw = node_add_catalog_assets_draw;
+  return type;
+}
+
+MenuType add_root_catalogs_menu_type()
+{
+  MenuType type{};
+  BLI_strncpy(type.idname, "NODE_MT_node_add_root_catalogs", sizeof(type.idname));
+  type.poll = node_add_menu_poll;
+  type.draw = add_root_catalogs_draw;
+  return type;
+}
+
+}  // namespace blender::ed::space_node
+
+/* Note: This is only necessary because Python can't set an asset catalog path context item. */
+void uiTemplateNodeAssetMenuItems(uiLayout *layout, bContext *C, const char *catalog_path)
+{
+  using namespace blender;
+  using namespace blender::ed::space_node;
+  bScreen &screen = *CTX_wm_screen(C);
+  SpaceNode &snode = *CTX_wm_space_node(C);
+  AssetItemTree &tree = *snode.runtime->assets_for_menu;
+  const bke::AssetCatalogTreeItem *item = tree.catalogs.find_root_item(catalog_path);
+  if (!item) {
+    return;
+  }
+  const bke::AssetCatalogPath &path = tree.full_catalog_per_tree_item.lookup(item);
+  PointerRNA path_ptr{
+      &screen.id, &RNA_AssetCatalogPath, const_cast<bke::AssetCatalogPath *>(&path)};
+  uiItemS(layout);
+  uiLayout *col = uiLayoutColumn(layout, false);
+  uiLayoutSetContextPointer(col, "asset_catalog_path", &path_ptr);
+  uiItemMContents(col, "NODE_MT_node_add_catalog_assets");
+}