2 files changed, 85 insertions, 0 deletions

diff --git a/source/blender/functions/FN_multi_function_network.hh b/source/blender/functions/FN_multi_function_network.hh
index e1d5ce8bd69..bb0c870746b 100644
--- a/source/blender/functions/FN_multi_function_network.hh
+++ b/source/blender/functions/FN_multi_function_network.hh
@@ -156,6 +156,15 @@ class MFSocket : NonCopyable, NonMovable {
 
   MFNode &node();
   const MFNode &node() const;
+
+  bool is_input() const;
+  bool is_output() const;
+
+  MFInputSocket &as_input();
+  const MFInputSocket &as_input() const;
+
+  MFOutputSocket &as_output();
+  const MFOutputSocket &as_output() const;
 };
 
 class MFInputSocket : public MFSocket {
@@ -205,6 +214,10 @@ class MFNetwork : NonCopyable, NonMovable {
   MFOutputSocket &add_input(StringRef name, MFDataType data_type);
   MFInputSocket &add_output(StringRef name, MFDataType data_type);
 
+  void relink(MFOutputSocket &old_output, MFOutputSocket &new_output);
+
+  void remove(MFNode &node);
+
   uint max_socket_id() const;
 
   std::string to_dot() const;
@@ -405,6 +418,40 @@ inline const MFNode &MFSocket::node() const
   return *m_node;
 }
 
+inline bool MFSocket::is_input() const
+{
+  return !m_is_output;
+}
+
+inline bool MFSocket::is_output() const
+{
+  return m_is_output;
+}
+
+inline MFInputSocket &MFSocket::as_input()
+{
+  BLI_assert(this->is_input());
+  return *(MFInputSocket *)this;
+}
+
+inline const MFInputSocket &MFSocket::as_input() const
+{
+  BLI_assert(this->is_input());
+  return *(const MFInputSocket *)this;
+}
+
+inline MFOutputSocket &MFSocket::as_output()
+{
+  BLI_assert(this->is_output());
+  return *(MFOutputSocket *)this;
+}
+
+inline const MFOutputSocket &MFSocket::as_output() const
+{
+  BLI_assert(this->is_output());
+  return *(const MFOutputSocket *)this;
+}
+
 /* --------------------------------------------------------------------
  * MFInputSocket inline methods.
  */
diff --git a/source/blender/functions/intern/multi_function_network.cc b/source/blender/functions/intern/multi_function_network.cc
index 269a32b73bd..93d062f3e5c 100644
--- a/source/blender/functions/intern/multi_function_network.cc
+++ b/source/blender/functions/intern/multi_function_network.cc
@@ -193,6 +193,44 @@ MFInputSocket &MFNetwork::add_output(StringRef name, MFDataType data_type)
   return this->add_dummy(name, {data_type}, {}, {name}, {}).input(0);
 }
 
+void MFNetwork::relink(MFOutputSocket &old_output, MFOutputSocket &new_output)
+{
+  BLI_assert(&old_output != &new_output);
+  for (MFInputSocket *input : old_output.targets()) {
+    input->m_origin = &new_output;
+  }
+  new_output.m_targets.extend(old_output.m_targets);
+  old_output.m_targets.clear();
+}
+
+void MFNetwork::remove(MFNode &node)
+{
+  for (MFInputSocket *socket : node.m_inputs) {
+    if (socket->m_origin != nullptr) {
+      socket->m_origin->m_targets.remove_first_occurrence_and_reorder(socket);
+    }
+    m_socket_or_null_by_id[socket->m_id] = nullptr;
+  }
+  for (MFOutputSocket *socket : node.m_outputs) {
+    for (MFInputSocket *other : socket->m_targets) {
+      other->m_origin = nullptr;
+    }
+    m_socket_or_null_by_id[socket->m_id] = nullptr;
+  }
+  node.destruct_sockets();
+  if (node.is_dummy()) {
+    MFDummyNode &dummy_node = node.as_dummy();
+    dummy_node.~MFDummyNode();
+    m_dummy_nodes.remove_contained(&dummy_node);
+  }
+  else {
+    MFFunctionNode &function_node = node.as_function();
+    function_node.~MFFunctionNode();
+    m_function_nodes.remove_contained(&function_node);
+  }
+  m_node_or_null_by_id[node.m_id] = nullptr;
+}
+
 std::string MFNetwork::to_dot() const
 {
   dot::DirectedGraph digraph;