|
| 1 | +export FixedSizeArray, FixedSizeVector, FixedSizeMatrix |
| 2 | + |
| 3 | +""" |
| 4 | + Internal() |
| 5 | +
|
| 6 | +Implementation detail. Do not use. |
| 7 | +""" |
| 8 | +struct Internal end |
| 9 | + |
| 10 | +struct FixedSizeArray{T,N,Mem<:GenericMemory{<:Any,T}} <: DenseArray{T,N} |
| 11 | + mem::Mem |
| 12 | + size::NTuple{N,Int} |
| 13 | + function FixedSizeArray{T,N,M}(::Internal, mem::M, size::NTuple{N,Int}) where {T,N,M<:GenericMemory{<:Any,T}} |
| 14 | + new{T,N,M}(mem, size) |
| 15 | + end |
| 16 | +end |
| 17 | + |
| 18 | +function Base.propertynames( |
| 19 | + # the `unused` is here because of https://github.com/JuliaLang/julia/issues/44428 |
| 20 | + (@nospecialize unused::FixedSizeArray), |
| 21 | + ::Bool = false, |
| 22 | +) |
| 23 | + () |
| 24 | +end |
| 25 | + |
| 26 | +const FixedSizeVector{T} = FixedSizeArray{T,1} |
| 27 | +const FixedSizeMatrix{T} = FixedSizeArray{T,2} |
| 28 | + |
| 29 | +const default_underlying_storage_type = Memory |
| 30 | + |
| 31 | +function FixedSizeArray{T,N,V}(::UndefInitializer, size::NTuple{N,Int}) where {T,N,V} |
| 32 | + FixedSizeArray{T,N,V}(Internal(), V(undef, checked_dims(size))::V, size) |
| 33 | +end |
| 34 | +function FixedSizeArray{T,N,V}(::UndefInitializer, size::NTuple{N,Integer}) where {T,N,V} |
| 35 | + ints = map(Int, size)::NTuple{N,Int} # prevent infinite recursion |
| 36 | + FixedSizeArray{T,N,V}(undef, ints) |
| 37 | +end |
| 38 | +function FixedSizeArray{T,N,V}(::UndefInitializer, size::Vararg{Integer,N}) where {T,N,V} |
| 39 | + FixedSizeArray{T,N,V}(undef, size) |
| 40 | +end |
| 41 | +function FixedSizeArray{T,<:Any,V}(::UndefInitializer, size::NTuple{N,Integer}) where {T,N,V} |
| 42 | + FixedSizeArray{T,N,V}(undef, size) |
| 43 | +end |
| 44 | +function FixedSizeArray{T,<:Any,V}(::UndefInitializer, size::Vararg{Integer,N}) where {T,N,V} |
| 45 | + FixedSizeArray{T,N,V}(undef, size) |
| 46 | +end |
| 47 | +function FixedSizeArray{T,N}(::UndefInitializer, size::NTuple{N,Integer}) where {T,N} |
| 48 | + FixedSizeArray{T,N,default_underlying_storage_type{T}}(undef, size) |
| 49 | +end |
| 50 | +function FixedSizeArray{T,N}(::UndefInitializer, size::Vararg{Integer,N}) where {T,N} |
| 51 | + FixedSizeArray{T,N}(undef, size) |
| 52 | +end |
| 53 | +function FixedSizeArray{T}(::UndefInitializer, size::Vararg{Integer,N}) where {T,N} |
| 54 | + FixedSizeArray{T,N}(undef, size) |
| 55 | +end |
| 56 | +function FixedSizeArray{T}(::UndefInitializer, size::NTuple{N,Integer}) where {T,N} |
| 57 | + FixedSizeArray{T,N}(undef, size) |
| 58 | +end |
| 59 | + |
| 60 | +Base.IndexStyle(::Type{<:FixedSizeArray}) = IndexLinear() |
| 61 | +Base.@propagate_inbounds Base.getindex(A::FixedSizeArray, i::Int) = A.mem[i] |
| 62 | +Base.@propagate_inbounds Base.@assume_effects :noub_if_noinbounds function Base.setindex!(A::FixedSizeArray{T}, x, i::Int) where {T} |
| 63 | + @boundscheck checkbounds(A, i) |
| 64 | + @inbounds A.mem[i] = x |
| 65 | + return A |
| 66 | +end |
| 67 | + |
| 68 | +Base.size(a::FixedSizeArray) = a.size |
| 69 | + |
| 70 | +function Base.similar(::T, ::Type{E}, size::NTuple{N,Int}) where {T<:FixedSizeArray,E,N} |
| 71 | + with_replaced_parameters(DenseArray, T, Val(E), Val(N))(undef, size) |
| 72 | +end |
| 73 | + |
| 74 | +Base.isassigned(a::FixedSizeArray, i::Int) = isassigned(a.mem, i) |
| 75 | + |
| 76 | +# safe product of a tuple of integers, for calculating dimensions size |
| 77 | + |
| 78 | +checked_dims_impl(a::Int, ::Tuple{}, have_overflow::Bool) = (a, have_overflow) |
| 79 | +function checked_dims_impl(a::Int, t::Tuple{Int,Vararg{Int,N}}, have_overflow::Bool) where {N} |
| 80 | + b = first(t) |
| 81 | + (m, o) = Base.Checked.mul_with_overflow(a, b) |
| 82 | + r = Base.tail(t)::NTuple{N,Int} |
| 83 | + checked_dims_impl(m, r, have_overflow | o)::Tuple{Int,Bool} |
| 84 | +end |
| 85 | + |
| 86 | +checked_dims(::Tuple{}) = 1 |
| 87 | +function checked_dims(t::Tuple{Int,Vararg{Int,N}}) where {N} |
| 88 | + any_is_zero = any(iszero, t)::Bool |
| 89 | + any_is_negative = any((x -> x < false), t)::Bool |
| 90 | + any_is_typemax = any((x -> x == typemax(x)), t)::Bool |
| 91 | + a = first(t) |
| 92 | + r = Base.tail(t)::NTuple{N,Int} |
| 93 | + (product, have_overflow) = checked_dims_impl(a, r, false)::Tuple{Int,Bool} |
| 94 | + if any_is_negative |
| 95 | + throw(ArgumentError("array dimension size can't be negative")) |
| 96 | + end |
| 97 | + if any_is_typemax |
| 98 | + throw(ArgumentError("array dimension size can't be the maximum representable value")) |
| 99 | + end |
| 100 | + if have_overflow & !any_is_zero |
| 101 | + throw(ArgumentError("array dimensions too great, can't represent length")) |
| 102 | + end |
| 103 | + product |
| 104 | +end |
| 105 | + |
| 106 | +# broadcasting |
| 107 | + |
| 108 | +function Base.BroadcastStyle(::Type{T}) where {T<:FixedSizeArray} |
| 109 | + Broadcast.ArrayStyle{stripped_type(DenseArray, T)}() |
| 110 | +end |
| 111 | + |
| 112 | +function Base.similar( |
| 113 | + bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{S}}, |
| 114 | + ::Type{E}, |
| 115 | +) where {S<:FixedSizeArray,E} |
| 116 | + similar(S{E}, axes(bc)) |
| 117 | +end |
| 118 | + |
| 119 | +# helper functions |
| 120 | + |
| 121 | +normalized_type(::Type{T}) where {T} = T |
| 122 | + |
| 123 | +function stripped_type_unchecked(::Type{DenseVector}, ::Type{<:GenericMemory{K,<:Any,AS}}) where {K,AS} |
| 124 | + GenericMemory{K,<:Any,AS} |
| 125 | +end |
| 126 | + |
| 127 | +Base.@assume_effects :consistent function stripped_type_unchecked( |
| 128 | + ::Type{DenseArray}, ::Type{<:FixedSizeArray{<:Any,<:Any,V}}, |
| 129 | +) where {V} |
| 130 | + U = stripped_type(DenseVector, V) |
| 131 | + FixedSizeArray{E,N,U{E}} where {E,N} |
| 132 | +end |
| 133 | + |
| 134 | +function stripped_type(::Type{T}, ::Type{S}) where {T,S<:T} |
| 135 | + ret = stripped_type_unchecked(T, S)::Type{<:T}::UnionAll |
| 136 | + S::Type{<:ret} |
| 137 | + normalized_type(ret) # ensure `UnionAll` type variable order is normalized |
| 138 | +end |
| 139 | + |
| 140 | +function with_replaced_parameters(::Type{T}, ::Type{S}, ::Val{P1}, ::Val{P2}) where {T,S<:T,P1,P2} |
| 141 | + t = T{P1,P2}::Type{<:T} |
| 142 | + s = stripped_type(T, S) |
| 143 | + S::Type{<:s} |
| 144 | + s{P1,P2}::Type{<:s}::Type{<:T}::Type{<:t} |
| 145 | +end |
| 146 | + |
| 147 | +dimension_count_of(::Base.SizeUnknown) = 1 |
| 148 | +dimension_count_of(::Base.HasLength) = 1 |
| 149 | +dimension_count_of(::Base.HasShape{N}) where {N} = convert(Int, N)::Int |
| 150 | + |
| 151 | +struct LengthIsUnknown end |
| 152 | +struct LengthIsKnown end |
| 153 | +length_status(::Base.SizeUnknown) = LengthIsUnknown() |
| 154 | +length_status(::Base.HasLength) = LengthIsKnown() |
| 155 | +length_status(::Base.HasShape) = LengthIsKnown() |
| 156 | + |
| 157 | +function check_count_value(n::Int) |
| 158 | + if n < 0 |
| 159 | + throw(ArgumentError("count can't be negative")) |
| 160 | + end |
| 161 | +end |
| 162 | +function check_count_value(n) |
| 163 | + throw(ArgumentError("count must be an `Int`")) |
| 164 | +end |
| 165 | + |
| 166 | +# TODO: use `SpecFSA` for implementing each `FixedSizeArray` constructor? |
| 167 | +struct SpecFSA{N,Mem<:GenericMemory} end |
| 168 | +function fsa_spec_from_type(::Type{FixedSizeArray}) |
| 169 | + SpecFSA{nothing,default_underlying_storage_type}() |
| 170 | +end |
| 171 | +function fsa_spec_from_type(::Type{FixedSizeArray{<:Any,M}}) where {M} |
| 172 | + check_count_value(M) |
| 173 | + SpecFSA{M,default_underlying_storage_type}() |
| 174 | +end |
| 175 | +function fsa_spec_from_type(::Type{FixedSizeArray{E}}) where {E} |
| 176 | + E::Type |
| 177 | + SpecFSA{nothing,default_underlying_storage_type{E}}() |
| 178 | +end |
| 179 | +function fsa_spec_from_type(::Type{FixedSizeArray{E,M}}) where {E,M} |
| 180 | + check_count_value(M) |
| 181 | + E::Type |
| 182 | + SpecFSA{M,default_underlying_storage_type{E}}() |
| 183 | +end |
| 184 | +function fsa_spec_from_type(::Type{FixedSizeArray{E,<:Any,V}}) where {E,V} |
| 185 | + E::Type |
| 186 | + V::Type{<:DenseVector{E}} |
| 187 | + SpecFSA{nothing,V}() |
| 188 | +end |
| 189 | +function fsa_spec_from_type(::Type{FixedSizeArray{E,M,V}}) where {E,M,V} |
| 190 | + check_count_value(M) |
| 191 | + E::Type |
| 192 | + V::Type{<:DenseVector{E}} |
| 193 | + SpecFSA{M,V}() |
| 194 | +end |
| 195 | +for V ∈ (Memory, AtomicMemory) |
| 196 | + T = FixedSizeArray{E,M,V{E}} where {E,M} |
| 197 | + @eval begin |
| 198 | + function fsa_spec_from_type(::Type{$T}) |
| 199 | + SpecFSA{nothing,$V}() |
| 200 | + end |
| 201 | + function fsa_spec_from_type(::Type{($T){<:Any,M}}) where {M} |
| 202 | + check_count_value(M) |
| 203 | + SpecFSA{M,$V}() |
| 204 | + end |
| 205 | + end |
| 206 | +end |
| 207 | + |
| 208 | +parent_type(::Type{<:FixedSizeArray{<:Any,<:Any,T}}) where {T} = T |
| 209 | + |
| 210 | +underlying_storage(m) = m |
| 211 | +underlying_storage(f::FixedSizeArray) = f.mem |
| 212 | + |
| 213 | +axes_are_one_based(axes) = all(isone ∘ first, axes) |
| 214 | + |
| 215 | +# converting constructors for copying other array types |
| 216 | + |
| 217 | +function FixedSizeArray{T,N,V}(src::AbstractArray{S,N}) where {T,N,V,S} |
| 218 | + axs = axes(src) |
| 219 | + if !axes_are_one_based(axs) |
| 220 | + throw(DimensionMismatch("source array has a non-one-based indexing axis")) |
| 221 | + end |
| 222 | + # Can't use `Base.size` because, according to it's doc string, it's not |
| 223 | + # available for all `AbstractArray` types. |
| 224 | + size = map(length, axs) |
| 225 | + dst = FixedSizeArray{T,N,V}(undef, size) |
| 226 | + copyto!(dst.mem, src) |
| 227 | + dst |
| 228 | +end |
| 229 | + |
| 230 | +FixedSizeArray{T,<:Any,V}(a::AbstractArray{<:Any,N}) where {V,T,N} = FixedSizeArray{T,N,V}(a) |
| 231 | + |
| 232 | +FixedSizeArray{T,N}(a::AbstractArray{<:Any,N}) where {T,N} = FixedSizeArray{T,N,default_underlying_storage_type{T}}(a) |
| 233 | +FixedSizeArray{T}(a::AbstractArray{<:Any,N}) where {T,N} = FixedSizeArray{T,N}(a) |
| 234 | +FixedSizeArray{<:Any,N}(a::AbstractArray{T,N}) where {T,N} = FixedSizeArray{T,N}(a) |
| 235 | +FixedSizeArray(a::AbstractArray{T,N}) where {T,N} = FixedSizeArray{T,N}(a) |
| 236 | + |
| 237 | +# conversion |
| 238 | + |
| 239 | +Base.convert(::Type{T}, a::T) where {T<:FixedSizeArray} = a |
| 240 | +Base.convert(::Type{T}, a::AbstractArray) where {T<:FixedSizeArray} = T(a)::T |
| 241 | + |
| 242 | +# `copyto!` |
| 243 | + |
| 244 | +Base.@propagate_inbounds function copyto5!(dst, doff, src, soff, n) |
| 245 | + copyto!(underlying_storage(dst), doff, underlying_storage(src), soff, n) |
| 246 | + dst |
| 247 | +end |
| 248 | + |
| 249 | +Base.@propagate_inbounds function copyto2!(dst, src) |
| 250 | + copyto!(underlying_storage(dst), underlying_storage(src)) |
| 251 | + dst |
| 252 | +end |
| 253 | + |
| 254 | +Base.@propagate_inbounds Base.copyto!(dst::FixedSizeArray, doff::Integer, src::FixedSizeArray, soff::Integer, n::Integer) = copyto5!(dst, doff, src, soff, n) |
| 255 | +Base.@propagate_inbounds Base.copyto!(dst::FixedSizeArray, src::FixedSizeArray) = copyto2!(dst, src) |
| 256 | + |
| 257 | +for A ∈ (Array, GenericMemory) # Add more? Too bad we have to hardcode to avoid ambiguity. |
| 258 | + @eval begin |
| 259 | + Base.@propagate_inbounds Base.copyto!(dst::FixedSizeArray, doff::Integer, src::$A, soff::Integer, n::Integer) = copyto5!(dst, doff, src, soff, n) |
| 260 | + Base.@propagate_inbounds Base.copyto!(dst::$A, doff::Integer, src::FixedSizeArray, soff::Integer, n::Integer) = copyto5!(dst, doff, src, soff, n) |
| 261 | + |
| 262 | + Base.@propagate_inbounds Base.copyto!(dst::FixedSizeArray, src::$A ) = copyto2!(dst, src) |
| 263 | + Base.@propagate_inbounds Base.copyto!(dst::$A, src::FixedSizeArray) = copyto2!(dst, src) |
| 264 | + end |
| 265 | +end |
| 266 | + |
| 267 | +# unsafe: the native address of the array's storage |
| 268 | + |
| 269 | +Base.cconvert(::Type{<:Ptr}, a::FixedSizeArray) = a.mem |
| 270 | + |
| 271 | +# `elsize`: part of the strided arrays interface, used for `pointer` |
| 272 | + |
| 273 | +Base.elsize(::Type{A}) where {A<:FixedSizeArray} = Base.elsize(parent_type(A)) |
| 274 | + |
| 275 | +# `reshape`: specializing it to ensure it returns a `FixedSizeArray` |
| 276 | + |
| 277 | +function Base.reshape(a::FixedSizeArray{T,<:Any,V}, size::NTuple{N,Int}) where {V,T,N} |
| 278 | + len = checked_dims(size) |
| 279 | + if length(a) != len |
| 280 | + throw(DimensionMismatch("new shape not consistent with existing array length")) |
| 281 | + end |
| 282 | + FixedSizeArray{T,N,V}(Internal(), a.mem, size) |
| 283 | +end |
0 commit comments