resolve ambiguities with StaticArrays (#304)

JuliaArrays · Apr 6, 2024 · e52b3cc · e52b3cc
1 parent b1836a6
commit e52b3cc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/ext/StructArraysStaticArraysExt.jl b/ext/StructArraysStaticArraysExt.jl
@@ -1,7 +1,7 @@
 module StructArraysStaticArraysExt
 
 using StructArrays
-using StaticArrays: StaticArray, FieldArray, tuple_prod, SVector, MVector
+using StaticArrays: StaticArray, FieldArray, tuple_prod, SVector, MVector, SOneTo
 
 """
     StructArrays.staticschema(::Type{<:StaticArray{S, T}}) where {S, T}
@@ -40,6 +40,10 @@ end
 StructArrays.component(s::FieldArray, i) = invoke(StructArrays.component, Tuple{Any, Any}, s, i)
 StructArrays.createinstance(T::Type{<:FieldArray}, args...) = invoke(StructArrays.createinstance, Tuple{Type{<:Any}, Vararg}, T, args...)
 
+# disambiguation
+Base.similar(s::StructArray, S::Type, sz::Tuple{Union{Integer, Base.OneTo, SOneTo}, Vararg{Union{Union{Integer, Base.OneTo, SOneTo}}}}) = StructArrays._similar(s, S, sz)
+Base.reshape(s::StructArray{T}, d::Tuple{SOneTo, Vararg{SOneTo}}) where {T} = StructArray{T}(map(x -> reshape(x, d), StructArrays.components(s)))
+
 # Broadcast overload
 using StaticArrays: StaticArrayStyle, similar_type, Size, SOneTo
 using StaticArrays: broadcast_flatten, broadcast_sizes, first_statictype

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1444,6 +1444,14 @@ end
         @test x .+ y == StructArray([StaticArrayType{Tuple{1,2}}(3*ones(1,2) .+ 2*i) for i = 0:1])
     end
 
+    let
+        # potential ambiguities
+        x = StructArray((SA[1,2], SA[1,2]))
+        @test eltype(similar(x)::StructArray) == eltype(x)
+        @test map(x->x, x)::StructArray == x
+        @test size(reshape(x, (SOneTo(2), SOneTo(1)))::StructArray) == (2, 1)
+    end
+
     # test FieldVector constructor (see https://github.com/JuliaArrays/StructArrays.jl/issues/205)
     struct FlippedVec2D <: FieldVector{2,Float64}
         x::Float64