metagraphoptimization.jl/src/operations/apply.jl

# functions that apply graph operations

# applies all unapplied operations in the DAG
function apply_all!(graph::DAG)
   while !isempty(graph.operationsToApply)
      # get next operation to apply from front of the deque
      op = popfirst!(graph.operationsToApply)

      # apply it
      appliedOp = apply_operation!(graph, op)

      # push to the end of the appliedOperations deque
      push!(graph.appliedOperations, appliedOp)
   end
   return nothing
end

function apply_operation!(graph::DAG, operation::Operation)
   error("Unknown operation type!")
end

function apply_operation!(graph::DAG, operation::NodeFusion)
   diff = node_fusion!(graph, operation.input[1], operation.input[2], operation.input[3])
   return AppliedNodeFusion(operation, diff)
end

function apply_operation!(graph::DAG, operation::NodeReduction)
   diff = node_reduction!(graph, operation.input)
   return AppliedNodeReduction(operation, diff)
end

function apply_operation!(graph::DAG, operation::NodeSplit)
   diff = node_split!(graph, operation.input)
   return AppliedNodeSplit(operation, diff)
end


function revert_operation!(graph::DAG, operation::AppliedOperation)
   error("Unknown operation type!")
end

function revert_operation!(graph::DAG, operation::AppliedNodeFusion)
   revert_diff!(graph, operation.diff)
   return operation.operation
end

function revert_operation!(graph::DAG, operation::AppliedNodeReduction)
   revert_diff!(graph, operation.diff)
   return operation.operation
end

function revert_operation!(graph::DAG, operation::AppliedNodeSplit)
   revert_diff!(graph, operation.diff)
   return operation.operation
end


function revert_diff!(graph::DAG, diff)
   # add removed nodes, remove added nodes, same for edges
   # note the order
   for edge in diff.addedEdges
      remove_edge!(graph, edge, false)
   end
   for node in diff.addedNodes
      remove_node!(graph, node, false)
   end

   for node in diff.removedNodes
      insert_node!(graph, node, false)
   end
   for edge in diff.removedEdges
      insert_edge!(graph, edge, false)
   end
end

# Fuse nodes n1 -> n2 -> n3 together into one node, return the applied difference to the graph
function node_fusion!(graph::DAG, n1::ComputeTaskNode, n2::DataTaskNode, n3::ComputeTaskNode)
   # clear snapshot
   get_snapshot_diff(graph)

   if !(n1 in graph) || !(n2 in graph) || !(n3 in graph)
      error("[Node Fusion] The given nodes are not part of the given graph")
   end

   if !is_child(n1, n2) || !is_child(n2, n3) || !is_parent(n3, n2) || !is_parent(n2, n1)
      # the checks are redundant but maybe a good sanity check
      error("[Node Fusion] The given nodes are not connected by edges which is required for node fusion")
   end

   # save children and parents
   n1_children = children(n1)
   n3_parents = parents(n3)
   n3_children = children(n3)

   if length(n2.parents) > 1
      error("[Node Fusion] The given data node has more than one parent")
   end
   if length(n2.children) > 1
      error("[Node Fusion] The given data node has more than one child")
   end
   if length(n1.parents) > 1
      error("[Node Fusion] The given n1 has more than one parent")
   end

   required_edge1 = make_edge(n1, n2)
   required_edge2 = make_edge(n2, n3)

   # remove the edges and nodes that will be replaced by the fused node
   remove_edge!(graph, required_edge1)
   remove_edge!(graph, required_edge2)
   remove_node!(graph, n1)
   remove_node!(graph, n2)

   # get n3's children now so it automatically excludes n2
   n3_children = children(n3)
   remove_node!(graph, n3)

   # create new node with the fused compute task
   new_node = ComputeTaskNode(FusedComputeTask{typeof(n1.task),typeof(n3.task)}())
   insert_node!(graph, new_node)

   # use a set for combined children of n1 and n3 to not get duplicates
   n1and3_children = Set{Node}()

   # remove edges from n1 children to n1
   for child in n1_children
      remove_edge!(graph, make_edge(child, n1))
      push!(n1and3_children, child)
   end

   # remove edges from n3 children to n3
   for child in n3_children
      remove_edge!(graph, make_edge(child, n3))
      push!(n1and3_children, child)
   end

   for child in n1and3_children
      insert_edge!(graph, make_edge(child, new_node))
   end

   # "repoint" parents of n3 from new node
   for parent in n3_parents
      remove_edge!(graph, make_edge(n3, parent))
      insert_edge!(graph, make_edge(new_node, parent))
   end

   return get_snapshot_diff(graph)
end

function node_reduction!(graph::DAG, nodes::Vector{Node})
   # clear snapshot
   get_snapshot_diff(graph)

   t = typeof(nodes[1].task)
   for n in nodes
      if n ∉ graph
         error("[Node Reduction] The given nodes are not part of the given graph")
      end

      if typeof(n.task) != t
         error("[Node Reduction] The given nodes are not of the same type")
      end
   end


   n1 = nodes[1]
   n1_children = children(n1)
   for n in nodes
      if Set(n1_children) != Set(n.children)
         error("[Node Reduction] The given nodes do not have equal prerequisite nodes which is required for node reduction")
      end
   end

   n1_parents = Set(n1.parents)
   new_parents = Set{Node}()

   # remove all of the nodes' parents and children and the nodes themselves (except for first node)
   for i in 2:length(nodes)
      n = nodes[i]
      for child in n1_children
         remove_edge!(graph, make_edge(child, n))
      end

      for parent in parents(n)
         remove_edge!(graph, make_edge(n, parent))

         # collect all parents
         push!(new_parents, parent)
      end

      remove_node!(graph, n)
   end

   setdiff!(new_parents, n1_parents)

   for parent in new_parents
      # now add parents of n2 to n1 without duplicates
      insert_edge!(graph, make_edge(n1, parent))
   end

   return get_snapshot_diff(graph)
end

function node_split!(graph::DAG, n1::Node)
   # clear snapshot
   get_snapshot_diff(graph)

   #=if !(n1 in graph)
      error("[Node Split] The given node is not part of the given graph")
   end=#

   n1_parents = parents(n1)
   n1_children = children(n1)

   #=if length(n1_parents) <= 1
      error("[Node Split] The given node does not have multiple parents which is required for node split")
   end=#

   for parent in n1_parents
      remove_edge!(graph, make_edge(n1, parent))
   end
   for child in n1_children
      remove_edge!(graph, make_edge(child, n1))
   end
   remove_node!(graph, n1)

   for parent in n1_parents
      n_copy = copy(n1)
      insert_node!(graph, n_copy)
      insert_edge!(graph, make_edge(n_copy, parent))

      for child in n1_children
         insert_edge!(graph, make_edge(child, n_copy))
      end
   end

   return get_snapshot_diff(graph)
end