Merge pull request #19 from mkitti/mkitti/unsafe_wrap

Implement Base.unsafe_wrap(MallocArray, ...)

Author: brenhinkeller

src/mallocarray.jl

@@ -165,6 +165,24 @@
         return c
     end
 
+    # Avoid Base.unsafe_wrap stack overflow
+    @inline _unsafe_wrap(::Type{MallocArray{T,N}}, ptr::Ptr, dims::Dims{N}, own::Bool = false) where {T,N} =
+        own ?
+            error("A MallocArray cannot be owned via `unsafe_wrap`. Set the keyword `own` to be false.") :
+            MallocArray{T,N}(ptr, prod(dims), dims)
+
+    # Overload Base.unsafe_wrap
+    """
+    ```julia
+    unsafe_wrap(MallocArray, ptr::Ptr{T}, dims)
+    ```
+    Create a `MallocArray{T}` wrapping around `ptr`
+    """
+    @inline Base.unsafe_wrap(::Union{Type{MallocArray},Type{MallocArray{T}},Type{MallocArray{T,N}}}, ptr::Ptr{T}, dims::Dims{N}; own = false) where {T,N} =
+        _unsafe_wrap(MallocArray{T,N}, ptr, dims, own)
+    @inline Base.unsafe_wrap(::Union{Type{MallocArray},Type{MallocArray{T}},Type{MallocArray{T,1}}}, ptr::Ptr{T}, dim::Integer; own = false) where {T} =
+        _unsafe_wrap(MallocArray{T,1}, ptr, (Int(dim),), own)
+
     # Other custom constructors
     """
     ```julia

test/testmallocarray.jl

@@ -178,6 +178,28 @@
     free(C)
     free(D)
 
+## ---
+
+## -- test Base.unsafe_wrap
+
+    ptr = StaticTools.malloc(64)
+    A = unsafe_wrap(MallocArray, Ptr{Int8}(ptr), 64)
+    A .= 1:64
+    @test typeof(A) == MallocArray{Int8, 1}
+    @test A == 1:64
+
+    B = unsafe_wrap(MallocArray{Int64}, Ptr{Int64}(ptr), (2,4))
+    B[:] .= 1:8
+    @test typeof(B) == MallocArray{Int64,2}
+    @test B == reshape(1:8, (2,4))
+
+    C = unsafe_wrap(MallocArray{Float64,3}, Ptr{Float64}(ptr), (2,2,2))
+    C[:] .= 1.0:1.0:8.0
+    @test typeof(C) == MallocArray{Float64,3}
+    @test C == reshape(1.0:1.0:8.0, (2,2,2))
+
+    free(A)
+
 ## ---
 
     A = mones(11,10)