metagraphoptimization.jl/src/graph/mute.jl

# for graph mutating functions we need to do a few things
# 1: mute the graph (duh)
# 2: keep track of what was changed for the diff (if track == true)
# 3: invalidate operation caches

"""
    insert_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)

Insert the node into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`remove_node!`](@ref), [`insert_edge!`](@ref), [`remove_edge!`](@ref)
"""
function insert_node!(
    graph::DAG,
    node::Node,
    track = true,
    invalidate_cache = true,
)
    # 1: mute
    push!(graph.nodes, node)

    # 2: keep track
    if (track)
        push!(graph.diff.addedNodes, node)
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return node
    end
    push!(graph.dirtyNodes, node)

    return node
end

"""
    insert_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)

Insert the edge between node1 (child) and node2 (parent) into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`remove_node!`](@ref), [`remove_edge!`](@ref)
"""
function insert_edge!(
    graph::DAG,
    node1::Node,
    node2::Node,
    track = true,
    invalidate_cache = true,
)
    # @assert (node2 ∉ node1.parents) && (node1 ∉ node2.children) "Edge to insert already exists"

    # 1: mute
    # edge points from child to parent
    push!(node1.parents, node2)
    push!(node2.children, node1)

    # 2: keep track
    if (track)
        push!(graph.diff.addedEdges, make_edge(node1, node2))
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node1)
    invalidate_operation_caches!(graph, node2)

    push!(graph.dirtyNodes, node1)
    push!(graph.dirtyNodes, node2)

    return nothing
end

"""
    remove_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)

Remove the node from the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`insert_edge!`](@ref), [`remove_edge!`](@ref)
"""
function remove_node!(
    graph::DAG,
    node::Node,
    track = true,
    invalidate_cache = true,
)
    # @assert node in graph.nodes "Trying to remove a node that's not in the graph"

    # 1: mute
    delete!(graph.nodes, node)

    # 2: keep track
    if (track)
        push!(graph.diff.removedNodes, node)
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node)
    delete!(graph.dirtyNodes, node)

    return nothing
end

"""
    remove_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)

Remove the edge between node1 (child) and node2 (parent) into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`remove_node!`](@ref), [`insert_edge!`](@ref)
"""
function remove_edge!(
    graph::DAG,
    node1::Node,
    node2::Node,
    track = true,
    invalidate_cache = true,
)
    # 1: mute
    pre_length1 = length(node1.parents)
    pre_length2 = length(node2.children)
    filter!(x -> x != node2, node1.parents)
    filter!(x -> x != node1, node2.children)

    #=@assert begin
       removed = pre_length1 - length(node1.parents)
       removed <= 1
    end "removed more than one node from node1's parents"=#

    #=@assert begin
       removed = pre_length2 - length(node2.children)
       removed <= 1
    end "removed more than one node from node2's children"=#

    # 2: keep track
    if (track)
        push!(graph.diff.removedEdges, make_edge(node1, node2))
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node1)
    invalidate_operation_caches!(graph, node2)
    if (node1 in graph)
        push!(graph.dirtyNodes, node1)
    end
    if (node2 in graph)
        push!(graph.dirtyNodes, node2)
    end

    return nothing
end

"""
    get_snapshot_diff(graph::DAG)

Return the graph's [`Diff`](@ref) since last time this function was called.

See also: [`revert_diff`](@ref), [`AppliedOperation`](@ref) and [`revert_operation`](@ref)
"""
function get_snapshot_diff(graph::DAG)
    return swapfield!(graph, :diff, Diff())
end

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

Invalidate the operation caches for a given [`NodeFusion`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeFusion)
    delete!(graph.possibleOperations, operation)

    # delete the operation from all caches of nodes involved in the operation
    filter!(!=(operation), operation.input[1].nodeFusions)
    filter!(!=(operation), operation.input[3].nodeFusions)

    operation.input[2].nodeFusion = missing

    return nothing
end

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

Invalidate the operation caches for a given [`NodeReduction`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeReduction)
    delete!(graph.possibleOperations, operation)

    for node in operation.input
        node.nodeReduction = missing
    end

    return nothing
end

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

Invalidate the operation caches for a given [`NodeSplit`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeSplit)
    delete!(graph.possibleOperations, operation)

    # delete the operation from all caches of nodes involved in the operation
    # for node split there is only one node
    operation.input.nodeSplit = missing

    return nothing
end

"""
    invalidate_operation_caches!(graph::DAG, node::ComputeTaskNode)

Invalidate the operation caches of the given node through calls to the respective [`invalidate_caches!`](@ref) functions.
"""
function invalidate_operation_caches!(graph::DAG, node::ComputeTaskNode)
    if !ismissing(node.nodeReduction)
        invalidate_caches!(graph, node.nodeReduction)
    end
    if !ismissing(node.nodeSplit)
        invalidate_caches!(graph, node.nodeSplit)
    end
    while !isempty(node.nodeFusions)
        invalidate_caches!(graph, pop!(node.nodeFusions))
    end
    return nothing
end

"""
    invalidate_operation_caches!(graph::DAG, node::DataTaskNode)

Invalidate the operation caches of the given node through calls to the respective [`invalidate_caches!`](@ref) functions.
"""
function invalidate_operation_caches!(graph::DAG, node::DataTaskNode)
    if !ismissing(node.nodeReduction)
        invalidate_caches!(graph, node.nodeReduction)
    end
    if !ismissing(node.nodeSplit)
        invalidate_caches!(graph, node.nodeSplit)
    end
    if !ismissing(node.nodeFusion)
        invalidate_caches!(graph, node.nodeFusion)
    end
    return nothing
end