metagraphoptimization.jl/src/operation/apply.jl

"""
    apply_all!(graph::DAG)

Apply all unapplied operations in the DAG. Is automatically called in all functions that require the latest state of the [`DAG`](@ref).
"""
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

"""
    apply_operation!(graph::DAG, operation::Operation)

Fallback implementation of apply_operation! for unimplemented operation types, throwing an error.
"""
function apply_operation!(graph::DAG, operation::Operation)
    return error("Unknown operation type!")
end

"""
    apply_operation!(graph::DAG, operation::NodeFusion)

Apply the given [`NodeFusion`](@ref) to the graph. Generic wrapper around [`node_fusion!`](@ref).

Return an [`AppliedNodeFusion`](@ref) object generated from the graph's [`Diff`](@ref).
"""
function apply_operation!(graph::DAG, operation::NodeFusion)
    diff = node_fusion!(
        graph,
        operation.input[1],
        operation.input[2],
        operation.input[3],
    )

    graph.properties += GraphProperties(diff)

    return AppliedNodeFusion(operation, diff)
end

"""
    apply_operation!(graph::DAG, operation::NodeReduction)

Apply the given [`NodeReduction`](@ref) to the graph. Generic wrapper around [`node_reduction!`](@ref).

Return an [`AppliedNodeReduction`](@ref) object generated from the graph's [`Diff`](@ref).
"""
function apply_operation!(graph::DAG, operation::NodeReduction)
    diff = node_reduction!(graph, operation.input)

    graph.properties += GraphProperties(diff)

    return AppliedNodeReduction(operation, diff)
end

"""
    apply_operation!(graph::DAG, operation::NodeSplit)

Apply the given [`NodeSplit`](@ref) to the graph. Generic wrapper around [`node_split!`](@ref).

Return an [`AppliedNodeSplit`](@ref) object generated from the graph's [`Diff`](@ref).
"""
function apply_operation!(graph::DAG, operation::NodeSplit)
    diff = node_split!(graph, operation.input)

    graph.properties += GraphProperties(diff)

    return AppliedNodeSplit(operation, diff)
end

"""
    revert_operation!(graph::DAG, operation::AppliedOperation)

Fallback implementation of operation reversion for unimplemented operation types, throwing an error.
"""
function revert_operation!(graph::DAG, operation::AppliedOperation)
    return error("Unknown operation type!")
end

"""
    revert_operation!(graph::DAG, operation::AppliedNodeFusion)

Revert the applied node fusion on the graph. Return the original [`NodeFusion`](@ref) operation.
"""
function revert_operation!(graph::DAG, operation::AppliedNodeFusion)
    revert_diff!(graph, operation.diff)
    return operation.operation
end

"""
    revert_operation!(graph::DAG, operation::AppliedNodeReduction)

Revert the applied node fusion on the graph. Return the original [`NodeReduction`](@ref) operation.
"""
function revert_operation!(graph::DAG, operation::AppliedNodeReduction)
    revert_diff!(graph, operation.diff)
    return operation.operation
end

"""
    revert_operation!(graph::DAG, operation::AppliedNodeSplit)

Revert the applied node fusion on the graph. Return the original [`NodeSplit`](@ref) operation.
"""
function revert_operation!(graph::DAG, operation::AppliedNodeSplit)
    revert_diff!(graph, operation.diff)
    return operation.operation
end

"""
    revert_diff!(graph::DAG, diff::Diff)

Revert the given diff on the graph. Used to revert the individual [`AppliedOperation`](@ref)s with [`revert_operation!`](@ref).
"""
function revert_diff!(graph::DAG, diff::Diff)
    # add removed nodes, remove added nodes, same for edges
    # note the order
    for edge in diff.addedEdges
        remove_edge!(graph, edge.edge[1], edge.edge[2], 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.edge[1], edge.edge[2], false)
    end

    graph.properties -= GraphProperties(diff)

    return nothing
end

"""
    node_fusion!(graph::DAG, n1::ComputeTaskNode, n2::DataTaskNode, n3::ComputeTaskNode)

Fuse nodes n1 -> n2 -> n3 together into one node, return the applied difference to the graph.

For details see [`NodeFusion`](@ref).
"""
function node_fusion!(
    graph::DAG,
    n1::ComputeTaskNode,
    n2::DataTaskNode,
    n3::ComputeTaskNode,
)
    # @assert is_valid_node_fusion_input(graph, n1, n2, n3)

    # clear snapshot
    get_snapshot_diff(graph)


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

    # remove the edges and nodes that will be replaced by the fused node
    remove_edge!(graph, n1, n2)
    remove_edge!(graph, n2, n3)
    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, child, n1)
        push!(n1and3_children, child)
    end

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

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

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

    return get_snapshot_diff(graph)
end

"""
    node_reduction!(graph::DAG, nodes::Vector{Node})

Reduce the given nodes together into one node, return the applied difference to the graph.

For details see [`NodeReduction`](@ref).
"""
function node_reduction!(graph::DAG, nodes::Vector{Node})
    # @assert is_valid_node_reduction_input(graph, nodes)

    # clear snapshot
    get_snapshot_diff(graph)

    n1 = nodes[1]
    n1_children = children(n1)

    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, child, n)
        end

        for parent in parents(n)
            remove_edge!(graph, 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 all input nodes to n1 without duplicates
        insert_edge!(graph, n1, parent)
    end

    return get_snapshot_diff(graph)
end

"""
    node_split!(graph::DAG, n1::Node)

Split the given node into one node per parent, return the applied difference to the graph.

For details see [`NodeSplit`](@ref).
"""
function node_split!(graph::DAG, n1::Node)
    # @assert is_valid_node_split_input(graph, n1)

    # clear snapshot
    get_snapshot_diff(graph)

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

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

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

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

    return get_snapshot_diff(graph)
end