coloring.jl

"""
    partial_distance2_coloring(bg::BipartiteGraph, ::Val{side}, order::AbstractOrder)

Compute a distance-2 coloring of the given `side` (`1` or `2`) in the bipartite graph `bg` and return a vector of integer colors.

A _distance-2 coloring_ is such that two vertices have different colors if they are at distance at most 2.

The vertices are colored in a greedy fashion, following the `order` supplied.

# See also

- [`BipartiteGraph`](@ref)
- [`AbstractOrder`](@ref)

# References

> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005), Algorithm 3.2
"""
function partial_distance2_coloring(
    bg::BipartiteGraph{T}, ::Val{side}, order::AbstractOrder
) where {T,side}
    color = Vector{T}(undef, nb_vertices(bg, Val(side)))
    forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side)))
    vertices_in_order = vertices(bg, Val(side), order)
    partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order)
    return color
end

function partial_distance2_coloring!(
    color::AbstractVector{<:Integer},
    forbidden_colors::AbstractVector{<:Integer},
    bg::BipartiteGraph,
    ::Val{side},
    vertices_in_order::AbstractVector{<:Integer},
) where {side}
    color .= 0
    forbidden_colors .= 0
    other_side = 3 - side
    for v in vertices_in_order
        for w in neighbors(bg, Val(side), v)
            for x in neighbors(bg, Val(other_side), w)
                if !iszero(color[x])
                    forbidden_colors[color[x]] = v
                end
            end
        end
        for i in eachindex(forbidden_colors)
            if forbidden_colors[i] != v
                color[v] = i
                break
            end
        end
    end
end

"""
    star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::Bool)

Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where

- `color` is the vector of integer colors
- `star_set` is a [`StarSet`](@ref) encoding the set of 2-colored stars

A _star coloring_ is a distance-1 coloring such that every path on 4 vertices uses at least 3 colors.

The vertices are colored in a greedy fashion, following the `order` supplied.

If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.

# See also

- [`AdjacencyGraph`](@ref)
- [`AbstractOrder`](@ref)

# References

> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
"""
function star_coloring(
    g::AdjacencyGraph{T}, order::AbstractOrder, postprocessing::Bool
) where {T<:Integer}
    # Initialize data structures
    nv = nb_vertices(g)
    ne = nb_edges(g)
    color = zeros(T, nv)
    forbidden_colors = zeros(T, nv)
    first_neighbor = fill((zero(T), zero(T), zero(T)), nv)  # at first no neighbors have been encountered
    treated = zeros(T, nv)
    star = Vector{T}(undef, ne)
    hub = T[]  # one hub for each star, including the trivial ones
    vertices_in_order = vertices(g, order)

    for v in vertices_in_order
        for (w, index_vw) in neighbors_with_edge_indices(g, v)
            !has_diagonal(g) || (v == w && continue)
            iszero(color[w]) && continue
            forbidden_colors[color[w]] = v
            (p, q, _) = first_neighbor[color[w]]
            if p == v  # Case 1
                if treated[q] != v
                    # forbid colors of neighbors of q
                    _treat!(treated, forbidden_colors, g, v, q, color)
                end
                # forbid colors of neighbors of w
                _treat!(treated, forbidden_colors, g, v, w, color)
            else
                first_neighbor[color[w]] = (v, w, index_vw)
                for (x, index_wx) in neighbors_with_edge_indices(g, w)
                    !has_diagonal(g) || (w == x && continue)
                    (x == v || iszero(color[x])) && continue
                    if x == hub[star[index_wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
                        forbidden_colors[color[x]] = v
                    end
                end
            end
        end
        for i in eachindex(forbidden_colors)
            if forbidden_colors[i] != v
                color[v] = i
                break
            end
        end
        _update_stars!(star, hub, g, v, color, first_neighbor)
    end
    star_set = StarSet(star, hub)
    if postprocessing
        # Reuse the vector forbidden_colors to compute offsets during post-processing
        offsets = forbidden_colors
        postprocess!(color, star_set, g, offsets)
    end
    return color, star_set
end

function _treat!(
    # modified
    treated::AbstractVector{<:Integer},
    forbidden_colors::AbstractVector{<:Integer},
    # not modified
    g::AdjacencyGraph,
    v::Integer,
    w::Integer,
    color::AbstractVector{<:Integer},
)
    for x in neighbors(g, w)
        !has_diagonal(g) || (w == x && continue)
        iszero(color[x]) && continue
        forbidden_colors[color[x]] = v
    end
    treated[w] = v
    return nothing
end

function _update_stars!(
    # modified
    star::AbstractVector{<:Integer},
    hub::AbstractVector{<:Integer},
    # not modified
    g::AdjacencyGraph,
    v::Integer,
    color::AbstractVector{<:Integer},
    first_neighbor::AbstractVector{<:Tuple},
)
    for (w, index_vw) in neighbors_with_edge_indices(g, v)
        !has_diagonal(g) || (v == w && continue)
        iszero(color[w]) && continue
        x_exists = false
        for (x, index_wx) in neighbors_with_edge_indices(g, w)
            !has_diagonal(g) || (w == x && continue)
            if x != v && color[x] == color[v]  # vw, wx ∈ E
                star_wx = star[index_wx]
                hub[star_wx] = w  # this may already be true
                star[index_vw] = star_wx
                x_exists = true
                break
            end
        end
        if !x_exists
            (p, q, index_pq) = first_neighbor[color[w]]
            if p == v && q != w  # vw, vq ∈ E and color[w] = color[q]
                star_vq = star[index_pq]
                hub[star_vq] = v  # this may already be true
                star[index_vw] = star_vq
            else  # vw forms a new star
                push!(hub, -max(v, w))  # star is trivial (composed only of two vertices) so we set the hub to a negative value, but it allows us to choose one of the two vertices
                star[index_vw] = length(hub)
            end
        end
    end
    return nothing
end

"""
    StarSet

Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algorithm.

# Fields

$TYPEDFIELDS
"""
struct StarSet{T}
    "a mapping from edges (pair of vertices) to their star index"
    star::Vector{T}
    "a mapping from star indices to their hub (undefined hubs for single-edge stars are the negative value of one of the vertices, picked arbitrarily)"
    hub::Vector{T}
end

"""
    acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::Bool)

Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a tuple `(color, tree_set)`, where

- `color` is the vector of integer colors
- `tree_set` is a [`TreeSet`](@ref) encoding the set of 2-colored trees

An _acyclic coloring_ is a distance-1 coloring with the further restriction that every cycle uses at least 3 colors.

The vertices are colored in a greedy fashion, following the `order` supplied.

If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.

# See also

- [`AdjacencyGraph`](@ref)
- [`AbstractOrder`](@ref)

# References

> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 3.1
"""
function acyclic_coloring(
    g::AdjacencyGraph{T}, order::AbstractOrder, postprocessing::Bool
) where {T<:Integer}
    # Initialize data structures
    nv = nb_vertices(g)
    ne = nb_edges(g)
    color = zeros(T, nv)
    forbidden_colors = zeros(T, nv)
    first_neighbor = fill((zero(T), zero(T), zero(T)), nv)  # at first no neighbors have been encountered
    first_visit_to_tree = fill((zero(T), zero(T)), ne)
    forest = Forest{T}(ne)
    vertices_in_order = vertices(g, order)

    for v in vertices_in_order
        for w in neighbors(g, v)
            !has_diagonal(g) || (v == w && continue)
            iszero(color[w]) && continue
            forbidden_colors[color[w]] = v
        end
        for w in neighbors(g, v)
            !has_diagonal(g) || (v == w && continue)
            iszero(color[w]) && continue
            for (x, index_wx) in neighbors_with_edge_indices(g, w)
                !has_diagonal(g) || (w == x && continue)
                iszero(color[x]) && continue
                if forbidden_colors[color[x]] != v
                    _prevent_cycle!(
                        v,
                        w,
                        x,
                        index_wx,
                        color,
                        first_visit_to_tree,
                        forbidden_colors,
                        forest,
                    )
                end
            end
        end
        for i in eachindex(forbidden_colors)
            if forbidden_colors[i] != v
                color[v] = i
                break
            end
        end
        for (w, index_vw) in neighbors_with_edge_indices(g, v)  # grow two-colored stars around the vertex v
            !has_diagonal(g) || (v == w && continue)
            iszero(color[w]) && continue
            _grow_star!(v, w, index_vw, color, first_neighbor, forest)
        end
        for (w, index_vw) in neighbors_with_edge_indices(g, v)
            !has_diagonal(g) || (v == w && continue)
            iszero(color[w]) && continue
            for (x, index_wx) in neighbors_with_edge_indices(g, w)
                !has_diagonal(g) || (w == x && continue)
                (x == v || iszero(color[x])) && continue
                if color[x] == color[v]
                    _merge_trees!(v, w, x, index_vw, index_wx, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
                end
            end
        end
    end

    buffer = forbidden_colors
    reverse_bfs_orders = first_visit_to_tree
    tree_set = TreeSet(g, forest, buffer, reverse_bfs_orders, ne)
    if postprocessing
        # Reuse the vector forbidden_colors to compute offsets during post-processing
        offsets = forbidden_colors
        postprocess!(color, tree_set, g, offsets)
    end
    return color, tree_set
end

function _prevent_cycle!(
    # not modified
    v::Integer,
    w::Integer,
    x::Integer,
    index_wx::Integer,
    color::AbstractVector{<:Integer},
    # modified
    first_visit_to_tree::AbstractVector{<:Tuple},
    forbidden_colors::AbstractVector{<:Integer},
    forest::Forest{<:Integer},
)
    root_wx = find_root!(forest, index_wx)  # root of the 2-colored tree T to which the edge wx belongs
    (p, q) = first_visit_to_tree[root_wx]
    if p != v  # T is being visited from vertex v for the first time
        first_visit_to_tree[root_wx] = (v, w)
    elseif q != w  # T is connected to vertex v via at least two edges
        forbidden_colors[color[x]] = v
    end
    return nothing
end

function _grow_star!(
    # not modified
    v::Integer,
    w::Integer,
    index_vw::Integer,
    color::AbstractVector{<:Integer},
    # modified
    first_neighbor::AbstractVector{<:Tuple},
    forest::Forest{<:Integer},
)
    # Create a new tree T_{vw} consisting only of edge vw
    (p, q, index_pq) = first_neighbor[color[w]]
    if p != v  # a neighbor of v with color[w] encountered for the first time
        first_neighbor[color[w]] = (v, w, index_vw)
    else  # merge T_{vw} with a two-colored star being grown around v
        root_vw = find_root!(forest, index_vw)
        root_pq = find_root!(forest, index_pq)
        root_union!(forest, root_vw, root_pq)
    end
    return nothing
end

function _merge_trees!(
    # not modified
    v::Integer,
    w::Integer,
    x::Integer,
    index_vw::Integer,
    index_wx::Integer,
    # modified
    forest::Forest{<:Integer},
)
    root_vw = find_root!(forest, index_vw)
    root_wx = find_root!(forest, index_wx)
    if root_vw != root_wx
        root_union!(forest, root_vw, root_wx)
    end
    return nothing
end

"""
    TreeSet

Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) algorithm.

# Fields

$TYPEDFIELDS
"""
struct TreeSet{T}
    reverse_bfs_orders::Vector{Tuple{T,T}}
    is_star::Vector{Bool}
    tree_edge_indices::Vector{T}
    nt::T
end

function TreeSet(
    g::AdjacencyGraph{T},
    forest::Forest{T},
    buffer::AbstractVector{T},
    reverse_bfs_orders::Vector{Tuple{T,T}},
    ne::Integer,
) where {T}
    S = pattern(g)
    edge_to_index = edge_indices(g)
    nv = nb_vertices(g)
    (; nt, ranks, parents) = forest

    # root_to_tree is a vector that maps a tree's root to the index of the tree
    # We can recycle the vector "ranks" because we don't need it anymore to merge trees
    root_to_tree = ranks
    fill!(root_to_tree, zero(T))

    # vector specifying the starting and ending indices of edges for each tree
    tree_edge_indices = zeros(T, nt + 1)

    # number of roots found
    nr = 0

    # determine the number of edges for each tree and map each root to a tree index
    for index_edge in 1:ne
        root = find_root!(forest, index_edge)

        # create a mapping between roots and tree indices
        if iszero(root_to_tree[root])
            nr += 1
            root_to_tree[root] = nr
        end

        # index of the tree that contains this edge
        index_tree = root_to_tree[root]

        # Update the number of edges for the current tree (shifted by 1 to facilitate the final cumsum)
        tree_edge_indices[index_tree + 1] += 1
    end

    # nvmax is the number of vertices in the largest tree of the forest
    # Note: the number of vertices in a tree is equal the number of edges plus one
    nvmax = maximum(tree_edge_indices) + one(T)

    # Vector containing the list of vertices, grouped by tree
    # Note: the total number of edges in the graph is "ne", so there are "ne + nt" vertices across all trees
    tree_vertices = Vector{T}(undef, ne + nt)

    # Provide the positions of the first and last neighbors for each vertex in "tree_vertices", within the tree to which the vertex belongs
    # These positions refer to indices in the vector "tree_neighbors"
    tree_neighbor_indices = zeros(T, ne + nt + 1)

    # Packed representation of the neighbors of each vertex in "tree_vertices"
    tree_neighbors = Vector{T}(undef, 2 * ne)

    # Track the positions for inserting vertices and neighbors per tree
    vertex_position = Vector{T}(undef, nt)
    neighbor_position = Vector{T}(undef, nt)

    # Compute starting positions for vertices and neighbors in each tree
    if nt > 0
        vertex_position[1] = zero(T)
        neighbor_position[1] = zero(T)
    end
    for k in 2:nt
        # Note: tree_edge_indices[k] is the number of edges in the tree k-1
        vertex_position[k] = vertex_position[k - 1] + tree_edge_indices[k] + 1
        neighbor_position[k] = neighbor_position[k - 1] + 2 * tree_edge_indices[k]
    end

    # found_in_tree indicates if a given vertex is in each tree
    found_in_tree = Vector{Bool}(undef, nt)

    rvS = rowvals(S)
    for j in axes(S, 2)
        # Reset the vector found_in_tree to indicate that the vertex j is not present in any tree yet
        fill!(found_in_tree, false)

        for pos in nzrange(S, j)
            i = rvS[pos]
            if i != j
                index_ij = edge_to_index[pos]

                # No need to call "find_root!" because paths have already been compressed
                root = parents[index_ij]

                # Index of the tree containing edge (i, j)
                index_tree = root_to_tree[root]

                # Position in tree_vertices where vertex j should be found or inserted
                vertex_index = vertex_position[index_tree]

                if !found_in_tree[index_tree]
                    # Mark that vertex j is present in the current tree
                    found_in_tree[index_tree] = true

                    # Insert j into tree_vertices
                    vertex_position[index_tree] += 1
                    vertex_index += 1
                    tree_vertices[vertex_index] = j
                end

                # Append neighbor i to the list of neighbors of j in the tree
                neighbor_position[index_tree] += 1
                neighbor_index = neighbor_position[index_tree]
                tree_neighbors[neighbor_index] = i

                # Increment neighbor count for j in the tree (shifted by 1 to facilitate the final cumsum)
                tree_neighbor_indices[vertex_index + 1] += 1
            end
        end
    end

    # Compute a shifted cumulative sum of tree_edge_indices, starting from one
    tree_edge_indices[1] = one(T)
    for k in 2:(nt + 1)
        tree_edge_indices[k] += tree_edge_indices[k - 1]
    end

    # Compute a shifted cumulative sum of tree_neighbor_indices, starting from one
    tree_neighbor_indices[1] = 1
    for k in 2:(ne + nt + 1)
        tree_neighbor_indices[k] += tree_neighbor_indices[k - 1]
    end

    # degrees is a vector of integers that stores the degree of each vertex in a tree
    degrees = buffer

    # For each vertex in the current tree, reverse_mapping will hold its corresponding index in tree_vertices
    reverse_mapping = Vector{T}(undef, nv)

    # Create a queue with a fixed size nvmax
    queue = Vector{T}(undef, nvmax)

    # Determine if each tree in the forest is a star
    # In a star, at most one vertex has a degree strictly greater than one
    is_star = found_in_tree

    # Number of edges treated
    num_edges_treated = zero(T)

    # reverse_bfs_orders contains the reverse breadth first (BFS) traversal order for each tree in the forest
    for k in 1:nt
        # Initialize the queue to store the leaves
        queue_start = 1
        queue_end = 0

        # Positions of the first and last vertices in the current tree
        # Note: tree_edge_indices contains the positions of the first and last edges,
        # so we add to add an offset k-1 between edge indices and vertex indices
        first_vertex = tree_edge_indices[k] + (k-1)
        last_vertex = tree_edge_indices[k + 1] + (k-1)

        # compute the degree of each vertex in the tree
        for index_vertex in first_vertex:last_vertex
            vertex = tree_vertices[index_vertex]
            degree =
                tree_neighbor_indices[index_vertex + 1] -
                tree_neighbor_indices[index_vertex]
            degrees[vertex] = degree

            # store a reverse mapping to get the position of the vertex in tree_vertices
            reverse_mapping[vertex] = index_vertex

            # the vertex is a leaf
            if degree == 1
                queue_end += 1
                queue[queue_end] = vertex
            end
        end

        # number of vertices in the tree
        nv_tree = tree_edge_indices[k + 1] - tree_edge_indices[k] + 1

        # Check that no more than one vertex has a degree strictly greater than one
        # "queue_end" currently represents the number of vertices considered as leaves in the tree before any pruning
        is_star[k] = queue_end >= nv_tree - 1

        # continue until all leaves are treated
        while queue_start <= queue_end
            leaf = queue[queue_start]
            queue_start += 1

            # Mark the vertex as removed
            degrees[leaf] = 0

            # Position of the leaf in tree_vertices
            index_leaf = reverse_mapping[leaf]

            # Positions of the first and last neighbors of the leaf in the current tree
            first_neighbor = tree_neighbor_indices[index_leaf]
            last_neighbor = tree_neighbor_indices[index_leaf + 1] - 1

            # Iterate over all neighbors of the leaf to be pruned
            for index_neighbor in first_neighbor:last_neighbor
                neighbor = tree_neighbors[index_neighbor]

                # Check if neighbor is the parent of the leaf or if it was a child before the tree was pruned
                if degrees[neighbor] != 0
                    # (leaf, neighbor) represents the next edge to visit during decompression
                    num_edges_treated += 1
                    reverse_bfs_orders[num_edges_treated] = (leaf, neighbor)

                    # reduce the degree of the neighbor
                    degrees[neighbor] -= 1

                    # check if the neighbor is now a leaf
                    if degrees[neighbor] == 1
                        queue_end += 1
                        queue[queue_end] = neighbor
                    end
                end
            end
        end
    end

    return TreeSet(reverse_bfs_orders, is_star, tree_edge_indices, nt)
end

## Postprocessing, mirrors decompression code

function postprocess!(
    color::AbstractVector{<:Integer},
    star_or_tree_set::Union{StarSet,TreeSet},
    g::AdjacencyGraph,
    offsets::AbstractVector{<:Integer},
)
    S = pattern(g)
    edge_to_index = edge_indices(g)
    # flag which colors are actually used during decompression
    nb_colors = maximum(color)
    color_used = zeros(Bool, nb_colors)

    # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
    if has_diagonal(g)
        for i in axes(S, 1)
            if !iszero(S[i, i])
                color_used[color[i]] = true
            end
        end
    end

    if star_or_tree_set isa StarSet
        # only the colors of the hubs are used
        (; star, hub) = star_or_tree_set
        nb_trivial_stars = 0

        # Iterate through all non-trivial stars
        for s in eachindex(hub)
            h = hub[s]
            if h > 0
                color_used[color[h]] = true
            else
                nb_trivial_stars += 1
            end
        end

        # Process the trivial stars (if any)
        if nb_trivial_stars > 0
            rvS = rowvals(S)
            for j in axes(S, 2)
                for k in nzrange(S, j)
                    i = rvS[k]
                    if i > j
                        index_ij = edge_to_index[k]
                        s = star[index_ij]
                        h = hub[s]
                        if h < 0
                            h = abs(h)
                            spoke = h == j ? i : j
                            if color_used[color[spoke]]
                                # Switch the hub and the spoke to possibly avoid adding one more used color
                                hub[s] = spoke
                            else
                                # Keep the current hub
                                color_used[color[h]] = true
                            end
                        end
                    end
                end
            end
        end
    else
        # only the colors of non-leaf vertices are used
        (; reverse_bfs_orders, is_star, tree_edge_indices, nt) = star_or_tree_set
        nb_trivial_trees = 0

        # Iterate through all non-trivial trees
        for k in 1:nt
            # Position of the first edge in the tree
            first = tree_edge_indices[k]

            # Total number of edges in the tree
            ne_tree = tree_edge_indices[k + 1] - first

            # Check if we have more than one edge in the tree (non-trivial tree)
            if ne_tree > 1
                # Determine if the tree is a star
                if is_star[k]
                    # It is a non-trivial star and only the color of the hub is needed
                    (_, hub) = reverse_bfs_orders[first]
                    color_used[color[hub]] = true
                else
                    # It is not a star and both colors are needed during the decompression
                    (i, j) = reverse_bfs_orders[first]
                    color_used[color[i]] = true
                    color_used[color[j]] = true
                end
            else
                nb_trivial_trees += 1
            end
        end

        # Process the trivial trees (if any)
        if nb_trivial_trees > 0
            for k in 1:nt
                # Position of the first edge in the tree
                first = tree_edge_indices[k]

                # Total number of edges in the tree
                ne_tree = tree_edge_indices[k + 1] - first

                # Check if we have exactly one edge in the tree
                if ne_tree == 1
                    (i, j) = reverse_bfs_orders[first]
                    if color_used[color[i]]
                        # Make i the root to avoid possibly adding one more used color
                        # Switch it with the (only) leaf
                        reverse_bfs_orders[first] = (j, i)
                    else
                        # Keep j as the root
                        color_used[color[j]] = true
                    end
                end
            end
        end
    end

    # if at least one of the colors is useless, modify the color assignments of vertices
    if any(!, color_used)
        num_colors_useless = 0

        # determine what are the useless colors and compute the offsets
        for ci in 1:nb_colors
            if color_used[ci]
                offsets[ci] = num_colors_useless
            else
                num_colors_useless += 1
            end
        end

        # assign the neutral color to every vertex with a useless color and remap the colors
        for i in eachindex(color)
            ci = color[i]
            if !color_used[ci]
                # assign the neutral color
                color[i] = 0
            else
                # remap the color to not have any gap
                color[i] -= offsets[ci]
            end
        end
    end
    return color
end