Take into account the comments of Guillaume

Author: amontoison

src/decompression.jl

@@ -530,29 +530,16 @@ end
 
 ## TreeSetColoringResult
 
-function compute_tree_value(is_star::Val{false}, B::AbstractMatrix, i::Integer, j::Integer, color::AbstractVector{<:Integer}, buffer::AbstractVector{<:Real})
-    # The tree is not a star
-    val = B[i, color[j]] - buffer[i]
-    buffer[j] = buffer[j] + val
-    return val
-end
-
-function compute_tree_value(is_star::Val{true}, B::AbstractMatrix, i::Integer, j::Integer, color::AbstractVector{<:Integer}, buffer::AbstractVector{<:Real})
-    # The tree is a star (trivial or non-trivial)
-    val = B[i, color[j]]
-    return val
-end
-
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; ag, color, reverse_bfs_orders, is_star, buffer) = result
+    (; ag, color, reverse_bfs_orders, buffer) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
     fill!(A, zero(R))
 
-    if eltype(buffer) == R || isempty(buffer)
+    if eltype(buffer) == R
         buffer_right_type = buffer
     else
         buffer_right_type = similar(buffer, R)
@@ -569,22 +556,17 @@ function decompress!(
 
     # Recover the off-diagonal coefficients of A
     for k in eachindex(reverse_bfs_orders)
-        is_star_k = is_star[k]
-        val_is_star_k = Val(is_star_k)
-
-        # We need the buffer only when the tree is not a star (trivial or non-trivial)
-        if !is_star_k
-            # Reset the buffer to zero for all vertices in a tree (except the root)
-            for (vertex, _) in reverse_bfs_orders[k]
-                buffer_right_type[vertex] = zero(R)
-            end
-            # Reset the buffer to zero for the root vertex
-            (_, root) = reverse_bfs_orders[k][end]
-            buffer_right_type[root] = zero(R)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
+            buffer_right_type[vertex] = zero(R)
         end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
 
         for (i, j) in reverse_bfs_orders[k]
-            val = compute_tree_value(val_is_star_k, B, i, j, color, buffer_right_type)
+            val = B[i, color[j]] - buffer_right_type[i]
+            buffer_right_type[j] = buffer_right_type[j] + val
 
             if in_triangle(i, j, uplo)
                 A[i, j] = val
@@ -607,7 +589,6 @@ function decompress!(
         ag,
         color,
         reverse_bfs_orders,
-        is_star,
         diagonal_indices,
         diagonal_nzind,
         lower_triangle_offsets,
@@ -619,7 +600,7 @@ function decompress!(
     nzA = nonzeros(A)
     uplo == :F && check_same_pattern(A, S)
 
-    if eltype(buffer) == R || isempty(buffer)
+    if eltype(buffer) == R
         buffer_right_type = buffer
     else
         buffer_right_type = similar(buffer, R)
@@ -652,23 +633,18 @@ function decompress!(
 
     # Recover the off-diagonal coefficients of A
     for k in eachindex(reverse_bfs_orders)
-        is_star_k = is_star[k]
-        val_is_star_k = Val(is_star_k)
-
-        # We need the buffer only when the tree is not a star (trivial or non-trivial)
-        if !is_star_k
-            # Reset the buffer to zero for all vertices in a tree (except the root)
-            for (vertex, _) in reverse_bfs_orders[k]
-                buffer_right_type[vertex] = zero(R)
-            end
-            # Reset the buffer to zero for the root vertex
-            (_, root) = reverse_bfs_orders[k][end]
-            buffer_right_type[root] = zero(R)
+        # Reset the buffer to zero for all vertices in a tree (except the root)
+        for (vertex, _) in reverse_bfs_orders[k]
+            buffer_right_type[vertex] = zero(R)
         end
+        # Reset the buffer to zero for the root vertex
+        (_, root) = reverse_bfs_orders[k][end]
+        buffer_right_type[root] = zero(R)
 
         for (i, j) in reverse_bfs_orders[k]
             counter += 1
-            compute_tree_value(val_is_star_k, B, i, j, color, buffer_right_type)
+            val = B[i, color[j]] - buffer_right_type[i]
+            buffer_right_type[j] = buffer_right_type[j] + val
 
             #! format: off
             # A[i,j] is in the lower triangular part of A

src/result.jl

@@ -279,7 +279,6 @@ struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
     color::Vector{Int}
     group::V
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
-    is_star::Vector{Bool}
     diagonal_indices::Vector{Int}
     diagonal_nzind::Vector{Int}
     lower_triangle_offsets::Vector{Int}
@@ -294,7 +293,7 @@ function TreeSetColoringResult(
     tree_set::TreeSet,
     decompression_eltype::Type{R},
 ) where {R}
-    (; reverse_bfs_orders, is_star) = tree_set
+    (; reverse_bfs_orders) = tree_set
     (; S) = ag
     nvertices = length(color)
     group = group_by_color(color)
@@ -362,16 +361,14 @@ function TreeSetColoringResult(
 
     # buffer holds the sum of edge values for subtrees in a tree.
     # For each vertex i, buffer[i] is the sum of edge values in the subtree rooted at i.
-    # Note that we don't need a buffer is all trees are stars.
-    buffer = all(is_star) ? R[] : Vector{R}(undef, nvertices)
+    buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
         A,
         ag,
         color,
         group,
         reverse_bfs_orders,
-        is_star,
         diagonal_indices,
         diagonal_nzind,
         lower_triangle_offsets,

test/allocations.jl

@@ -19,7 +19,7 @@ end
 
 @testset "Distance-2 coloring" begin
     test_noallocs_distance2_coloring(1000)
-end;
+end
 
 function test_noallocs_sparse_decompression(
     n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
@@ -121,7 +121,7 @@ end
     ]
         test_noallocs_sparse_decompression(1000; structure, partition, decompression)
     end
-end;
+end
 
 @testset "Structured decompression" begin
     @testset "$structure - $partition - $decompression" for (
@@ -131,22 +131,4 @@ end;
     ]
         test_noallocs_structured_decompression(1000; structure, partition, decompression)
     end
-end;
-
-@testset "Multi-precision acyclic decompression" begin
-    @testset "$format" for format in ("dense", "sparse")
-        A = [0 0 1; 0 1 0; 1 0 0]
-        if format == "sparse"
-            A = sparse(A)
-        end
-        problem = ColoringProblem(; structure=:symmetric, partition=:column)
-        result = coloring(A, problem, GreedyColoringAlgorithm{:substitution}())
-        @test isempty(result.buffer)
-        for T in (Float32, Float64)
-            C = rand(T) * T.(A)
-            B = compress(C, result)
-            bench_multiprecision = @be decompress!(C, B, result)
-            @test minimum(bench_multiprecision).allocs == 0
-        end
-    end
 end