From b6c38d998f531b248c7322a452d0122651eb1ed3 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Mon, 12 Aug 2024 20:33:40 +0200
Subject: [PATCH] Make decompression more efficient (#71)

* Make decompression more efficient

* Add skipped test for decompression into dense matrix
---
 docs/src/dev.md              |   4 +-
 src/SparseMatrixColorings.jl |   3 +-
 src/decompression.jl         | 192 +++++++++++-----------------
 src/graph.jl                 |   6 +-
 src/interface.jl             | 108 +++++++++-------
 src/result.jl                | 240 ++++++++++++++++++++++++++++-------
 src/sparsematrixcsc.jl       | 107 ----------------
 test/graph.jl                |   4 +-
 test/performance.jl          |  92 ++++++++++----
 test/random.jl               |  27 ++--
 test/small.jl                |  28 ++--
 test/utils.jl                |  16 ++-
 12 files changed, 430 insertions(+), 397 deletions(-)
 delete mode 100644 src/sparsematrixcsc.jl

diff --git a/docs/src/dev.md b/docs/src/dev.md
index f2f0a05..bad5217 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -34,10 +34,10 @@ SparseMatrixColorings.TreeSet
 ## Concrete coloring results
 
 ```@docs
-SparseMatrixColorings.DefaultColoringResult
+SparseMatrixColorings.NonSymmetricColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
-SparseMatrixColorings.DirectSparseColoringResult
+SparseMatrixColorings.LinearSystemColoringResult
 ```
 
 ## Testing
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 39586fa..10552b4 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -20,7 +20,9 @@ using LinearAlgebra:
     Transpose,
     adjoint,
     checksquare,
+    factorize,
     issymmetric,
+    ldiv!,
     parent,
     transpose
 using Random: AbstractRNG, default_rng, randperm
@@ -45,7 +47,6 @@ include("matrices.jl")
 include("interface.jl")
 include("decompression.jl")
 include("check.jl")
-include("sparsematrixcsc.jl")
 include("examples.jl")
 
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
diff --git a/src/decompression.jl b/src/decompression.jl
index 6392cae..55962e6 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -7,7 +7,7 @@ Compress `A` given a coloring `result` of the sparsity pattern of `A`.
 - If `result` comes from a `:bidirectional` partition, the output is a tuple of matrices `(Br, Bc)`, where `Br` is compressed by row and `Bc` by column.
 
 Compression means summing either the columns or the rows of `A` which share the same color.
-It is undone by calling [`decompress`](@ref).
+It is undone by calling [`decompress`](@ref) or [`decompress!`](@ref).
 
 !!! warning
     At the moment, `:bidirectional` partitions are not implemented.
@@ -67,6 +67,7 @@ end
     decompress(B::AbstractMatrix, result::AbstractColoringResult)
 
 Decompress `B` into a new matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+The in-place alternative is [`decompress!`](@ref).
 
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
@@ -127,10 +128,14 @@ end
     )
 
 Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
+The out-of-place alternative is [`decompress`](@ref).
 
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
 
+!!! note
+    In-place decompression is faster when `A isa SparseMatrixCSC`.
+
 # Example
 
 ```jldoctest
@@ -190,10 +195,10 @@ function decompress!(
     return decompress_aux!(A, B, result)
 end
 
+## NonSymmetricColoringResult
+
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:nonsymmetric,:column,:direct},
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -209,9 +214,7 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:nonsymmetric,:row,:direct},
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -227,24 +230,31 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:symmetric,:column,:direct},
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
-    A .= zero(R)
-    S = get_matrix(result)
-    color = column_colors(result)
-    group = column_groups(result)
-    for ij in findall(!iszero, S)
-        i, j = Tuple(ij)
-        k, l = symmetric_coefficient(i, j, color, group, S)
-        A[i, j] = B[k, l]
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
+    end
+    return A
+end
+
+function decompress_aux!(
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
+) where {R<:Real}
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
     end
     return A
 end
 
+## StarSetColoringResult
+
 function decompress_aux!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult{:column}
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -258,117 +268,35 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:symmetric,:column,:substitution},
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
-    # build T such that T * strict_upper_nonzeros(A) = B
-    # and solve a linear least-squares problem
-    # only consider the strict upper triangle of A because of symmetry
-    # TODO: make more efficient
-    A .= zero(R)
-    S = sparse(get_matrix(result))
-    color = column_colors(result)
-
-    n = checksquare(S)
-    strict_upper_nonzero_inds = Tuple{Int,Int}[]
-    I, J, _ = findnz(S)
-    for (i, j) in zip(I, J)
-        (i < j) && push!(strict_upper_nonzero_inds, (i, j))
-        (i == j) && (A[i, i] = B[i, color[i]])
-    end
-
-    T = spzeros(float(R), length(B), length(strict_upper_nonzero_inds))
-    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
-        ci = color[i]
-        cj = color[j]
-        ki = (ci - 1) * n + j  # A[i, j] appears in B[j, ci]
-        kj = (cj - 1) * n + i  # A[i, j] appears in B[i, cj]
-        T[ki, l] = one(float(R))
-        T[kj, l] = one(float(R))
-    end
-
-    strict_upper_nonzeros_A = T \ vec(B)
-    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
-        A[i, j] = strict_upper_nonzeros_A[l]
-        A[j, i] = strict_upper_nonzeros_A[l]
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
     end
     return A
 end
 
+## TreeSetColoringResult
+
+# TODO: add method for A::SparseMatrixCSC
+
 function decompress_aux!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
 ) where {R<:Real}
-    n = checksquare(A)
     A .= zero(R)
     S = get_matrix(result)
     color = column_colors(result)
-
-    # forest is a structure DisjointSets from DataStructures.jl
-    # - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
-    # - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
-    # - forest.internal.ngroups: the number of trees in the forest
-    forest = result.tree_set.forest
-    ntrees = forest.internal.ngroups
-
-    # vector of trees where each tree contains the indices of its edges
-    trees = [Int[] for i in 1:ntrees]
-
-    # dictionary that maps a tree's root to the index of the tree
-    roots = Dict{Int,Int}()
-
-    k = 0
-    for edge in forest.revmap
-        root_edge = find_root!(forest, edge)
-        root = forest.intmap[root_edge]
-        if !haskey(roots, root)
-            k += 1
-            roots[root] = k
-        end
-        index_tree = roots[root]
-        push!(trees[index_tree], forest.intmap[edge])
-    end
-
-    # vector of dictionaries where each dictionary stores the degree of each vertex in a tree
-    degrees = [Dict{Int,Int}() for k in 1:ntrees]
-    for k in 1:ntrees
-        tree = trees[k]
-        degree = degrees[k]
-        for edge_index in tree
-            i, j = forest.revmap[edge_index]
-            !haskey(degree, i) && (degree[i] = 0)
-            !haskey(degree, j) && (degree[j] = 0)
-            degree[i] += 1
-            degree[j] += 1
-        end
-    end
-
-    # depth-first search (DFS) traversal order for each tree in the forest
-    dfs_orders = [Vector{Tuple{Int,Int}}() for k in 1:ntrees]
-    for k in 1:ntrees
-        tree = trees[k]
-        degree = degrees[k]
-        while sum(values(degree)) != 0
-            for (t, edge_index) in enumerate(tree)
-                if edge_index != 0
-                    i, j = forest.revmap[edge_index]
-                    if (degree[i] == 1) || (degree[j] == 1)  # leaf vertex
-                        if degree[i] > degree[j]  # vertex i is the parent of vertex j
-                            i, j = j, i  # ensure that i always denotes a leaf vertex
-                        end
-                        degree[i] -= 1  # decrease the degree of vertex i
-                        degree[j] -= 1  # decrease the degree of vertex j
-                        tree[t] = 0  # remove the edge (i,j)
-                        push!(dfs_orders[k], (i, j))
-                    end
-                end
-            end
-        end
-    end
+    @compat (; degrees, dfs_orders, stored_values) = result
 
     # stored_values holds the sum of edge values for subtrees in a tree.
     # For each vertex i, stored_values[i] is the sum of edge values in the subtree rooted at i.
-    stored_values = Vector{R}(undef, n)
+    stored_values_right_type = if R == Float64
+        stored_values
+    else
+        similar(stored_values, R)
+    end
 
     # Recover the diagonal coefficients of A
     for i in axes(A, 1)
@@ -378,19 +306,43 @@ function decompress_aux!(
     end
 
     # Recover the off-diagonal coefficients of A
-    for k in 1:ntrees
+    for k in eachindex(degrees, dfs_orders)
         vertices = keys(degrees[k])
         for vertex in vertices
-            stored_values[vertex] = zero(R)
+            stored_values_right_type[vertex] = zero(R)
         end
 
         tree = dfs_orders[k]
         for (i, j) in tree
-            val = B[i, color[j]] - stored_values[i]
-            stored_values[j] = stored_values[j] + val
+            val = B[i, color[j]] - stored_values_right_type[i]
+            stored_values_right_type[j] = stored_values_right_type[j] + val
             A[i, j] = val
             A[j, i] = val
         end
     end
     return A
 end
+
+## MatrixInverseColoringResult
+
+function decompress_aux!(
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult
+) where {R<:Real}
+    S = get_matrix(result)
+    color = column_colors(result)
+    @compat (; strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) = result
+
+    # TODO: for some reason I cannot use ldiv! with a sparse QR
+    strict_upper_nonzeros_A = T_factorization \ vec(B)
+    A .= zero(R)
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
+        A[i, j] = strict_upper_nonzeros_A[l]
+        A[j, i] = strict_upper_nonzeros_A[l]
+    end
+    return A
+end
diff --git a/src/graph.jl b/src/graph.jl
index ffbc61d..c1932c2 100644
--- a/src/graph.jl
+++ b/src/graph.jl
@@ -16,7 +16,6 @@ struct Graph{T<:Integer}
 end
 
 Graph(S::SparseMatrixCSC) = Graph(S.colptr, S.rowval)
-Graph(S::AbstractMatrix) = Graph(sparse(S))
 
 Base.length(g::Graph) = length(g.colptr) - 1
 SparseArrays.nnz(g::Graph) = length(g.rowval)
@@ -91,7 +90,6 @@ The adjacency graph of a symmetrix matric `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 adjacency_graph(H::SparseMatrixCSC) = Graph(H - Diagonal(H))
-adjacency_graph(H::AbstractMatrix) = adjacency_graph(sparse(H))
 
 """
     bipartite_graph(J::AbstractMatrix)
@@ -109,9 +107,7 @@ The bipartite graph of a matrix `A ∈ ℝ^{m × n}` is `Gb(A) = (V₁, V₂, E)
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 function bipartite_graph(J::SparseMatrixCSC)
-    g1 = Graph(transpose(J))  # rows to columns
+    g1 = Graph(sparse(transpose(J)))  # rows to columns
     g2 = Graph(J)  # columns to rows
     return BipartiteGraph(g1, g2)
 end
-
-bipartite_graph(J::AbstractMatrix) = bipartite_graph(sparse(J))
diff --git a/src/interface.jl b/src/interface.jl
index 7866dec..250811c 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -1,5 +1,5 @@
 """
-    ColoringProblem{structure,partition,decompression}
+    ColoringProblem{structure,partition}
 
 Selector type for the coloring problem to solve, enabling multiple dispatch.
 
@@ -7,17 +7,10 @@ It is passed as an argument to the main function [`coloring`](@ref).
 
 # Constructor
 
-    ColoringProblem(;
-        structure::Symbol=:nonsymmetric,
-        partition::Symbol=:column,
-        decompression::Symbol=:direct,
-    )
-
-# Type parameters
+    ColoringProblem(; structure::Symbol=:nonsymmetric, partition::Symbol=:column)
 
 - `structure::Symbol`: either `:nonsymmetric` or `:symmetric`
 - `partition::Symbol`: either `:column`, `:row` or `:bidirectional`
-- `decompression::Symbol`: either `:direct` or `:substitution`
 
 #  Link to automatic differentiation
 
@@ -30,24 +23,22 @@ Matrix coloring is often used in automatic differentiation, and here is the tran
 | Jacobian | forward + reverse    | `:nonsymmetric` | `:bidirectional` |
 | Hessian  | any                  | `:symmetric`    | `:column`        |
 
+!!! warning
+    With a `:symmetric` structure, you have to use a `:column` partition.
+
 !!! warning
     At the moment, `:bidirectional` partitions are not implemented.
 """
-struct ColoringProblem{structure,partition,decompression} end
+struct ColoringProblem{structure,partition} end
 
-function ColoringProblem(;
-    structure::Symbol=:nonsymmetric,
-    partition::Symbol=:column,
-    decompression::Symbol=:direct,
-)
+function ColoringProblem(; structure::Symbol=:nonsymmetric, partition::Symbol=:column)
     @assert structure in (:nonsymmetric, :symmetric)
     @assert partition in (:column, :row, :bidirectional)
-    @assert decompression in (:direct, :substitution)
-    return ColoringProblem{structure,partition,decompression}()
+    return ColoringProblem{structure,partition}()
 end
 
 """
-    GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm
+    GreedyColoringAlgorithm{decompression} <: ADTypes.AbstractColoringAlgorithm
 
 Greedy coloring algorithm for sparse matrices which colors columns or rows one after the other, following a configurable order.
 
@@ -55,10 +46,13 @@ It is passed as an argument to the main function [`coloring`](@ref).
 
 # Constructor
 
-    GreedyColoringAlgorithm(order::AbstractOrder=NaturalOrder())
+    GreedyColoringAlgorithm(
+        order::AbstractOrder=NaturalOrder();
+        decompression::Symbol=:direct
+    )
 
-The choice of [`AbstractOrder`](@ref) impacts the resulting number of colors.
-It defaults to [`NaturalOrder`](@ref) for reproducibility, but [`LargestFirst`](@ref) can sometimes be a better option.
+- `order::AbstractOrder`: the order in which the columns or rows are colored, which can impact the number of colors.
+- `decompression::Symbol`: either `:direct` or `:substitution`. Usually `:substitution` leads to fewer colors, at the cost of a more expensive coloring (and decompression). When `:substitution` is not applicable, it falls back on `:direct` decompression.
 
 # ADTypes coloring interface
 
@@ -69,22 +63,37 @@ It defaults to [`NaturalOrder`](@ref) for reproducibility, but [`LargestFirst`](
 - [`ADTypes.symmetric_coloring`](@extref ADTypes.symmetric_coloring)
 
 See their respective docstrings for details.
+
+# See also
+
+- [`AbstractOrder`](@ref)
+- [`decompress`](@ref)
 """
-struct GreedyColoringAlgorithm{O<:AbstractOrder} <: ADTypes.AbstractColoringAlgorithm
+struct GreedyColoringAlgorithm{decompression,O<:AbstractOrder} <:
+       ADTypes.AbstractColoringAlgorithm
     order::O
 end
 
-GreedyColoringAlgorithm() = GreedyColoringAlgorithm(NaturalOrder())
+function GreedyColoringAlgorithm(
+    order::AbstractOrder=NaturalOrder(); decompression::Symbol=:direct
+)
+    @assert decompression in (:direct, :substitution)
+    return GreedyColoringAlgorithm{decompression,typeof(order)}(order)
+end
 
 """
     coloring(
         S::AbstractMatrix,
         problem::ColoringProblem,
-        algo::GreedyColoringAlgorithm
+        algo::GreedyColoringAlgorithm;
+        [decompression_eltype=Float64]
     )
 
 Solve a [`ColoringProblem`](@ref) on the matrix `S` with a [`GreedyColoringAlgorithm`](@ref) and return an [`AbstractColoringResult`](@ref).
 
+The result can be used to [`compress`](@ref) and [`decompress`](@ref) a matrix `A` with the same sparsity pattern as `S`.
+If `eltype(A) == decompression_eltype`, decompression might be faster.
+
 # Example
 
 ```jldoctest
@@ -124,64 +133,77 @@ julia> column_groups(result)
 - [`ColoringProblem`](@ref)
 - [`GreedyColoringAlgorithm`](@ref)
 - [`AbstractColoringResult`](@ref)
+- [`compress`](@ref)
+- [`decompress`](@ref)
 """
 function coloring end
 
 function coloring(
-    S::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:column,:direct},
-    algo::GreedyColoringAlgorithm,
+    A::AbstractMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    decompression_eltype=Float64,
 )
+    S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(2), algo.order)
-    return DefaultColoringResult{:nonsymmetric,:column,:direct}(S, color)
+    return NonSymmetricColoringResult{:column}(S, color)
 end
 
 function coloring(
-    S::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:row,:direct},
-    algo::GreedyColoringAlgorithm,
+    A::AbstractMatrix,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    decompression_eltype=Float64,
 )
+    S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(1), algo.order)
-    return DefaultColoringResult{:nonsymmetric,:row,:direct}(S, color)
+    return NonSymmetricColoringResult{:row}(S, color)
 end
 
 function coloring(
-    S::AbstractMatrix,
-    ::ColoringProblem{:symmetric,:column,:direct},
-    algo::GreedyColoringAlgorithm,
+    A::AbstractMatrix,
+    ::ColoringProblem{:symmetric,:column},
+    algo::GreedyColoringAlgorithm{:direct};
+    decompression_eltype=Float64,
 )
+    S = sparse(A)
     ag = adjacency_graph(S)
     color, star_set = star_coloring(ag, algo.order)
-    return StarSetColoringResult{:column}(S, color, star_set)
+    return StarSetColoringResult(S, color, star_set)
 end
 
 function coloring(
-    S::AbstractMatrix,
-    ::ColoringProblem{:symmetric,:column,:substitution},
-    algo::GreedyColoringAlgorithm,
+    A::AbstractMatrix,
+    ::ColoringProblem{:symmetric,:column},
+    algo::GreedyColoringAlgorithm{:substitution};
+    decompression_eltype=Float64,
 )
+    S = sparse(A)
     ag = adjacency_graph(S)
     color, tree_set = acyclic_coloring(ag, algo.order)
-    return TreeSetColoringResult{:column}(S, color, tree_set)
+    return TreeSetColoringResult(S, color, tree_set, decompression_eltype)
 end
 
 ## ADTypes interface
 
-function ADTypes.column_coloring(S::AbstractMatrix, algo::GreedyColoringAlgorithm)
+function ADTypes.column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+    S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(2), algo.order)
     return color
 end
 
-function ADTypes.row_coloring(S::AbstractMatrix, algo::GreedyColoringAlgorithm)
+function ADTypes.row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+    S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(1), algo.order)
     return color
 end
 
-function ADTypes.symmetric_coloring(S::AbstractMatrix, algo::GreedyColoringAlgorithm)
+function ADTypes.symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+    S = sparse(A)
     ag = adjacency_graph(S)
     color, star_set = star_coloring(ag, algo.order)
     return color
diff --git a/src/result.jl b/src/result.jl
index 39a124a..6acf93b 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -9,7 +9,7 @@ It is the supertype of the object returned by the main function [`coloring`](@re
 
 # Type parameters
 
-Same as those of the [`ColoringProblem`](@ref) that was solved to obtain the result:
+Combination between the type parameters of [`ColoringProblem`](@ref) and [`GreedyColoringAlgorithm`](@ref):
 
 - `structure::Symbol`: either `:nonsymmetric` or `:symmetric`
 - `partition::Symbol`: either `:column`, `:row` or `:bidirectional`
@@ -19,12 +19,12 @@ Same as those of the [`ColoringProblem`](@ref) that was solved to obtain the res
 
 - [`column_colors`](@ref) and [`column_groups`](@ref) (for a `:column` or `:bidirectional` partition) 
 - [`row_colors`](@ref) and [`row_groups`](@ref) (for a `:row` or `:bidirectional` partition)
-- [`decompress`](@ref) and [`decompress!`](@ref)
+- [`compress`](@ref), [`decompress`](@ref) and [`decompress!`](@ref)
 
 !!! warning
     Unlike the methods above, the concrete subtypes of `AbstractColoringResult` are not part of the public API and may change without notice.
 """
-abstract type AbstractColoringResult{structure,partition,decompression,M<:AbstractMatrix} end
+abstract type AbstractColoringResult{structure,partition,decompression,M<:SparseMatrixCSC} end
 
 """
     get_matrix(result::AbstractColoringResult)
@@ -78,7 +78,7 @@ function group_by_color(color::AbstractVector{<:Integer})
     return group
 end
 
-get_matrix(result::AbstractColoringResult) = result.matrix
+get_matrix(result::AbstractColoringResult) = result.S
 
 column_colors(result::AbstractColoringResult{s,:column}) where {s} = result.color
 column_groups(result::AbstractColoringResult{s,:column}) where {s} = result.group
@@ -91,7 +91,7 @@ row_groups(result::AbstractColoringResult{s,:row}) where {s} = result.group
 """
 $TYPEDEF
 
-Default storage for the result of a coloring algorithm, containing minimal information.
+Storage for the result of a nonsymmetric coloring with direct decompression.
 
 # Fields
 
@@ -101,46 +101,88 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct DefaultColoringResult{structure,partition,decompression,M} <:
-       AbstractColoringResult{structure,partition,decompression,M}
+struct NonSymmetricColoringResult{partition,M} <:
+       AbstractColoringResult{:nonsymmetric,partition,:direct,M}
     "matrix that was colored"
-    matrix::M
+    S::M
     "one integer color for each column or row (depending on `partition`)"
     color::Vector{Int}
     "color groups for columns or rows (depending on `partition`)"
     group::Vector{Vector{Int}}
+    "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A` when `A isa SparseMatrixCSC`. They satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
+    compressed_indices::Vector{Int}
 end
 
-function DefaultColoringResult{structure,partition,decompression}(
-    matrix::M, color::Vector{Int}
-) where {structure,partition,decompression,M}
-    return DefaultColoringResult{structure,partition,decompression,M}(
-        matrix, color, group_by_color(color)
+function NonSymmetricColoringResult{:column}(S::SparseMatrixCSC, color::Vector{Int})
+    group = group_by_color(color)
+    n = size(S, 1)
+    I, J, _ = findnz(S)
+    compressed_indices = zeros(Int, nnz(S))
+    for k in eachindex(I, J, compressed_indices)
+        i, j = I[k], J[k]
+        c = color[j]
+        # A[i, j] = B[i, c]
+        compressed_indices[k] = (c - 1) * n + i
+    end
+    return NonSymmetricColoringResult{:column,typeof(S)}(
+        S, color, group, compressed_indices
     )
 end
 
-function DefaultColoringResult(
-    result::AbstractColoringResult{structure,:column,decompression}
-) where {structure,decompression}
-    return DefaultColoringResult{structure,:column,decompression}(
-        get_matrix(result), column_colors(result)
-    )
+function NonSymmetricColoringResult{:row}(S::SparseMatrixCSC, color::Vector{Int})
+    group = group_by_color(color)
+    C = maximum(color)
+    I, J, _ = findnz(S)
+    compressed_indices = zeros(Int, nnz(S))
+    for k in eachindex(I, J, compressed_indices)
+        i, j = I[k], J[k]
+        c = color[i]
+        # A[i, j] = B[c, j]
+        compressed_indices[k] = (j - 1) * C + c
+    end
+    return NonSymmetricColoringResult{:row,typeof(S)}(S, color, group, compressed_indices)
 end
 
-function DefaultColoringResult(
-    result::AbstractColoringResult{structure,:row,decompression}
-) where {structure,decompression}
-    return DefaultColoringResult{structure,:row,decompression}(
-        get_matrix(result), row_colors(result)
-    )
+"""
+$TYPEDEF
+
+Storage for the result of a symmetric coloring with direct decompression.
+
+# Fields
+
+$TYPEDFIELDS
+
+# See also
+
+- [`AbstractColoringResult`](@ref)
+- [`NonSymmetricColoringResult`](@ref)
+"""
+struct StarSetColoringResult{M} <: AbstractColoringResult{:symmetric,:column,:direct,M}
+    S::M
+    color::Vector{Int}
+    group::Vector{Vector{Int}}
+    star_set::StarSet
+    compressed_indices::Vector{Int}
+end
+
+function StarSetColoringResult(S::SparseMatrixCSC, color::Vector{Int}, star_set::StarSet)
+    group = group_by_color(color)
+    n = size(S, 1)
+    I, J, _ = findnz(S)
+    compressed_indices = zeros(Int, nnz(S))
+    for k in eachindex(I, J, compressed_indices)
+        i, j = I[k], J[k]
+        l, c = symmetric_coefficient(i, j, color, star_set)
+        # A[i, j] = B[l, c]
+        compressed_indices[k] = (c - 1) * n + l
+    end
+    return StarSetColoringResult(S, color, group, star_set, compressed_indices)
 end
 
 """
 $TYPEDEF
 
-Storage for the result of a symmetric coloring algorithm with direct decompression.
-
-Similar to [`DefaultColoringResult`](@ref) but contains an additional [`StarSet`](@ref).
+Storage for the result of a symmetric coloring with decompression by substitution.
 
 # Fields
 
@@ -149,29 +191,100 @@ $TYPEDFIELDS
 # See also
 
 - [`AbstractColoringResult`](@ref)
+- [`NonSymmetricColoringResult`](@ref)
 """
-struct StarSetColoringResult{partition,M} <:
-       AbstractColoringResult{:symmetric,partition,:direct,M}
-    matrix::M
+struct TreeSetColoringResult{M,R} <:
+       AbstractColoringResult{:symmetric,:column,:substitution,M}
+    S::M
     color::Vector{Int}
     group::Vector{Vector{Int}}
-    star_set::StarSet
+    tree_set::TreeSet
+    degrees::Vector{Dict{Int,Int}}
+    dfs_orders::Vector{Vector{Tuple{Int,Int}}}
+    stored_values::Vector{R}
 end
 
-function StarSetColoringResult{partition}(
-    matrix::M, color::Vector{Int}, star_set::StarSet
-) where {partition,M}
-    return StarSetColoringResult{partition,M}(
-        matrix, color, group_by_color(color), star_set
+function TreeSetColoringResult(
+    S::SparseMatrixCSC, color::Vector{Int}, tree_set::TreeSet, decompression_eltype::Type{R}
+) where {R}
+    group = group_by_color(color)
+
+    # forest is a structure DisjointSets from DataStructures.jl
+    # - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
+    # - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
+    # - forest.internal.ngroups: the number of trees in the forest
+    forest = tree_set.forest
+    ntrees = forest.internal.ngroups
+
+    # vector of trees where each tree contains the indices of its edges
+    trees = [Int[] for i in 1:ntrees]
+
+    # dictionary that maps a tree's root to the index of the tree
+    roots = Dict{Int,Int}()
+
+    k = 0
+    for edge in forest.revmap
+        root_edge = find_root!(forest, edge)
+        root = forest.intmap[root_edge]
+        if !haskey(roots, root)
+            k += 1
+            roots[root] = k
+        end
+        index_tree = roots[root]
+        push!(trees[index_tree], forest.intmap[edge])
+    end
+
+    # vector of dictionaries where each dictionary stores the degree of each vertex in a tree
+    degrees = [Dict{Int,Int}() for k in 1:ntrees]
+    for k in 1:ntrees
+        tree = trees[k]
+        degree = degrees[k]
+        for edge_index in tree
+            i, j = forest.revmap[edge_index]
+            !haskey(degree, i) && (degree[i] = 0)
+            !haskey(degree, j) && (degree[j] = 0)
+            degree[i] += 1
+            degree[j] += 1
+        end
+    end
+
+    # depth-first search (DFS) traversal order for each tree in the forest
+    dfs_orders = [Vector{Tuple{Int,Int}}() for k in 1:ntrees]
+    for k in 1:ntrees
+        tree = trees[k]
+        degree = degrees[k]
+        while sum(values(degree)) != 0
+            for (t, edge_index) in enumerate(tree)
+                if edge_index != 0
+                    i, j = forest.revmap[edge_index]
+                    if (degree[i] == 1) || (degree[j] == 1)  # leaf vertex
+                        if degree[i] > degree[j]  # vertex i is the parent of vertex j
+                            i, j = j, i  # ensure that i always denotes a leaf vertex
+                        end
+                        degree[i] -= 1  # decrease the degree of vertex i
+                        degree[j] -= 1  # decrease the degree of vertex j
+                        tree[t] = 0  # remove the edge (i,j)
+                        push!(dfs_orders[k], (i, j))
+                    end
+                end
+            end
+        end
+    end
+
+    n = checksquare(S)
+    stored_values = Vector{R}(undef, n)
+
+    return TreeSetColoringResult(
+        S, color, group, tree_set, degrees, dfs_orders, stored_values
     )
 end
 
+## LinearSystemColoringResult
+
 """
 $TYPEDEF
 
-Storage for the result of a symmetric coloring algorithm with decompression by substitution.
-
-Similar to [`DefaultColoringResult`](@ref) but contains an additional [`TreeSet`](@ref).
+Storage for the result of a symmetric coloring with any decompression.
 
 # Fields
 
@@ -180,19 +293,48 @@ $TYPEDFIELDS
 # See also
 
 - [`AbstractColoringResult`](@ref)
+- [`NonSymmetricColoringResult`](@ref)
 """
-struct TreeSetColoringResult{partition,M} <:
-       AbstractColoringResult{:symmetric,partition,:substitution,M}
-    matrix::M
+struct LinearSystemColoringResult{M,R,F} <:
+       AbstractColoringResult{:symmetric,:column,:substitution,M}
+    S::M
     color::Vector{Int}
     group::Vector{Vector{Int}}
-    tree_set::TreeSet
+    strict_upper_nonzero_inds::Vector{Tuple{Int,Int}}
+    strict_upper_nonzeros_A::Vector{R}  # TODO: adjust type
+    T_factorization::F  # TODO: adjust type
 end
 
-function TreeSetColoringResult{partition}(
-    matrix::M, color::Vector{Int}, tree_set::TreeSet
-) where {partition,M}
-    return TreeSetColoringResult{partition,M}(
-        matrix, color, group_by_color(color), tree_set
+function LinearSystemColoringResult(
+    S::SparseMatrixCSC, color::Vector{Int}, decompression_eltype::Type{R}
+) where {R}
+    group = group_by_color(color)
+    C = maximum(color)
+
+    # build T such that T * strict_upper_nonzeros(A) = B
+    # and solve a linear least-squares problem
+    # only consider the strict upper triangle of A because of symmetry
+    n = checksquare(S)
+    strict_upper_nonzero_inds = Tuple{Int,Int}[]
+    I, J, _ = findnz(S)
+    for (i, j) in zip(I, J)
+        (i < j) && push!(strict_upper_nonzero_inds, (i, j))
+    end
+
+    T = spzeros(Float64, n * C, length(strict_upper_nonzero_inds))
+    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
+        ci = color[i]
+        cj = color[j]
+        ki = (ci - 1) * n + j  # A[i, j] appears in B[j, ci]
+        kj = (cj - 1) * n + i  # A[i, j] appears in B[i, cj]
+        T[ki, l] = 1
+        T[kj, l] = 1
+    end
+    T_factorization = factorize(T)
+
+    strict_upper_nonzeros_A = Vector{R}(undef, size(T, 2))
+
+    return LinearSystemColoringResult(
+        S, color, group, strict_upper_nonzero_inds, strict_upper_nonzeros_A, T_factorization
     )
 end
diff --git a/src/sparsematrixcsc.jl b/src/sparsematrixcsc.jl
deleted file mode 100644
index c49c7fb..0000000
--- a/src/sparsematrixcsc.jl
+++ /dev/null
@@ -1,107 +0,0 @@
-## Result
-
-"""
-$TYPEDEF
-
-Storage for the result of a coloring algorithm when the decompression target is a `SparseMatrixCSC`.
-
-# Fields
-
-$TYPEDFIELDS
-
-# See also
-
-- [`AbstractColoringResult`](@ref)
-"""
-struct DirectSparseColoringResult{structure,partition,M} <:
-       AbstractColoringResult{structure,partition,:direct,M}
-    "matrix that was colored"
-    matrix::M
-    "one integer color for each column or row (depending on `partition`)"
-    color::Vector{Int}
-    "color groups for columns or rows (depending on `partition`)"
-    group::Vector{Vector{Int}}
-    "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A`: they satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
-    compressed_indices::Vector{Int}
-end
-
-function DirectSparseColoringResult{structure,partition}(
-    matrix::M, color::Vector{Int}, compressed_indices::Vector{Int}
-) where {structure,partition,M}
-    return DirectSparseColoringResult{structure,partition,M}(
-        matrix, color, group_by_color(color), compressed_indices
-    )
-end
-
-## Coloring
-
-function coloring(
-    S::SparseMatrixCSC,
-    ::ColoringProblem{:nonsymmetric,:column,:direct},
-    algo::GreedyColoringAlgorithm,
-)
-    bg = bipartite_graph(S)
-    color = partial_distance2_coloring(bg, Val(2), algo.order)
-    n = size(S, 1)
-    I, J, _ = findnz(S)
-    compressed_indices = zeros(Int, nnz(S))
-    for k in eachindex(I, J, compressed_indices)
-        i, j = I[k], J[k]
-        c = color[j]
-        # A[i, j] = B[i, c]
-        compressed_indices[k] = (c - 1) * n + i
-    end
-    return DirectSparseColoringResult{:nonsymmetric,:column}(S, color, compressed_indices)
-end
-
-function coloring(
-    S::SparseMatrixCSC,
-    ::ColoringProblem{:nonsymmetric,:row,:direct},
-    algo::GreedyColoringAlgorithm,
-)
-    bg = bipartite_graph(S)
-    color = partial_distance2_coloring(bg, Val(1), algo.order)
-    n = size(S, 1)
-    C = maximum(color)
-    I, J, _ = findnz(S)
-    compressed_indices = zeros(Int, nnz(S))
-    for k in eachindex(I, J, compressed_indices)
-        i, j = I[k], J[k]
-        c = color[i]
-        # A[i, j] = B[c, j]
-        compressed_indices[k] = (j - 1) * C + c
-    end
-    return DirectSparseColoringResult{:nonsymmetric,:row}(S, color, compressed_indices)
-end
-
-function coloring(
-    S::SparseMatrixCSC,
-    ::ColoringProblem{:symmetric,:column,:direct},
-    algo::GreedyColoringAlgorithm,
-)
-    ag = adjacency_graph(S)
-    color, star_set = star_coloring(ag, algo.order)
-    n = size(S, 1)
-    I, J, _ = findnz(S)
-    compressed_indices = zeros(Int, nnz(S))
-    for k in eachindex(I, J, compressed_indices)
-        i, j = I[k], J[k]
-        l, c = symmetric_coefficient(i, j, color, star_set)
-        # A[i, j] = B[l, c]
-        compressed_indices[k] = (c - 1) * n + l
-    end
-    return DirectSparseColoringResult{:symmetric,:column}(S, color, compressed_indices)
-end
-
-## Decompression
-
-function decompress_aux!(
-    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::DirectSparseColoringResult
-) where {R<:Real}
-    nzA = nonzeros(A)
-    ind = result.compressed_indices
-    for i in eachindex(nzA, ind)
-        nzA[i] = B[ind[i]]
-    end
-    return A
-end
diff --git a/test/graph.jl b/test/graph.jl
index 6a46984..509193d 100644
--- a/test/graph.jl
+++ b/test/graph.jl
@@ -6,11 +6,11 @@ using Test
 ## Standard graph
 
 @testset "Graph" begin
-    g = Graph([
+    g = Graph(sparse([
         1 0 1
         1 1 0
         0 0 0
-    ])
+    ]))
 
     @test length(g) == 3
     @test neighbors(g, 1) == [1, 2]
diff --git a/test/performance.jl b/test/performance.jl
index 3a4e73b..3c4657e 100644
--- a/test/performance.jl
+++ b/test/performance.jl
@@ -14,45 +14,87 @@ rng = StableRNG(63)
 @testset "Coloring - type stability" begin
     n = 10
     A = sprand(rng, Bool, n, n, 3 / n)
-    algo = GreedyColoringAlgorithm()
-    @test_opt target_modules = (SparseMatrixColorings,) column_coloring(A, algo)
-    @test_opt target_modules = (SparseMatrixColorings,) row_coloring(A, algo)
+
+    # ADTypes
+    @test_opt target_modules = (SparseMatrixColorings,) column_coloring(
+        A, GreedyColoringAlgorithm()
+    )
+    @test_opt target_modules = (SparseMatrixColorings,) row_coloring(
+        A, GreedyColoringAlgorithm()
+    )
     @test_opt target_modules = (SparseMatrixColorings,) symmetric_coloring(
-        Symmetric(A), algo
+        Symmetric(A), GreedyColoringAlgorithm()
+    )
+
+    # Native
+    @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        A,
+        ColoringProblem(; structure=:nonsymmetric, partition=:column),
+        GreedyColoringAlgorithm(; decompression=:direct),
+    )
+    @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        A,
+        ColoringProblem(; structure=:nonsymmetric, partition=:row),
+        GreedyColoringAlgorithm(; decompression=:direct),
+    )
+    @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        A,
+        ColoringProblem(; structure=:symmetric, partition=:column),
+        GreedyColoringAlgorithm(; decompression=:direct),
+    )
+    @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        A,
+        ColoringProblem(; structure=:symmetric, partition=:column),
+        GreedyColoringAlgorithm(; decompression=:substitution),
     )
-    @test_opt target_modules = (SparseMatrixColorings,) coloring(A, ColoringProblem(), algo)
 end
 
-function benchmark_distance2_coloring(n)
-    return @be (;
+function test_noallocs_distance2_coloring(n)
+    bench = @be (;
         bg=bipartite_graph(sprand(Bool, n, n, 3 / n)),
         color=Vector{Int}(undef, n),
         forbidden_colors=Vector{Int}(undef, n),
     ) partial_distance2_coloring!(_.color, _.forbidden_colors, _.bg, Val(1), 1:n) evals = 1
+    @test minimum(bench).allocs == 0
 end
 
 @testset "Coloring - allocations" begin
-    @test minimum(benchmark_distance2_coloring(10)).allocs == 0
-    @test minimum(benchmark_distance2_coloring(100)).allocs == 0
-    @test minimum(benchmark_distance2_coloring(1000)).allocs == 0
+    test_noallocs_distance2_coloring(1000)
 end
 
-function benchmark_sparse_decompression(n)
-    A = sprand(n, 2n, 5 / n)
-    result = coloring(
-        A,
-        ColoringProblem(; structure=:nonsymmetric, partition=:column),
-        GreedyColoringAlgorithm(),
-    )
-    group = column_groups(result)
-    B = stack(group; dims=2) do g
-        dropdims(sum(A[:, g]; dims=2); dims=2)
+function test_noallocs_sparse_decompression(
+    n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
+)
+    @testset "$structure - $partition - $decompression" begin
+        A = if structure == :nonsymmetric
+            sprand(n, 2n, 5 / n)
+        elseif structure == :symmetric
+            sparse(Symmetric(sprand(n, n, 5 / n)))
+        end
+        result = coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+        B = compress(A, result)
+        bench1 = @be respectful_similar(A) decompress!(_, B, result) evals = 1
+        @test minimum(bench1).allocs == 0
+        bench2 = @be similar(Matrix(A)) decompress!(_, B, result) evals = 1
+        @test_skip minimum(bench2).allocs == 0
     end
-    return @be respectful_similar(A) decompress!(_, B, result) evals = 1
 end
 
-@testset "SparseMatrixCSC decompression - allocations" begin
-    @test minimum(benchmark_sparse_decompression(10)).allocs == 0
-    @test minimum(benchmark_sparse_decompression(100)).allocs == 0
-    @test minimum(benchmark_sparse_decompression(1000)).allocs == 0
+@testset "Decompression - allocations" begin
+    test_noallocs_sparse_decompression(
+        1000; structure=:nonsymmetric, partition=:column, decompression=:direct
+    )
+    test_noallocs_sparse_decompression(
+        1000; structure=:nonsymmetric, partition=:row, decompression=:direct
+    )
+    test_noallocs_sparse_decompression(
+        1000; structure=:symmetric, partition=:column, decompression=:direct
+    )
+    test_noallocs_sparse_decompression(
+        1000; structure=:symmetric, partition=:column, decompression=:substitution
+    )
 end
diff --git a/test/random.jl b/test/random.jl
index cae7302..7dfcaa9 100644
--- a/test/random.jl
+++ b/test/random.jl
@@ -5,19 +5,14 @@ using LinearAlgebra: I, Symmetric
 using SparseArrays
 using SparseMatrixColorings
 using SparseMatrixColorings:
-    DefaultColoringResult,
     structurally_orthogonal_columns,
     symmetrically_orthogonal_columns,
-    directly_recoverable_columns,
-    matrix_versions,
-    respectful_similar
+    directly_recoverable_columns
 using StableRNGs
 using Test
 
 rng = StableRNG(63)
 
-algo = GreedyColoringAlgorithm()
-
 asymmetric_params = vcat(
     [(10, 20, p) for p in (0.1:0.1:0.5)],
     [(20, 10, p) for p in (0.1:0.1:0.5)],
@@ -31,9 +26,8 @@ symmetric_params = vcat(
 )
 
 @testset "Column coloring & decompression" begin
-    problem = ColoringProblem(;
-        structure=:nonsymmetric, partition=:column, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
     @testset "Size ($m, $n) - sparsity $p" for (m, n, p) in asymmetric_params
         A0 = sprand(rng, m, n, p)
         color0 = column_coloring(A0, algo)
@@ -44,9 +38,8 @@ symmetric_params = vcat(
 end;
 
 @testset "Row coloring & decompression" begin
-    problem = ColoringProblem(;
-        structure=:nonsymmetric, partition=:row, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
     @testset "Size ($m, $n) - sparsity $p" for (m, n, p) in asymmetric_params
         A0 = sprand(rng, m, n, p)
         color0 = row_coloring(A0, algo)
@@ -57,9 +50,8 @@ end;
 end;
 
 @testset "Symmetric coloring & direct decompression" begin
-    problem = ColoringProblem(;
-        structure=:symmetric, partition=:column, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
     @testset "Size ($n, $n) - sparsity $p" for (n, p) in symmetric_params
         A0 = Symmetric(sprand(rng, n, n, p))
         color0 = symmetric_coloring(A0, algo)
@@ -70,9 +62,8 @@ end;
 end;
 
 @testset "Symmetric coloring & substitution decompression" begin
-    problem = ColoringProblem(;
-        structure=:symmetric, partition=:column, decompression=:substitution
-    )
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:substitution)
     @testset "Size ($n, $n) - sparsity $p" for (n, p) in symmetric_params
         A0 = Symmetric(sprand(rng, n, n, p))
         color0 = column_colors(coloring(A0, problem, algo))
diff --git a/test/small.jl b/test/small.jl
index 6f055a2..db449ac 100644
--- a/test/small.jl
+++ b/test/small.jl
@@ -5,25 +5,18 @@ using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
 using SparseMatrixColorings:
-    DefaultColoringResult,
-    group_by_color,
     structurally_orthogonal_columns,
     symmetrically_orthogonal_columns,
     directly_recoverable_columns,
-    matrix_versions,
-    respectful_similar,
     what_fig_41,
     what_fig_61,
     efficient_fig_1,
     efficient_fig_4
 using Test
 
-algo = GreedyColoringAlgorithm()
-
 @testset "Column coloring & decompression" begin
-    problem = ColoringProblem(;
-        structure=:nonsymmetric, partition=:column, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
     A0 = sparse([
         1 0 2
         0 3 4
@@ -41,9 +34,8 @@ algo = GreedyColoringAlgorithm()
 end;
 
 @testset "Row coloring & decompression" begin
-    problem = ColoringProblem(;
-        structure=:nonsymmetric, partition=:row, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
     A0 = sparse([
         1 0 3
         0 2 0
@@ -60,9 +52,9 @@ end;
 end;
 
 @testset "Symmetric coloring & direct decompression" begin
-    problem = ColoringProblem(;
-        structure=:symmetric, partition=:column, decompression=:direct
-    )
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:direct)
+
     @testset "Fig 4.1 from 'What color is your Jacobian'" begin
         example = what_fig_41()
         A0, B0, color0 = example.A, example.B, example.color
@@ -81,9 +73,9 @@ end;
 end;
 
 @testset "Symmetric coloring & substitution decompression" begin
-    problem = ColoringProblem(;
-        structure=:symmetric, partition=:column, decompression=:substitution
-    )
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:substitution)
+
     @testset "Fig 6.1 from 'What color is your Jacobian'" begin
         example = what_fig_61()
         A0, B0, color0 = example.A, example.B, example.color
diff --git a/test/utils.jl b/test/utils.jl
index 717bb9d..6cf92ad 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -1,18 +1,17 @@
 using SparseMatrixColorings
-using SparseMatrixColorings: ColoringProblem, DefaultColoringResult
+using SparseMatrixColorings: LinearSystemColoringResult, matrix_versions, respectful_similar
 using Test
 
 function test_coloring_decompression(
     A0::AbstractMatrix,
-    problem::ColoringProblem{structure,partition,decompression},
-    algo::GreedyColoringAlgorithm;
+    problem::ColoringProblem{structure,partition},
+    algo::GreedyColoringAlgorithm{decompression};
     B0=nothing,
     color0=nothing,
 ) where {structure,partition,decompression}
     color_vec = Vector{Int}[]
     @testset "A::$(typeof(A))" for A in matrix_versions(A0)
-        result = coloring(A, problem, algo)
-        default_result = DefaultColoringResult(result)
+        result = coloring(A, problem, algo; decompression_eltype=eltype(A))
         color = if partition == :column
             column_colors(result)
         elseif partition == :row
@@ -22,12 +21,15 @@ function test_coloring_decompression(
         B = compress(A, result)
         !isnothing(color0) && @test color == color0
         !isnothing(B0) && @test B == B0
-        @test decompress(B, default_result) ≈ A0
         @test decompress(B, result) ≈ A0
         @test decompress(B, result) ≈ A0  # check result wasn't modified
-        @test decompress!(respectful_similar(A), B, default_result) ≈ A0
         @test decompress!(respectful_similar(A), B, result) ≈ A0
         @test decompress!(respectful_similar(A), B, result) ≈ A0
+        if structure == :symmetric
+            linresult = LinearSystemColoringResult(sparse(A), color, eltype(A))
+            @test decompress(B, linresult) ≈ A0
+            @test decompress!(respectful_similar(A), B, linresult) ≈ A0
+        end
     end
     @test all(color_vec .== Ref(color_vec[1]))
 end