diff --git a/docs/src/dev.md b/docs/src/dev.md
index 9fff4e65..2f253dbe 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -23,7 +23,6 @@ SparseMatrixColorings.bidirectional_pattern
 
 ```@docs
 SparseMatrixColorings.partial_distance2_coloring
-SparseMatrixColorings.symmetric_coefficient
 SparseMatrixColorings.star_coloring
 SparseMatrixColorings.acyclic_coloring
 SparseMatrixColorings.group_by_color
diff --git a/src/coloring.jl b/src/coloring.jl
index 3c1d49a9..587edc1c 100644
--- a/src/coloring.jl
+++ b/src/coloring.jl
@@ -82,18 +82,18 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
     ne = nb_edges(g)
     color = zeros(Int, nv)
     forbidden_colors = zeros(Int, nv)
-    first_neighbor = fill((0, 0), nv)  # at first no neighbors have been encountered
+    first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     treated = zeros(Int, nv)
-    star = Dict{Tuple{Int,Int},Int}()
-    sizehint!(star, ne)
+    star = Vector{Int}(undef, ne)
     hub = Int[]  # one hub for each star, including the trivial ones
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
-        for w in neighbors(g, v)
+        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]]
+            (p, q, _) = first_neighbor[color[w]]
             if p == v  # Case 1
                 if treated[q] != v
                     # forbid colors of neighbors of q
@@ -102,11 +102,11 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
                 # forbid colors of neighbors of w
                 _treat!(treated, forbidden_colors, g, v, w, color)
             else
-                first_neighbor[color[w]] = (v, w)
-                for x in neighbors(g, w)
+                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
-                    wx = _sort(w, x)
-                    if x == hub[star[wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
+                    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
@@ -123,29 +123,12 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::
     star_set = StarSet(star, hub)
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
-        postprocess!(color, star_set, g, forbidden_colors)
+        offsets = forbidden_colors
+        postprocess!(color, star_set, g, offsets)
     end
     return color, star_set
 end
 
-"""
-    StarSet
-
-Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algorithm.
-
-# Fields
-
-$TYPEDFIELDS
-"""
-struct StarSet
-    "a mapping from edges (pair of vertices) to their star index"
-    star::Dict{Tuple{Int,Int},Int}
-    "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{Int}
-end
-
-_sort(u, v) = (min(u, v), max(u, v))
-
 function _treat!(
     # modified
     treated::AbstractVector{<:Integer},
@@ -157,6 +140,7 @@ function _treat!(
     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
@@ -166,7 +150,7 @@ end
 
 function _update_stars!(
     # modified
-    star::Dict{<:Tuple,<:Integer},
+    star::AbstractVector{<:Integer},
     hub::AbstractVector{<:Integer},
     # not modified
     g::AdjacencyGraph,
@@ -174,30 +158,29 @@ function _update_stars!(
     color::AbstractVector{<:Integer},
     first_neighbor::AbstractVector{<:Tuple},
 )
-    for w in neighbors(g, v)
+    for (w, index_vw) in neighbors_with_edge_indices(g, v)
+        !has_diagonal(g) || (v == w && continue)
         iszero(color[w]) && continue
-        vw = _sort(v, w)
         x_exists = false
-        for x in neighbors(g, w)
+        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
-                wx = _sort(w, x)
-                star_wx = star[wx]
+                star_wx = star[index_wx]
                 hub[star_wx] = w  # this may already be true
-                star[vw] = star_wx
+                star[index_vw] = star_wx
                 x_exists = true
                 break
             end
         end
         if !x_exists
-            (p, q) = first_neighbor[color[w]]
+            (p, q, index_pq) = first_neighbor[color[w]]
             if p == v && q != w  # vw, vq ∈ E and color[w] = color[q]
-                vq = _sort(v, q)
-                star_vq = star[vq]
+                star_vq = star[index_pq]
                 hub[star_vq] = v  # this may already be true
-                star[vw] = star_vq
+                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[vw] = length(hub)
+                star[index_vw] = length(hub)
             end
         end
     end
@@ -205,41 +188,19 @@ function _update_stars!(
 end
 
 """
-    symmetric_coefficient(
-        i::Integer, j::Integer,
-        color::AbstractVector{<:Integer},
-        star_set::StarSet
-    )
-
-Return the indices `(k, c)` such that `A[i, j] = B[k, c]`, where `A` is the uncompressed symmetric matrix and `B` is the column-compressed matrix.
+    StarSet
 
-This function corresponds to algorithm `DirectRecover2` in the paper.
+Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algorithm.
 
-# References
+# Fields
 
-> [_Efficient Computation of Sparse Hessians Using Coloring and Automatic Differentiation_](https://pubsonline.informs.org/doi/abs/10.1287/ijoc.1080.0286), Gebremedhin et al. (2009), Figure 3
+$TYPEDFIELDS
 """
-function symmetric_coefficient(
-    i::Integer, j::Integer, color::AbstractVector{<:Integer}, star_set::StarSet
-)
-    (; star, hub) = star_set
-    if i == j
-        # diagonal
-        return i, color[j]
-    end
-    if i > j  # keys of star are sorted tuples
-        # star only contains one triangle
-        i, j = j, i
-    end
-    star_id = star[i, j]
-    h = abs(hub[star_id])
-    if h == j
-        # i is the spoke
-        return i, color[h]
-    else
-        # j is the spoke
-        return j, color[h]
-    end
+struct StarSet
+    "a mapping from edges (pair of vertices) to their star index"
+    star::Vector{Int}
+    "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{Int}
 end
 
 """
@@ -271,23 +232,33 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessin
     ne = nb_edges(g)
     color = zeros(Int, nv)
     forbidden_colors = zeros(Int, nv)
-    first_neighbor = fill((0, 0), nv)  # at first no neighbors have been encountered
+    first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
     first_visit_to_tree = fill((0, 0), ne)
     forest = Forest{Int}(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 in neighbors(g, w)
+            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, color, first_visit_to_tree, forbidden_colors, forest
+                        v,
+                        w,
+                        x,
+                        index_wx,
+                        color,
+                        first_visit_to_tree,
+                        forbidden_colors,
+                        forest,
                     )
                 end
             end
@@ -298,25 +269,29 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessin
                 break
             end
         end
-        for w in neighbors(g, v)  # grow two-colored stars around the vertex v
+        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, color, first_neighbor, forest)
+            _grow_star!(v, w, index_vw, color, first_neighbor, forest)
         end
-        for w in neighbors(g, v)
+        for (w, index_vw) in neighbors_with_edge_indices(g, v)
+            !has_diagonal(g) || (v == w && continue)
             iszero(color[w]) && continue
-            for x in neighbors(g, w)
+            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, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
+                    _merge_trees!(v, w, x, index_vw, index_wx, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
                 end
             end
         end
     end
 
-    tree_set = TreeSet(forest, nb_vertices(g))
+    tree_set = TreeSet(g, forest)
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
-        postprocess!(color, tree_set, g, forbidden_colors)
+        offsets = forbidden_colors
+        postprocess!(color, tree_set, g, offsets)
     end
     return color, tree_set
 end
@@ -326,14 +301,14 @@ function _prevent_cycle!(
     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},
 )
-    wx = _sort(w, x)
-    id = find_root!(forest, wx)  # The edge wx belongs to the 2-colored tree T, represented by an edge with an integer ID
+    id = find_root!(forest, index_wx)  # The edge wx belongs to the 2-colored tree T, represented by an edge with an integer ID
     (p, q) = first_visit_to_tree[id]
     if p != v  # T is being visited from vertex v for the first time
         first_visit_to_tree[id] = (v, w)
@@ -347,21 +322,19 @@ function _grow_star!(
     # not modified
     v::Integer,
     w::Integer,
+    index_vw::Integer,
     color::AbstractVector{<:Integer},
     # modified
     first_neighbor::AbstractVector{<:Tuple},
     forest::Forest{<:Integer},
 )
-    vw = _sort(v, w)
-    push!(forest, vw)  # Create a new tree T_{vw} consisting only of edge vw
-    (p, q) = first_neighbor[color[w]]
+    # 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)
+        first_neighbor[color[w]] = (v, w, index_vw)
     else  # merge T_{vw} with a two-colored star being grown around v
-        vw = _sort(v, w)
-        pq = _sort(p, q)
-        root1 = find_root!(forest, vw)
-        root2 = find_root!(forest, pq)
+        root1 = find_root!(forest, index_vw)
+        root2 = find_root!(forest, index_pq)
         root_union!(forest, root1, root2)
     end
     return nothing
@@ -372,13 +345,13 @@ function _merge_trees!(
     v::Integer,
     w::Integer,
     x::Integer,
+    index_vw::Integer,
+    index_wx::Integer,
     # modified
     forest::Forest{<:Integer},
 )
-    vw = _sort(v, w)
-    wx = _sort(w, x)
-    root1 = find_root!(forest, vw)
-    root2 = find_root!(forest, wx)
+    root1 = find_root!(forest, index_vw)
+    root2 = find_root!(forest, index_wx)
     if root1 != root2
         root_union!(forest, root1, root2)
     end
@@ -399,10 +372,10 @@ struct TreeSet
     is_star::Vector{Bool}
 end
 
-function TreeSet(forest::Forest{Int}, nvertices::Int)
-    # Forest is a structure defined in forest.jl
-    # - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
-    # - forest.num_trees: the number of trees in the forest
+function TreeSet(g::AdjacencyGraph, forest::Forest{Int})
+    S = pattern(g)
+    edge_to_index = edge_indices(g)
+    nv = nb_vertices(g)
     nt = forest.num_trees
 
     # dictionary that maps a tree's root to the index of the tree
@@ -412,38 +385,45 @@ function TreeSet(forest::Forest{Int}, nvertices::Int)
     # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
     trees = [Dict{Int,Vector{Int}}() for i in 1:nt]
 
-    # counter of the number of roots found
-    k = 0
-    for edge in keys(forest.intmap)
-        i, j = edge
-        root = find_root!(forest, edge)
-
-        # Update roots
-        if !haskey(roots, root)
-            k += 1
-            roots[root] = k
-        end
+    # current number of roots found
+    nr = 0
+
+    rvS = rowvals(S)
+    for j in axes(S, 2)
+        for pos in nzrange(S, j)
+            i = rvS[pos]
+            if i > j
+                index_ij = edge_to_index[pos]
+                root = find_root!(forest, index_ij)
+
+                # Update roots
+                if !haskey(roots, root)
+                    nr += 1
+                    roots[root] = nr
+                end
 
-        # index of the tree T that contains this edge
-        index_tree = roots[root]
+                # index of the tree T that contains this edge
+                index_tree = roots[root]
 
-        # Update the neighbors of i in the tree T
-        if !haskey(trees[index_tree], i)
-            trees[index_tree][i] = [j]
-        else
-            push!(trees[index_tree][i], j)
-        end
+                # Update the neighbors of i in the tree T
+                if !haskey(trees[index_tree], i)
+                    trees[index_tree][i] = [j]
+                else
+                    push!(trees[index_tree][i], j)
+                end
 
-        # Update the neighbors of j in the tree T
-        if !haskey(trees[index_tree], j)
-            trees[index_tree][j] = [i]
-        else
-            push!(trees[index_tree][j], i)
+                # Update the neighbors of j in the tree T
+                if !haskey(trees[index_tree], j)
+                    trees[index_tree][j] = [i]
+                else
+                    push!(trees[index_tree][j], i)
+                end
+            end
         end
     end
 
     # degrees is a vector of integers that stores the degree of each vertex in a tree
-    degrees = Vector{Int}(undef, nvertices)
+    degrees = Vector{Int}(undef, nv)
 
     # reverse breadth first (BFS) traversal order for each tree in the forest
     reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:nt]
@@ -541,7 +521,8 @@ function postprocess!(
     g::AdjacencyGraph,
     offsets::Vector{Int},
 )
-    (; S) = g
+    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)
@@ -562,9 +543,9 @@ function postprocess!(
 
         # Iterate through all non-trivial stars
         for s in eachindex(hub)
-            j = hub[s]
-            if j > 0
-                color_used[color[j]] = true
+            h = hub[s]
+            if h > 0
+                color_used[color[h]] = true
             else
                 nb_trivial_stars += 1
             end
@@ -572,17 +553,25 @@ function postprocess!(
 
         # Process the trivial stars (if any)
         if nb_trivial_stars > 0
-            for ((i, j), s) in pairs(star)
-                h = hub[s]
-                if h < 0
-                    h = abs(h)
-                    spoke = i == h ? j : i
-                    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
+            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
diff --git a/src/forest.jl b/src/forest.jl
index 71e8bb00..42bc0208 100644
--- a/src/forest.jl
+++ b/src/forest.jl
@@ -10,12 +10,8 @@ Structure that provides fast union-find operations for constructing a forest dur
 $TYPEDFIELDS
 """
 mutable struct Forest{T<:Integer}
-    "current number of edges in the forest"
-    num_edges::T
     "current number of distinct trees in the forest"
     num_trees::T
-    "dictionary mapping each edge represented as a tuple of vertices to its unique integer index"
-    intmap::Dict{Tuple{T,T},T}
     "vector storing the index of a parent in the tree for each edge, used in union-find operations"
     parents::Vector{T}
     "vector approximating the depth of each tree to optimize path compression"
@@ -23,20 +19,10 @@ mutable struct Forest{T<:Integer}
 end
 
 function Forest{T}(n::Integer) where {T<:Integer}
-    num_edges = zero(T)
-    num_trees = zero(T)
-    intmap = Dict{Tuple{T,T},T}()
-    sizehint!(intmap, n)
+    num_trees = T(n)
     parents = collect(Base.OneTo(T(n)))
     ranks = zeros(T, T(n))
-    return Forest{T}(num_edges, num_trees, intmap, parents, ranks)
-end
-
-function Base.push!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
-    forest.num_edges += 1
-    forest.intmap[edge] = forest.num_edges
-    forest.num_trees += one(T)
-    return forest
+    return Forest{T}(num_trees, parents, ranks)
 end
 
 function _find_root!(parents::Vector{T}, index_edge::T) where {T<:Integer}
@@ -47,8 +33,8 @@ function _find_root!(parents::Vector{T}, index_edge::T) where {T<:Integer}
     return p
 end
 
-function find_root!(forest::Forest{T}, edge::Tuple{T,T}) where {T<:Integer}
-    return _find_root!(forest.parents, forest.intmap[edge])
+function find_root!(forest::Forest{T}, index_edge::T) where {T<:Integer}
+    return _find_root!(forest.parents, index_edge)
 end
 
 function root_union!(forest::Forest{T}, index_edge1::T, index_edge2::T) where {T<:Integer}
diff --git a/src/graph.jl b/src/graph.jl
index 2cd9c7fd..047fdae4 100644
--- a/src/graph.jl
+++ b/src/graph.jl
@@ -98,19 +98,21 @@ end
 
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
-bidirectional_pattern(A::AbstractMatrix; symmetric_pattern) =
+bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) =
     bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
 
-function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
+function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T}
     m, n = size(S)
     p = m + n
     nnzS = nnz(S)
-    rowval = Vector{Int}(undef, 2 * nnzS)
-    colptr = zeros(Int, p + 1)
+    rowval = Vector{T}(undef, 2 * nnzS)
+    colptr = zeros(T, p + 1)
+    edge_to_index = Vector{T}(undef, 2 * nnzS)
 
     # Update rowval and colptr for the block A
     for i in 1:nnzS
         rowval[i] = S.rowval[i] + n
+        edge_to_index[i] = i
     end
     for j in 1:n
         colptr[j] = S.colptr[j]
@@ -118,9 +120,16 @@ function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
 
     # Update rowval and colptr for the block Aᵀ
     if symmetric_pattern
+        # Use colptr[n+1:n+m] to store the offsets during the update of edge_to_index
+        offsets = colptr
+
         # We use the sparsity pattern of A for Aᵀ
-        for i in 1:nnzS
-            rowval[nnzS + i] = S.rowval[i]
+        for k in 1:nnzS
+            r = S.rowval[k]
+            rowval[nnzS + k] = r
+            pos = S.colptr[r] + offsets[n + r]
+            edge_to_index[nnzS + pos] = edge_to_index[k]
+            offsets[n + r] += 1
         end
         # m and n are identical because symmetric_pattern is true
         for j in 1:m
@@ -147,6 +156,7 @@ function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
                 i = S.rowval[index]
                 pos = colptr[n + i]
                 rowval[nnzS + pos] = j
+                edge_to_index[nnzS + pos] = edge_to_index[index]
                 colptr[n + i] += 1
             end
         end
@@ -159,8 +169,32 @@ function bidirectional_pattern(S::SparsityPatternCSC; symmetric_pattern)
     end
 
     # Create the SparsityPatternCSC of the augmented adjacency matrix
-    S_and_Sᵀ = SparsityPatternCSC{Int}(p, p, colptr, rowval)
-    return S_and_Sᵀ
+    S_and_Sᵀ = SparsityPatternCSC{T}(p, p, colptr, rowval)
+    return S_and_Sᵀ, edge_to_index
+end
+
+function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}
+    # edge_to_index gives an index for each edge
+    edge_to_index = Vector{T}(undef, nnz(S))
+    offsets = zeros(T, S.n)
+    counter = 0
+    rvS = rowvals(S)
+    for j in axes(S, 2)
+        for k in nzrange(S, j)
+            i = rvS[k]
+            if i > j
+                counter += 1
+                edge_to_index[k] = counter
+                k2 = S.colptr[i] + offsets[i]
+                edge_to_index[k2] = counter
+                offsets[i] += 1
+            elseif i == j
+                # this should never be used, make sure it errors
+                edge_to_index[k] = 0
+            end
+        end
+    end
+    return edge_to_index
 end
 
 ## Adjacency graph
@@ -177,11 +211,12 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 # Constructors
 
-    AdjacencyGraph(A::SparseMatrixCSC)
+    AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
 
 # Fields
 
 - `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, whose diagonal is empty whenever `has_diagonal` is `false`
+- `edge_to_index::Vector{T}`: A vector mapping each nonzero of `S` to a unique edge index (ignoring diagonal and accounting for symmetry, so that `(i, j)` and `(j, i)` get the same index)
 
 # References
 
@@ -189,10 +224,15 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 """
 struct AdjacencyGraph{T,has_diagonal}
     S::SparsityPatternCSC{T}
+    edge_to_index::Vector{T}
 end
 
-function AdjacencyGraph(S::SparsityPatternCSC{T}; has_diagonal::Bool=true) where {T}
-    return AdjacencyGraph{T,has_diagonal}(S)
+function AdjacencyGraph(
+    S::SparsityPatternCSC{T},
+    edge_to_index::Vector{T}=build_edge_to_index(S);
+    has_diagonal::Bool=true,
+) where {T}
+    return AdjacencyGraph{T,has_diagonal}(S, edge_to_index)
 end
 
 function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
@@ -204,24 +244,28 @@ function AdjacencyGraph(A::AbstractMatrix; has_diagonal::Bool=true)
 end
 
 pattern(g::AdjacencyGraph) = g.S
+edge_indices(g::AdjacencyGraph) = g.edge_to_index
 nb_vertices(g::AdjacencyGraph) = pattern(g).n
 vertices(g::AdjacencyGraph) = 1:nb_vertices(g)
+has_diagonal(::AdjacencyGraph{T,hd}) where {T,hd} = hd
 
-has_diagonal(::AdjacencyGraph{T,hl}) where {T,hl} = hl
-
-function neighbors(g::AdjacencyGraph{T,true}, v::Integer) where {T}
+function neighbors(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
     neighbors_v = view(rowvals(S), nzrange(S, v))
-    return Iterators.filter(!=(v), neighbors_v)  # TODO: optimize
+    return neighbors_v
 end
 
-function neighbors(g::AdjacencyGraph{T,false}, v::Integer) where {T}
+function neighbors_with_edge_indices(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
-    neighbors_v = view(rowvals(S), nzrange(S, v))
-    return neighbors_v
+    range_v = nzrange(S, v)
+    neighbors_v = view(rowvals(S), range_v)
+    edges_indices_v = view(edge_indices(g), range_v)
+    return zip(neighbors_v, edges_indices_v)
 end
 
-function degree(g::AdjacencyGraph, v::Integer)
+degree(g::AdjacencyGraph{T,false}, v::Integer) where {T} = g.S.colptr[v + 1] - g.S.colptr[v]
+
+function degree(g::AdjacencyGraph{T,true}, v::Integer) where {T}
     d = 0
     for u in neighbors(g, v)
         if u != v
@@ -231,12 +275,12 @@ function degree(g::AdjacencyGraph, v::Integer)
     return d
 end
 
-function nb_edges(g::AdjacencyGraph)
+nb_edges(g::AdjacencyGraph{T,false}) where {T} = nnz(g.S) ÷ 2
+
+function nb_edges(g::AdjacencyGraph{T,true}) where {T}
     ne = 0
     for v in vertices(g)
-        for u in neighbors(g, v)
-            ne += 1
-        end
+        ne += degree(g, v)
     end
     return ne ÷ 2
 end
@@ -246,6 +290,7 @@ minimum_degree(g::AdjacencyGraph) = minimum(Base.Fix1(degree, g), vertices(g))
 
 function has_neighbor(g::AdjacencyGraph, v::Integer, u::Integer)
     for w in neighbors(g, v)
+        !has_diagonal(g) || (v == w && continue)
         if w == u
             return true
         end
@@ -280,7 +325,7 @@ A `BipartiteGraph` has two sets of vertices, one for the rows of `A` (which we c
 
 # Constructors
 
-    BipartiteGraph(A::SparseMatrixCSC; symmetric_pattern=false)
+    BipartiteGraph(A::SparseMatrixCSC; symmetric_pattern::Bool=false)
 
 When `symmetric_pattern` is `true`, this construction is more efficient.
 
@@ -293,7 +338,7 @@ When `symmetric_pattern` is `true`, this construction is more efficient.
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct BipartiteGraph{T<:Integer}
+struct BipartiteGraph{T}
     S1::SparsityPatternCSC{T}
     S2::SparsityPatternCSC{T}
 end
diff --git a/src/interface.jl b/src/interface.jl
index e1c1b0d2..3997993d 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -265,7 +265,7 @@ function _coloring(
     decompression_eltype::Type,
     symmetric_pattern::Bool,
 )
-    ag = AdjacencyGraph(A)
+    ag = AdjacencyGraph(A; has_diagonal=true)
     color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         return StarSetColoringResult(A, ag, color, star_set)
@@ -282,7 +282,7 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    ag = AdjacencyGraph(A)
+    ag = AdjacencyGraph(A; has_diagonal=true)
     color, tree_set = acyclic_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         return TreeSetColoringResult(A, ag, color, tree_set, R)
@@ -299,8 +299,8 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
+    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
     color, star_set = star_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
@@ -319,8 +319,8 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    A_and_Aᵀ = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ; has_diagonal=false)
+    A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
     color, tree_set = acyclic_coloring(ag, algo.order, algo.postprocessing)
     if speed_setting isa WithResult
         symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
diff --git a/src/order.jl b/src/order.jl
index 1f79ba1a..88bc4f96 100644
--- a/src/order.jl
+++ b/src/order.jl
@@ -262,6 +262,7 @@ function vertices(
         u = pop_next_candidate!(db; direction)
         direction == :low2high ? push!(π, u) : pushfirst!(π, u)
         for v in neighbors(g, u)
+            !has_diagonal(g) || (u == v && continue)
             already_ordered(db, v) && continue
             update_bucket!(db, v; degtype, direction)
         end
diff --git a/src/result.jl b/src/result.jl
index 8bbefddf..ae8d7ce1 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -241,19 +241,40 @@ end
 function StarSetColoringResult(
     A::AbstractMatrix, ag::AdjacencyGraph, color::Vector{Int}, star_set::StarSet
 )
-    S = ag.S
+    (; star, hub) = star_set
+    S = pattern(ag)
+    edge_to_index = edge_indices(ag)
+    n = S.n
     group = group_by_color(color)
-    n = size(S, 1)
-    rv = rowvals(S)
-    compressed_indices = zeros(Int, nnz(S))
+    rvS = rowvals(S)
+    compressed_indices = zeros(Int, nnz(S))  # needs to be independent from the storage in the graph, in case the graph gets reused
     for j in axes(S, 2)
         for k in nzrange(S, j)
-            i = rv[k]
-            l, c = symmetric_coefficient(i, j, color, star_set)
-            # A[i, j] = B[l, c]
-            compressed_indices[k] = (c - 1) * n + l
+            i = rvS[k]
+            if i == j
+                # diagonal coefficients
+                c = color[i]
+                compressed_indices[k] = (c - 1) * n + i
+            else
+                # off-diagonal coefficients
+                index_ij = edge_to_index[k]
+                s = star[index_ij]
+                h = abs(hub[s])
+
+                # Assign the non-hub vertex (spoke) to the correct position in spokes
+                if i == h
+                    # i is the hub and j is the spoke
+                    c = color[i]
+                    compressed_indices[k] = (c - 1) * n + j
+                else  # j == h
+                    # j is the hub and i is the spoke
+                    c = color[j]
+                    compressed_indices[k] = (c - 1) * n + i
+                end
+            end
         end
     end
+
     return StarSetColoringResult(A, ag, color, group, star_set, compressed_indices)
 end
 
diff --git a/test/forest.jl b/test/forest.jl
index dc5e6bf0..5a3be8d2 100644
--- a/test/forest.jl
+++ b/test/forest.jl
@@ -3,74 +3,41 @@ using Test
 
 @testset "Constructor Forest" begin
     forest = Forest{Int}(5)
-
-    @test forest.num_edges == 0
-    @test forest.num_trees == 0
-    @test length(forest.intmap) == 0
+    @test forest.num_trees == 5
     @test length(forest.parents) == 5
     @test all(forest.parents .== 1:5)
     @test all(forest.ranks .== 0)
 end
 
-@testset "Push edge" begin
-    forest = Forest{Int}(5)
-
-    push!(forest, (1, 2))
-    @test forest.num_edges == 1
-    @test forest.num_trees == 1
-    @test haskey(forest.intmap, (1, 2))
-    @test forest.intmap[(1, 2)] == 1
-    @test forest.num_trees == 1
-
-    push!(forest, (3, 4))
-    @test forest.num_edges == 2
-    @test forest.num_trees == 2
-    @test haskey(forest.intmap, (3, 4))
-    @test forest.intmap[(3, 4)] == 2
-    @test forest.num_trees == 2
-end
-
 @testset "Find root" begin
     forest = Forest{Int}(5)
-    push!(forest, (1, 2))
-    push!(forest, (3, 4))
-
-    @test find_root!(forest, (1, 2)) == 1
-    @test find_root!(forest, (3, 4)) == 2
+    @test find_root!(forest, 3) == 3
+    @test find_root!(forest, 5) == 5
 end
 
 @testset "Root union" begin
     forest = Forest{Int}(5)
-    push!(forest, (1, 2))
-    push!(forest, (4, 5))
-    push!(forest, (2, 4))
-    @test forest.num_trees == 3
-
-    root1 = find_root!(forest, (1, 2))
-    root3 = find_root!(forest, (2, 4))
+    root1 = find_root!(forest, 1)
+    root3 = find_root!(forest, 3)
     @test root1 != root3
 
     root_union!(forest, root1, root3)
-    @test find_root!(forest, (2, 4)) == 1
+    @test find_root!(forest, 3) == 1
     @test forest.parents[1] == 1
     @test forest.parents[3] == 1
     @test forest.ranks[1] == 1
     @test forest.ranks[3] == 0
-    @test forest.num_trees == 2
+    @test forest.num_trees == 4
 
-    root1 = find_root!(forest, (1, 2))
-    root2 = find_root!(forest, (4, 5))
+    root1 = find_root!(forest, 1)
+    root2 = find_root!(forest, 2)
     @test root1 != root2
+
     root_union!(forest, root1, root2)
-    @test find_root!(forest, (4, 5)) == 1
+    @test find_root!(forest, 2) == 1
     @test forest.parents[1] == 1
     @test forest.parents[2] == 1
     @test forest.ranks[1] == 1
     @test forest.ranks[2] == 0
-    @test forest.num_trees == 1
-
-    push!(forest, (1, 4))
-    @test forest.num_trees == 2
-    @test forest.intmap[(1, 4)] == 4
-    @test forest.parents[4] == 4
+    @test forest.num_trees == 3
 end
diff --git a/test/graph.jl b/test/graph.jl
index 71e96792..477ebce4 100644
--- a/test/graph.jl
+++ b/test/graph.jl
@@ -16,32 +16,47 @@ using Test
 
 @testset "SparsityPatternCSC" begin
     @testset "Transpose" begin
-        @test all(1:1000) do _
+        for _ in 1:1000
             m, n = rand(100:1000), rand(100:1000)
             p = 0.05 * rand()
             A = sprand(m, n, p)
             S = SparsityPatternCSC(A)
             Sᵀ = transpose(S)
             Sᵀ_true = SparsityPatternCSC(sparse(transpose(A)))
-            Sᵀ.colptr == Sᵀ_true.colptr && Sᵀ.rowval == Sᵀ_true.rowval
+            @test Sᵀ.colptr == Sᵀ_true.colptr
+            @test Sᵀ.rowval == Sᵀ_true.rowval
         end
     end
     @testset "Bidirectional" begin
-        @test all(1:1000) do _
-            m, n = rand(100:1000), rand(100:1000)
-            p = 0.05 * rand()
-            A = sprand(Bool, m, n, p)
-            A_and_Aᵀ = [spzeros(Bool, n, n) transpose(A); A spzeros(Bool, m, m)]
-            S_and_Sᵀ = bidirectional_pattern(A; symmetric_pattern=false)
-            S_and_Sᵀ.colptr == A_and_Aᵀ.colptr && S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+        @testset "symmetric_pattern = false" begin
+            for _ in 1:1000
+                m, n = rand(100:1000), rand(100:1000)
+                p = 0.05 * rand()
+                A = sprand(Bool, m, n, p)
+                A_and_Aᵀ = [spzeros(Bool, n, n) transpose(A); A spzeros(Bool, m, m)]
+                S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=false)
+                @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr
+                @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+                M = SparseMatrixCSC(
+                    m + n, m + n, S_and_Sᵀ.colptr, S_and_Sᵀ.rowval, edge_to_index
+                )
+                @test issymmetric(M)
+            end
         end
-        @test all(1:1000) do _
-            m = rand(100:1000)
-            p = 0.05 * rand()
-            A = sparse(Symmetric(sprand(Bool, m, m, p)))
-            A_and_Aᵀ = [spzeros(Bool, m, m) transpose(A); A spzeros(Bool, m, m)]
-            S_and_Sᵀ = bidirectional_pattern(A; symmetric_pattern=true)
-            S_and_Sᵀ.colptr == A_and_Aᵀ.colptr && S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+        @testset "symmetric_pattern = true" begin
+            for _ in 1:1000
+                m = rand(100:1000)
+                p = 0.05 * rand()
+                A = sparse(Symmetric(sprand(Bool, m, m, p)))
+                A_and_Aᵀ = [spzeros(Bool, m, m) transpose(A); A spzeros(Bool, m, m)]
+                S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=true)
+                @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr
+                @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval
+                M = SparseMatrixCSC(
+                    2 * m, 2 * m, S_and_Sᵀ.colptr, S_and_Sᵀ.rowval, edge_to_index
+                )
+                @test issymmetric(M)
+            end
         end
     end
     @testset "size" begin
@@ -127,15 +142,35 @@ end;
         0 0 0 1 1 1 1 0
     ])
 
-    B = transpose(A) * A
-    g = AdjacencyGraph(B - Diagonal(B))
+    g = AdjacencyGraph(transpose(A) * A)
     @test nb_vertices(g) == 8
-    @test collect(neighbors(g, 1)) == [6, 7, 8]
-    @test collect(neighbors(g, 2)) == [5, 7, 8]
-    @test collect(neighbors(g, 3)) == [5, 6, 8]
-    @test collect(neighbors(g, 4)) == [5, 6, 7]
-    @test collect(neighbors(g, 5)) == [2, 3, 4, 6, 7, 8]
-    @test collect(neighbors(g, 6)) == [1, 3, 4, 5, 7, 8]
-    @test collect(neighbors(g, 7)) == [1, 2, 4, 5, 6, 8]
-    @test collect(neighbors(g, 8)) == [1, 2, 3, 5, 6, 7]
+    # wrong neighbors, it's okay they are filtered after
+    @test collect(neighbors(g, 1)) == sort(vcat(1, [6, 7, 8]))
+    @test collect(neighbors(g, 2)) == sort(vcat(2, [5, 7, 8]))
+    @test collect(neighbors(g, 3)) == sort(vcat(3, [5, 6, 8]))
+    @test collect(neighbors(g, 4)) == sort(vcat(4, [5, 6, 7]))
+    @test collect(neighbors(g, 5)) == sort(vcat(5, [2, 3, 4, 6, 7, 8]))
+    @test collect(neighbors(g, 6)) == sort(vcat(6, [1, 3, 4, 5, 7, 8]))
+    @test collect(neighbors(g, 7)) == sort(vcat(7, [1, 2, 4, 5, 6, 8]))
+    @test collect(neighbors(g, 8)) == sort(vcat(8, [1, 2, 3, 5, 6, 7]))
+    # right degree
+    @test degree(g, 1) == 3
+    @test degree(g, 2) == 3
+    @test degree(g, 3) == 3
+    @test degree(g, 4) == 3
+    @test degree(g, 5) == 6
+    @test degree(g, 6) == 6
+    @test degree(g, 7) == 6
+    @test degree(g, 8) == 6
+
+    g = AdjacencyGraph(transpose(A) * A; has_diagonal=false)
+    # wrong degree
+    @test degree(g, 1) == 4
+    @test degree(g, 2) == 4
+    @test degree(g, 3) == 4
+    @test degree(g, 4) == 4
+    @test degree(g, 5) == 7
+    @test degree(g, 6) == 7
+    @test degree(g, 7) == 7
+    @test degree(g, 8) == 7
 end
diff --git a/test/order.jl b/test/order.jl
index efd030ed..789157b5 100644
--- a/test/order.jl
+++ b/test/order.jl
@@ -61,14 +61,15 @@ end;
 end;
 
 @testset "LargestFirst" begin
-    A = sparse([
-        0 1 0
-        1 0 1
-        0 1 0
-    ])
-    ag = AdjacencyGraph(A)
-
-    @test vertices(ag, LargestFirst()) == [2, 1, 3]
+    for has_diagonal in (false, true)
+        A = sparse([
+            0 1 0
+            1 0 1
+            0 1 0
+        ])
+        ag = AdjacencyGraph(A; has_diagonal)
+        @test vertices(ag, LargestFirst()) == [2, 1, 3]
+    end
 
     A = sparse([
         1 1 0 0