diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 98af2ddd..a6f556db 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -675,6 +675,15 @@ function Base.view(A::AbstractVectorOfArray{T, N, <:AbstractVector{T}}, end function Base.view(A::AbstractVectorOfArray, I::Vararg{Any, M}) where {M} @inline + # Special handling for heterogeneous arrays when viewing a single column + # The issue is that to_indices uses axes, which is based on the first element's size + # For heterogeneous arrays, we need to use the actual size of the specific column + if length(I) == 2 && I[1] == Colon() && I[2] isa Int + @boundscheck checkbounds(A.u, I[2]) + # Use the actual size of the specific column instead of relying on axes/to_indices + J = (Base.OneTo(length(A.u[I[2]])), I[2]) + return SubArray(A, J) + end J = map(i -> Base.unalias(A, i), to_indices(A, I)) @boundscheck checkbounds(A, J...) SubArray(A, J) @@ -826,7 +835,12 @@ end function Base.fill!(VA::AbstractVectorOfArray, x) for i in 1:length(VA.u) if VA[:, i] isa AbstractArray - fill!(VA[:, i], x) + if ArrayInterface.ismutable(VA.u[i]) || VA.u[i] isa AbstractVectorOfArray + fill!(VA[:, i], x) + else + # For immutable arrays like SVector, create a new filled array + VA.u[i] = fill(x, StaticArraysCore.similar_type(VA.u[i])) + end else VA[:, i] = x end diff --git a/test/basic_indexing.jl b/test/basic_indexing.jl index 442e761f..7fe8b80f 100644 --- a/test/basic_indexing.jl +++ b/test/basic_indexing.jl @@ -145,6 +145,31 @@ diffeq = DiffEqArray(recs, t) @test diffeq[:, 1] == recs[1] @test diffeq[1:2, 1:2] == [1 3; 2 5] +# Test views of heterogeneous arrays (issue #453) +f = VectorOfArray([[1.0], [2.0, 3.0]]) +@test length(view(f, :, 1)) == 1 +@test length(view(f, :, 2)) == 2 +@test view(f, :, 1) == [1.0] +@test view(f, :, 2) == [2.0, 3.0] +@test collect(view(f, :, 1)) == f[:, 1] +@test collect(view(f, :, 2)) == f[:, 2] + +f2 = VectorOfArray([[1.0, 2.0], [3.0]]) +@test length(view(f2, :, 1)) == 2 +@test length(view(f2, :, 2)) == 1 +@test view(f2, :, 1) == [1.0, 2.0] +@test view(f2, :, 2) == [3.0] +@test collect(view(f2, :, 1)) == f2[:, 1] +@test collect(view(f2, :, 2)) == f2[:, 2] + +# Test that views can be modified +f3 = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) +v = view(f3, :, 2) +@test length(v) == 3 +v[1] = 10.0 +@test f3[1, 2] == 10.0 +@test f3.u[2][1] == 10.0 + t = 1:5 recs = [rand(2, 2) for i in 1:5] testva = VectorOfArray(recs) diff --git a/test/utils_test.jl b/test/utils_test.jl index 63f78238..9b9dbded 100644 --- a/test/utils_test.jl +++ b/test/utils_test.jl @@ -139,6 +139,35 @@ end @test u1.u[2] isa SVector end +# Test recursivefill! with immutable StaticArrays (issue #461) +@testset "recursivefill! with immutable StaticArrays (issue #461)" begin + # Test with only immutable SVectors + x = VectorOfArray([SVector{2}(ones(2)), SVector{2}(ones(2))]) + recursivefill!(x, 0.0) + @test all(x.u[i] == SVector{2}(zeros(2)) for i in 1:2) + @test all(x.u[i] isa SVector for i in 1:2) + + # Test with mixed immutable and mutable StaticArrays + x = VectorOfArray([SVector{2}(ones(2)), MVector{2}(ones(2))]) + recursivefill!(x, 0.0) + @test all(x.u[i] == [0.0, 0.0] for i in 1:2) + @test x.u[1] isa SVector + @test x.u[2] isa MVector + + # Test fill! on VectorOfArray with immutable SVectors + x = VectorOfArray([SVector{2}(ones(2)), SVector{2}(ones(2))]) + fill!(x, 0.0) + @test all(x.u[i] == SVector{2}(zeros(2)) for i in 1:2) + @test all(x.u[i] isa SVector for i in 1:2) + + # Test fill! on VectorOfArray with mixed types + x = VectorOfArray([SVector{2}(ones(2)), MVector{2}(ones(2))]) + fill!(x, 0.0) + @test all(x.u[i] == [0.0, 0.0] for i in 1:2) + @test x.u[1] isa SVector + @test x.u[2] isa MVector +end + import KernelAbstractions: get_backend @testset "KernelAbstractions" begin v = VectorOfArray([randn(2) for i in 1:10])