EnzymeTestUtils: fix recursive data structures

[wsmoses]

fixes #3116

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/EnzymeTestUtils/src/test_approx.jl b/lib/EnzymeTestUtils/src/test_approx.jl
index 2a5684b6..62e3f96d 100644
--- a/lib/EnzymeTestUtils/src/test_approx.jl
+++ b/lib/EnzymeTestUtils/src/test_approx.jl
@@ -1,9 +1,9 @@
 test_approx(x, y; kwargs...) = test_approx(x, y, ""; kwargs...)
-function test_approx(x::Number, y::Number, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::Number, y::Number, msg; seen = Base.IdSet(), kwargs...)
     @test_msg msg isapprox(x, y; kwargs...)
     return nothing
 end
-function test_approx(x::Array{<:Number}, y::Array{<:Number}, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::Array{<:Number}, y::Array{<:Number}, msg; seen = Base.IdSet(), kwargs...)
     @test_msg msg isapprox(x, y; kwargs...)
     return nothing
 end
@@ -13,7 +13,7 @@ function zero_copy(x)
     copyto!(x.uplo == 'U' ? UpperTriangular(y) : LowerTriangular(y), x.uplo == 'U' ? UpperTriangular(x) : LowerTriangular(x))
     y
 end
-function test_approx(x::LinearAlgebra.HermOrSym{<:Number}, y::LinearAlgebra.HermOrSym{<:Number}, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::LinearAlgebra.HermOrSym{<:Number}, y::LinearAlgebra.HermOrSym{<:Number}, msg; seen = Base.IdSet(), kwargs...)
     x in seen && return nothing
     push!(seen, x)
     # on 1.11+ similar of a hermitian or symmetric leaves the other part of the array undefined
@@ -24,13 +24,13 @@ function test_approx(x::LinearAlgebra.HermOrSym{<:Number}, y::LinearAlgebra.Herm
     test_approx(x2, y2, msg; seen, kwargs...)
     return nothing
 end
-function test_approx(x::AbstractArray{<:Number}, y::AbstractArray{<:Number}, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::AbstractArray{<:Number}, y::AbstractArray{<:Number}, msg; seen = Base.IdSet(), kwargs...)
     @test_msg msg isapprox(x, y; kwargs...)
     # for custom array types, fields should also match
     _test_fields_approx(x, y, msg; seen, kwargs...)
     return nothing
 end
-function test_approx(x::AbstractArray, y::AbstractArray, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::AbstractArray, y::AbstractArray, msg; seen = Base.IdSet(), kwargs...)
     x in seen && return nothing
     push!(seen, x)
     @test_msg "$msg: indices must match" eachindex(x) == eachindex(y)
@@ -40,7 +40,7 @@ function test_approx(x::AbstractArray, y::AbstractArray, msg; seen=Base.IdSet(),
     end
     return nothing
 end
-function test_approx(x::Tuple, y::Tuple, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::Tuple, y::Tuple, msg; seen = Base.IdSet(), kwargs...)
     @test_msg "$msg: lengths must match" length(x) == length(y)
     for i in eachindex(x)
         msg_new = "$msg: ::$(typeof(x))[$i]"
@@ -48,7 +48,7 @@ function test_approx(x::Tuple, y::Tuple, msg; seen=Base.IdSet(), kwargs...)
     end
     return nothing
 end
-function test_approx(x::Dict, y::Dict, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::Dict, y::Dict, msg; seen = Base.IdSet(), kwargs...)
     x in seen && return nothing
     push!(seen, x)
     @test_msg "$msg: keys must match" issetequal(keys(x), keys(y))
@@ -58,13 +58,13 @@ function test_approx(x::Dict, y::Dict, msg; seen=Base.IdSet(), kwargs...)
     end
     return nothing
 end
-function test_approx(x::Type, y::Type, msg; seen=Base.IdSet(), kwargs...)
+function test_approx(x::Type, y::Type, msg; seen = Base.IdSet(), kwargs...)
     @test_msg "$msg: types must match" x === y
     return nothing
 end
-test_approx(x, y, msg; seen=Base.IdSet(), kwargs...) = _test_fields_approx(x, y, msg; seen, kwargs...)
+test_approx(x, y, msg; seen = Base.IdSet(), kwargs...) = _test_fields_approx(x, y, msg; seen, kwargs...)
 
-function _test_fields_approx(x, y, msg; seen=Base.IdSet(), kwargs...)
+function _test_fields_approx(x, y, msg; seen = Base.IdSet(), kwargs...)
     @test_msg "$msg: types must match" typeof(x) == typeof(y)
     if ismutable(x)
         x in seen && return nothing
diff --git a/lib/EnzymeTestUtils/src/to_vec.jl b/lib/EnzymeTestUtils/src/to_vec.jl
index 47180127..a94c5644 100644
--- a/lib/EnzymeTestUtils/src/to_vec.jl
+++ b/lib/EnzymeTestUtils/src/to_vec.jl
@@ -193,7 +193,7 @@ function to_vec(x::GenericMemory, seen_vecs::AliasDict)
     if has_seen || is_const
         x_vec = Float32[]
     else
-        seen_vecs[x] = Float32[]
+            seen_vecs[x] = Float32[]
         from_vecs = []
         subvec_inds = UnitRange{Int}[]
         l = 0
@@ -220,14 +220,14 @@ function to_vec(x::GenericMemory, seen_vecs::AliasDict)
         has_seen && return reshape(seen_xs[x], size(x))
         is_const && return x
         x_new = typeof(x)(undef, size(x))
-        seen_xs[x] = x_new
+            seen_xs[x] = x_new
         k = 1
         for i in eachindex(x)
             isassigned(x, i) || continue
             xi = from_vecs[k](@view(x_vec_new[subvec_inds[k]]), seen_xs)
             x_new[i] = xi
             k += 1
-        end
+            end
         return x_new
     end
     return x_vec, Memory_from_vec

[github-actions]

Benchmark Results

	main	74137d9...	main / 74137d9...
basics/make_zero/namedtuple	0.0553 ± 0.0041 μs	0.0557 ± 0.0041 μs	0.994 ± 0.1
basics/make_zero/struct	0.256 ± 0.0083 μs	0.254 ± 0.0069 μs	1.01 ± 0.043
basics/overhead	5.24 ± 0.01 ns	4.54 ± 0.01 ns	1.15 ± 0.0034
basics/remake_zero!/namedtuple	0.231 ± 0.0087 μs	0.231 ± 0.0092 μs	1 ± 0.055
basics/remake_zero!/struct	0.231 ± 0.011 μs	0.23 ± 0.01 μs	1 ± 0.064
fold_broadcast/multidim_sum_bcast/1D	10.4 ± 0.36 μs	10.3 ± 0.41 μs	1.01 ± 0.053
fold_broadcast/multidim_sum_bcast/2D	12.7 ± 0.36 μs	12.7 ± 0.37 μs	1 ± 0.041
time_to_load	1.1 ± 0.0012 s	1.1 ± 0.0041 s	0.991 ± 0.0038

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/27852966542/artifacts/7759894083.