Functions: extend multi-function network api

3 files changed, 57 insertions, 10 deletions

diff --git a/source/blender/functions/FN_multi_function_network.hh b/source/blender/functions/FN_multi_function_network.hh
index a9d8508cdb8..413f68a8531 100644
--- a/source/blender/functions/FN_multi_function_network.hh
+++ b/source/blender/functions/FN_multi_function_network.hh
@@ -216,10 +216,18 @@ class MFNetwork : NonCopyable, NonMovable {
   void relink(MFOutputSocket &old_output, MFOutputSocket &new_output);
 
   void remove(MFNode &node);
+  void remove(Span<MFNode *> nodes);
 
-  uint max_socket_id() const;
+  uint socket_id_amount() const;
+  uint node_id_amount() const;
 
-  std::string to_dot() const;
+  Span<MFDummyNode *> dummy_nodes();
+  Span<MFFunctionNode *> function_nodes();
+
+  MFNode *node_or_null_by_id(uint id);
+  const MFNode *node_or_null_by_id(uint id) const;
+
+  std::string to_dot(Span<const MFNode *> marked_nodes = {}) const;
 };
 
 /* --------------------------------------------------------------------
@@ -325,7 +333,7 @@ inline Span<const MFOutputSocket *> MFNode::outputs() const
   return outputs_;
 }
 
-template<typename FuncT> void MFNode::foreach_origin_socket(const FuncT &func) const
+template<typename FuncT> inline void MFNode::foreach_origin_socket(const FuncT &func) const
 {
   for (const MFInputSocket *socket : inputs_) {
     const MFOutputSocket *origin = socket->origin();
@@ -483,9 +491,34 @@ inline Span<const MFInputSocket *> MFOutputSocket::targets() const
  * MFNetwork inline methods.
  */
 
-inline uint MFNetwork::max_socket_id() const
+inline Span<MFDummyNode *> MFNetwork::dummy_nodes()
+{
+  return dummy_nodes_;
+}
+
+inline Span<MFFunctionNode *> MFNetwork::function_nodes()
+{
+  return function_nodes_;
+}
+
+inline MFNode *MFNetwork::node_or_null_by_id(uint id)
+{
+  return node_or_null_by_id_[id];
+}
+
+inline const MFNode *MFNetwork::node_or_null_by_id(uint id) const
+{
+  return node_or_null_by_id_[id];
+}
+
+inline uint MFNetwork::socket_id_amount() const
+{
+  return socket_or_null_by_id_.size();
+}
+
+inline uint MFNetwork::node_id_amount() const
 {
-  return socket_or_null_by_id_.size() - 1;
+  return node_or_null_by_id_.size();
 }
 
 }  // namespace blender::fn
diff --git a/source/blender/functions/intern/multi_function_network.cc b/source/blender/functions/intern/multi_function_network.cc
index 5df70d92a4e..fd141a48bea 100644
--- a/source/blender/functions/intern/multi_function_network.cc
+++ b/source/blender/functions/intern/multi_function_network.cc
@@ -230,7 +230,14 @@ void MFNetwork::remove(MFNode &node)
   node_or_null_by_id_[node.id_] = nullptr;
 }
 
-std::string MFNetwork::to_dot() const
+void MFNetwork::remove(Span<MFNode *> nodes)
+{
+  for (MFNode *node : nodes) {
+    this->remove(*node);
+  }
+}
+
+std::string MFNetwork::to_dot(Span<const MFNode *> marked_nodes) const
 {
   dot::DirectedGraph digraph;
   digraph.set_rankdir(dot::Attr_rankdir::LeftToRight);
@@ -256,6 +263,13 @@ std::string MFNetwork::to_dot() const
     dot_nodes.add_new(node, dot_node_ref);
   }
 
+  for (const MFDummyNode *node : dummy_nodes_) {
+    dot_nodes.lookup(node).node().set_background_color("#77EE77");
+  }
+  for (const MFNode *node : marked_nodes) {
+    dot_nodes.lookup(node).node().set_background_color("#7777EE");
+  }
+
   for (const MFNode *to_node : all_nodes) {
     dot::NodeWithSocketsRef to_dot_node = dot_nodes.lookup(to_node);
 
diff --git a/source/blender/functions/intern/multi_function_network_evaluation.cc b/source/blender/functions/intern/multi_function_network_evaluation.cc
index 08a254dc300..36d363dc56c 100644
--- a/source/blender/functions/intern/multi_function_network_evaluation.cc
+++ b/source/blender/functions/intern/multi_function_network_evaluation.cc
@@ -58,7 +58,7 @@ class MFNetworkEvaluationStorage {
   uint min_array_size_;
 
  public:
-  MFNetworkEvaluationStorage(IndexMask mask, uint max_socket_id);
+  MFNetworkEvaluationStorage(IndexMask mask, uint socket_id_amount);
   ~MFNetworkEvaluationStorage();
 
   /* Add the values that have been provided by the caller of the multi-function network. */
@@ -142,7 +142,7 @@ void MFNetworkEvaluator::call(IndexMask mask, MFParams params, MFContext context
   }
 
   const MFNetwork &network = outputs_[0]->node().network();
-  Storage storage(mask, network.max_socket_id());
+  Storage storage(mask, network.socket_id_amount());
 
   Vector<const MFInputSocket *> outputs_to_initialize_in_the_end;
 
@@ -507,9 +507,9 @@ struct OwnVectorValue : public Value {
 /** \name Storage methods
  * \{ */
 
-MFNetworkEvaluationStorage::MFNetworkEvaluationStorage(IndexMask mask, uint max_socket_id)
+MFNetworkEvaluationStorage::MFNetworkEvaluationStorage(IndexMask mask, uint socket_id_amount)
     : mask_(mask),
-      value_per_output_id_(max_socket_id + 1, nullptr),
+      value_per_output_id_(socket_id_amount, nullptr),
       min_array_size_(mask.min_array_size())
 {
 }