Improve ArrayInterface compatibility (#28)

* improve ArrayInterface compatibility

* a bit safer implementation

* one more restriction

* restrict a test to newer versions of Julia

* strides improvements

Author: mateuszbaran

Project.toml

@@ -1,7 +1,7 @@
 name = "HybridArrays"
 uuid = "1baab800-613f-4b0a-84e4-9cd3431bfbb9"
 authors = ["Mateusz Baran <[email protected]>"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/HybridArrays.jl

@@ -10,6 +10,8 @@ import Base: convert,
     getindex,
     dataids,
     promote_rule,
+    pointer,
+    strides,
     setindex!,
     size,
     length,

src/abstractarray.jl

@@ -1,10 +1,14 @@
 
 Base.dataids(sa::HybridArray) = Base.dataids(sa.data)
 
-@inline size(sa::HybridArray{S}) where S = size(sa.data)
+@inline size(sa::HybridArray{S,T,N,N}) where {S,T,N} = size(sa.data)
 
 @inline length(sa::HybridArray) = length(sa.data)
 
+@inline strides(sa::HybridArray{S,T,N,N}) where {S,T,N} = strides(sa.data)
+
+@inline pointer(sa::HybridArray) = pointer(sa.data)
+
 @generated function _sized_abstract_array_axes(::Type{S}, ax::Tuple) where S<:Tuple
     exprs = Any[]
     map(enumerate(S.parameters)) do (i, si)

src/array_interface_compat.jl

@@ -14,3 +14,47 @@ end
 function ArrayInterface.restructure(x::HybridArray{S}, y) where {S}
     return HybridArray{S}(reshape(convert(Array, y), size(x)...))
 end
+
+function ArrayInterface.strides(x::HybridArray)
+    return ArrayInterface.strides(x.data)
+end
+
+@generated function ArrayInterface.strides(x::HybridArray{S,T,N,N,Array{T,N}}) where {S,T,N}
+    collected_strides = []
+    i = 1
+    for (argnum, Sarg) in enumerate(S.parameters)
+        if i > 0
+            push!(collected_strides, ArrayInterface.StaticInt(i))
+        else
+            push!(collected_strides, :(datastrides[$argnum]))
+        end
+        if Sarg isa Integer
+            i *= Sarg
+        else
+            i = -1
+        end
+    end
+    return quote
+        datastrides = strides(x.data)
+        return tuple($(collected_strides...))
+    end
+end
+
+@generated function ArrayInterface.size(x::HybridArray{S}) where {S}
+    collected_sizes = []
+    for (argnum, Sarg) in enumerate(S.parameters)
+        if Sarg isa Integer
+            push!(collected_sizes, ArrayInterface.StaticInt(Sarg))
+        else
+            push!(collected_sizes, :(datasize[$argnum]))
+        end
+    end
+    return quote
+        datasize = ArrayInterface.size(x.data)
+        return tuple($(collected_sizes...))
+    end
+end
+
+ArrayInterface.contiguous_axis(::Type{HybridArray{S,T,N,N,TData}}) where {S,T,N,TData} = ArrayInterface.contiguous_axis(TData)
+ArrayInterface.contiguous_batch_size(::Type{HybridArray{S,T,N,N,TData}}) where {S,T,N,TData} = ArrayInterface.contiguous_batch_size(TData)
+ArrayInterface.stride_rank(::Type{HybridArray{S,T,N,N,TData}}) where {S,T,N,TData} = ArrayInterface.stride_rank(TData)

src/convert.jl

@@ -65,3 +65,5 @@ end
 @inline function convert(::Type{HybridArray{S}}, a::TData) where {S,T,M,TData<:AbstractArray{T,M}}
     convert(HybridArray{S,T}, a)
 end
+
+@inline Base.unsafe_convert(::Type{Ptr{T}}, A::HybridArray{S,T}) where {S,T} = pointer(A.data)

test/abstractarray.jl

@@ -78,5 +78,27 @@ using StaticArrays, HybridArrays, Test, LinearAlgebra
         @test_throws TypeError HybridArrays.new_out_size_nongen(Size{Tuple{1,2}}, 'a')
     end
 
+    @testset "strides" begin
+        M = HybridMatrix{2, StaticArrays.Dynamic(), Int}([1 2; 3 4])
+
+        @test strides(M) == strides(M.data)
+    end
+
+    @testset "pointer" begin
+        M = HybridMatrix{2, StaticArrays.Dynamic(), Int}([1 2; 3 4])
+        MT = HybridMatrix{2, StaticArrays.Dynamic(), Int}([1 2; 3 4]')
+
+        @test pointer(M) == pointer(M.data)
+        if VERSION >= v"1.5"
+            # pointer on Adjoint is not available on earilier versions of Julia
+            @test pointer(MT) == pointer(MT.data)
+        end
+    end
+
+    @testset "unsafe_convert" begin
+        M = HybridMatrix{2, StaticArrays.Dynamic(), Int}([1 2; 3 4])
+        @test Base.unsafe_convert(Ptr{Int}, M) === pointer(M.data)
+    end
+
     @test HybridArrays._totally_linear() === true
 end

test/array_interface_compat.jl

@@ -1,11 +1,25 @@
 
 using HybridArrays, ArrayInterface, Test, StaticArrays
+using ArrayInterface: StaticInt
 
 @testset "ArrayInterface compatibility" begin
     M = HybridMatrix{2, StaticArrays.Dynamic()}([1 2; 4 5])
+    MV = HybridMatrix{3, StaticArrays.Dynamic()}(view(randn(10, 10), 1:2:5, 1:3:12))
     @test ArrayInterface.ismutable(M)
     @test ArrayInterface.can_setindex(M)
     @test ArrayInterface.parent_type(M) === Matrix{Int}
     @test ArrayInterface.restructure(M, [2, 4, 6, 8]) == HybridMatrix{2, StaticArrays.Dynamic()}([2 6; 4 8])
     @test isa(ArrayInterface.restructure(M, [2, 4, 6, 8]), HybridMatrix{2, StaticArrays.Dynamic()})
+
+    M2 = HybridArray{Tuple{2, 3, StaticArrays.Dynamic(), StaticArrays.Dynamic()}}(randn(2, 3, 5, 7))
+    @test (@inferred ArrayInterface.strides(M2)) === (StaticInt(1), StaticInt(2), StaticInt(6), 30)
+    @test (@inferred ArrayInterface.strides(MV)) === (2, 30)
+    @test (@inferred ArrayInterface.size(M2)) === (StaticInt(2), StaticInt(3), 5, 7)
+
+    @test ArrayInterface.contiguous_axis(typeof(M2)) === ArrayInterface.contiguous_axis(typeof(M2.data))
+    @test ArrayInterface.contiguous_batch_size(typeof(M2)) === ArrayInterface.contiguous_batch_size(typeof(M2.data))
+    @test ArrayInterface.stride_rank(typeof(M2)) === ArrayInterface.stride_rank(typeof(M2.data))
+    @test ArrayInterface.contiguous_axis(typeof(M')) === ArrayInterface.contiguous_axis(typeof(M.data'))
+    @test ArrayInterface.contiguous_batch_size(typeof(M')) === ArrayInterface.contiguous_batch_size(typeof(M.data'))
+    @test ArrayInterface.stride_rank(typeof(M')) === ArrayInterface.stride_rank(typeof(M.data'))
 end