test2/source/blender/editors/space_node/link_drag_search.cc

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

#include "AS_asset_library.hh"
#include "AS_asset_representation.hh"

#include "BLI_listbase.h"

#include "DNA_space_types.h"

#include "BKE_asset.hh"
#include "BKE_context.hh"
#include "BKE_idprop.hh"
#include "BKE_lib_id.hh"
#include "BKE_node_runtime.hh"
#include "BKE_node_tree_update.hh"
#include "BKE_screen.hh"

#include "UI_string_search.hh"

#include "NOD_socket.hh"
#include "NOD_socket_search_link.hh"

#include "BLT_translation.hh"

#include "WM_api.hh"

#include "DEG_depsgraph_build.hh"

#include "ED_asset.hh"
#include "ED_node.hh"

#include "node_intern.hh"

using blender::nodes::SocketLinkOperation;

namespace blender::ed::space_node {

struct LinkDragSearchStorage {
  bNode &from_node;
  bNodeSocket &from_socket;
  float2 cursor;
  Vector<SocketLinkOperation> search_link_ops;
  char search[256];
  bool update_items_tag = true;

  eNodeSocketInOut in_out() const
  {
    return static_cast<eNodeSocketInOut>(from_socket.in_out);
  }
};

static void link_drag_search_listen_fn(const wmRegionListenerParams *params, void *arg)
{
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg);
  const wmNotifier *wmn = params->notifier;

  switch (wmn->category) {
    case NC_ASSET:
      if (wmn->data == ND_ASSET_LIST_READING) {
        storage.update_items_tag = true;
      }
      break;
  }
}

static void add_reroute_node_fn(nodes::LinkSearchOpParams &params)
{
  bNode &reroute = params.add_node("NodeReroute");
  if (params.socket.in_out == SOCK_IN) {
    nodeAddLink(&params.node_tree,
                &reroute,
                static_cast<bNodeSocket *>(reroute.outputs.first),
                &params.node,
                &params.socket);
  }
  else {
    nodeAddLink(&params.node_tree,
                &params.node,
                &params.socket,
                &reroute,
                static_cast<bNodeSocket *>(reroute.inputs.first));
  }
}

static void add_group_input_node_fn(nodes::LinkSearchOpParams &params)
{
  /* Add a group input based on the connected socket, and add a new group input node. */
  bNodeTreeInterfaceSocket *socket_iface = params.node_tree.tree_interface.add_socket(
      params.socket.name,
      params.socket.description,
      params.socket.typeinfo->idname,
      NODE_INTERFACE_SOCKET_INPUT,
      nullptr);
  socket_iface->init_from_socket_instance(&params.socket);
  params.node_tree.tree_interface.active_item_set(&socket_iface->item);

  bNode &group_input = params.add_node("NodeGroupInput");

  /* This is necessary to create the new sockets in the other input nodes. */
  ED_node_tree_propagate_change(&params.C, CTX_data_main(&params.C), &params.node_tree);

  /* Hide the new input in all other group input nodes, to avoid making them taller. */
  for (bNode *node : params.node_tree.all_nodes()) {
    if (node->type == NODE_GROUP_INPUT) {
      bNodeSocket *new_group_input_socket = nodeFindSocket(
          node, SOCK_OUT, socket_iface->identifier);
      if (new_group_input_socket) {
        new_group_input_socket->flag |= SOCK_HIDDEN;
      }
    }
  }

  /* Hide all existing inputs in the new group input node, to only display the new one. */
  LISTBASE_FOREACH (bNodeSocket *, socket, &group_input.outputs) {
    socket->flag |= SOCK_HIDDEN;
  }

  bNodeSocket *socket = nodeFindSocket(&group_input, SOCK_OUT, socket_iface->identifier);
  if (socket) {
    /* Unhide the socket for the new input in the new node and make a connection to it. */
    socket->flag &= ~SOCK_HIDDEN;
    nodeAddLink(&params.node_tree, &group_input, socket, &params.node, &params.socket);

    bke::node_socket_move_default_value(
        *CTX_data_main(&params.C), params.node_tree, params.socket, *socket);
  }
}

static void add_existing_group_input_fn(nodes::LinkSearchOpParams &params,
                                        const bNodeTreeInterfaceSocket &interface_socket)
{
  const eNodeSocketInOut in_out = eNodeSocketInOut(params.socket.in_out);
  NodeTreeInterfaceSocketFlag flag = NodeTreeInterfaceSocketFlag(0);
  SET_FLAG_FROM_TEST(flag, in_out & SOCK_IN, NODE_INTERFACE_SOCKET_INPUT);
  SET_FLAG_FROM_TEST(flag, in_out & SOCK_OUT, NODE_INTERFACE_SOCKET_OUTPUT);

  bNode &group_input = params.add_node("NodeGroupInput");

  LISTBASE_FOREACH (bNodeSocket *, socket, &group_input.outputs) {
    socket->flag |= SOCK_HIDDEN;
  }

  bNodeSocket *socket = nodeFindSocket(&group_input, SOCK_OUT, interface_socket.identifier);
  if (socket != nullptr) {
    socket->flag &= ~SOCK_HIDDEN;
    nodeAddLink(&params.node_tree, &group_input, socket, &params.node, &params.socket);
  }
}

/**
 * \note This could use #search_link_ops_for_socket_templates, but we have to store the inputs and
 * outputs as IDProperties for assets because of asset indexing, so that's all we have without
 * loading the file.
 */
static void search_link_ops_for_asset_metadata(const bNodeTree &node_tree,
                                               const bNodeSocket &socket,
                                               const asset_system::AssetRepresentation &asset,
                                               Vector<SocketLinkOperation> &search_link_ops)
{
  const AssetMetaData &asset_data = asset.get_metadata();
  const IDProperty *tree_type = BKE_asset_metadata_idprop_find(&asset_data, "type");
  if (tree_type == nullptr || IDP_Int(tree_type) != node_tree.type) {
    return;
  }

  const bNodeTreeType &node_tree_type = *node_tree.typeinfo;
  const eNodeSocketInOut in_out = socket.in_out == SOCK_OUT ? SOCK_IN : SOCK_OUT;

  const IDProperty *sockets = BKE_asset_metadata_idprop_find(
      &asset_data, in_out == SOCK_IN ? "inputs" : "outputs");

  int weight = -1;
  Set<StringRef> socket_names;
  LISTBASE_FOREACH (IDProperty *, socket_property, &sockets->data.group) {
    if (socket_property->type != IDP_STRING) {
      continue;
    }
    const char *socket_idname = IDP_String(socket_property);
    const bNodeSocketType *socket_type = nodeSocketTypeFind(socket_idname);
    if (socket_type == nullptr) {
      continue;
    }
    eNodeSocketDatatype from = (eNodeSocketDatatype)socket.type;
    eNodeSocketDatatype to = (eNodeSocketDatatype)socket_type->type;
    if (socket.in_out == SOCK_OUT) {
      std::swap(from, to);
    }
    if (node_tree_type.validate_link && !node_tree_type.validate_link(from, to)) {
      continue;
    }
    if (!socket_names.add(socket_property->name)) {
      /* See comment in #search_link_ops_for_declarations. */
      continue;
    }

    const StringRef asset_name = asset.get_name();
    const StringRef socket_name = socket_property->name;

    search_link_ops.append(
        {asset_name + " " + UI_MENU_ARROW_SEP + socket_name,
         [&asset, socket_property, in_out](nodes::LinkSearchOpParams &params) {
           Main &bmain = *CTX_data_main(&params.C);

           bNode &node = params.add_node(params.node_tree.typeinfo->group_idname);
           node.flag &= ~NODE_OPTIONS;

           node.id = asset::asset_local_id_ensure_imported(bmain, asset);
           id_us_plus(node.id);
           BKE_ntree_update_tag_node_property(&params.node_tree, &node);
           DEG_relations_tag_update(&bmain);

           /* Create the inputs and outputs on the new node. */
           nodes::update_node_declaration_and_sockets(params.node_tree, node);

           bNodeSocket *new_node_socket = bke::node_find_enabled_socket(
               node, in_out, socket_property->name);
           if (new_node_socket != nullptr) {
             /* Rely on the way #nodeAddLink switches in/out if necessary. */
             nodeAddLink(&params.node_tree, &params.node, &params.socket, &node, new_node_socket);
           }
         },
         weight});

    weight--;
  }
}

static void gather_search_link_ops_for_all_assets(const bContext &C,
                                                  const bNodeTree &node_tree,
                                                  const bNodeSocket &socket,
                                                  Vector<SocketLinkOperation> &search_link_ops)
{
  const AssetLibraryReference library_ref = asset_system::all_library_reference();
  asset::AssetFilterSettings filter_settings{};
  filter_settings.id_types = FILTER_ID_NT;

  asset::list::storage_fetch(&library_ref, &C);
  asset::list::iterate(library_ref, [&](asset_system::AssetRepresentation &asset) {
    if (!asset::filter_matches_asset(&filter_settings, asset)) {
      return true;
    }
    search_link_ops_for_asset_metadata(node_tree, socket, asset, search_link_ops);
    return true;
  });
}

/**
 * Call the callback to gather compatible socket connections for all node types, and the operations
 * that will actually make the connections. Also add some custom operations like connecting a group
 * output node.
 */
static void gather_socket_link_operations(const bContext &C,
                                          bNodeTree &node_tree,
                                          const bNodeSocket &socket,
                                          Vector<SocketLinkOperation> &search_link_ops)
{
  const SpaceNode &snode = *CTX_wm_space_node(&C);
  NODE_TYPES_BEGIN (node_type) {
    const char *disabled_hint;
    if (node_type->poll && !node_type->poll(node_type, &node_tree, &disabled_hint)) {
      continue;
    }
    if (node_type->add_ui_poll && !node_type->add_ui_poll(&C)) {
      continue;
    }
    if (StringRefNull(node_type->ui_name).endswith("(Legacy)")) {
      continue;
    }
    if (node_type->gather_link_search_ops) {
      nodes::GatherLinkSearchOpParams params{
          *node_type, snode, node_tree, socket, search_link_ops};
      node_type->gather_link_search_ops(params);
    }
  }
  NODE_TYPES_END;

  search_link_ops.append({IFACE_("Reroute"), add_reroute_node_fn});

  const bool is_node_group = !(node_tree.id.flag & LIB_EMBEDDED_DATA);

  if (is_node_group && socket.in_out == SOCK_IN) {
    search_link_ops.append({IFACE_("Group Input"), add_group_input_node_fn});

    int weight = -1;
    node_tree.tree_interface.foreach_item([&](const bNodeTreeInterfaceItem &item) {
      if (item.item_type != NODE_INTERFACE_SOCKET) {
        return true;
      }
      const bNodeTreeInterfaceSocket &interface_socket =
          reinterpret_cast<const bNodeTreeInterfaceSocket &>(item);
      {
        const bNodeSocketType *from_typeinfo = nodeSocketTypeFind(interface_socket.socket_type);
        const eNodeSocketDatatype from = from_typeinfo ? eNodeSocketDatatype(from_typeinfo->type) :
                                                         SOCK_CUSTOM;
        const eNodeSocketDatatype to = eNodeSocketDatatype(socket.typeinfo->type);
        if (node_tree.typeinfo->validate_link && !node_tree.typeinfo->validate_link(from, to)) {
          return true;
        }
      }
      search_link_ops.append({std::string(IFACE_("Group Input")) + " " + UI_MENU_ARROW_SEP +
                                  (interface_socket.name ? interface_socket.name : ""),
                              [interface_socket](nodes::LinkSearchOpParams &params) {
                                add_existing_group_input_fn(params, interface_socket);
                              },
                              weight});
      weight--;
      return true;
    });
  }

  gather_search_link_ops_for_all_assets(C, node_tree, socket, search_link_ops);
}

static void link_drag_search_update_fn(
    const bContext *C, void *arg, const char *str, uiSearchItems *items, const bool is_first)
{
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg);
  if (storage.update_items_tag) {
    bNodeTree *node_tree = CTX_wm_space_node(C)->edittree;
    storage.search_link_ops.clear();
    gather_socket_link_operations(*C, *node_tree, storage.from_socket, storage.search_link_ops);
    storage.update_items_tag = false;
  }

  ui::string_search::StringSearch<SocketLinkOperation> search{
      string_search::MainWordsHeuristic::All};

  for (SocketLinkOperation &op : storage.search_link_ops) {
    search.add(op.name, &op, op.weight);
  }

  /* Don't filter when the menu is first opened, but still run the search
   * so the items are in the same order they will appear in while searching. */
  const char *string = is_first ? "" : str;
  const Vector<SocketLinkOperation *> filtered_items = search.query(string);

  for (SocketLinkOperation *item : filtered_items) {
    if (!UI_search_item_add(items, item->name.c_str(), item, ICON_NONE, 0, 0)) {
      break;
    }
  }
}

static void link_drag_search_exec_fn(bContext *C, void *arg1, void *arg2)
{
  Main &bmain = *CTX_data_main(C);
  SpaceNode &snode = *CTX_wm_space_node(C);
  bNodeTree &node_tree = *snode.edittree;
  LinkDragSearchStorage &storage = *static_cast<LinkDragSearchStorage *>(arg1);
  SocketLinkOperation *item = static_cast<SocketLinkOperation *>(arg2);
  if (item == nullptr) {
    return;
  }

  node_deselect_all(node_tree);

  Vector<bNode *> new_nodes;
  nodes::LinkSearchOpParams params{
      *C, node_tree, storage.from_node, storage.from_socket, new_nodes};
  item->fn(params);
  if (new_nodes.is_empty()) {
    return;
  }

  /* For now, assume that only one node is created by the callback. */
  BLI_assert(new_nodes.size() == 1);
  bNode *new_node = new_nodes.first();

  new_node->locx = storage.cursor.x / UI_SCALE_FAC;
  new_node->locy = storage.cursor.y / UI_SCALE_FAC + 20;
  if (storage.in_out() == SOCK_IN) {
    new_node->locx -= new_node->width;
  }

  nodeSetSelected(new_node, true);
  nodeSetActive(&node_tree, new_node);

  /* Ideally it would be possible to tag the node tree in some way so it updates only after the
   * translate operation is finished, but normally moving nodes around doesn't cause updates. */
  ED_node_tree_propagate_change(C, &bmain, &node_tree);

  /* Start translation operator with the new node. */
  wmOperatorType *ot = WM_operatortype_find("NODE_OT_translate_attach_remove_on_cancel", true);
  BLI_assert(ot);
  PointerRNA ptr;
  WM_operator_properties_create_ptr(&ptr, ot);
  WM_operator_name_call_ptr(C, ot, WM_OP_INVOKE_DEFAULT, &ptr, nullptr);
  WM_operator_properties_free(&ptr);
}

static void link_drag_search_free_fn(void *arg)
{
  LinkDragSearchStorage *storage = static_cast<LinkDragSearchStorage *>(arg);
  delete storage;
}

static uiBlock *create_search_popup_block(bContext *C, ARegion *region, void *arg_op)
{
  LinkDragSearchStorage &storage = *(LinkDragSearchStorage *)arg_op;

  uiBlock *block = UI_block_begin(C, region, "_popup", UI_EMBOSS);
  UI_block_flag_enable(block, UI_BLOCK_LOOP | UI_BLOCK_MOVEMOUSE_QUIT | UI_BLOCK_SEARCH_MENU);
  UI_block_theme_style_set(block, UI_BLOCK_THEME_STYLE_POPUP);

  uiBut *but = uiDefSearchBut(block,
                              storage.search,
                              0,
                              ICON_VIEWZOOM,
                              sizeof(storage.search),
                              storage.in_out() == SOCK_OUT ? 10 : 10 - UI_searchbox_size_x(),
                              10,
                              UI_searchbox_size_x(),
                              UI_UNIT_Y,
                              "");
  UI_but_func_search_set_sep_string(but, UI_MENU_ARROW_SEP);
  UI_but_func_search_set_listen(but, link_drag_search_listen_fn);
  UI_but_func_search_set(but,
                         nullptr,
                         link_drag_search_update_fn,
                         &storage,
                         false,
                         link_drag_search_free_fn,
                         link_drag_search_exec_fn,
                         nullptr);
  UI_but_flag_enable(but, UI_BUT_ACTIVATE_ON_INIT);

  /* Fake button to hold space for the search items. */
  uiDefBut(block,
           UI_BTYPE_LABEL,
           0,
           "",
           storage.in_out() == SOCK_OUT ? 10 : 10 - UI_searchbox_size_x(),
           10 - UI_searchbox_size_y(),
           UI_searchbox_size_x(),
           UI_searchbox_size_y(),
           nullptr,
           0,
           0,
           nullptr);

  const int2 offset = {0, -UI_UNIT_Y};
  UI_block_bounds_set_popup(block, 0.3f * U.widget_unit, offset);
  return block;
}

void invoke_node_link_drag_add_menu(bContext &C,
                                    bNode &node,
                                    bNodeSocket &socket,
                                    const float2 &cursor)
{
  LinkDragSearchStorage *storage = new LinkDragSearchStorage{node, socket, cursor};
  /* Use the "_ex" variant with `can_refresh` false to avoid a double free when closing Blender. */
  UI_popup_block_invoke_ex(&C, create_search_popup_block, storage, nullptr, false);
}

}  // namespace blender::ed::space_node