Implement Base.unsafe_wrap for OffsetArrays (pull request) (#290)

Author: mkitti

src/OffsetArrays.jl

@@ -623,6 +623,37 @@ function Base.replace_in_print_matrix(A::OffsetArray{<:Any,1}, i::Integer, j::In
     Base.replace_in_print_matrix(parent(A), ip, j, s)
 end
 
+# Actual unsafe_wrap implementation
+@inline function _unsafe_wrap(pointer::Ptr{T}, inds::NTuple{N}; own = false, kw...) where {T,N}
+    _checkindices(N, inds, "indices")
+    AA = Base.unsafe_wrap(Array, pointer, map(_indexlength, inds); own=own)
+    OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds); kw...)
+end
+const OffsetArrayUnion{T,N} = Union{Type{OffsetArray}, Type{OffsetArray{T}}, Type{OffsetArray{T,N}}, Type{OffsetArray{T1, N} where T1}} where {T,N}
+
+"""
+    Base.unsafe_wrap(OffsetArray, pointer::Ptr{T}, inds...; own=false, kw...)
+    Base.unsafe_wrap(OffsetArray, pointer::Ptr{T}, inds::Tuple; own=false, kw...)
+
+Construct an `OffsetArray` around a pointer with axes defined by `inds` which may be Integers or AbstractUnitRanges.
+If the keyword `own` is true, then Julia will free the pointer when the array is garbage collected.
+Other keywords are forwarded to the `OffsetArray` constructor.
+"""
+@inline function Base.unsafe_wrap(::OffsetArrayUnion{T,N}, pointer::Ptr{T}, inds::NTuple{N, OffsetAxisKnownLength}; kw...) where {T,N}
+    _unsafe_wrap(pointer, inds; kw...)
+end
+# Avoid ambiguity
+@inline function Base.unsafe_wrap(::OffsetArrayUnion{T,N}, pointer::Ptr{T}, inds::NTuple{N, <:Integer}; kw...) where {T,N}
+    _unsafe_wrap(pointer, inds; kw...)
+end
+@inline function Base.unsafe_wrap(::OffsetArrayUnion{T,N}, pointer::Ptr{T}, inds::Vararg{OffsetAxisKnownLength,N}; kw...) where {T,N}
+    _unsafe_wrap(pointer, inds; kw...)
+end
+# Avoid ambiguity
+@inline function Base.unsafe_wrap(::OffsetArrayUnion{T,N}, pointer::Ptr{T}, inds::Vararg{<:Integer,N}; kw...) where {T,N}
+    _unsafe_wrap(pointer, inds; kw...)
+end
+
 """
     no_offset_view(A)
 

test/runtests.jl

@@ -2683,3 +2683,19 @@ end
     A = reshape(collect(1:9), 3, 3)
     @test OffsetArrays.centered(A, RoundDown) == OffsetArrays.centered(A, RoundUp)
 end
+@testset "unsafe_wrap" begin
+    p = Ptr{UInt16}(Libc.malloc(2*3*4*2))
+    @test unsafe_wrap(OffsetArray, p, 2, 3, 4) isa OffsetArray{UInt16, 3}
+    @test unsafe_wrap(OffsetArray, p, (2, 3, 4)) isa OffsetArray{UInt16, 3}
+    @test unsafe_wrap(OffsetVector, p, 2*3*4) isa OffsetVector{UInt16}
+    @test unsafe_wrap(OffsetMatrix, p, 2*3, 4) isa OffsetMatrix{UInt16}
+    @test unsafe_wrap(OffsetArray{UInt16}, p, 2, 3, 4) isa OffsetArray{UInt16, 3}
+    @test unsafe_wrap(OffsetArray{UInt16}, p, (2, 3, 4)) isa OffsetArray{UInt16, 3}
+    @test unsafe_wrap(OffsetVector{UInt16}, p, 2*3*4) isa OffsetVector{UInt16}
+    @test unsafe_wrap(OffsetMatrix{UInt16}, p, 2*3, 4) isa OffsetMatrix{UInt16}
+    p = Ptr{UInt8}(p)
+    @test unsafe_wrap(OffsetArray, p, 2:3, 3:5, 4:7) isa OffsetArray{UInt8, 3}
+    @test unsafe_wrap(OffsetArray, p, (2:3, 3:5, 4:7)) isa OffsetArray{UInt8, 3}
+    @test unsafe_wrap(OffsetVector, p, 1:(2*3*4) .- 1) isa OffsetVector{UInt8}
+    @test unsafe_wrap(OffsetMatrix, p, 1:(2*3) .+ 6, 4:7; own = true) isa OffsetMatrix{UInt8}
+end