From 889f1953c2e43e07ef95c3db468992b499461464 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 12:16:11 +1300
Subject: [PATCH 01/17] Initial code by Claude

---
 Project.toml                     |   2 +
 src/ArrayDiff.jl                 |   4 +-
 src/Coloring/topological_sort.jl | 106 ------------------
 src/coloring_compat.jl           | 184 +++++++++++++++++++++++++++++++
 src/forward_over_reverse.jl      |   4 +-
 src/graph_tools.jl               |   8 +-
 src/indexed_set.jl               |  62 +++++++++++
 src/mathoptinterface_api.jl      |   2 +-
 src/types.jl                     |  10 +-
 test/ArrayDiff.jl                |   1 -
 test/ReverseAD.jl                | 138 ++++++++++++-----------
 11 files changed, 334 insertions(+), 187 deletions(-)
 delete mode 100644 src/Coloring/topological_sort.jl
 create mode 100644 src/coloring_compat.jl
 create mode 100644 src/indexed_set.jl

diff --git a/Project.toml b/Project.toml
index 014259e..117cc9e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,6 +9,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 
 [compat]
 DataStructures = "0.18, 0.19"
@@ -16,4 +17,5 @@ ForwardDiff = "1"
 MathOptInterface = "1.40"
 NaNMath = "1"
 SparseArrays = "1.10"
+SparseMatrixColorings = "1"
 julia = "1.10"
diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 1f08c1e..588efce 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -10,6 +10,7 @@ import ForwardDiff
 import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 import SparseArrays
+import SparseMatrixColorings
 
 """
     Mode() <: AbstractAutomaticDifferentiation
@@ -48,7 +49,8 @@ import NaNMath:
     pow,
     sqrt
 
-include("Coloring/Coloring.jl")
+include("indexed_set.jl")
+include("coloring_compat.jl")
 include("graph_tools.jl")
 include("sizes.jl")
 include("types.jl")
diff --git a/src/Coloring/topological_sort.jl b/src/Coloring/topological_sort.jl
deleted file mode 100644
index 0b8d0ca..0000000
--- a/src/Coloring/topological_sort.jl
+++ /dev/null
@@ -1,106 +0,0 @@
-# Modified from Graphs.jl.
-# Workaround for edge_index not being O(1) on SimpleGraph.
-# edge_index was only needed to test for cycles, so
-# this implementation skips that check.
-
-# Graphs.jl is licensed under the MIT License:
-#
-# Copyright (c) 2012: John Myles White and other contributors.
-#
-# Permission is hereby granted, free of charge, to any person obtaining
-# a copy of this software and associated documentation files (the
-# "Software"), to deal in the Software without restriction, including
-# without limitation the rights to use, copy, modify, merge, publish,
-# distribute, sublicense, and/or sell copies of the Software, and to
-# permit persons to whom the Software is furnished to do so, subject to
-# the following conditions:
-#
-# The above copyright notice and this permission notice shall be
-# included in all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
-# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
-# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
-# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
-# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
-
-mutable struct _TopologicalSortVisitor
-    vertices::Vector{Int}
-    parents::Vector{Int}
-
-    function _TopologicalSortVisitor(n::Int)
-        vs = Int[]
-        sizehint!(vs, n)
-        return new(vs, zeros(Int, n))
-    end
-end
-
-function _traverse_graph(
-    adjlist,
-    offsets,
-    s,
-    visitor::_TopologicalSortVisitor,
-    vertexcolormap,
-    vertex_stack,
-    index_stack,
-)
-    vertexcolormap[s] = 1
-    resize!(vertex_stack, 1)
-    vertex_stack[1] = s
-    resize!(index_stack, 1)
-    index_stack[1] = 1
-    while !isempty(vertex_stack)
-        u = pop!(vertex_stack)
-        out_idx = pop!(index_stack)
-        len_uegs = offsets[u+1] - offsets[u]
-        found_new_vertex = false
-        while out_idx <= len_uegs && !found_new_vertex
-            v = adjlist[offsets[u]+out_idx-1]
-            out_idx += 1
-            v_color = vertexcolormap[v]
-            if v_color == 0
-                found_new_vertex = true
-                vertexcolormap[v] = 1
-                push!(vertex_stack, u)
-                push!(index_stack, out_idx)
-                visitor.parents[v] = u
-                push!(vertex_stack, v)
-                push!(index_stack, 1)
-            end
-        end
-        if !found_new_vertex
-            push!(visitor.vertices, u)
-            vertexcolormap[u] = 2
-        end
-    end
-    return
-end
-
-function reverse_topological_sort_by_dfs(
-    adjlist,
-    offsets,
-    num_vertices,
-    cmap::Vector{Int},
-    vertex_stack::Vector{Int} = Int[],
-    index_stack::Vector{Int} = Int[],
-)
-    @assert length(cmap) == num_vertices
-    fill!(cmap, 0)
-    visitor = _TopologicalSortVisitor(num_vertices)
-    for s in 1:num_vertices
-        if cmap[s] == 0
-            _traverse_graph(
-                adjlist,
-                offsets,
-                s,
-                visitor,
-                cmap,
-                vertex_stack,
-                index_stack,
-            )
-        end
-    end
-    return visitor.vertices, visitor.parents
-end
diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
new file mode 100644
index 0000000..02f90e8
--- /dev/null
+++ b/src/coloring_compat.jl
@@ -0,0 +1,184 @@
+# Copyright (c) 2017: Miles Lubin and contributors
+# Copyright (c) 2017: Google Inc.
+#
+# Use of this source code is governed by an MIT-style license that can be found
+# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
+
+"""
+    struct ColoringResult
+        result::SparseMatrixColorings.TreeSetColoringResult
+        local_indices::Vector{Int}  # map from local to global indices
+    end
+
+Wrapper around TreeSetColoringResult that also stores local_indices mapping.
+"""
+struct ColoringResult
+    result::SparseMatrixColorings.TreeSetColoringResult
+    local_indices::Vector{Int}  # map from local to global indices
+end
+
+"""
+    _hessian_color_preprocess(
+        edgelist,
+        num_total_var,
+        seen_idx = IndexedSet(0),
+    )
+
+`edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the
+diagonal.
+
+Returns `(I, J, result)` where `I` and `J` are the row and column indices
+of the Hessian structure, and `result` is a `TreeSetColoringResult` from
+SparseMatrixColorings.
+"""
+function _hessian_color_preprocess(
+    edgelist,
+    num_total_var,
+    seen_idx = IndexedSet(0),
+)
+    resize!(seen_idx, num_total_var)
+    I, J = Int[], Int[]
+    for (i, j) in edgelist
+        push!(seen_idx, i)
+        push!(seen_idx, j)
+        push!(I, i)
+        push!(J, j)
+    end
+    local_indices = sort!(collect(seen_idx))
+    empty!(seen_idx)
+    global_to_local_idx = seen_idx.nzidx # steal for storage
+    for k in eachindex(local_indices)
+        global_to_local_idx[local_indices[k]] = k
+    end
+    # only do the coloring on the local indices
+    for k in eachindex(I)
+        I[k] = global_to_local_idx[I[k]]
+        J[k] = global_to_local_idx[J[k]]
+    end
+    
+    # Create sparsity pattern matrix
+    n = length(local_indices)
+    S = SparseArrays.spzeros(Bool, n, n)
+    for k in eachindex(I)
+        i, j = I[k], J[k]
+        S[i, j] = true
+        S[j, i] = true  # symmetric
+    end
+    
+    # Perform coloring using SparseMatrixColorings
+    problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
+    tree_result = SparseMatrixColorings.coloring(S, problem, algo)
+    
+    # Convert back to global indices
+    for k in eachindex(I)
+        I[k] = local_indices[I[k]]
+        J[k] = local_indices[J[k]]
+    end
+    
+    # Wrap result with local_indices
+    result = ColoringResult(tree_result, local_indices)
+    return I, J, result
+end
+
+"""
+    _seed_matrix(result::ColoringResult)
+
+Allocate a seed matrix for the coloring result.
+"""
+function _seed_matrix(result::ColoringResult)
+    n = length(result.local_indices)
+    ncolors = SparseMatrixColorings.ncolors(result.result)
+    return Matrix{Float64}(undef, n, ncolors)
+end
+
+"""
+    _prepare_seed_matrix!(R, result::ColoringResult)
+
+Prepare the seed matrix R for Hessian computation.
+"""
+function _prepare_seed_matrix!(R, result::ColoringResult)
+    color = SparseMatrixColorings.column_colors(result.result)
+    N = length(result.local_indices)
+    @assert N == size(R, 1)
+    @assert size(R, 2) == SparseMatrixColorings.ncolors(result.result)
+    fill!(R, 0.0)
+    for i in 1:N
+        if color[i] > 0
+            R[i, color[i]] = 1
+        end
+    end
+    return
+end
+
+"""
+    _recover_from_matmat!(
+        V::AbstractVector{T},
+        R::AbstractMatrix{T},
+        result::ColoringResult,
+        stored_values::AbstractVector{T},
+    ) where {T}
+
+Recover the Hessian values from the Hessian-matrix product H*R_seed.
+R is the result of H*R_seed where R_seed is the seed matrix.
+`stored_values` is a temporary vector.
+"""
+function _recover_from_matmat!(
+    V::AbstractVector{T},
+    R::AbstractMatrix{T},
+    result::ColoringResult,
+    stored_values::AbstractVector{T},
+) where {T}
+    tree_result = result.result
+    color = SparseMatrixColorings.column_colors(tree_result)
+    N = length(result.local_indices)
+    S = tree_result.ag.S
+    # Compute number of off-diagonal nonzeros from the length of V
+    # V contains N diagonal elements + nnz_offdiag off-diagonal elements
+    nnz_offdiag = length(V) - N
+    @assert length(stored_values) >= N
+    
+    # Recover diagonal elements
+    k = 0
+    for i in 1:N
+        k += 1
+        if color[i] > 0
+            V[k] = R[i, color[i]]
+        else
+            V[k] = zero(T)
+        end
+    end
+    
+    # Recover off-diagonal elements using tree structure
+    (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
+    fill!(stored_values, zero(T))
+    
+    for tree_idx in 1:nt
+        first = tree_edge_indices[tree_idx]
+        last = tree_edge_indices[tree_idx + 1] - 1
+        
+        # Reset stored_values for vertices in this tree
+        for pos in first:last
+            (vertex, _) = reverse_bfs_orders[pos]
+            stored_values[vertex] = zero(T)
+        end
+        (_, root) = reverse_bfs_orders[last]
+        stored_values[root] = zero(T)
+        
+        # Recover edge values
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
+            if color[j] > 0
+                value = R[i, color[j]] - stored_values[i]
+            else
+                value = zero(T)
+            end
+            stored_values[j] += value
+            k += 1
+            V[k] = value
+        end
+    end
+    
+    @assert k == length(V)
+    return
+end
diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl
index ea4c7e2..c69854a 100644
--- a/src/forward_over_reverse.jl
+++ b/src/forward_over_reverse.jl
@@ -43,7 +43,7 @@ function _eval_hessian(
         return 0
     end
     chunk = min(size(ex.seed_matrix, 2), d.max_chunk)
-    Coloring.prepare_seed_matrix!(ex.seed_matrix, ex.rinfo)
+    _prepare_seed_matrix!(ex.seed_matrix, ex.rinfo)
     # Compute hessian-vector products
     num_products = size(ex.seed_matrix, 2) # number of hessian-vector products
     num_chunks = div(num_products, chunk)
@@ -66,7 +66,7 @@ function _eval_hessian(
     end
     # TODO(odow): consider reverting to a view.
     output_slice = _UnsafeVectorView(nzcount, length(ex.hess_I), pointer(H))
-    Coloring.recover_from_matmat!(
+    _recover_from_matmat!(
         output_slice,
         ex.seed_matrix,
         ex.rinfo,
diff --git a/src/graph_tools.jl b/src/graph_tools.jl
index 7e9c366..b91e0d7 100644
--- a/src/graph_tools.jl
+++ b/src/graph_tools.jl
@@ -151,7 +151,7 @@ end
 
 """
     _compute_gradient_sparsity!(
-        indices::Coloring.IndexedSet,
+        indices::IndexedSet,
         nodes::Vector{Nonlinear.Node},
     )
 
@@ -159,7 +159,7 @@ Compute the sparsity pattern of the gradient of an expression (that is, a list o
 which variable indices are present).
 """
 function _compute_gradient_sparsity!(
-    indices::Coloring.IndexedSet,
+    indices::IndexedSet,
     nodes::Vector{Nonlinear.Node},
 )
     for node in nodes
@@ -180,7 +180,7 @@ end
         nodes::Vector{Nonlinear.Node},
         adj,
         input_linearity::Vector{Linearity},
-        indexedset::Coloring.IndexedSet,
+        indexedset::IndexedSet,
         subexpression_edgelist::Vector{Set{Tuple{Int,Int}}},
         subexpression_variables::Vector{Vector{Int}},
     )
@@ -207,7 +207,7 @@ function _compute_hessian_sparsity(
     nodes::Vector{Nonlinear.Node},
     adj,
     input_linearity::Vector{Linearity},
-    indexedset::Coloring.IndexedSet,
+    indexedset::IndexedSet,
     subexpression_edgelist::Vector{Set{Tuple{Int,Int}}},
     subexpression_variables::Vector{Vector{Int}},
 )
diff --git a/src/indexed_set.jl b/src/indexed_set.jl
new file mode 100644
index 0000000..b2bb631
--- /dev/null
+++ b/src/indexed_set.jl
@@ -0,0 +1,62 @@
+# Copyright (c) 2017: Miles Lubin and contributors
+# Copyright (c) 2017: Google Inc.
+#
+# Use of this source code is governed by an MIT-style license that can be found
+# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
+
+"""
+    mutable struct IndexedSet
+        nzidx::Vector{Int}
+        empty::BitVector
+        nnz::Int
+    end
+
+Represent the set `nzidx[1:nnz]` by additionally setting `empty[i]` to `false`
+for each element `i` of the set for fast membership check.
+"""
+mutable struct IndexedSet
+    nzidx::Vector{Int}
+    empty::BitVector
+    nnz::Int
+end
+
+IndexedSet(n::Integer) = IndexedSet(zeros(Int, n), trues(n), 0)
+
+function Base.push!(v::IndexedSet, i::Integer)
+    if v.empty[i]  # new index
+        v.nzidx[v.nnz+=1] = i
+        v.empty[i] = false
+    end
+    return
+end
+
+function Base.empty!(v::IndexedSet)
+    for i in 1:v.nnz
+        v.empty[v.nzidx[i]] = true
+    end
+    v.nnz = 0
+    return v
+end
+
+# Returns the maximum index that the set can contain,
+# not the cardinality of the set like `length(::Base.Set)`
+Base.length(v::IndexedSet) = length(v.nzidx)
+
+function Base.resize!(v::IndexedSet, n::Integer)
+    if n > length(v)
+        @assert v.nnz == 0 # only resize empty vector
+        resize!(v.nzidx, n)
+        resize!(v.empty, n)
+        fill!(v.empty, true)
+    end
+    return
+end
+
+Base.collect(v::IndexedSet) = v.nzidx[1:v.nnz]
+
+function Base.union!(v::IndexedSet, s)
+    for x in s
+        push!(v, x)
+    end
+    return
+end
diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl
index 21c6804..31995ca 100644
--- a/src/mathoptinterface_api.jl
+++ b/src/mathoptinterface_api.jl
@@ -42,7 +42,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     d.last_x = fill(NaN, N)
     d.want_hess = :Hess in requested_features
     want_hess_storage = (:HessVec in requested_features) || d.want_hess
-    coloring_storage = Coloring.IndexedSet(N)
+    coloring_storage = IndexedSet(N)
     max_expr_length = 0
     max_expr_with_sub_length = 0
     #
diff --git a/src/types.jl b/src/types.jl
index 4f1f19a..5edae64 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -69,7 +69,7 @@ struct _FunctionStorage
     # Nonzero pattern of Hessian matrix
     hess_I::Vector{Int}
     hess_J::Vector{Int}
-    rinfo::Coloring.RecoveryInfo # coloring info for hessians
+    rinfo::Union{ColoringResult,Nothing} # coloring info for hessians
     seed_matrix::Matrix{Float64}
     # subexpressions which this function depends on, ordered for forward pass.
     dependent_subexpressions::Vector{Int}
@@ -77,7 +77,7 @@ struct _FunctionStorage
     function _FunctionStorage(
         expr::_SubexpressionStorage,
         num_variables,
-        coloring_storage::Coloring.IndexedSet,
+        coloring_storage::IndexedSet,
         want_hess::Bool,
         subexpressions::Vector{_SubexpressionStorage},
         dependent_subexpressions,
@@ -104,12 +104,12 @@ struct _FunctionStorage
                 subexpression_edgelist,
                 subexpression_variables,
             )
-            hess_I, hess_J, rinfo = Coloring.hessian_color_preprocess(
+            hess_I, hess_J, rinfo = _hessian_color_preprocess(
                 edgelist,
                 num_variables,
                 coloring_storage,
             )
-            seed_matrix = Coloring.seed_matrix(rinfo)
+            seed_matrix = _seed_matrix(rinfo)
             return new(
                 expr,
                 grad_sparsity,
@@ -125,7 +125,7 @@ struct _FunctionStorage
                 grad_sparsity,
                 Int[],
                 Int[],
-                Coloring.RecoveryInfo(),
+                nothing,
                 Array{Float64}(undef, 0, 0),
                 dependent_subexpressions,
             )
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 4b90f3b..5e894ec 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -8,7 +8,6 @@ import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 
 import ArrayDiff
-const Coloring = ArrayDiff.Coloring
 
 function runtests()
     for name in names(@__MODULE__; all = true)
diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl
index 552ff3f..60e06d3 100644
--- a/test/ReverseAD.jl
+++ b/test/ReverseAD.jl
@@ -14,7 +14,7 @@ import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 
 import ArrayDiff
-const Coloring = ArrayDiff.Coloring
+# Coloring submodule removed - using SparseMatrixColorings instead
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -414,79 +414,83 @@ struct _ColoringGraph
     edges::Vector{Tuple{Int,Int}}
 end
 
-function to_adjlist(graph::_ColoringGraph)
-    I = [i for (i, _) in graph.edges]
-    J = [j for (_, j) in graph.edges]
-    return Coloring.UndirectedGraph(I, J, graph.num_vertices)
-end
-
-function test_coloring_edge_free_graph()
-    graph = _ColoringGraph(10, [])
-    _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 1
-    return
-end
-
-function test_coloring_one_edge_graph()
-    graph = _ColoringGraph(10, [(2, 4)])
-    color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 2
-    @test color[2] != color[4]
-    return
-end
-
-function test_coloring_two_edge_graph()
-    graph = _ColoringGraph(10, [(2, 4), (2, 3)])
-    color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 2
-    @test color[3] == color[4]
-    return
-end
-
-function test_coloring_three_edge_graph()
-    graph = _ColoringGraph(10, [(2, 4), (2, 3), (3, 4)])
-    color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 3
-    # TODO: What is this testing?
-    Coloring.recovery_preprocess(to_adjlist(graph), color, numcolors, Int[])
-    return
-end
-
-function test_coloring_two_edge_three_vertex_graph()
-    graph = _ColoringGraph(3, [(1, 3), (2, 3)])
-    _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 2
-    return
-end
-
-function test_coloring_four_edge_four_vertex_graph()
-    graph = _ColoringGraph(4, [(1, 2), (2, 3), (3, 4), (4, 1)])
-    _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
-    @test numcolors == 3
-    return
-end
-
-function test_coloring_topological_sort()
-    # graph = _ColoringGraph(6, [(1, 2), (1, 3), (1, 6), (2, 4), (2, 5)])
-    vec = [3, 6, 2, 1, 4, 5, 1, 2, 2, 1]
-    offset = [1, 4, 7, 8, 9, 10, 11]
-    v = Coloring.reverse_topological_sort_by_dfs(vec, offset, 6, zeros(Int, 6))
-    @test v[1] == [3, 6, 4, 5, 2, 1]
-    @test v[2] == [0, 1, 1, 2, 2, 1]
-    return
-end
+# Tests for internal Coloring submodule removed - now using SparseMatrixColorings
+# The coloring functionality is tested indirectly through the Hessian evaluation tests
+
+# function to_adjlist(graph::_ColoringGraph)
+#     I = [i for (i, _) in graph.edges]
+#     J = [j for (_, j) in graph.edges]
+#     return Coloring.UndirectedGraph(I, J, graph.num_vertices)
+# end
+
+# function test_coloring_edge_free_graph()
+#     graph = _ColoringGraph(10, [])
+#     _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 1
+#     return
+# end
+
+# function test_coloring_one_edge_graph()
+#     graph = _ColoringGraph(10, [(2, 4)])
+#     color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 2
+#     @test color[2] != color[4]
+#     return
+# end
+
+# function test_coloring_two_edge_graph()
+#     graph = _ColoringGraph(10, [(2, 4), (2, 3)])
+#     color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 2
+#     @test color[3] == color[4]
+#     return
+# end
+
+# function test_coloring_three_edge_graph()
+#     graph = _ColoringGraph(10, [(2, 4), (2, 3), (3, 4)])
+#     color, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 3
+#     # TODO: What is this testing?
+#     Coloring.recovery_preprocess(to_adjlist(graph), color, numcolors, Int[])
+#     return
+# end
+
+# function test_coloring_two_edge_three_vertex_graph()
+#     graph = _ColoringGraph(3, [(1, 3), (2, 3)])
+#     _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 2
+#     return
+# end
+
+# function test_coloring_four_edge_four_vertex_graph()
+#     graph = _ColoringGraph(4, [(1, 2), (2, 3), (3, 4), (4, 1)])
+#     _, numcolors = Coloring.acyclic_coloring(to_adjlist(graph))
+#     @test numcolors == 3
+#     return
+# end
+
+# function test_coloring_topological_sort()
+#     # graph = _ColoringGraph(6, [(1, 2), (1, 3), (1, 6), (2, 4), (2, 5)])
+#     vec = [3, 6, 2, 1, 4, 5, 1, 2, 2, 1]
+#     offset = [1, 4, 7, 8, 9, 10, 11]
+#     v = Coloring.reverse_topological_sort_by_dfs(vec, offset, 6, zeros(Int, 6))
+#     @test v[1] == [3, 6, 4, 5, 2, 1]
+#     @test v[2] == [0, 1, 1, 2, 2, 1]
+#     return
+# end
 
 function test_coloring_end_to_end_hessian_coloring_and_recovery()
-    I, J, rinfo = Coloring.hessian_color_preprocess(Set([(1, 2)]), 2)
-    R = Coloring.seed_matrix(rinfo)
-    Coloring.prepare_seed_matrix!(R, rinfo)
+    # Test the new coloring API through the compatibility layer
+    I, J, rinfo = ArrayDiff._hessian_color_preprocess(Set([(1, 2)]), 2, ArrayDiff.IndexedSet(0))
+    R = ArrayDiff._seed_matrix(rinfo)
+    ArrayDiff._prepare_seed_matrix!(R, rinfo)
     @test I == [1, 2, 2]
     @test J == [1, 2, 1]
     @test R == [1.0 0.0; 0.0 1.0]
     hess = [3.4 2.1; 2.1 1.3]
     matmat = hess * R
     V = zeros(3)
-    Coloring.recover_from_matmat!(V, matmat, rinfo, zeros(3))
+    ArrayDiff._recover_from_matmat!(V, matmat, rinfo, zeros(3))
     @test V == [3.4, 1.3, 2.1]
     return
 end
@@ -547,7 +551,7 @@ function test_linearity()
         nodes = ArrayDiff._replace_moi_variables(expr.nodes, variables)
         ret = ArrayDiff._classify_linearity(nodes, adj, ArrayDiff.Linearity[])
         @test ret[1] == test_value
-        indexed_set = Coloring.IndexedSet(100)
+        indexed_set = ArrayDiff.IndexedSet(100)
         edge_list = ArrayDiff._compute_hessian_sparsity(
             nodes,
             adj,

From e5866eb0ac6c4143f3de4b90f58e0b84d4a1b28e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 12:24:08 +1300
Subject: [PATCH 02/17] Manual fix

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 117cc9e..8623c57 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,5 +17,5 @@ ForwardDiff = "1"
 MathOptInterface = "1.40"
 NaNMath = "1"
 SparseArrays = "1.10"
-SparseMatrixColorings = "1"
+SparseMatrixColorings = "0.4"
 julia = "1.10"

From 9e56de64fff231fd9d08a4e102593dfc725a8b63 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 12:24:51 +1300
Subject: [PATCH 03/17] Fix empty cases

---
 src/coloring_compat.jl | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index 02f90e8..af0350b 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -46,6 +46,35 @@ function _hessian_color_preprocess(
     end
     local_indices = sort!(collect(seen_idx))
     empty!(seen_idx)
+    
+    # Handle empty case (no edges in Hessian)
+    if isempty(local_indices)
+        # Return empty structure - no variables to color
+        # We still need to return a valid ColoringResult, but with empty local_indices
+        # The I and J vectors are already empty, which is correct
+        # For the result, we'll create a minimal valid structure with a diagonal element
+        # Note: This case should rarely occur in practice
+        S = SparseArrays.spdiagm(0 => [true])
+        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
+        tree_result = SparseMatrixColorings.coloring(S, problem, algo)
+        result = ColoringResult(tree_result, Int[])
+        return I, J, result
+    end
+    
+    # Also handle case where we have vertices but no edges (diagonal-only Hessian)
+    if isempty(I)
+        # Create identity matrix pattern (diagonal only)
+        n = length(local_indices)
+        S = SparseArrays.spdiagm(0 => trues(n))
+        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
+        tree_result = SparseMatrixColorings.coloring(S, problem, algo)
+        result = ColoringResult(tree_result, local_indices)
+        # I and J are already empty, which is correct for no off-diagonal elements
+        return I, J, result
+    end
+    
     global_to_local_idx = seen_idx.nzidx # steal for storage
     for k in eachindex(local_indices)
         global_to_local_idx[local_indices[k]] = k

From 82e0c4a5f116594e8180511e9a9eae342e331728 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 12:28:57 +1300
Subject: [PATCH 04/17] Fixes

---
 src/coloring_compat.jl | 48 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index af0350b..f08fa58 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -99,15 +99,53 @@ function _hessian_color_preprocess(
     algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
     tree_result = SparseMatrixColorings.coloring(S, problem, algo)
     
-    # Convert back to global indices
-    for k in eachindex(I)
-        I[k] = local_indices[I[k]]
-        J[k] = local_indices[J[k]]
+    # Reconstruct I and J from the tree structure (matching original _indirect_recover_structure)
+    # First add all diagonal elements
+    N = length(local_indices)
+    
+    # Count off-diagonal elements from tree structure
+    (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
+    nnz_offdiag = 0
+    for tree_idx in 1:nt
+        first = tree_edge_indices[tree_idx]
+        last = tree_edge_indices[tree_idx + 1] - 1
+        nnz_offdiag += (last - first + 1)
+    end
+    
+    I_new = Vector{Int}(undef, N + nnz_offdiag)
+    J_new = Vector{Int}(undef, N + nnz_offdiag)
+    k = 0
+    
+    # Add all diagonal elements
+    for i in 1:N
+        k += 1
+        I_new[k] = local_indices[i]
+        J_new[k] = local_indices[i]
     end
     
+    # Then add off-diagonal elements from the tree structure
+    for tree_idx in 1:nt
+        first = tree_edge_indices[tree_idx]
+        last = tree_edge_indices[tree_idx + 1] - 1
+        for pos in first:last
+            (i_local, j_local) = reverse_bfs_orders[pos]
+            # Convert from local to global indices and normalize (lower triangle)
+            i_global = local_indices[i_local]
+            j_global = local_indices[j_local]
+            if j_global > i_global
+                i_global, j_global = j_global, i_global
+            end
+            k += 1
+            I_new[k] = i_global
+            J_new[k] = j_global
+        end
+    end
+    
+    @assert k == length(I_new)
+    
     # Wrap result with local_indices
     result = ColoringResult(tree_result, local_indices)
-    return I, J, result
+    return I_new, J_new, result
 end
 
 """

From 4b71e3ba435c8bef1961b5102eff73061dd08c11 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 15:11:23 +1300
Subject: [PATCH 05/17] Fix type stability

---
 src/coloring_compat.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index f08fa58..a48d765 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -12,8 +12,8 @@
 
 Wrapper around TreeSetColoringResult that also stores local_indices mapping.
 """
-struct ColoringResult
-    result::SparseMatrixColorings.TreeSetColoringResult
+struct ColoringResult{M<:AbstractMatrix,T<:Integer,G<:SparseMatrixColorings.AdjacencyGraph{T},GT<:SparseMatrixColorings.AbstractGroups{T},R}
+    result::SparseMatrixColorings.TreeSetColoringResult{M,T,G,GT,R}
     local_indices::Vector{Int}  # map from local to global indices
 end
 

From b5da94dd29085fb93d06919aaf663b9283b66341 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 17:36:52 +1300
Subject: [PATCH 06/17] Fix type stability

---
 src/coloring_compat.jl      |  9 ++++-----
 src/forward_over_reverse.jl |  2 +-
 src/mathoptinterface_api.jl | 12 ++++++------
 src/types.jl                | 38 ++++++++++++++++++++++++-------------
 test/ReverseAD.jl           |  4 ++--
 5 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index a48d765..8a0cd80 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -12,8 +12,8 @@
 
 Wrapper around TreeSetColoringResult that also stores local_indices mapping.
 """
-struct ColoringResult{M<:AbstractMatrix,T<:Integer,G<:SparseMatrixColorings.AdjacencyGraph{T},GT<:SparseMatrixColorings.AbstractGroups{T},R}
-    result::SparseMatrixColorings.TreeSetColoringResult{M,T,G,GT,R}
+struct ColoringResult{R<:SparseMatrixColorings.AbstractColoringResult}
+    result::R
     local_indices::Vector{Int}  # map from local to global indices
 end
 
@@ -21,6 +21,7 @@ end
     _hessian_color_preprocess(
         edgelist,
         num_total_var,
+        algo::SparseMatrixColorings.GreedyColoringAlgorithm,
         seen_idx = IndexedSet(0),
     )
 
@@ -34,6 +35,7 @@ SparseMatrixColorings.
 function _hessian_color_preprocess(
     edgelist,
     num_total_var,
+    algo::SparseMatrixColorings.GreedyColoringAlgorithm,
     seen_idx = IndexedSet(0),
 )
     resize!(seen_idx, num_total_var)
@@ -56,7 +58,6 @@ function _hessian_color_preprocess(
         # Note: This case should rarely occur in practice
         S = SparseArrays.spdiagm(0 => [true])
         problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
-        algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, Int[])
         return I, J, result
@@ -68,7 +69,6 @@ function _hessian_color_preprocess(
         n = length(local_indices)
         S = SparseArrays.spdiagm(0 => trues(n))
         problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
-        algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, local_indices)
         # I and J are already empty, which is correct for no off-diagonal elements
@@ -96,7 +96,6 @@ function _hessian_color_preprocess(
     
     # Perform coloring using SparseMatrixColorings
     problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
-    algo = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
     tree_result = SparseMatrixColorings.coloring(S, problem, algo)
     
     # Reconstruct I and J from the tree structure (matching original _indirect_recover_structure)
diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl
index c69854a..5221fb5 100644
--- a/src/forward_over_reverse.jl
+++ b/src/forward_over_reverse.jl
@@ -65,7 +65,7 @@ function _eval_hessian(
         )
     end
     # TODO(odow): consider reverting to a view.
-    output_slice = _UnsafeVectorView(nzcount, length(ex.hess_I), pointer(H))
+    output_slice = _UnsafeVectorView{Float64}(nzcount, length(ex.hess_I), pointer(H))::_UnsafeVectorView{Float64}
     _recover_from_matmat!(
         output_slice,
         ex.seed_matrix,
diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl
index 31995ca..ef35af9 100644
--- a/src/mathoptinterface_api.jl
+++ b/src/mathoptinterface_api.jl
@@ -19,7 +19,7 @@ function MOI.features_available(d::NLPEvaluator)
     return [:Grad, :Jac, :JacVec, :Hess, :HessVec]
 end
 
-function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
+function MOI.initialize(d::NLPEvaluator{R}, requested_features::Vector{Symbol}) where {R}
     # Check that we support the features requested by the user.
     available_features = MOI.features_available(d)
     for feature in requested_features
@@ -38,7 +38,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     d.objective = nothing
     d.user_output_buffer = zeros(largest_user_input_dimension)
     d.jac_storage = zeros(max(N, largest_user_input_dimension))
-    d.constraints = _FunctionStorage[]
+    d.constraints = _FunctionStorage{R}[]
     d.last_x = fill(NaN, N)
     d.want_hess = :Hess in requested_features
     want_hess_storage = (:HessVec in requested_features) || d.want_hess
@@ -111,11 +111,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
             shared_partials_storage_ϵ,
             d,
         )
-        objective = _FunctionStorage(
+        objective = _FunctionStorage{R}(
             subexpr,
             N,
             coloring_storage,
-            d.want_hess,
+            d.want_hess ? d.coloring_algorithm : nothing,
             d.subexpressions,
             individual_order[1],
             subexpression_edgelist,
@@ -137,11 +137,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
         )
         push!(
             d.constraints,
-            _FunctionStorage(
+            _FunctionStorage{R}(
                 subexpr,
                 N,
                 coloring_storage,
-                d.want_hess,
+                d.want_hess ? d.coloring_algorithm : nothing,
                 d.subexpressions,
                 individual_order[idx],
                 subexpression_edgelist,
diff --git a/src/types.jl b/src/types.jl
index 5edae64..8493f03 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -63,28 +63,28 @@ function _subexpression_and_linearity(
     linearity
 end
 
-struct _FunctionStorage
+struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
     expr::_SubexpressionStorage
     grad_sparsity::Vector{Int}
     # Nonzero pattern of Hessian matrix
     hess_I::Vector{Int}
     hess_J::Vector{Int}
-    rinfo::Union{ColoringResult,Nothing} # coloring info for hessians
+    rinfo::Union{Nothing,ColoringResult{R}}
     seed_matrix::Matrix{Float64}
     # subexpressions which this function depends on, ordered for forward pass.
     dependent_subexpressions::Vector{Int}
 
-    function _FunctionStorage(
+    function _FunctionStorage{R}(
         expr::_SubexpressionStorage,
         num_variables,
         coloring_storage::IndexedSet,
-        want_hess::Bool,
+        coloring_algorithm::Union{Nothing,SparseMatrixColorings.GreedyColoringAlgorithm},
         subexpressions::Vector{_SubexpressionStorage},
         dependent_subexpressions,
         subexpression_edgelist,
         subexpression_variables,
         linearity::Vector{Linearity},
-    )
+    ) where {R}
         empty!(coloring_storage)
         _compute_gradient_sparsity!(coloring_storage, expr.nodes)
         for k in dependent_subexpressions
@@ -95,7 +95,7 @@ struct _FunctionStorage
         end
         grad_sparsity = sort!(collect(coloring_storage))
         empty!(coloring_storage)
-        if want_hess
+        if !isnothing(coloring_algorithm)
             edgelist = _compute_hessian_sparsity(
                 expr.nodes,
                 expr.adj,
@@ -107,10 +107,11 @@ struct _FunctionStorage
             hess_I, hess_J, rinfo = _hessian_color_preprocess(
                 edgelist,
                 num_variables,
+                coloring_algorithm,
                 coloring_storage,
             )
             seed_matrix = _seed_matrix(rinfo)
-            return new(
+            return new{R}(
                 expr,
                 grad_sparsity,
                 hess_I,
@@ -120,7 +121,7 @@ struct _FunctionStorage
                 dependent_subexpressions,
             )
         else
-            return new(
+            return new{R}(
                 expr,
                 grad_sparsity,
                 Int[],
@@ -137,6 +138,7 @@ end
     NLPEvaluator(
         model::Nonlinear.Model,
         ordered_variables::Vector{MOI.VariableIndex},
+        coloring_algorithm::SparseMatrixColorings.AbstractColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
     )
 
 Return an `NLPEvaluator` object that implements the `MOI.AbstractNLPEvaluator`
@@ -145,12 +147,13 @@ interface.
 !!! warning
     Before using, you must initialize the evaluator using `MOI.initialize`.
 """
-mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator
+mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.AbstractNLPEvaluator
     data::Nonlinear.Model
     ordered_variables::Vector{MOI.VariableIndex}
+    coloring_algorithm::C
 
-    objective::Union{Nothing,_FunctionStorage}
-    constraints::Vector{_FunctionStorage}
+    objective::Union{Nothing,_FunctionStorage{R}}
+    constraints::Vector{_FunctionStorage{R}}
     subexpressions::Vector{_SubexpressionStorage}
     subexpression_order::Vector{Int}
     # Storage for the subexpressions in reverse-mode automatic differentiation.
@@ -183,8 +186,17 @@ mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator
 
     function NLPEvaluator(
         data::Nonlinear.Model,
-        ordered_variables::Vector{MOI.VariableIndex},
+        ordered_variables::Vector{MOI.VariableIndex};
+        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
     )
-        return new(data, ordered_variables)
+        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        C = typeof(coloring_algorithm)
+        R = Base.promote_op(
+            SparseMatrixColorings.coloring,
+            SparseArrays.SparseMatrixCSC{Bool,Int},
+            typeof(problem),
+            C,
+        )
+        return new{R,C}(data, ordered_variables, coloring_algorithm)
     end
 end
diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl
index 60e06d3..cba3c74 100644
--- a/test/ReverseAD.jl
+++ b/test/ReverseAD.jl
@@ -14,7 +14,6 @@ import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 
 import ArrayDiff
-# Coloring submodule removed - using SparseMatrixColorings instead
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -481,7 +480,8 @@ end
 
 function test_coloring_end_to_end_hessian_coloring_and_recovery()
     # Test the new coloring API through the compatibility layer
-    I, J, rinfo = ArrayDiff._hessian_color_preprocess(Set([(1, 2)]), 2, ArrayDiff.IndexedSet(0))
+    coloring_algorithm = ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
+    I, J, rinfo = ArrayDiff._hessian_color_preprocess(Set([(1, 2)]), 2, coloring_algorithm, ArrayDiff.IndexedSet(0))
     R = ArrayDiff._seed_matrix(rinfo)
     ArrayDiff._prepare_seed_matrix!(R, rinfo)
     @test I == [1, 2, 2]

From 0c96ed31c0a11899bcd167203b192a139be2cfda Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 20 Nov 2025 17:42:23 +1300
Subject: [PATCH 07/17] Add option

---
 src/ArrayDiff.jl | 12 ++++++++----
 src/types.jl     |  2 +-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 588efce..1ddf715 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -13,20 +13,24 @@ import SparseArrays
 import SparseMatrixColorings
 
 """
-    Mode() <: AbstractAutomaticDifferentiation
+    Mode(coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm) <: AbstractAutomaticDifferentiation
 
 Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 """
-struct Mode <: MOI.Nonlinear.AbstractAutomaticDifferentiation end
+struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.Nonlinear.AbstractAutomaticDifferentiation
+    coloring_algorithm::C
+end
+
+Mode() = Mode(SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution))
 
 function MOI.Nonlinear.Evaluator(
     model::MOI.Nonlinear.Model,
-    ::Mode,
+    mode::Mode,
     ordered_variables::Vector{MOI.VariableIndex},
 )
     return MOI.Nonlinear.Evaluator(
         model,
-        NLPEvaluator(model, ordered_variables),
+        NLPEvaluator(model, ordered_variables, mode.coloring_algorithm),
     )
 end
 
diff --git a/src/types.jl b/src/types.jl
index 8493f03..544068e 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -186,7 +186,7 @@ mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm}
 
     function NLPEvaluator(
         data::Nonlinear.Model,
-        ordered_variables::Vector{MOI.VariableIndex};
+        ordered_variables::Vector{MOI.VariableIndex},
         coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
     )
         problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)

From f138f8b24fa229446c11bf36ce729c4d97d85d21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 21 Nov 2025 09:37:52 +1300
Subject: [PATCH 08/17] Fix format

---
 src/ArrayDiff.jl            | 11 +++++--
 src/coloring_compat.jl      | 57 +++++++++++++++++++++----------------
 src/forward_over_reverse.jl | 13 ++++-----
 src/mathoptinterface_api.jl |  5 +++-
 src/types.jl                | 19 ++++++++++---
 test/ReverseAD.jl           | 12 ++++++--
 6 files changed, 77 insertions(+), 40 deletions(-)

diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 1ddf715..98b3c79 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -17,11 +17,18 @@ import SparseMatrixColorings
 
 Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 """
-struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.Nonlinear.AbstractAutomaticDifferentiation
+struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <:
+       MOI.Nonlinear.AbstractAutomaticDifferentiation
     coloring_algorithm::C
 end
 
-Mode() = Mode(SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution))
+function Mode()
+    return Mode(
+        SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        ),
+    )
+end
 
 function MOI.Nonlinear.Evaluator(
     model::MOI.Nonlinear.Model,
diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index 8a0cd80..889258a 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -48,7 +48,7 @@ function _hessian_color_preprocess(
     end
     local_indices = sort!(collect(seen_idx))
     empty!(seen_idx)
-    
+
     # Handle empty case (no edges in Hessian)
     if isempty(local_indices)
         # Return empty structure - no variables to color
@@ -57,24 +57,30 @@ function _hessian_color_preprocess(
         # For the result, we'll create a minimal valid structure with a diagonal element
         # Note: This case should rarely occur in practice
         S = SparseArrays.spdiagm(0 => [true])
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, Int[])
         return I, J, result
     end
-    
+
     # Also handle case where we have vertices but no edges (diagonal-only Hessian)
     if isempty(I)
         # Create identity matrix pattern (diagonal only)
         n = length(local_indices)
         S = SparseArrays.spdiagm(0 => trues(n))
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         tree_result = SparseMatrixColorings.coloring(S, problem, algo)
         result = ColoringResult(tree_result, local_indices)
         # I and J are already empty, which is correct for no off-diagonal elements
         return I, J, result
     end
-    
+
     global_to_local_idx = seen_idx.nzidx # steal for storage
     for k in eachindex(local_indices)
         global_to_local_idx[local_indices[k]] = k
@@ -84,7 +90,7 @@ function _hessian_color_preprocess(
         I[k] = global_to_local_idx[I[k]]
         J[k] = global_to_local_idx[J[k]]
     end
-    
+
     # Create sparsity pattern matrix
     n = length(local_indices)
     S = SparseArrays.spzeros(Bool, n, n)
@@ -93,39 +99,42 @@ function _hessian_color_preprocess(
         S[i, j] = true
         S[j, i] = true  # symmetric
     end
-    
+
     # Perform coloring using SparseMatrixColorings
-    problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+    problem = SparseMatrixColorings.ColoringProblem(;
+        structure = :symmetric,
+        partition = :column,
+    )
     tree_result = SparseMatrixColorings.coloring(S, problem, algo)
-    
+
     # Reconstruct I and J from the tree structure (matching original _indirect_recover_structure)
     # First add all diagonal elements
     N = length(local_indices)
-    
+
     # Count off-diagonal elements from tree structure
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
     nnz_offdiag = 0
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
+        last = tree_edge_indices[tree_idx+1] - 1
         nnz_offdiag += (last - first + 1)
     end
-    
+
     I_new = Vector{Int}(undef, N + nnz_offdiag)
     J_new = Vector{Int}(undef, N + nnz_offdiag)
     k = 0
-    
+
     # Add all diagonal elements
     for i in 1:N
         k += 1
         I_new[k] = local_indices[i]
         J_new[k] = local_indices[i]
     end
-    
+
     # Then add off-diagonal elements from the tree structure
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
+        last = tree_edge_indices[tree_idx+1] - 1
         for pos in first:last
             (i_local, j_local) = reverse_bfs_orders[pos]
             # Convert from local to global indices and normalize (lower triangle)
@@ -139,9 +148,9 @@ function _hessian_color_preprocess(
             J_new[k] = j_global
         end
     end
-    
+
     @assert k == length(I_new)
-    
+
     # Wrap result with local_indices
     result = ColoringResult(tree_result, local_indices)
     return I_new, J_new, result
@@ -203,7 +212,7 @@ function _recover_from_matmat!(
     # V contains N diagonal elements + nnz_offdiag off-diagonal elements
     nnz_offdiag = length(V) - N
     @assert length(stored_values) >= N
-    
+
     # Recover diagonal elements
     k = 0
     for i in 1:N
@@ -214,15 +223,15 @@ function _recover_from_matmat!(
             V[k] = zero(T)
         end
     end
-    
+
     # Recover off-diagonal elements using tree structure
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_result
     fill!(stored_values, zero(T))
-    
+
     for tree_idx in 1:nt
         first = tree_edge_indices[tree_idx]
-        last = tree_edge_indices[tree_idx + 1] - 1
-        
+        last = tree_edge_indices[tree_idx+1] - 1
+
         # Reset stored_values for vertices in this tree
         for pos in first:last
             (vertex, _) = reverse_bfs_orders[pos]
@@ -230,7 +239,7 @@ function _recover_from_matmat!(
         end
         (_, root) = reverse_bfs_orders[last]
         stored_values[root] = zero(T)
-        
+
         # Recover edge values
         for pos in first:last
             (i, j) = reverse_bfs_orders[pos]
@@ -244,7 +253,7 @@ function _recover_from_matmat!(
             V[k] = value
         end
     end
-    
+
     @assert k == length(V)
     return
 end
diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl
index 5221fb5..5f3be96 100644
--- a/src/forward_over_reverse.jl
+++ b/src/forward_over_reverse.jl
@@ -65,13 +65,12 @@ function _eval_hessian(
         )
     end
     # TODO(odow): consider reverting to a view.
-    output_slice = _UnsafeVectorView{Float64}(nzcount, length(ex.hess_I), pointer(H))::_UnsafeVectorView{Float64}
-    _recover_from_matmat!(
-        output_slice,
-        ex.seed_matrix,
-        ex.rinfo,
-        d.output_ϵ,
-    )
+    output_slice = _UnsafeVectorView{Float64}(
+        nzcount,
+        length(ex.hess_I),
+        pointer(H),
+    )::_UnsafeVectorView{Float64}
+    _recover_from_matmat!(output_slice, ex.seed_matrix, ex.rinfo, d.output_ϵ)
     for i in 1:length(output_slice)
         output_slice[i] *= scale
     end
diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl
index ef35af9..3ff81a5 100644
--- a/src/mathoptinterface_api.jl
+++ b/src/mathoptinterface_api.jl
@@ -19,7 +19,10 @@ function MOI.features_available(d::NLPEvaluator)
     return [:Grad, :Jac, :JacVec, :Hess, :HessVec]
 end
 
-function MOI.initialize(d::NLPEvaluator{R}, requested_features::Vector{Symbol}) where {R}
+function MOI.initialize(
+    d::NLPEvaluator{R},
+    requested_features::Vector{Symbol},
+) where {R}
     # Check that we support the features requested by the user.
     available_features = MOI.features_available(d)
     for feature in requested_features
diff --git a/src/types.jl b/src/types.jl
index 544068e..d0402a5 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -78,7 +78,10 @@ struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
         expr::_SubexpressionStorage,
         num_variables,
         coloring_storage::IndexedSet,
-        coloring_algorithm::Union{Nothing,SparseMatrixColorings.GreedyColoringAlgorithm},
+        coloring_algorithm::Union{
+            Nothing,
+            SparseMatrixColorings.GreedyColoringAlgorithm,
+        },
         subexpressions::Vector{_SubexpressionStorage},
         dependent_subexpressions,
         subexpression_edgelist,
@@ -147,7 +150,10 @@ interface.
 !!! warning
     Before using, you must initialize the evaluator using `MOI.initialize`.
 """
-mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm} <: MOI.AbstractNLPEvaluator
+mutable struct NLPEvaluator{
+    R,
+    C<:SparseMatrixColorings.GreedyColoringAlgorithm,
+} <: MOI.AbstractNLPEvaluator
     data::Nonlinear.Model
     ordered_variables::Vector{MOI.VariableIndex}
     coloring_algorithm::C
@@ -187,9 +193,14 @@ mutable struct NLPEvaluator{R,C<:SparseMatrixColorings.GreedyColoringAlgorithm}
     function NLPEvaluator(
         data::Nonlinear.Model,
         ordered_variables::Vector{MOI.VariableIndex},
-        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
+        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        ),
     )
-        problem = SparseMatrixColorings.ColoringProblem(; structure=:symmetric, partition=:column)
+        problem = SparseMatrixColorings.ColoringProblem(;
+            structure = :symmetric,
+            partition = :column,
+        )
         C = typeof(coloring_algorithm)
         R = Base.promote_op(
             SparseMatrixColorings.coloring,
diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl
index cba3c74..4c10399 100644
--- a/test/ReverseAD.jl
+++ b/test/ReverseAD.jl
@@ -480,8 +480,16 @@ end
 
 function test_coloring_end_to_end_hessian_coloring_and_recovery()
     # Test the new coloring API through the compatibility layer
-    coloring_algorithm = ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution)
-    I, J, rinfo = ArrayDiff._hessian_color_preprocess(Set([(1, 2)]), 2, coloring_algorithm, ArrayDiff.IndexedSet(0))
+    coloring_algorithm =
+        ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(;
+            decompression = :substitution,
+        )
+    I, J, rinfo = ArrayDiff._hessian_color_preprocess(
+        Set([(1, 2)]),
+        2,
+        coloring_algorithm,
+        ArrayDiff.IndexedSet(0),
+    )
     R = ArrayDiff._seed_matrix(rinfo)
     ArrayDiff._prepare_seed_matrix!(R, rinfo)
     @test I == [1, 2, 2]

From 090384ff09a341e9f4b321fb8670b80cb2b590a3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 21 Nov 2025 09:39:28 +1300
Subject: [PATCH 09/17] Remove uncovered lines

---
 src/coloring_compat.jl | 15 ---------------
 1 file changed, 15 deletions(-)

diff --git a/src/coloring_compat.jl b/src/coloring_compat.jl
index 889258a..fd69d41 100644
--- a/src/coloring_compat.jl
+++ b/src/coloring_compat.jl
@@ -66,21 +66,6 @@ function _hessian_color_preprocess(
         return I, J, result
     end
 
-    # Also handle case where we have vertices but no edges (diagonal-only Hessian)
-    if isempty(I)
-        # Create identity matrix pattern (diagonal only)
-        n = length(local_indices)
-        S = SparseArrays.spdiagm(0 => trues(n))
-        problem = SparseMatrixColorings.ColoringProblem(;
-            structure = :symmetric,
-            partition = :column,
-        )
-        tree_result = SparseMatrixColorings.coloring(S, problem, algo)
-        result = ColoringResult(tree_result, local_indices)
-        # I and J are already empty, which is correct for no off-diagonal elements
-        return I, J, result
-    end
-
     global_to_local_idx = seen_idx.nzidx # steal for storage
     for k in eachindex(local_indices)
         global_to_local_idx[local_indices[k]] = k

From f2dff9b70c158e614151f84419ebf1afd1d183ea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 21 Nov 2025 09:40:00 +1300
Subject: [PATCH 10/17] Rename coloring file

---
 src/ArrayDiff.jl                        | 2 +-
 src/{coloring_compat.jl => coloring.jl} | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename src/{coloring_compat.jl => coloring.jl} (100%)

diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 98b3c79..439996b 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -61,7 +61,7 @@ import NaNMath:
     sqrt
 
 include("indexed_set.jl")
-include("coloring_compat.jl")
+include("coloring.jl")
 include("graph_tools.jl")
 include("sizes.jl")
 include("types.jl")
diff --git a/src/coloring_compat.jl b/src/coloring.jl
similarity index 100%
rename from src/coloring_compat.jl
rename to src/coloring.jl

From c2a0d38bf3749a4e3aae6ab16a9bbab296cd02a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 11:36:51 +0100
Subject: [PATCH 11/17] Remove Coloring

---
 src/Coloring/Coloring.jl | 574 ---------------------------------------
 1 file changed, 574 deletions(-)
 delete mode 100644 src/Coloring/Coloring.jl

diff --git a/src/Coloring/Coloring.jl b/src/Coloring/Coloring.jl
deleted file mode 100644
index 1ef8dc5..0000000
--- a/src/Coloring/Coloring.jl
+++ /dev/null
@@ -1,574 +0,0 @@
-# Copyright (c) 2017: Miles Lubin and contributors
-# Copyright (c) 2017: Google Inc.
-#
-# Use of this source code is governed by an MIT-style license that can be found
-# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
-
-module Coloring
-
-import DataStructures
-
-include("topological_sort.jl")
-
-"""
-    mutable struct IndexedSet
-        nzidx::Vector{Int}
-        empty::BitVector
-        nnz::Int
-    end
-
-Represent the set `nzidx[1:nnz]` by additionally setting `empty[i]` to `false`
-for each element `i` of the set for fast membership check.
-"""
-mutable struct IndexedSet
-    nzidx::Vector{Int}
-    empty::BitVector
-    nnz::Int
-end
-
-IndexedSet(n::Integer) = IndexedSet(zeros(Int, n), trues(n), 0)
-
-function Base.push!(v::IndexedSet, i::Integer)
-    if v.empty[i]  # new index
-        v.nzidx[v.nnz+=1] = i
-        v.empty[i] = false
-    end
-    return
-end
-
-function Base.empty!(v::IndexedSet)
-    for i in 1:v.nnz
-        v.empty[v.nzidx[i]] = true
-    end
-    v.nnz = 0
-    return v
-end
-
-# Returns the maximum index that the set can contain,
-# not the cardinality of the set like `length(::Base.Set)`
-Base.length(v::IndexedSet) = length(v.nzidx)
-
-function Base.resize!(v::IndexedSet, n::Integer)
-    if n > length(v)
-        @assert v.nnz == 0 # only resize empty vector
-        resize!(v.nzidx, n)
-        resize!(v.empty, n)
-        fill!(v.empty, true)
-    end
-    return
-end
-
-Base.collect(v::IndexedSet) = v.nzidx[1:v.nnz]
-
-function Base.union!(v::IndexedSet, s)
-    for x in s
-        push!(v, x)
-    end
-    return
-end
-
-"""
-    struct UndirectedGraph
-        adjlist::Vector{Int}
-        edgeindex::Vector{Int}
-        offsets::Vector{Int}
-        edges::Vector{Tuple{Int,Int}}
-    end
-
-Compact storage for an undirected graph. The number of nodes is given by
-`length(offsets) - 1`. The edges of node `u` are given by `edges[e]` for
-`e in edgeindex[offsets[u]:(offsets[u+1] - 1)]`. The neighbors are also
-stored at `adjlist[offsets[u]:(offsets[u+1] - 1)]`.
-"""
-struct UndirectedGraph
-    adjlist::Vector{Int}
-    edgeindex::Vector{Int} # corresponding edge number, indexed as adjlist
-    offsets::Vector{Int}
-    edges::Vector{Tuple{Int,Int}}
-end
-
-_num_vertices(g::UndirectedGraph) = length(g.offsets) - 1
-
-_num_edges(g::UndirectedGraph) = length(g.edges)
-
-_num_neighbors(i::Int, g::UndirectedGraph) = g.offsets[i+1] - g.offsets[i]
-
-_start_neighbors(i::Int, g::UndirectedGraph) = g.offsets[i]
-
-function UndirectedGraph(I, J, nel)
-    adjcount = zeros(Int, nel)
-    n_edges = 0
-    for k in eachindex(I)
-        i = I[k]
-        j = J[k]
-        if i == j
-            continue
-        end
-        n_edges += 1
-        adjcount[i] += 1
-        adjcount[j] += 1
-    end
-    offsets = Vector{Int}(undef, nel + 1)
-    offsets[1] = 1
-    for k in 1:nel
-        offsets[k+1] = offsets[k] + adjcount[k]
-    end
-    fill!(adjcount, 0)
-    edges = Vector{Tuple{Int,Int}}(undef, n_edges)
-    adjlist = Vector{Int}(undef, offsets[nel+1] - 1)
-    edgeindex = Vector{Int}(undef, length(adjlist))
-    edge_count = 0
-    for k in eachindex(I)
-        i = I[k]
-        j = J[k]
-        if i == j
-            continue
-        end
-        edge_count += 1
-        adjlist[offsets[i]+adjcount[i]] = j
-        edgeindex[offsets[i]+adjcount[i]] = edge_count
-        adjcount[i] += 1
-        adjlist[offsets[j]+adjcount[j]] = i
-        edgeindex[offsets[j]+adjcount[j]] = edge_count
-        adjcount[j] += 1
-        edges[edge_count] = (i, j)
-    end
-    @assert edge_count == n_edges
-    return UndirectedGraph(adjlist, edgeindex, offsets, edges)
-end
-
-struct _Edge
-    index::Int
-    source::Int
-    target::Int
-end
-
-function _prevent_cycle(
-    v,
-    w,
-    x,
-    e_idx1,
-    e_idx2,
-    S,
-    firstVisitToTree,
-    forbiddenColors,
-    color,
-)
-    er = DataStructures.find_root!(S, e_idx2)
-    @inbounds first = firstVisitToTree[er]
-    p = first.source # but this depends on the order?
-    q = first.target
-    @inbounds if p != v
-        firstVisitToTree[er] = _Edge(e_idx1, v, w)
-    elseif q != w
-        forbiddenColors[color[x]] = v
-    end
-    return
-end
-
-function _grow_star(v, w, e_idx, firstNeighbor, color, S)
-    @inbounds e = firstNeighbor[color[w]]
-    p = e.source
-    @inbounds if p != v
-        firstNeighbor[color[w]] = _Edge(e_idx, v, w)
-    else
-        union!(S, e_idx, e.index)
-    end
-    return
-end
-
-function _merge_trees(eg, eg1, S)
-    e1 = DataStructures.find_root!(S, eg)
-    e2 = DataStructures.find_root!(S, eg1)
-    if e1 != e2
-        union!(S, eg, eg1)
-    end
-    return
-end
-
-"""
-    acyclic_coloring(g::UndirectedGraph)
-
-Implement the acyclic coloring algorithm of Gebremdehin, Tarafdar, Manne, and
-Pothen, "New Acyclic and Star Coloring Algorithms with Application to Computing
-Hessians." SIAM J. Sci. Comput. 2007.
-
-Returns `Tuple{Vector{Int},Int}` giving the color index of each node in `g`, as
-well as the total number of colors used.
-"""
-function acyclic_coloring(g::UndirectedGraph)
-    if _num_edges(g) == 0
-        return fill(1, _num_vertices(g)), 1
-    end
-    num_colors = 0
-    forbiddenColors = Int[]
-    firstNeighbor = _Edge[]
-    firstVisitToTree = fill(_Edge(0, 0, 0), _num_edges(g))
-    color = fill(0, _num_vertices(g))
-    # disjoint set forest of edges in the graph
-    S = DataStructures.IntDisjointSet{Int}(_num_edges(g))
-    @inbounds for v in 1:_num_vertices(g)
-        n_neighbor = _num_neighbors(v, g)
-        start_neighbor = _start_neighbors(v, g)
-        for k in 0:(n_neighbor-1)
-            w = g.adjlist[start_neighbor+k]
-            if color[w] == 0
-                continue
-            end
-            forbiddenColors[color[w]] = v
-        end
-        for k in 0:(n_neighbor-1)
-            w = g.adjlist[start_neighbor+k]
-            e_idx = g.edgeindex[start_neighbor+k]
-            if color[w] == 0
-                continue
-            end
-            n_neighbor_w = _num_neighbors(w, g)
-            start_neighbor_w = _start_neighbors(w, g)
-            for k2 in 0:(n_neighbor_w-1)
-                x = g.adjlist[start_neighbor_w+k2]
-                e2_idx = g.edgeindex[start_neighbor_w+k2]
-                if color[x] == 0
-                    continue
-                end
-                if forbiddenColors[color[x]] != v
-                    _prevent_cycle(
-                        v,
-                        w,
-                        x,
-                        e_idx,
-                        e2_idx,
-                        S,
-                        firstVisitToTree,
-                        forbiddenColors,
-                        color,
-                    )
-                end
-            end
-        end
-        # find feasible color
-        found = false
-        for k in 1:num_colors
-            if forbiddenColors[k] != v
-                color[v] = k
-                found = true
-                break
-            end
-        end
-        if !found
-            num_colors += 1
-            push!(forbiddenColors, 0)
-            push!(firstNeighbor, _Edge(0, 0, 0))
-            color[v] = num_colors
-        end
-        for k in 0:(n_neighbor-1)
-            w = g.adjlist[start_neighbor+k]
-            e_idx = g.edgeindex[start_neighbor+k]
-            if color[w] == 0
-                continue
-            end
-            _grow_star(v, w, e_idx, firstNeighbor, color, S)
-        end
-        for k in 0:(n_neighbor-1)
-            w = g.adjlist[start_neighbor+k]
-            e_idx = g.edgeindex[start_neighbor+k]
-            if color[w] == 0
-                continue
-            end
-            n_neighbor_w = _num_neighbors(w, g)
-            start_neighbor_w = _start_neighbors(w, g)
-            for k2 in 0:(n_neighbor_w-1)
-                x = g.adjlist[start_neighbor_w+k2]
-                e2_idx = g.edgeindex[start_neighbor_w+k2]
-                if color[x] == 0 || x == v
-                    continue
-                end
-                if color[x] == color[v]
-                    _merge_trees(e_idx, e2_idx, S)
-                end
-            end
-        end
-    end
-    return color, num_colors
-end
-
-struct RecoveryInfo
-    vertexmap::Vector{Vector{Int}}
-    postorder::Vector{Vector{Int}}
-    parents::Vector{Vector{Int}}
-    color::Vector{Int}
-    num_colors::Int
-    nnz::Int # number of off-diagonal nonzeros
-    local_indices::Vector{Int} # map back to global indices
-end
-
-function RecoveryInfo()
-    return RecoveryInfo(
-        Vector{Vector{Int}}(undef, 0),
-        Vector{Vector{Int}}(undef, 0),
-        Vector{Vector{Int}}(undef, 0),
-        Vector{Int}(undef, 0),
-        0,
-        0,
-        Vector{Int}(undef, 0),
-    )
-end
-
-function recovery_preprocess(
-    g::UndirectedGraph,
-    color,
-    num_colors,
-    local_indices,
-)
-    # represent two-color subgraph as:
-    # list of vertices (with map to global indices)
-    # adjacency list in a single vector (with list of offsets)
-    # linear index of pair of colors
-    twocolorindex = zeros(Int32, num_colors, num_colors)
-    seen_twocolors = 0
-    # count of edges in each subgraph
-    edge_count = Int[]
-    for k in eachindex(g.edges)
-        u, v = g.edges[k]
-        i = min(color[u], color[v])
-        j = max(color[u], color[v])
-        if twocolorindex[i, j] == 0
-            seen_twocolors += 1
-            twocolorindex[i, j] = seen_twocolors
-            push!(edge_count, 0)
-        end
-        idx = twocolorindex[i, j]
-        edge_count[idx] += 1
-    end
-    # edges sorted by twocolor subgraph
-    sorted_edges = Vector{Vector{Tuple{Int,Int}}}(undef, seen_twocolors)
-    for idx in 1:seen_twocolors
-        sorted_edges[idx] = Tuple{Int,Int}[]
-        sizehint!(sorted_edges[idx], edge_count[idx])
-    end
-    for k in eachindex(g.edges)
-        u, v = g.edges[k]
-        i = min(color[u], color[v])
-        j = max(color[u], color[v])
-        idx = twocolorindex[i, j]
-        push!(sorted_edges[idx], (u, v))
-    end
-    # list of unique vertices in each twocolor subgraph
-    vertexmap = Vector{Vector{Int}}(undef, seen_twocolors)
-    postorder = Vector{Vector{Int}}(undef, seen_twocolors)
-    parents = Vector{Vector{Int}}(undef, seen_twocolors)
-    # temporary lookup map from global index to subgraph index
-    revmap = zeros(Int, _num_vertices(g))
-    adjcount = zeros(Int, _num_vertices(g))
-    cmap = zeros(Int, 0) # shared storage for DFS
-    for idx in 1:seen_twocolors
-        my_edges = sorted_edges[idx]
-        vlist = Int[]
-        # build up the vertex list and adjacency count
-        for k in eachindex(my_edges)
-            u, v = my_edges[k]
-            # seen these vertices yet?
-            if revmap[u] == 0
-                push!(vlist, u)
-                revmap[u] = length(vlist)
-                adjcount[u] = 0
-            end
-            if revmap[v] == 0
-                push!(vlist, v)
-                revmap[v] = length(vlist)
-                adjcount[v] = 0
-            end
-            adjcount[u] += 1
-            adjcount[v] += 1
-        end
-        # set up offsets for adjlist
-        offset = Vector{Int}(undef, length(vlist) + 1)
-        offset[1] = 1
-        for k in eachindex(vlist)
-            offset[k+1] = offset[k] + adjcount[vlist[k]]
-            adjcount[vlist[k]] = 0
-        end
-        # adjlist for node u in twocolor idx starts at
-        # vec[offset[u]]
-        # u has global index vlist[u]
-        vec = Vector{Int}(undef, offset[length(vlist)+1] - 1)
-        # now fill in
-        for k in eachindex(my_edges)
-            u, v = my_edges[k]
-            u_rev = revmap[u] # indices in the subgraph
-            v_rev = revmap[v]
-            vec[offset[u_rev]+adjcount[u]] = v_rev
-            vec[offset[v_rev]+adjcount[v]] = u_rev
-            adjcount[u] += 1
-            adjcount[v] += 1
-        end
-        resize!(cmap, length(vlist))
-        order, parent =
-            reverse_topological_sort_by_dfs(vec, offset, length(vlist), cmap)
-        for k in eachindex(vlist)
-            revmap[vlist[k]] = 0  # clear for reuse
-        end
-        postorder[idx] = order
-        parents[idx] = parent
-        vertexmap[idx] = vlist
-    end
-    return RecoveryInfo(
-        vertexmap,
-        postorder,
-        parents,
-        color,
-        num_colors,
-        _num_edges(g),
-        local_indices,
-    )
-end
-
-_normalize(i, j) = (j > i) ? (j, i) : (i, j)
-
-function _indirect_recover_structure(rinfo::RecoveryInfo)
-    N = length(rinfo.color)
-    I = zeros(Int, rinfo.nnz + N)
-    J = zeros(Int, rinfo.nnz + N)
-    k = 0
-    for i in 1:N
-        k += 1
-        I[k] = J[k] = i
-    end
-    for t in eachindex(rinfo.postorder)
-        vmap = rinfo.vertexmap[t]
-        order = rinfo.postorder[t]
-        parent = rinfo.parents[t]
-        for z in eachindex(order)
-            v = order[z]
-            p = parent[v]
-            if p == 0
-                continue
-            end
-            k += 1
-            I[k], J[k] = _normalize(vmap[v], vmap[p])
-        end
-    end
-    @assert k == rinfo.nnz + N
-    return I, J
-end
-
-"""
-    hessian_color_preprocess(
-        edgelist,
-        num_total_var,
-        seen_idx = IndexedSet(0),
-    )
-
-`edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the
-diagonal.
-"""
-function hessian_color_preprocess(
-    edgelist,
-    num_total_var,
-    seen_idx = IndexedSet(0),
-)
-    resize!(seen_idx, num_total_var)
-    I, J = Int[], Int[]
-    for (i, j) in edgelist
-        push!(seen_idx, i)
-        push!(seen_idx, j)
-        push!(I, i)
-        push!(J, j)
-    end
-    local_indices = sort!(collect(seen_idx))
-    empty!(seen_idx)
-    global_to_local_idx = seen_idx.nzidx # steal for storage
-    for k in eachindex(local_indices)
-        global_to_local_idx[local_indices[k]] = k
-    end
-    # only do the coloring on the local indices
-    for k in eachindex(I)
-        I[k] = global_to_local_idx[I[k]]
-        J[k] = global_to_local_idx[J[k]]
-    end
-    g = UndirectedGraph(I, J, length(local_indices))
-    color, num_colors = acyclic_coloring(g)
-    @assert length(color) == _num_vertices(g)
-    rinfo = recovery_preprocess(g, color, num_colors, local_indices)
-    I, J = _indirect_recover_structure(rinfo)
-    # convert back to global indices
-    for k in eachindex(I)
-        I[k] = local_indices[I[k]]
-        J[k] = local_indices[J[k]]
-    end
-    return I, J, rinfo
-end
-
-"""
-    seed_matrix(rinfo::RecoveryInfo)
-
-Allocate a seed matrix for the Coloring.
-"""
-function seed_matrix(rinfo::RecoveryInfo)
-    return Matrix{Float64}(undef, length(rinfo.local_indices), rinfo.num_colors)
-end
-
-"""
-    prepare_seed_matrix!(R, rinfo::RecoveryInfo)
-"""
-function prepare_seed_matrix!(R, rinfo::RecoveryInfo)
-    N = length(rinfo.color)
-    @assert N == size(R, 1) == length(rinfo.local_indices)
-    @assert size(R, 2) == rinfo.num_colors
-    fill!(R, 0.0)
-    for i in 1:N
-        R[i, rinfo.color[i]] = 1
-    end
-    return
-end
-
-"""
-    recover_from_matmat!(
-        V::AbstractVector{T},
-        R::AbstractMatrix{T},
-        rinfo::RecoveryInfo,
-        stored_values::AbstractVector{T},
-    ) where {T}
-
-recover the hessian values from the hessian-matrix solution
-stored_values is a temporary vector of length >= length(rinfo.local_indices)
-"""
-function recover_from_matmat!(
-    V::AbstractVector{T},
-    R::AbstractMatrix{T},
-    rinfo::RecoveryInfo,
-    stored_values::AbstractVector{T},
-) where {T}
-    N = length(rinfo.color) # number of local variables
-    @assert length(V) == rinfo.nnz + N
-    @assert length(stored_values) >= length(rinfo.local_indices)
-    k = 0
-    for i in 1:N
-        k += 1
-        V[k] = R[i, rinfo.color[i]]
-    end
-    for t in eachindex(rinfo.vertexmap)
-        vmap = rinfo.vertexmap[t]
-        order = rinfo.postorder[t]
-        parent = rinfo.parents[t]
-        for z in eachindex(order)
-            @inbounds stored_values[z] = 0.0
-        end
-        @inbounds for z in eachindex(order)
-            v = order[z]
-            p = parent[v]
-            if p == 0
-                continue
-            end
-            i, j = vmap[v], vmap[p]
-            value = R[i, rinfo.color[j]] - stored_values[v]
-            stored_values[p] += value
-            k += 1
-            V[k] = value
-        end
-    end
-    @assert k == rinfo.nnz + N
-    return
-end
-
-end  # module

From 53468ce2db0982b0792965c066cc3f94423aa96a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 11:40:33 +0100
Subject: [PATCH 12/17] Remove indexed set

---
 src/coloring.jl    |  6 ++---
 src/indexed_set.jl | 62 ----------------------------------------------
 2 files changed, 2 insertions(+), 66 deletions(-)
 delete mode 100644 src/indexed_set.jl

diff --git a/src/coloring.jl b/src/coloring.jl
index fd69d41..c222bd6 100644
--- a/src/coloring.jl
+++ b/src/coloring.jl
@@ -22,7 +22,7 @@ end
         edgelist,
         num_total_var,
         algo::SparseMatrixColorings.GreedyColoringAlgorithm,
-        seen_idx = IndexedSet(0),
+        seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0),
     )
 
 `edgelist` contains the nonzeros in the Hessian, *including* nonzeros on the
@@ -36,7 +36,7 @@ function _hessian_color_preprocess(
     edgelist,
     num_total_var,
     algo::SparseMatrixColorings.GreedyColoringAlgorithm,
-    seen_idx = IndexedSet(0),
+    seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0),
 )
     resize!(seen_idx, num_total_var)
     I, J = Int[], Int[]
@@ -192,10 +192,8 @@ function _recover_from_matmat!(
     tree_result = result.result
     color = SparseMatrixColorings.column_colors(tree_result)
     N = length(result.local_indices)
-    S = tree_result.ag.S
     # Compute number of off-diagonal nonzeros from the length of V
     # V contains N diagonal elements + nnz_offdiag off-diagonal elements
-    nnz_offdiag = length(V) - N
     @assert length(stored_values) >= N
 
     # Recover diagonal elements
diff --git a/src/indexed_set.jl b/src/indexed_set.jl
deleted file mode 100644
index b2bb631..0000000
--- a/src/indexed_set.jl
+++ /dev/null
@@ -1,62 +0,0 @@
-# Copyright (c) 2017: Miles Lubin and contributors
-# Copyright (c) 2017: Google Inc.
-#
-# Use of this source code is governed by an MIT-style license that can be found
-# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
-
-"""
-    mutable struct IndexedSet
-        nzidx::Vector{Int}
-        empty::BitVector
-        nnz::Int
-    end
-
-Represent the set `nzidx[1:nnz]` by additionally setting `empty[i]` to `false`
-for each element `i` of the set for fast membership check.
-"""
-mutable struct IndexedSet
-    nzidx::Vector{Int}
-    empty::BitVector
-    nnz::Int
-end
-
-IndexedSet(n::Integer) = IndexedSet(zeros(Int, n), trues(n), 0)
-
-function Base.push!(v::IndexedSet, i::Integer)
-    if v.empty[i]  # new index
-        v.nzidx[v.nnz+=1] = i
-        v.empty[i] = false
-    end
-    return
-end
-
-function Base.empty!(v::IndexedSet)
-    for i in 1:v.nnz
-        v.empty[v.nzidx[i]] = true
-    end
-    v.nnz = 0
-    return v
-end
-
-# Returns the maximum index that the set can contain,
-# not the cardinality of the set like `length(::Base.Set)`
-Base.length(v::IndexedSet) = length(v.nzidx)
-
-function Base.resize!(v::IndexedSet, n::Integer)
-    if n > length(v)
-        @assert v.nnz == 0 # only resize empty vector
-        resize!(v.nzidx, n)
-        resize!(v.empty, n)
-        fill!(v.empty, true)
-    end
-    return
-end
-
-Base.collect(v::IndexedSet) = v.nzidx[1:v.nnz]
-
-function Base.union!(v::IndexedSet, s)
-    for x in s
-        push!(v, x)
-    end
-    return
-end

From 06ed8bcc59d153de6623d0dddc161c345b8ff873 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 11:51:40 +0100
Subject: [PATCH 13/17] Remove IndexedSet

---
 src/ArrayDiff.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 439996b..4bfff5f 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -60,7 +60,6 @@ import NaNMath:
     pow,
     sqrt
 
-include("indexed_set.jl")
 include("coloring.jl")
 include("graph_tools.jl")
 include("sizes.jl")

From 79af9b5503b402e7aa8fc520e88e1fed2b04d1d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 11:54:24 +0100
Subject: [PATCH 14/17] Fixes

---
 src/graph_tools.jl          | 4 ++--
 src/mathoptinterface_api.jl | 2 +-
 src/types.jl                | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/graph_tools.jl b/src/graph_tools.jl
index b91e0d7..f813016 100644
--- a/src/graph_tools.jl
+++ b/src/graph_tools.jl
@@ -159,7 +159,7 @@ Compute the sparsity pattern of the gradient of an expression (that is, a list o
 which variable indices are present).
 """
 function _compute_gradient_sparsity!(
-    indices::IndexedSet,
+    indices::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet,
     nodes::Vector{Nonlinear.Node},
 )
     for node in nodes
@@ -207,7 +207,7 @@ function _compute_hessian_sparsity(
     nodes::Vector{Nonlinear.Node},
     adj,
     input_linearity::Vector{Linearity},
-    indexedset::IndexedSet,
+    indexedset::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet,
     subexpression_edgelist::Vector{Set{Tuple{Int,Int}}},
     subexpression_variables::Vector{Vector{Int}},
 )
diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl
index 3ff81a5..943ea61 100644
--- a/src/mathoptinterface_api.jl
+++ b/src/mathoptinterface_api.jl
@@ -45,7 +45,7 @@ function MOI.initialize(
     d.last_x = fill(NaN, N)
     d.want_hess = :Hess in requested_features
     want_hess_storage = (:HessVec in requested_features) || d.want_hess
-    coloring_storage = IndexedSet(N)
+    coloring_storage = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(N)
     max_expr_length = 0
     max_expr_with_sub_length = 0
     #
diff --git a/src/types.jl b/src/types.jl
index d0402a5..776ad83 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -77,7 +77,7 @@ struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
     function _FunctionStorage{R}(
         expr::_SubexpressionStorage,
         num_variables,
-        coloring_storage::IndexedSet,
+        coloring_storage::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet,
         coloring_algorithm::Union{
             Nothing,
             SparseMatrixColorings.GreedyColoringAlgorithm,

From 85444366fdbb6ffded9b1b70b99b99f4ac415c82 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 11:57:53 +0100
Subject: [PATCH 15/17] Fixes

---
 test/ReverseAD.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl
index 4c10399..129bcf8 100644
--- a/test/ReverseAD.jl
+++ b/test/ReverseAD.jl
@@ -488,7 +488,7 @@ function test_coloring_end_to_end_hessian_coloring_and_recovery()
         Set([(1, 2)]),
         2,
         coloring_algorithm,
-        ArrayDiff.IndexedSet(0),
+        MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(0),
     )
     R = ArrayDiff._seed_matrix(rinfo)
     ArrayDiff._prepare_seed_matrix!(R, rinfo)
@@ -559,7 +559,7 @@ function test_linearity()
         nodes = ArrayDiff._replace_moi_variables(expr.nodes, variables)
         ret = ArrayDiff._classify_linearity(nodes, adj, ArrayDiff.Linearity[])
         @test ret[1] == test_value
-        indexed_set = ArrayDiff.IndexedSet(100)
+        indexed_set = MOI.Nonlinear.ReverseAD.Coloring.IndexedSet(100)
         edge_list = ArrayDiff._compute_hessian_sparsity(
             nodes,
             adj,

From 07051a1f978a4bfc0661ccf9a9d93ee6daf59e65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 12:31:52 +0100
Subject: [PATCH 16/17] Faster sparse matrix construction

---
 Project.toml      |  2 +-
 src/ArrayDiff.jl  | 11 ++++++-----
 src/coloring.jl   | 41 +++++++++++++++++++++--------------------
 src/types.jl      | 16 ++++++++--------
 test/ReverseAD.jl |  2 +-
 5 files changed, 37 insertions(+), 35 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8623c57..5b92302 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ArrayDiff"
 uuid = "c45fa1ca-6901-44ac-ae5b-5513a4852d50"
-authors = ["Benoît Legat <benoit.legat@gmail.com>"]
 version = "0.1.0"
+authors = ["Benoît Legat <benoit.legat@gmail.com>"]
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 4bfff5f..8b8af5a 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -6,25 +6,26 @@
 
 module ArrayDiff
 
+import LinearAlgebra as LA
+import SparseArrays
+import SparseMatrixColorings as SMC
 import ForwardDiff
 import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
-import SparseArrays
-import SparseMatrixColorings
 
 """
-    Mode(coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm) <: AbstractAutomaticDifferentiation
+    Mode(coloring_algorithm::SMC.GreedyColoringAlgorithm) <: AbstractAutomaticDifferentiation
 
 Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
 """
-struct Mode{C<:SparseMatrixColorings.GreedyColoringAlgorithm} <:
+struct Mode{C<:SMC.GreedyColoringAlgorithm} <:
        MOI.Nonlinear.AbstractAutomaticDifferentiation
     coloring_algorithm::C
 end
 
 function Mode()
     return Mode(
-        SparseMatrixColorings.GreedyColoringAlgorithm(;
+        SMC.GreedyColoringAlgorithm(;
             decompression = :substitution,
         ),
     )
diff --git a/src/coloring.jl b/src/coloring.jl
index c222bd6..e23c3fa 100644
--- a/src/coloring.jl
+++ b/src/coloring.jl
@@ -6,13 +6,13 @@
 
 """
     struct ColoringResult
-        result::SparseMatrixColorings.TreeSetColoringResult
+        result::SMC.TreeSetColoringResult
         local_indices::Vector{Int}  # map from local to global indices
     end
 
 Wrapper around TreeSetColoringResult that also stores local_indices mapping.
 """
-struct ColoringResult{R<:SparseMatrixColorings.AbstractColoringResult}
+struct ColoringResult{R<:SMC.AbstractColoringResult}
     result::R
     local_indices::Vector{Int}  # map from local to global indices
 end
@@ -21,7 +21,7 @@ end
     _hessian_color_preprocess(
         edgelist,
         num_total_var,
-        algo::SparseMatrixColorings.GreedyColoringAlgorithm,
+        algo::SMC.GreedyColoringAlgorithm,
         seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0),
     )
 
@@ -35,7 +35,7 @@ SparseMatrixColorings.
 function _hessian_color_preprocess(
     edgelist,
     num_total_var,
-    algo::SparseMatrixColorings.GreedyColoringAlgorithm,
+    algo::SMC.GreedyColoringAlgorithm,
     seen_idx = MOI.Nonlinear.Coloring.IndexedSet(0),
 )
     resize!(seen_idx, num_total_var)
@@ -45,6 +45,10 @@ function _hessian_color_preprocess(
         push!(seen_idx, j)
         push!(I, i)
         push!(J, j)
+        if i != j
+            push!(I, j)
+            push!(J, i)
+        end
     end
     local_indices = sort!(collect(seen_idx))
     empty!(seen_idx)
@@ -56,12 +60,12 @@ function _hessian_color_preprocess(
         # The I and J vectors are already empty, which is correct
         # For the result, we'll create a minimal valid structure with a diagonal element
         # Note: This case should rarely occur in practice
-        S = SparseArrays.spdiagm(0 => [true])
-        problem = SparseMatrixColorings.ColoringProblem(;
+        S = SMC.SparsityPatternCSC(SparseArrays.spdiagm(0 => [true]))
+        problem = SMC.ColoringProblem(;
             structure = :symmetric,
             partition = :column,
         )
-        tree_result = SparseMatrixColorings.coloring(S, problem, algo)
+        tree_result = SMC.coloring(S, problem, algo)
         result = ColoringResult(tree_result, Int[])
         return I, J, result
     end
@@ -78,19 +82,16 @@ function _hessian_color_preprocess(
 
     # Create sparsity pattern matrix
     n = length(local_indices)
-    S = SparseArrays.spzeros(Bool, n, n)
-    for k in eachindex(I)
-        i, j = I[k], J[k]
-        S[i, j] = true
-        S[j, i] = true  # symmetric
-    end
+    S = SMC.SparsityPatternCSC(
+        SparseArrays.sparse(I, J, trues(length(I)), n, n, &)
+    )
 
-    # Perform coloring using SparseMatrixColorings
-    problem = SparseMatrixColorings.ColoringProblem(;
+    # Perform coloring using SMC
+    problem = SMC.ColoringProblem(;
         structure = :symmetric,
         partition = :column,
     )
-    tree_result = SparseMatrixColorings.coloring(S, problem, algo)
+    tree_result = SMC.coloring(S, problem, algo)
 
     # Reconstruct I and J from the tree structure (matching original _indirect_recover_structure)
     # First add all diagonal elements
@@ -148,7 +149,7 @@ Allocate a seed matrix for the coloring result.
 """
 function _seed_matrix(result::ColoringResult)
     n = length(result.local_indices)
-    ncolors = SparseMatrixColorings.ncolors(result.result)
+    ncolors = SMC.ncolors(result.result)
     return Matrix{Float64}(undef, n, ncolors)
 end
 
@@ -158,10 +159,10 @@ end
 Prepare the seed matrix R for Hessian computation.
 """
 function _prepare_seed_matrix!(R, result::ColoringResult)
-    color = SparseMatrixColorings.column_colors(result.result)
+    color = SMC.column_colors(result.result)
     N = length(result.local_indices)
     @assert N == size(R, 1)
-    @assert size(R, 2) == SparseMatrixColorings.ncolors(result.result)
+    @assert size(R, 2) == SMC.ncolors(result.result)
     fill!(R, 0.0)
     for i in 1:N
         if color[i] > 0
@@ -190,7 +191,7 @@ function _recover_from_matmat!(
     stored_values::AbstractVector{T},
 ) where {T}
     tree_result = result.result
-    color = SparseMatrixColorings.column_colors(tree_result)
+    color = SMC.column_colors(tree_result)
     N = length(result.local_indices)
     # Compute number of off-diagonal nonzeros from the length of V
     # V contains N diagonal elements + nnz_offdiag off-diagonal elements
diff --git a/src/types.jl b/src/types.jl
index 776ad83..d7373f9 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -63,7 +63,7 @@ function _subexpression_and_linearity(
     linearity
 end
 
-struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
+struct _FunctionStorage{R<:SMC.AbstractColoringResult}
     expr::_SubexpressionStorage
     grad_sparsity::Vector{Int}
     # Nonzero pattern of Hessian matrix
@@ -80,7 +80,7 @@ struct _FunctionStorage{R<:SparseMatrixColorings.AbstractColoringResult}
         coloring_storage::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet,
         coloring_algorithm::Union{
             Nothing,
-            SparseMatrixColorings.GreedyColoringAlgorithm,
+            SMC.GreedyColoringAlgorithm,
         },
         subexpressions::Vector{_SubexpressionStorage},
         dependent_subexpressions,
@@ -141,7 +141,7 @@ end
     NLPEvaluator(
         model::Nonlinear.Model,
         ordered_variables::Vector{MOI.VariableIndex},
-        coloring_algorithm::SparseMatrixColorings.AbstractColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(; decompression=:substitution),
+        coloring_algorithm::SMC.AbstractColoringAlgorithm = SMC.GreedyColoringAlgorithm(; decompression=:substitution),
     )
 
 Return an `NLPEvaluator` object that implements the `MOI.AbstractNLPEvaluator`
@@ -152,7 +152,7 @@ interface.
 """
 mutable struct NLPEvaluator{
     R,
-    C<:SparseMatrixColorings.GreedyColoringAlgorithm,
+    C<:SMC.GreedyColoringAlgorithm,
 } <: MOI.AbstractNLPEvaluator
     data::Nonlinear.Model
     ordered_variables::Vector{MOI.VariableIndex}
@@ -193,18 +193,18 @@ mutable struct NLPEvaluator{
     function NLPEvaluator(
         data::Nonlinear.Model,
         ordered_variables::Vector{MOI.VariableIndex},
-        coloring_algorithm::SparseMatrixColorings.GreedyColoringAlgorithm = SparseMatrixColorings.GreedyColoringAlgorithm(;
+        coloring_algorithm::SMC.GreedyColoringAlgorithm = SMC.GreedyColoringAlgorithm(;
             decompression = :substitution,
         ),
     )
-        problem = SparseMatrixColorings.ColoringProblem(;
+        problem = SMC.ColoringProblem(;
             structure = :symmetric,
             partition = :column,
         )
         C = typeof(coloring_algorithm)
         R = Base.promote_op(
-            SparseMatrixColorings.coloring,
-            SparseArrays.SparseMatrixCSC{Bool,Int},
+            SMC.coloring,
+            SMC.SparsityPatternCSC{Int},
             typeof(problem),
             C,
         )
diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl
index 129bcf8..e2605e0 100644
--- a/test/ReverseAD.jl
+++ b/test/ReverseAD.jl
@@ -481,7 +481,7 @@ end
 function test_coloring_end_to_end_hessian_coloring_and_recovery()
     # Test the new coloring API through the compatibility layer
     coloring_algorithm =
-        ArrayDiff.SparseMatrixColorings.GreedyColoringAlgorithm(;
+        ArrayDiff.SMC.GreedyColoringAlgorithm(;
             decompression = :substitution,
         )
     I, J, rinfo = ArrayDiff._hessian_color_preprocess(

From ffb58d6e464860d4fe53339b589dd8001236e7b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 9 Jan 2026 16:45:30 +0100
Subject: [PATCH 17/17] Remove LA import

---
 src/ArrayDiff.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 8b8af5a..2b78193 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -6,7 +6,6 @@
 
 module ArrayDiff
 
-import LinearAlgebra as LA
 import SparseArrays
 import SparseMatrixColorings as SMC
 import ForwardDiff