Make reshape and view on Memory produce Arrays and delete wrap (#53896)

- Make reshape and view with one based indexing on Memory produce Arrays
- delete wrap

Implements
#53552 (comment)

---------

Co-authored-by: Jameson Nash <[email protected]>
(cherry picked from commit 273d91e)

Author: 2 people

NEWS.md

@@ -87,7 +87,6 @@ New library functions
 * `Sys.username()` can be used to return the current user's username ([#51897]).
 * `Sys.isreadable(), Sys.iswritable()` can be used to check if the current user has access permissions
   that permit reading and writing, respectively. ([#53320]).
-* `wrap(Array, m::Union{MemoryRef{T}, Memory{T}}, dims)` is the safe counterpart to `unsafe_wrap` ([#52049]).
 * `GC.logging_enabled()` can be used to test whether GC logging has been enabled via `GC.enable_logging` ([#51647]).
 * `IdSet` is now exported from Base and considered public ([#53262]).
 * `@time` now reports a count of any lock conflicts where a `ReentrantLock` had to wait, plus a new macro

base/array.jl

@@ -3067,54 +3067,3 @@ intersect(r::AbstractRange, v::AbstractVector) = intersect(v, r)
         _getindex(v, i)
     end
 end
-
-"""
-    wrap(Array, m::Union{Memory{T}, MemoryRef{T}}, dims)
-
-Create an array of size `dims` using `m` as the underlying memory. This can be thought of as a safe version
-of [`unsafe_wrap`](@ref) utilizing `Memory` or `MemoryRef` instead of raw pointers.
-"""
-function wrap end
-
-# validity checking for _wrap calls, separate from allocation of Array so that it can be more likely to inline into the caller
-function _wrap(ref::MemoryRef{T}, dims::NTuple{N, Int}) where {T, N}
-    mem = ref.mem
-    mem_len = length(mem) + 1 - memoryrefoffset(ref)
-    len = Core.checked_dims(dims...)
-    @boundscheck mem_len >= len || invalid_wrap_err(mem_len, dims, len)
-    if N != 1 && !(ref === GenericMemoryRef(mem) && len === mem_len)
-        mem = ccall(:jl_genericmemory_slice, Memory{T}, (Any, Ptr{Cvoid}, Int), mem, ref.ptr_or_offset, len)
-        ref = MemoryRef(mem)
-    end
-    return ref
-end
-
-@noinline invalid_wrap_err(len, dims, proddims) = throw(DimensionMismatch(
-    "Attempted to wrap a MemoryRef of length $len with an Array of size dims=$dims, which is invalid because prod(dims) = $proddims > $len, so that the array would have more elements than the underlying memory can store."))
-
-@eval @propagate_inbounds function wrap(::Type{Array}, m::MemoryRef{T}, dims::NTuple{N, Integer}) where {T, N}
-    dims = convert(Dims, dims)
-    ref = _wrap(m, dims)
-    $(Expr(:new, :(Array{T, N}), :ref, :dims))
-end
-
-@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}, dims::NTuple{N, Integer}) where {T, N}
-    dims = convert(Dims, dims)
-    ref = _wrap(MemoryRef(m), dims)
-    $(Expr(:new, :(Array{T, N}), :ref, :dims))
-end
-@eval @propagate_inbounds function wrap(::Type{Array}, m::MemoryRef{T}, l::Integer) where {T}
-    dims = (Int(l),)
-    ref = _wrap(m, dims)
-    $(Expr(:new, :(Array{T, 1}), :ref, :dims))
-end
-@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}, l::Integer) where {T}
-    dims = (Int(l),)
-    ref = _wrap(MemoryRef(m), (l,))
-    $(Expr(:new, :(Array{T, 1}), :ref, :dims))
-end
-@eval @propagate_inbounds function wrap(::Type{Array}, m::Memory{T}) where {T}
-    ref = MemoryRef(m)
-    dims = (length(m),)
-    $(Expr(:new, :(Array{T, 1}), :ref, :dims))
-end

base/exports.jl

@@ -460,7 +460,6 @@ export
     vcat,
     vec,
     view,
-    wrap,
     zeros,
 
 # search, find, match and related functions

base/genericmemory.jl

@@ -288,3 +288,27 @@ function indcopy(sz::Dims, I::GenericMemory)
     src = eltype(I)[I[i][_findin(I[i], i < n ? (1:sz[i]) : (1:s))] for i = 1:n]
     dst, src
 end
+
+# Wrapping a memory region in an Array
+@eval begin # @eval for the Array construction. Block for the docstring.
+    function reshape(m::GenericMemory{M, T}, dims::Vararg{Int, N}) where {M, T, N}
+        len = Core.checked_dims(dims...)
+        length(m) == len || throw(DimensionMismatch("parent has $(length(m)) elements, which is incompatible with size $(dims)"))
+        ref = MemoryRef(m)
+        $(Expr(:new, :(Array{T, N}), :ref, :dims))
+    end
+
+    """
+        view(m::GenericMemory{M, T}, inds::Union{UnitRange, OneTo})
+
+    Create a vector `v::Vector{T}` backed by the specified indices of `m`. It is only safe to
+    resize `v` if `m` is subseqently not used.
+    """
+    function view(m::GenericMemory{M, T}, inds::Union{UnitRange, OneTo}) where {M, T}
+        isempty(inds) && return T[] # needed to allow view(Memory{T}(undef, 0), 2:1)
+        @boundscheck checkbounds(m, inds)
+        ref = MemoryRef(m, first(inds)) # @inbounds would be safe here but does not help performance.
+        dims = (length(inds),)
+        $(Expr(:new, :(Array{T, 1}), :ref, :dims))
+    end
+end

base/iobuffer.jl

@@ -42,7 +42,7 @@ end
 
 # allocate Vector{UInt8}s for IOBuffer storage that can efficiently become Strings
 StringMemory(n::Integer) = unsafe_wrap(Memory{UInt8}, _string_n(n))
-StringVector(n::Integer) = wrap(Array, StringMemory(n))
+StringVector(n::Integer) = view(StringMemory(n), 1:n)::Vector{UInt8}
 
 # IOBuffers behave like Files. They are typically readable and writable. They are seekable. (They can be appendable).
 
@@ -456,7 +456,7 @@ function take!(io::IOBuffer)
         if nbytes == 0 || io.reinit
             data = StringVector(0)
         elseif io.writable
-            data = wrap(Array, MemoryRef(io.data, io.offset + 1), nbytes)
+            data = view(io.data, io.offset+1:nbytes+io.offset)
         else
             data = copyto!(StringVector(io.size), 1, io.data, io.offset + 1, nbytes)
         end
@@ -465,7 +465,7 @@ function take!(io::IOBuffer)
         if nbytes == 0
             data = StringVector(0)
         elseif io.writable
-            data = wrap(Array, MemoryRef(io.data, io.ptr), nbytes)
+            data = view(io.data, io.ptr:io.ptr+nbytes-1)
         else
             data = read!(io, data)
         end
@@ -491,11 +491,7 @@ state.  This should only be used internally for performance-critical
 It might save an allocation compared to `take!` (if the compiler elides the
 Array allocation), as well as omits some checks.
 """
-_unsafe_take!(io::IOBuffer) =
-    wrap(Array, io.size == io.offset ?
-        MemoryRef(Memory{UInt8}()) :
-        MemoryRef(io.data, io.offset + 1),
-        io.size - io.offset)
+_unsafe_take!(io::IOBuffer) = view(io.data, io.offset+1:io.size)
 
 function write(to::IO, from::GenericIOBuffer)
     written::Int = bytesavailable(from)

base/strings/string.jl

@@ -117,7 +117,10 @@ String(s::AbstractString) = print_to_string(s)
 @assume_effects :total String(s::Symbol) = unsafe_string(unsafe_convert(Ptr{UInt8}, s))
 
 unsafe_wrap(::Type{Memory{UInt8}}, s::String) = ccall(:jl_string_to_genericmemory, Ref{Memory{UInt8}}, (Any,), s)
-unsafe_wrap(::Type{Vector{UInt8}}, s::String) = wrap(Array, unsafe_wrap(Memory{UInt8}, s))
+function unsafe_wrap(::Type{Vector{UInt8}}, s::String)
+    mem = unsafe_wrap(Memory{UInt8}, s)
+    view(mem, eachindex(mem))
+end
 
 Vector{UInt8}(s::CodeUnits{UInt8,String}) = copyto!(Vector{UInt8}(undef, length(s)), s)
 Vector{UInt8}(s::String) = Vector{UInt8}(codeunits(s))

doc/src/base/arrays.md

@@ -139,7 +139,6 @@ Base.reshape
 Base.dropdims
 Base.vec
 Base.SubArray
-Base.wrap
 ```
 
 ## Concatenation and permutation

test/arrayops.jl

@@ -3187,23 +3187,37 @@ end
     end
 end
 
-@testset "Wrapping Memory into Arrays" begin
-    mem = Memory{Int}(undef, 10) .= 1
-    memref = MemoryRef(mem)
-    @test_throws DimensionMismatch wrap(Array, mem, (10, 10))
-    @test wrap(Array, mem, (5,)) == ones(Int, 5)
-    @test wrap(Array, mem, 2) == ones(Int, 2)
-    @test wrap(Array, memref, 10) == ones(Int, 10)
-    @test wrap(Array, memref, (2,2,2)) == ones(Int,2,2,2)
-    @test wrap(Array, mem, (5, 2)) == ones(Int, 5, 2)
-
-    memref2 = MemoryRef(mem, 3)
-    @test wrap(Array, memref2, (5,)) == ones(Int, 5)
-    @test wrap(Array, memref2, 2) == ones(Int, 2)
-    @test wrap(Array, memref2, (2,2,2)) == ones(Int,2,2,2)
-    @test wrap(Array, memref2, (3, 2)) == ones(Int, 3, 2)
-    @test_throws DimensionMismatch wrap(Array, memref2, 9)
-    @test_throws DimensionMismatch wrap(Array, memref2, 10)
+@testset "Wrapping Memory into Arrays with view and reshape" begin
+    mem::Memory{Int} = Memory{Int}(undef, 10) .= 11:20
+
+    @test_throws DimensionMismatch reshape(mem, 10, 10)
+    @test_throws DimensionMismatch reshape(mem, 5)
+    @test_throws BoundsError view(mem, 1:10, 1:10)
+    @test_throws BoundsError view(mem, 1:11)
+    @test_throws BoundsError view(mem, 3:11)
+    @test_throws BoundsError view(mem, 0:4)
+
+    @test @inferred(view(mem, 1:5))::Vector{Int} == 11:15
+    @test @inferred(view(mem, 1:2))::Vector{Int} == 11:12
+    @test @inferred(view(mem, 1:10))::Vector{Int} == 11:20
+    @test @inferred(view(mem, 3:8))::Vector{Int} == 13:18
+    @test @inferred(view(mem, 20:19))::Vector{Int} == []
+    @test @inferred(view(mem, -5:-7))::Vector{Int} == []
+    @test @inferred(reshape(mem, 5, 2))::Matrix{Int} == reshape(11:20, 5, 2)
+
+    empty_mem = Memory{Module}(undef, 0)
+    @test_throws BoundsError view(empty_mem, 0:1)
+    @test_throws BoundsError view(empty_mem, 1:2)
+    @test_throws DimensionMismatch reshape(empty_mem, 1)
+    @test_throws DimensionMismatch reshape(empty_mem, 1, 2, 3)
+    @test_throws ArgumentError reshape(empty_mem, 2^16, 2^16, 2^16, 2^16)
+
+    @test @inferred(view(empty_mem, 1:0))::Vector{Module} == []
+    @test @inferred(view(empty_mem, 10:3))::Vector{Module} == []
+    @test isempty(@inferred(reshape(empty_mem, 0, 7, 1))::Array{Module, 3})
+
+    offset_inds = OffsetArrays.IdOffsetRange(values=3:6, indices=53:56)
+    @test @inferred(view(collect(mem), offset_inds)) == view(mem, offset_inds)
 end
 
 @testset "Memory size" begin