diff --git a/.gitignore b/.gitignore index d2c6291..86ff35b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ *.jl.mem deps/deps.jl .DS_Store +Manifest.toml diff --git a/Project.toml b/Project.toml index 5fa2aa3..ba1557e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "QuasiArrays" uuid = "c4ea9172-b204-11e9-377d-29865faadc5c" authors = ["Sheehan Olver "] -version = "0.1.1" +version = "0.2" [deps] ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" @@ -9,8 +9,8 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [compat] -ArrayLayouts = "0.1.4" -LazyArrays = "0.14.7, 0.15" +ArrayLayouts = "0.1.5" +LazyArrays = "0.15" julia = "1.3" [extras] diff --git a/src/QuasiArrays.jl b/src/QuasiArrays.jl index 47ab28f..4f3e386 100644 --- a/src/QuasiArrays.jl +++ b/src/QuasiArrays.jl @@ -13,7 +13,7 @@ import Base: @_inline_meta, DimOrInd, OneTo, @_propagate_inbounds_meta, @_noinli import Base: ViewIndex, Slice, IdentityUnitRange, ScalarIndex, RangeIndex, view, viewindexing, ensure_indexable, index_dimsum, check_parent_index_match, reindex, _isdisjoint, unsafe_indices, _unsafe_ind2sub, _ind2sub, _sub2ind, _ind2sub_recurse, _lookup, SubArray, - parentindices, reverse, ndims, checkbounds, + parentindices, reverse, ndims, checkbounds, uncolon, promote_shape, maybeview, checkindex, checkbounds_indices, throw_boundserror, rdims, replace_in_print_matrix, show, hcat, vcat, hvcat @@ -47,7 +47,7 @@ import Base.IteratorsMD export AbstractQuasiArray, AbstractQuasiMatrix, AbstractQuasiVector, materialize, QuasiArray, QuasiMatrix, QuasiVector, QuasiDiagonal, Inclusion, QuasiAdjoint, QuasiTranspose, ApplyQuasiArray, ApplyQuasiMatrix, ApplyQuasiVector, - BroadcastQuasiArray, BroadcastQuasiMatrix, BroadcastQuasiVector + BroadcastQuasiArray, BroadcastQuasiMatrix, BroadcastQuasiVector, indextype if VERSION < v"1.3-" """ diff --git a/src/abstractquasiarray.jl b/src/abstractquasiarray.jl index 6937678..a22066d 100644 --- a/src/abstractquasiarray.jl +++ b/src/abstractquasiarray.jl @@ -42,7 +42,7 @@ convert(::Type{AbstractArray{T,N}}, a::AbstractQuasiArray{<:Any,N}) where {T,N} convert(::Type{AbstractMatrix}, a::AbstractQuasiMatrix) = convert(AbstractArray, a) convert(::Type{AbstractVector}, a::AbstractQuasiVector) = convert(AbstractArray, a) - +indextype(A::AbstractQuasiArray) = Tuple{map(eltype, axes(A))...} """ @@ -218,8 +218,8 @@ end keys(s::IndexStyle, A::AbstractQuasiArray, B::AbstractQuasiArray...) = eachindex(s, A, B...) """ - lastindex(collection) -> Number - lastindex(collection, d) -> Number + lastindex(collection) -> index + lastindex(collection, d) -> index Return the last index of `collection`. If `d` is given, return the last index of `collection` along dimension `d`. @@ -239,8 +239,8 @@ lastindex(a::AbstractQuasiArray) = (@_inline_meta; last(eachindex(IndexLinear(), lastindex(a::AbstractQuasiArray, d) = (@_inline_meta; last(axes(a, d))) """ - firstindex(collection) -> Number - firstindex(collection, d) -> Number + firstindex(collection) -> index + firstindex(collection, d) -> index Return the first index of `collection`. If `d` is given, return the first index of `collection` along dimension `d`. @@ -259,7 +259,7 @@ firstindex(a::AbstractQuasiArray, d) = (@_inline_meta; first(axes(a, d))) first(a::AbstractQuasiArray) = a[first(eachindex(a))] stride(A::AbstractQuasiArray, k::Integer) = strides(A)[k] -function isassigned(a::AbstractQuasiArray, i::Number...) +function isassigned(a::AbstractQuasiArray, i...) try a[i...] true @@ -277,11 +277,6 @@ function checkbounds(::Type{Bool}, A::AbstractQuasiArray, I...) checkbounds_indices(Bool, axes(A), I) end -# Linear indexing is explicitly allowed when there is only one (non-cartesian) index -function checkbounds(::Type{Bool}, A::AbstractQuasiArray, i) - @_inline_meta - checkindex(Bool, eachindex(IndexLinear(), A), i) -end # As a special extension, allow using logical arrays that match the source array exactly function checkbounds(::Type{Bool}, A::AbstractQuasiArray{<:Any,N}, I::AbstractQuasiArray{Bool,N}) where N @_inline_meta @@ -291,7 +286,7 @@ end function checkbounds(A::AbstractQuasiArray, I...) @_inline_meta - checkbounds(Bool, A, I...) || throw_boundserror(A, I) + checkbounds(Bool, A, to_indices(A,I)...) || throw_boundserror(A, I) nothing end @@ -302,11 +297,11 @@ similar(a::AbstractQuasiArray, ::Type{T}) where {T} = simila similar(a::AbstractQuasiArray{T}, dims::Tuple) where {T} = similar(a, T, dims) similar(a::AbstractQuasiArray{T}, dims::QuasiDimOrInd...) where {T} = similar(a, T, dims) similar(a::AbstractQuasiArray, ::Type{T}, dims::QuasiDimOrInd...) where {T} = similar(a, T, dims) -similar(::Type{<:AbstractQuasiArray{T}}, shape::NTuple{N,AbstractQuasiOrVector{<:Number}}) where {N,T} = +similar(::Type{<:AbstractQuasiArray{T}}, shape::NTuple{N,AbstractQuasiOrVector}) where {N,T} = QuasiArray{T,N}(undef, convert.(AbstractVector, shape)) -similar(a::AbstractQuasiArray, ::Type{T}, dims::NTuple{N,AbstractQuasiOrVector{<:Number}}) where {T,N} = +similar(a::AbstractQuasiArray, ::Type{T}, dims::NTuple{N,AbstractQuasiOrVector}) where {T,N} = QuasiArray{T,N}(undef, convert.(AbstractVector, dims)) -similar(a::AbstractQuasiArray, ::Type{T}, dims::Vararg{AbstractQuasiOrVector{<:Number},N}) where {T,N} = +similar(a::AbstractQuasiArray, ::Type{T}, dims::Vararg{AbstractQuasiOrVector,N}) where {T,N} = QuasiArray{T,N}(undef, convert.(AbstractVector, dims)) similar(a::AbstractQuasiArray{T}, m::Int) where {T} = Vector{T}(undef, m) @@ -376,10 +371,12 @@ end isempty(a::AbstractQuasiArray) = (length(a) == 0) -function getindex(A::AbstractQuasiArray, I...) +getindex(A::AbstractQuasiArray, I...) = _getindex(indextype(A), A, I) + +function _getindex(::Type{IND}, A::AbstractQuasiArray, I) where IND @_propagate_inbounds_meta - error_if_canonical_getindex(IndexStyle(A), A, I...) - _getindex(IndexStyle(A), A, to_indices(A, I)...) + error_if_canonical_getindex(IndexStyle(A), A, I) + _getindex(IND, IndexStyle(A), A, to_indices(A, I)) end function unsafe_getindex(A::AbstractQuasiArray, I...) @_inline_meta @@ -387,45 +384,32 @@ function unsafe_getindex(A::AbstractQuasiArray, I...) r end -error_if_canonical_getindex(::IndexLinear, A::AbstractQuasiArray, ::Number) = - error("getindex not defined for ", typeof(A)) -error_if_canonical_getindex(::IndexCartesian, A::AbstractQuasiArray{T,N}, ::Vararg{Number,N}) where {T,N} = - error("getindex not defined for ", typeof(A)) -error_if_canonical_getindex(::IndexStyle, ::AbstractQuasiArray, ::Any...) = nothing - ## Internal definitions -_getindex(::IndexStyle, A::AbstractQuasiArray, I...) = lazy_getindex(A, I...) +_getindex(_, ::IndexStyle, A::AbstractQuasiArray, I) = lazy_getindex(A, I...) -## IndexLinear Scalar indexing: canonical method is one Int -_getindex(::IndexLinear, A::AbstractQuasiArray, i::Number) = (@_propagate_inbounds_meta; getindex(A, i)) -function _getindex(::IndexLinear, A::AbstractQuasiArray, I::Vararg{Number,M}) where M - @_inline_meta - @boundscheck checkbounds(A, I...) # generally _to_linear_index requires bounds checking - @inbounds r = getindex(A, _to_linear_index(A, I...)) - r -end -_to_linear_index(A::AbstractQuasiArray, i::Number) = i -_to_linear_index(A::AbstractQuasiVector, i::Number, I::Number...) = i -_to_linear_index(A::AbstractQuasiArray) = 1 -_to_linear_index(A::AbstractQuasiArray, I::Number...) = (@_inline_meta; _sub2ind(A, I...)) -## IndexCartesian Scalar indexing: Canonical method is full dimensionality of Numbers -function _getindex(::IndexCartesian, A::AbstractQuasiArray, I::Vararg{Number,M}) where M +## IndexCartesian Scalar indexing: Canonical method is full dimensionality of indices +function _getindex(::Type{IND}, ::IndexCartesian, A::AbstractQuasiArray, I::IND) where {M,IND} @_inline_meta @boundscheck checkbounds(A, I...) # generally _to_subscript_indices requires bounds checking @inbounds r = getindex(A, _to_subscript_indices(A, I...)...) r end -function _getindex(::IndexCartesian, A::AbstractQuasiArray{T,N}, I::Vararg{Number, N}) where {T,N} - @_propagate_inbounds_meta - getindex(A, I...) -end -_to_subscript_indices(A::AbstractQuasiArray, i::Number) = (@_inline_meta; _unsafe_ind2sub(A, i)) + +error_if_canonical_getindex(::IndexCartesian, A::AbstractQuasiArray{T,N}, I::Tuple) where {T,N} = + _error_if_canonical_getindex(indextype(A), A, I) + +_error_if_canonical_getindex(::Type{IND}, A::AbstractQuasiArray{T,N}, I::IND) where {T,N,IND} = + error("getindex not defined for ", typeof(A)) + +_error_if_canonical_getindex(::Type, ::AbstractQuasiArray, ::Any...) = nothing + +_to_subscript_indices(A::AbstractQuasiArray, i) = (@_inline_meta; _unsafe_ind2sub(A, i)) _to_subscript_indices(A::AbstractQuasiArray{T,N}) where {T,N} = (@_inline_meta; fill_to_length((), 1, Val(N))) _to_subscript_indices(A::AbstractQuasiArray{T,0}) where {T} = () -_to_subscript_indices(A::AbstractQuasiArray{T,0}, i::Number) where {T} = () -_to_subscript_indices(A::AbstractQuasiArray{T,0}, I::Number...) where {T} = () -function _to_subscript_indices(A::AbstractQuasiArray{T,N}, I::Number...) where {T,N} +_to_subscript_indices(A::AbstractQuasiArray{T,0}, i) where {T} = () +_to_subscript_indices(A::AbstractQuasiArray{T,0}, I...) where {T} = () +function _to_subscript_indices(A::AbstractQuasiArray{T,N}, I...) where {T,N} @_inline_meta J, Jrem = IteratorsMD.split(I, Val(N)) _to_subscript_indices(A, J, Jrem) @@ -437,7 +421,7 @@ function __to_subscript_indices(A::AbstractQuasiArray, @_inline_meta (J..., map(first, tail(_remaining_size(J, axes(A))))...) end -_to_subscript_indices(A::AbstractQuasiArray{T,N}, I::Vararg{Number,N}) where {T,N} = I +_to_subscript_indices(A::AbstractQuasiArray{T,N}, I::Vararg{Any,N}) where {T,N} = I ## Setindex! is defined similarly. We first dispatch to an internal _setindex! # function that allows dispatch on array storage @@ -445,8 +429,8 @@ _to_subscript_indices(A::AbstractQuasiArray{T,N}, I::Vararg{Number,N}) where {T, function setindex!(A::AbstractQuasiArray, v, I...) @_propagate_inbounds_meta - error_if_canonical_setindex(IndexStyle(A), A, I...) - _setindex!(IndexStyle(A), A, v, to_indices(A, I)...) + error_if_canonical_setindex(IndexStyle(A), A, I) + _setindex!(indextype(A), IndexStyle(A), A, v, to_indices(A, I)) end function unsafe_setindex!(A::AbstractQuasiArray, v, I...) @_inline_meta @@ -454,34 +438,26 @@ function unsafe_setindex!(A::AbstractQuasiArray, v, I...) r end -error_if_canonical_setindex(::IndexLinear, A::AbstractQuasiArray, ::Number) = - error("setindex! not defined for ", typeof(A)) -error_if_canonical_setindex(::IndexCartesian, A::AbstractQuasiArray{T,N}, ::Vararg{Number,N}) where {T,N} = +error_if_canonical_setindex(::IndexCartesian, A::AbstractQuasiArray{T,N}, I::Tuple) where {T,N} = + _error_if_canonical_setindex(indextype(A), A, I) + +_error_if_canonical_setindex(::Type{IND}, A::AbstractQuasiArray{T,N}, I::IND) where {T,N,IND} = error("setindex! not defined for ", typeof(A)) -error_if_canonical_setindex(::IndexStyle, ::AbstractQuasiArray, ::Any...) = nothing +_error_if_canonical_setindex(::Type, ::AbstractQuasiArray, ::Any...) = nothing ## Internal definitions -_setindex!(::IndexStyle, A::AbstractQuasiArray, v, I...) = +_setindex!(::Type, ::IndexStyle, A::AbstractQuasiArray, v, I) = error("setindex! for $(typeof(A)) with types $(typeof(I)) is not supported") -## IndexLinear Scalar indexing -_setindex!(::IndexLinear, A::AbstractQuasiArray, v, i::Number) = (@_propagate_inbounds_meta; setindex!(A, v, i)) -function _setindex!(::IndexLinear, A::AbstractQuasiArray, v, I::Vararg{Number,M}) where M - @_inline_meta - @boundscheck checkbounds(A, I...) - @inbounds r = setindex!(A, v, _to_linear_index(A, I...)) - r -end - # IndexCartesian Scalar indexing -function _setindex!(::IndexCartesian, A::AbstractQuasiArray{T,N}, v, I::Vararg{Number, N}) where {T,N} +function _setindex!(::Type{IND}, ::IndexCartesian, A::AbstractQuasiArray{T,N}, v, I::NTuple{N,Any}) where {T,N,IND} @_propagate_inbounds_meta setindex!(A, v, I...) end -function _setindex!(::IndexCartesian, A::AbstractQuasiArray, v, I::Vararg{Number,M}) where M +function _setindex!(::Type{IND}, ::IndexCartesian, A::AbstractQuasiArray, v, I::NTuple{M,Any}) where {M,IND} @_inline_meta @boundscheck checkbounds(A, I...) - @inbounds r = setindex!(A, v, _to_subscript_indices(A, I...)...) + @inbounds r = setindex!(A, v, _to_subscript_indices(IND, A, I...)...) r end @@ -516,8 +492,8 @@ dataids(A::AbstractQuasiArray) = (UInt(objectid(A)),) ## structured matrix methods ## -replace_in_print_matrix(A::AbstractQuasiMatrix,i::Number,j::Number,s::AbstractString) = s -replace_in_print_matrix(A::AbstractQuasiVector,i::Number,j::Number,s::AbstractString) = s +replace_in_print_matrix(A::AbstractQuasiMatrix,i,j,s::AbstractString) = s +replace_in_print_matrix(A::AbstractQuasiVector,i,j,s::AbstractString) = s ## Concatenation ## eltypeof(x::AbstractQuasiArray) = eltype(x) @@ -565,45 +541,14 @@ function (==)(A::AbstractQuasiArray, B::AbstractQuasiArray) return anymissing ? missing : true end -# _sub2ind and _ind2sub -# fallbacks -function _sub2ind(A::AbstractQuasiArray, I...) - @_inline_meta - _sub2ind(axes(A), I...) -end - -function _ind2sub(A::AbstractQuasiArray, ind) - @_inline_meta - _ind2sub(axes(A), ind) -end - -# Vectorized forms -function _sub2ind(inds::Indices{1}, I1::AbstractQuasiVector{T}, I::AbstractQuasiVector{T}...) where T<:Number - throw(ArgumentError("Linear indexing is not defined for one-dimensional arrays")) -end -_sub2ind(inds::Tuple{OneTo}, I1::AbstractQuasiVector{T}, I::AbstractQuasiVector{T}...) where {T<:Number} = - _sub2ind_vecs(inds, I1, I...) -_sub2ind(inds::Union{DimsInteger,Indices}, I1::AbstractQuasiVector{T}, I::AbstractQuasiVector{T}...) where {T<:Number} = - _sub2ind_vecs(inds, I1, I...) -function _sub2ind_vecs(inds, I::AbstractQuasiVector...) - I1 = I[1] - Iinds = axes1(I1) - for j = 2:length(I) - axes1(I[j]) == Iinds || throw(DimensionMismatch("indices of I[1] ($(Iinds)) does not match indices of I[$j] ($(axes1(I[j])))")) - end - Iout = similar(I1) - _sub2ind!(Iout, inds, Iinds, I) - Iout -end - _lookup(ind, r::Inclusion) = ind -_ind2sub(dims::NTuple{N,Number}, ind::Number) where N = (@_inline_meta; _ind2sub_recurse(dims, ind-1)) -_ind2sub(inds::QuasiIndices, ind::Number) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) -_ind2sub(inds::Tuple{Inclusion{<:Number},AbstractUnitRange{<:Integer}}, ind::Number) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) -_ind2sub(inds::Tuple{AbstractUnitRange{<:Integer},Inclusion{<:Number}}, ind::Number) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) +_ind2sub(dims::NTuple{N,Any}, ind) where N = (@_inline_meta; _ind2sub_recurse(dims, ind-1)) +_ind2sub(inds::QuasiIndices, ind) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) +_ind2sub(inds::Tuple{Inclusion{<:Any},AbstractUnitRange{<:Integer}}, ind) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) +_ind2sub(inds::Tuple{AbstractUnitRange{<:Integer},Inclusion{<:Any}}, ind) = (@_inline_meta; _ind2sub_recurse(inds, ind-1)) -function _ind2sub(inds::Union{NTuple{N,Number},QuasiIndices{N}}, ind::AbstractQuasiVector{<:Number}) where N +function _ind2sub(inds::Union{NTuple{N,Any},QuasiIndices{N}}, ind::AbstractQuasiVector) where N M = length(ind) t = ntuple(n->similar(ind),Val(N)) for (i,idx) in pairs(IndexLinear(), ind) diff --git a/src/indices.jl b/src/indices.jl index acffa50..61b7082 100644 --- a/src/indices.jl +++ b/src/indices.jl @@ -100,9 +100,23 @@ to_index(I::AbstractQuasiArray) = I to_index(I::AbstractQuasiArray{<:Union{AbstractArray, Colon}}) = throw(ArgumentError("invalid index: $I of type $(typeof(I))")) -to_quasi_index(i::Number) = i -to_quasi_index(i) = Base.to_index(i) -to_index(A::AbstractQuasiArray, i) = to_quasi_index(i) +to_quasi_index(::Type{IND}, i) where IND = convert(IND, i) +to_quasi_index(::Type{IND}, I::AbstractArray) where IND<:AbstractArray = convert(IND, I) +to_quasi_index(::Type{IND}, I::AbstractQuasiArray) where IND<:AbstractQuasiArray = convert(IND, I) +to_quasi_index(::Type{IND}, I::AbstractArray) where IND = convert(AbstractArray{IND}, I) +to_quasi_index(::Type{IND}, I::AbstractQuasiArray) where IND = convert(AbstractQuasiArray{IND}, I) +to_quasi_index(::Type{IND}, I::AbstractArray{<:AbstractArray}) where IND<:AbstractArray = convert(AbstractArray{IND}, I) +to_quasi_index(::Type{IND}, I::AbstractQuasiArray{<:AbstractArray}) where IND<:AbstractArray = convert(AbstractQuasiArray{IND}, I) + +to_quasi_index(A, IND, i) = to_quasi_index(IND,i) + +to_indices(A::AbstractQuasiArray, inds, ::Tuple{}) = () +to_indices(A::AbstractQuasiArray, inds, I::Tuple{Any,Vararg{Any}}) = + (@_inline_meta; (to_quasi_index(A, eltype(inds[1]), I[1]), to_indices(A, _maybetail(inds), tail(I))...)) +@inline to_indices(A::AbstractQuasiArray, inds, I::Tuple{Colon, Vararg{Any}}) = + (uncolon(inds, I), to_indices(A, _maybetail(inds), tail(I))...) + + LinearIndices(A::AbstractQuasiArray) = LinearIndices(axes(A)) @@ -164,9 +178,9 @@ size(S::Inclusion) = (cardinality(S.domain),) length(S::Inclusion) = cardinality(S.domain) unsafe_length(S::Inclusion) = cardinality(S.domain) cardinality(S::Inclusion) = cardinality(S.domain) -getindex(S::Inclusion{T}, i::Number) where T = +getindex(S::Inclusion{T}, i::T) where T = (@_inline_meta; @boundscheck checkbounds(S, i); convert(T,i)) -getindex(S::Inclusion{T}, i::AbstractVector{<:Number}) where T = +getindex(S::Inclusion{T}, i::AbstractVector{T}) where T = (@_inline_meta; @boundscheck checkbounds(S, i); convert(AbstractVector{T},i)) getindex(S::Inclusion, i::Inclusion) = (@_inline_meta; @boundscheck checkbounds(S, i); copy(S)) @@ -176,10 +190,10 @@ iterate(S::Inclusion, s...) = iterate(S.domain, s...) in(x, S::Inclusion) = x in S.domain -checkindex(::Type{Bool}, inds::Inclusion, i::Number) = i ∈ inds.domain +checkindex(::Type{Bool}, inds::Inclusion, i) = i ∈ inds.domain checkindex(::Type{Bool}, inds::Inclusion, ::Colon) = true checkindex(::Type{Bool}, inds::Inclusion, ::Inclusion) = true -function checkindex(::Type{Bool}, inds::Inclusion, I::AbstractArray) +function __checkindex(::Type{Bool}, inds::Inclusion, I::AbstractArray) @_inline_meta b = true for i in I @@ -188,9 +202,16 @@ function checkindex(::Type{Bool}, inds::Inclusion, I::AbstractArray) b end -function checkindex(::Type{Bool}, inds::Inclusion, r::AbstractRange) +checkindex(::Type{Bool}, inds::Inclusion{T}, I::AbstractArray{T}) where T = + __checkindex(Bool, inds, I) +checkindex(::Type{Bool}, inds::Inclusion{T}, I::AbstractArray{T}) where T<:AbstractArray = + __checkindex(Bool, inds, I) +checkindex(::Type{Bool}, inds::Inclusion{T}, I::AbstractArray{<:AbstractArray}) where T<:AbstractArray = + __checkindex(Bool, inds, convert(AbstractArray{T}, I)) + +function checkindex(::Type{Bool}, inds::Inclusion{T}, r::AbstractRange) where T @_propagate_inbounds_meta - isempty(r) | (checkindex(Bool, inds, first(r)) & checkindex(Bool, inds, last(r))) + isempty(r) | (checkindex(Bool, inds, convert(T, first(r))) & checkindex(Bool, inds, last(r))) end checkindex(::Type{Bool}, indx::Inclusion, I::AbstractVector{Bool}) = indx == axes1(I) checkindex(::Type{Bool}, indx::Inclusion, I::AbstractArray{Bool}) = false diff --git a/src/multidimensional.jl b/src/multidimensional.jl index d8b1f72..c742f7f 100644 --- a/src/multidimensional.jl +++ b/src/multidimensional.jl @@ -58,20 +58,17 @@ module QuasiIteratorsMD 5 ``` """ + function _QuasiCartesianIndex end struct QuasiCartesianIndex{N,II<:Tuple} <: AbstractCartesianIndex{N} I::II - QuasiCartesianIndex{N,II}(index::II) where {N,II} = new(index) + global _QuasiCartesianIndex(index::NTuple{N,Any}) where N = new{N,typeof(index)}(index) end - QuasiCartesianIndex{N}(index::NTuple{N,Number}) where N = QuasiCartesianIndex{N,typeof(index)}(index) - QuasiCartesianIndex(index::NTuple{N,Number}) where {N} = QuasiCartesianIndex{N}(index) - QuasiCartesianIndex(index::Number...) = QuasiCartesianIndex(index) - QuasiCartesianIndex{N}(index::Vararg{Number,N}) where {N} = QuasiCartesianIndex{N}(index) + QuasiCartesianIndex(index...) = _QuasiCartesianIndex(flatten(index)) + QuasiCartesianIndex{N}(index::Vararg{Any,N}) where {N} = _QuasiCartesianIndex(index) # Allow passing tuples smaller than N - QuasiCartesianIndex{N}(index::Tuple) where {N} = QuasiCartesianIndex{N}(fill_to_length(index, 1, Val(N))) - QuasiCartesianIndex{N}(index::Number...) where {N} = QuasiCartesianIndex{N}(index) - QuasiCartesianIndex{N}() where {N} = QuasiCartesianIndex{N}(()) - # Un-nest passed QuasiCartesianIndexes - QuasiCartesianIndex(index::Union{Number, QuasiCartesianIndex}...) = QuasiCartesianIndex(flatten(index)) + QuasiCartesianIndex{N}(index...) where {N} = _QuasiCartesianIndex(fill_to_length(index, 1, Val(N))) + QuasiCartesianIndex{N}() where {N} = _QuasiCartesianIndex(fill_to_length((), 1, Val(N))) + QuasiCartesianIndex{N,II}(index...) where {N,II<:NTuple{N,Any}} = _QuasiCartesianIndex(convert(II, index)) flatten(I::Tuple{}) = I flatten(I::Tuple{Any}) = I flatten(I::Tuple{<:QuasiCartesianIndex}) = I[1].I @@ -79,11 +76,10 @@ module QuasiIteratorsMD @inline _flatten() = () @inline _flatten(i, I...) = (i, _flatten(I...)...) @inline _flatten(i::QuasiCartesianIndex, I...) = (i.I..., _flatten(I...)...) - QuasiCartesianIndex(index::Tuple{Vararg{Union{Number, QuasiCartesianIndex}}}) = QuasiCartesianIndex(index...) show(io::IO, i::QuasiCartesianIndex) = (print(io, "QuasiCartesianIndex"); show(io, i.I)) Base.convert(::Type{QuasiCartesianIndex{N,II}}, Q::QuasiCartesianIndex) where {N,II<:Tuple} = - QuasiCartesianIndex{N,II}(convert(II, Q.I)) + QuasiCartesianIndex{N,II}(convert(II, Q.I)...) # length length(::QuasiCartesianIndex{N}) where {N} = N @@ -131,7 +127,7 @@ module QuasiIteratorsMD icmp(a, b) = ifelse(isless(a,b), 1, ifelse(a==b, 0, -1)) # conversions - convert(::Type{T}, index::QuasiCartesianIndex{1}) where {T<:Number} = convert(T, index[1]) + convert(::Type{T}, index::QuasiCartesianIndex{1}) where {T} = convert(T, index[1]) convert(::Type{T}, index::QuasiCartesianIndex) where {T<:Tuple} = convert(T, index.I) # hashing @@ -246,18 +242,18 @@ module QuasiIteratorsMD For cartesian to linear index conversion, see [`LinearIndices`](@ref). """ - struct QuasiCartesianIndices{N,R<:NTuple{N,AbstractQuasiOrVector{<:Number}},RR<:NTuple{N,Number}} <: AbstractArray{QuasiCartesianIndex{N,RR},N} + struct QuasiCartesianIndices{N,R<:NTuple{N,AbstractQuasiOrVector},RR<:NTuple{N,Any}} <: AbstractArray{QuasiCartesianIndex{N,RR},N} indices::R end - QuasiCartesianIndices(nd::NTuple{N,AbstractVector{<:Number}}) where N = + QuasiCartesianIndices(nd::NTuple{N,AbstractVector}) where N = QuasiCartesianIndices{N,typeof(nd),Tuple{map(eltype,nd)...}}(nd) - QuasiCartesianIndices(nd::NTuple{N,AbstractQuasiOrVector{<:Number}}) where N = + QuasiCartesianIndices(nd::NTuple{N,AbstractQuasiOrVector}) where N = QuasiCartesianIndices(convert.(AbstractArray,nd)) QuasiCartesianIndices(::Tuple{}) = QuasiCartesianIndices{0,typeof(())}(()) QuasiCartesianIndices(index::QuasiCartesianIndex) = QuasiCartesianIndices(index.I) - QuasiCartesianIndices(sz::NTuple{N,<:Number}) where {N} = QuasiCartesianIndices(map(Base.OneTo, sz)) + QuasiCartesianIndices(sz::NTuple{N,Any}) where {N} = QuasiCartesianIndices(map(Base.OneTo, sz)) QuasiCartesianIndices(A::AbstractQuasiArray) = QuasiCartesianIndices(axes(A)) @@ -313,7 +309,7 @@ module QuasiIteratorsMD Base.IndexStyle(::Type{QuasiCartesianIndices{N,R}}) where {N,R} = IndexCartesian() @inline function Base.getindex(iter::QuasiCartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R} @boundscheck checkbounds(iter, I...) - QuasiCartesianIndex(getindex.(iter.indices, I)) + _QuasiCartesianIndex(getindex.(iter.indices, I)) end ndims(R::QuasiCartesianIndices) = ndims(typeof(R)) @@ -373,14 +369,6 @@ end @inline checkbounds_indices(::Type{Bool}, IA::Tuple, I::Tuple{QuasiCartesianIndex,Vararg{Any}}) = checkbounds_indices(Bool, IA, (I[1].I..., tail(I)...)) -# Indexing into Array with mixtures of Numbers and QuasiCartesianIndices is -# extremely performance-sensitive. While the abstract fallbacks support this, -# codegen has extra support for SIMDification that sub2ind doesn't (yet) support -@propagate_inbounds getindex(A::Array, i1::Union{Number, QuasiCartesianIndex}, I::Union{Number, QuasiCartesianIndex}...) = - A[to_indices(A, (i1, I...))...] -@propagate_inbounds setindex!(A::Array, v, i1::Union{Number, QuasiCartesianIndex}, I::Union{Number, QuasiCartesianIndex}...) = - (A[to_indices(A, (i1, I...))...] = v; A) - # Support indexing with an array of QuasiCartesianIndex{N}s # Here we try to consume N of the indices (if there are that many available) # The first two simply handle ambiguities @@ -422,13 +410,15 @@ end # In simple cases, we know that we don't need to use axes(A). Optimize those # until Julia gets smart enough to elide the call on its own: -@inline to_indices(A, I::Tuple{Vararg{Union{Number, QuasiCartesianIndex}}}) = to_indices(A, (), I) +@inline to_indices(A::AbstractQuasiArray, I::Tuple{Vararg{Union{Any, QuasiCartesianIndex}}}) = to_indices(A, axes(A), I) +to_indices(A::AbstractQuasiArray, I::Tuple{Any}) = (@_inline_meta; to_indices(A, axes(A), I)) # But some index types require more context spanning multiple indices # QuasiCartesianIndexes are simple; they just splat out -@inline to_indices(A, inds, I::Tuple{QuasiCartesianIndex, Vararg{Any}}) = +@inline to_indices(A::AbstractQuasiArray, inds, I::Tuple{QuasiCartesianIndex, Vararg{Any}}) = to_indices(A, inds, (I[1].I..., tail(I)...)) +@inline to_indices(A::AbstractQuasiArray, I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = to_indices(A, axes(A), I) # But for arrays of QuasiCartesianIndex, we just skip the appropriate number of inds -@inline function to_indices(A, inds, I::Tuple{AbstractArray{QuasiCartesianIndex{N}}, Vararg{Any}}) where N +@inline function to_indices(A::AbstractQuasiArray, inds, I::Tuple{AbstractArray{QuasiCartesianIndex{N}}, Vararg{Any}}) where N _, indstail = IteratorsMD.split(inds, Val(N)) (to_index(A, I[1]), to_indices(A, indstail, tail(I))...) end @@ -440,7 +430,7 @@ const CI0 = Union{QuasiCartesianIndex{0}, AbstractArray{QuasiCartesianIndex{0}}} getindex(x::Number, i::QuasiCartesianIndex{0}) = x getindex(t::Tuple, i::QuasiCartesianIndex{1}) = getindex(t, i.I[1]) -@inline function _getindex(l::IndexStyle, A::AbstractQuasiArray, I::Union{Number, AbstractArray}...) +@inline function _getindex(l::IndexStyle, A::AbstractQuasiArray, I::Union{Any, AbstractArray}...) @boundscheck checkbounds(A, I...) return _unsafe_getindex(l, _maybe_reshape(l, A, I...), I...) end @@ -450,7 +440,7 @@ end (ntuple(x->true, Val(N))..., index_dimsum(I...)...) end -Slice(d::AbstractQuasiVector{<:Number}) = Inclusion(d) +Slice(d::AbstractQuasiVector) = Inclusion(d) _maybe_reshape(::IndexLinear, A::AbstractQuasiArray, I...) = A @@ -459,11 +449,11 @@ _maybe_reshape(::IndexCartesian, A::AbstractQuasiVector, I...) = A @inline __maybe_reshape(A::AbstractQuasiArray{T,N}, ::NTuple{N,Any}) where {T,N} = A @inline __maybe_reshape(A::AbstractQuasiArray, ::NTuple{N,Any}) where {N} = reshape(A, Val(N)) -_unsafe_getindex(::IndexStyle, A::AbstractQuasiArray, I::Vararg{Union{Number, AbstractArray}, N}) where N = +_unsafe_getindex(::IndexStyle, A::AbstractQuasiArray, I::Vararg{Union{Any, AbstractArray}, N}) where N = lazy_getindex(A, I...) # Always index with the exactly indices provided. -@generated function _unsafe_getindex!(dest::Union{AbstractArray,AbstractQuasiArray}, src::AbstractQuasiArray, I::Vararg{Union{Number, AbstractArray}, N}) where N +@generated function _unsafe_getindex!(dest::Union{AbstractArray,AbstractQuasiArray}, src::AbstractQuasiArray, I::Vararg{Union{Any, AbstractArray}, N}) where N quote @_inline_meta D = eachindex(dest) diff --git a/src/quasiarray.jl b/src/quasiarray.jl index 8a5f968..a603925 100644 --- a/src/quasiarray.jl +++ b/src/quasiarray.jl @@ -1,31 +1,31 @@ -struct QuasiArray{T,N,AXES<:NTuple{N,AbstractVector{<:Number}}} <: AbstractQuasiArray{T,N} +struct QuasiArray{T,N,AXES<:NTuple{N,AbstractVector}} <: AbstractQuasiArray{T,N} parent::Array{T,N} axes::AXES - function QuasiArray{T,N,AXES}(par::AbstractArray{T,N}, axes::AXES) where {T,N,AXES<:NTuple{N,AbstractVector{<:Number}}} + function QuasiArray{T,N,AXES}(par::AbstractArray{T,N}, axes::AXES) where {T,N,AXES<:NTuple{N,AbstractVector}} size(par) == length.(axes) || throw(ArgumentError("Axes must be compatible with parent dimensions")) new{T,N,AXES}(convert(Array{T,N},par),axes) end end -const QuasiMatrix{T,AXES<:NTuple{2,AbstractVector{<:Number}}} = QuasiArray{T,2,AXES} -const QuasiVector{T,AXES<:Tuple{AbstractVector{<:Number}}} = QuasiArray{T,1,AXES} +const QuasiMatrix{T,AXES<:NTuple{2,AbstractVector}} = QuasiArray{T,2,AXES} +const QuasiVector{T,AXES<:Tuple{AbstractVector}} = QuasiArray{T,1,AXES} -QuasiArray{T,N}(::UndefInitializer, axes::NTuple{N,AbstractVector{<:Number}}) where {T,N} = +QuasiArray{T,N}(::UndefInitializer, axes::NTuple{N,AbstractVector}) where {T,N} = QuasiArray(Array{T}(undef, map(length,axes)), axes) -QuasiArray{T,N}(::UndefInitializer, axes::Vararg{AbstractVector{<:Number},N}) where {T,N} = +QuasiArray{T,N}(::UndefInitializer, axes::Vararg{AbstractVector,N}) where {T,N} = QuasiArray{T,N}(undef, axes) -QuasiVector(::UndefInitializer, axes::AbstractVector{<:Number}) where T = +QuasiVector(::UndefInitializer, axes::AbstractVector) where T = QuasiArray(Vector(undef,length(axes)), (axes,)) -QuasiMatrix(::UndefInitializer, ax1::AbstractVector{<:Number}, ax2::AbstractVector{<:Number}) where T = +QuasiMatrix(::UndefInitializer, ax1::AbstractVector, ax2::AbstractVector) where T = QuasiArray(Matrix(undef,length(ax1),length(ax2)), (ax1,ax2)) -QuasiArray(par::AbstractArray{T,N}, axes::NTuple{N,AbstractVector{<:Number}}) where {T,N} = +QuasiArray(par::AbstractArray{T,N}, axes::NTuple{N,AbstractVector}) where {T,N} = QuasiArray{T,N,typeof(axes)}(par, axes) -QuasiMatrix(par::AbstractMatrix{T}, axes::NTuple{2,AbstractVector{<:Number}}) where T = +QuasiMatrix(par::AbstractMatrix{T}, axes::NTuple{2,AbstractVector}) where T = QuasiArray{T,2,typeof(axes)}(par, axes) -QuasiVector(par::AbstractVector{T}, axes::Tuple{AbstractVector{<:Number}}) where T = +QuasiVector(par::AbstractVector{T}, axes::Tuple{AbstractVector}) where T = QuasiArray{T,1,typeof(axes)}(par, axes) QuasiVector(par::AbstractVector{T}, axes::AbstractArray) where {T} = @@ -49,7 +49,7 @@ _inclusion(d) = Inclusion(d) axes(A::QuasiArray) = _inclusion.(A.axes) parent(A::QuasiArray) = A.parent -@propagate_inbounds @inline function getindex(A::QuasiArray{<:Any,N}, I::Vararg{Number,N}) where N +@propagate_inbounds @inline function _getindex(::Type{IND}, A::QuasiArray{<:Any,N}, I::IND) where {N,IND} @boundscheck checkbounds(A, I...) A.parent[findfirst.(isequal.(I), A.axes)...] end diff --git a/src/quasibroadcast.jl b/src/quasibroadcast.jl index 75de787..1746d33 100644 --- a/src/quasibroadcast.jl +++ b/src/quasibroadcast.jl @@ -108,16 +108,16 @@ Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple{}) = () (ifelse(keep[1], I[1], Idefault[1]), _newindex(tail(I), tail(keep), tail(Idefault))...) @inline _newindex(I, keep::Tuple{}, Idefault) = () # truncate if keep is shorter than I # for now we assume indexing is simple -Base.@propagate_inbounds newindex(arg, I::QuasiCartesianIndex) = QuasiCartesianIndex(_newindex(axes(arg), I.I)) -@inline newindex(I::QuasiCartesianIndex, keep, Idefault) = QuasiCartesianIndex(_newindex(I.I, keep, Idefault)) +Base.@propagate_inbounds newindex(arg, I::QuasiCartesianIndex) = QuasiCartesianIndex(_newindex(axes(arg), I.I)...) +@inline newindex(I::QuasiCartesianIndex, keep, Idefault) = QuasiCartesianIndex(_newindex(I.I, keep, Idefault)...) @inline function Base.getindex(bc::Broadcasted, I::QuasiCartesianIndex) @boundscheck checkbounds(bc, I) @inbounds _broadcast_getindex(bc, I) end -Base.@propagate_inbounds Base.getindex(bc::Broadcasted{<:AbstractQuasiArrayStyle}, i1::Number, I::Number...) = bc[QuasiCartesianIndex((i1, I...))] -Base.@propagate_inbounds Base.getindex(bc::Broadcasted{<:AbstractQuasiArrayStyle}) = bc[QuasiCartesianIndex(())] +Base.@propagate_inbounds Base.getindex(bc::Broadcasted{<:AbstractQuasiArrayStyle}, i1::Number, I::Number...) = bc[QuasiCartesianIndex(i1, I...)] +Base.@propagate_inbounds Base.getindex(bc::Broadcasted{<:AbstractQuasiArrayStyle}) = bc[QuasiCartesianIndex()] @inline Base.checkbounds(bc::Broadcasted{<:AbstractQuasiArrayStyle}, I::Number) = Base.checkbounds_indices(Bool, axes(bc), (I,)) || Base.throw_boundserror(bc, (I,)) diff --git a/src/subquasiarray.jl b/src/subquasiarray.jl index 29c4d51..e24bc7f 100644 --- a/src/subquasiarray.jl +++ b/src/subquasiarray.jl @@ -71,9 +71,43 @@ end const QViewIndex = Union{ViewIndex,AbstractQuasiArray} +# This computes the linear indexing compatibility for a given tuple of indices +quasi_viewindexing(::Tuple{}, I::Tuple{}) = IndexLinear() +# Leading scalar indices simply increase the stride +quasi_viewindexing(axs::Tuple{AbstractQuasiVector{IND}, Vararg{Any}}, I::Tuple{IND, Vararg{Any}}) where IND = + (@_inline_meta; quasi_viewindexing(tail(axs), tail(I))) +quasi_viewindexing(axs::Tuple{AbstractVector{IND}, Vararg{Any}}, I::Tuple{IND, Vararg{Any}}) where IND = + (@_inline_meta; quasi_viewindexing(tail(axs), tail(I))) + +# Slices may begin a section which may be followed by any number of Slices +# quasi_viewindexing(axs, I::Tuple{Slice, Slice, Vararg{Any}}) = (@_inline_meta; quasi_viewindexing(tail(I))) +# # A UnitRange can follow Slices, but only if all other indices are scalar +# quasi_viewindexing(I::Tuple{Slice, AbstractUnitRange, Vararg{ScalarIndex}}) = IndexLinear() +# quasi_viewindexing(I::Tuple{Slice, Slice, Vararg{ScalarIndex}}) = IndexLinear() # disambiguate +# # In general, ranges are only fast if all other indices are scalar +# quasi_viewindexing(I::Tuple{AbstractRange, Vararg{ScalarIndex}}) = IndexLinear() +# # All other index combinations are slow +# quasi_viewindexing(I::Tuple{Vararg{Any}}) = IndexCartesian() +# # Of course, all other array types are slow +quasi_viewindexing(axs::Tuple{AbstractQuasiVector{IND}, Vararg{Any}}, I::Tuple{AbstractArray{IND}, Vararg{Any}}) where IND = IndexCartesian() +quasi_viewindexing(axs::Tuple{AbstractVector{IND}, Vararg{Any}}, I::Tuple{AbstractArray{IND}, Vararg{Any}}) where IND = IndexCartesian() + +# combined dimensionality of all indices +# rather than returning N, it returns an NTuple{N,Bool} so the result is inferrable +@inline quasi_index_dimsum(axs::Tuple{AbstractQuasiVector{IND},Vararg{Any}}, inds::Tuple{IND,Vararg{Any}}) where IND = + (quasi_index_dimsum(tail(axs), tail(inds))...,) +@inline quasi_index_dimsum(axs::Tuple{AbstractVector{IND},Vararg{Any}}, inds::Tuple{IND,Vararg{Any}}) where IND = + (quasi_index_dimsum(tail(axs), tail(inds))...,) +@inline quasi_index_dimsum(axs, inds::Tuple{Colon,Vararg{Any}}) = (true, quasi_index_dimsum(tail(axs), tail(inds))...) +@inline quasi_index_dimsum(axs::Tuple{AbstractQuasiVector{IND},Vararg{Any}}, inds::Tuple{AbstractArray{IND,N},Vararg{Any}}) where {N,IND} = + (ntuple(x->true, Val(N))..., quasi_index_dimsum(tail(axs), tail(inds))...) +@inline quasi_index_dimsum(axs::Tuple{AbstractVector{IND},Vararg{Any}}, inds::Tuple{AbstractArray{IND,N},Vararg{Any}}) where {N,IND} = + (ntuple(x->true, Val(N))..., quasi_index_dimsum(tail(axs), tail(inds))...) +quasi_index_dimsum(::Tuple{}, ::Tuple{}) = () + function SubArray(parent::AbstractQuasiArray, indices::Tuple) @_inline_meta - SubArray(IndexStyle(viewindexing(indices), IndexStyle(parent)), parent, ensure_indexable(indices), index_dimsum(indices...)) + SubArray(IndexStyle(quasi_viewindexing(axes(parent), indices), IndexStyle(parent)), parent, ensure_indexable(indices), quasi_index_dimsum(axes(parent), indices)) end function unsafe_view(A::AbstractQuasiArray, I::Vararg{ViewIndex,N}) where {N} @@ -108,20 +142,12 @@ _subarray(A::AbstractArray, idxs) = SubArray(A, idxs) _subarray(A::AbstractQuasiArray, idxs) = SubQuasiArray(A, idxs) _subarray(A::AbstractQuasiArray, idxs::NTuple{N,ViewIndex}) where {N} = SubArray(A, idxs) - -if VERSION < v"1.2-" - function _maybe_reindex(V, I, ::Tuple{}) - @_inline_meta - @inbounds idxs = to_indices(V.parent, reindex(V, V.indices, I)) - _subarray(V.parent, idxs) - end -else - function _maybe_reindex(V, I, ::Tuple{}) - @_inline_meta - @inbounds idxs = to_indices(V.parent, reindex(V.indices, I)) - _subarray(V.parent, idxs) - end +function _maybe_reindex(V, I, ::Tuple{}) + @_inline_meta + @inbounds idxs = to_indices(V.parent, reindex(V.indices, I)) + _subarray(V.parent, idxs) end + ## Re-indexing is the heart of a view, transforming A[i, j][x, y] to A[i[x], j[y]] # # Recursively look through the heads of the parent- and sub-indices, considering @@ -132,64 +158,73 @@ end AbstractZeroDimQuasiArray{T} = AbstractQuasiArray{T, 0} -if VERSION < v"1.2-" - # Re-index into parent vectors with one subindex - reindex(V, idxs::Tuple{AbstractQuasiVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) = - (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], reindex(V, tail(idxs), tail(subidxs))...)) - - # Parent matrices are re-indexed with two sub-indices - reindex(V, idxs::Tuple{AbstractQuasiMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) = - (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], reindex(V, tail(idxs), tail(tail(subidxs)))...)) - - # In general, we index N-dimensional parent arrays with N indices - @generated function reindex(V, idxs::Tuple{AbstractQuasiArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where {T,N} - if length(subidxs.parameters) >= N - subs = [:(subidxs[$d]) for d in 1:N] - tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)] - :(@_propagate_inbounds_meta; (idxs[1][$(subs...)], reindex(V, tail(idxs), ($(tail...),))...)) - else - :(throw(ArgumentError("cannot re-index $(ndims(V)) dimensional SubQuasiArray with fewer than $(ndims(V)) indices\nThis should not occur; please submit a bug report."))) - end - end -else - # Re-index into parent vectors with one subindex - reindex(idxs::Tuple{AbstractQuasiVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) = - (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], reindex(tail(idxs), tail(subidxs))...)) - - # Parent matrices are re-indexed with two sub-indices - reindex(idxs::Tuple{AbstractQuasiMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) = - (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], reindex(tail(idxs), tail(tail(subidxs)))...)) - - # In general, we index N-dimensional parent arrays with N indices - @generated function reindex(idxs::Tuple{AbstractQuasiArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where {T,N} - if length(subidxs.parameters) >= N - subs = [:(subidxs[$d]) for d in 1:N] - tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)] - :(@_propagate_inbounds_meta; (idxs[1][$(subs...)], reindex(tail(idxs), ($(tail...),))...)) - else - :(throw(ArgumentError("cannot re-index SubArray with fewer indices than dimensions\nThis should not occur; please submit a bug report."))) - end +# Re-index into parent vectors with one subindex +reindex(idxs::Tuple{AbstractQuasiVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) = + (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], reindex(tail(idxs), tail(subidxs))...)) + +# Parent matrices are re-indexed with two sub-indices +reindex(idxs::Tuple{AbstractQuasiMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) = + (@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], reindex(tail(idxs), tail(tail(subidxs)))...)) + +# In general, we index N-dimensional parent arrays with N indices +@generated function reindex(idxs::Tuple{AbstractQuasiArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where {T,N} + if length(subidxs.parameters) >= N + subs = [:(subidxs[$d]) for d in 1:N] + tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)] + :(@_propagate_inbounds_meta; (idxs[1][$(subs...)], reindex(tail(idxs), ($(tail...),))...)) + else + :(throw(ArgumentError("cannot re-index SubArray with fewer indices than dimensions\nThis should not occur; please submit a bug report."))) end end +# indices are taken from the range/vector +# Since bounds-checking is performance-critical and uses +# indices, it's worth optimizing these implementations thoroughly +axes(S::SubArray{T,N,<:AbstractQuasiArray}) where {T,N} = + _quasi_indices_sub(axes(parent(S)), S.indices) +_quasi_indices_sub(axs::Tuple{AbstractQuasiVector{IND},Vararg{Any}}, inds::Tuple{IND,Vararg{Any}}) where IND = + (@_inline_meta; _quasi_indices_sub(tail(axs), tail(inds))) +_quasi_indices_sub(axs::Tuple{AbstractVector{IND},Vararg{Any}}, inds::Tuple{IND,Vararg{Any}}) where IND = + (@_inline_meta; _quasi_indices_sub(tail(axs), tail(inds))) +_quasi_indices_sub(::Tuple{}, ::Tuple{}) = () +function _quasi_indices_sub(axs::Tuple{AbstractQuasiVector{IND},Vararg{Any}}, inds::Tuple{AbstractArray{IND},Vararg{Any}}) where IND + @_inline_meta + (unsafe_indices(inds[1])..., _quasi_indices_sub(tail(axs), tail(inds))...) +end +function _quasi_indices_sub(axs::Tuple{AbstractVector{IND},Vararg{Any}}, inds::Tuple{AbstractArray{IND},Vararg{Any}}) where IND + @_inline_meta + (unsafe_indices(inds[1])..., _quasi_indices_sub(tail(axs), tail(inds))...) +end + +quasi_reindex(axs::Tuple{AbstractQuasiVector{IND}, Vararg{Any}}, idxs::Tuple{IND, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where IND = + (@_propagate_inbounds_meta; (idxs[1], quasi_reindex(tail(axs), tail(idxs), subidxs)...)) +quasi_reindex(axs::Tuple{AbstractVector{IND}, Vararg{Any}}, idxs::Tuple{IND, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where IND = + (@_propagate_inbounds_meta; (idxs[1], quasi_reindex(tail(axs), tail(idxs), subidxs)...)) +quasi_reindex(axs::Tuple{AbstractQuasiVector{IND}, Vararg{Any}}, idxs::Tuple{AbstractVector{IND}, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) where IND = + (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], quasi_reindex(tail(axs), tail(idxs), tail(subidxs))...)) +quasi_reindex(axs::Tuple{AbstractVector{IND}, Vararg{Any}}, idxs::Tuple{AbstractVector{IND}, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) where IND = + (@_propagate_inbounds_meta; (idxs[1][subidxs[1]], quasi_reindex(tail(axs), tail(idxs), tail(subidxs))...)) + + + +quasi_reindex(::Tuple{}, ::Tuple{}, ::Tuple{}) = () + +function getindex(V::SubArray{T,N,<:AbstractQuasiArray}, I::Vararg{Int,N}) where {T,N} + @_inline_meta + @boundscheck checkbounds(V, I...) + @inbounds r = V.parent[quasi_reindex(axes(parent(V)), V.indices, I)...] + r +end + # In general, we simply re-index the parent indices by the provided ones SlowSubQuasiArray{T,N,P,I} = SubQuasiArray{T,N,P,I,false} -if VERSION < v"1.2-" - function getindex(V::SlowSubQuasiArray{T,N}, I::Vararg{Number,N}) where {T,N} - @_inline_meta - @boundscheck checkbounds(V, I...) - @inbounds r = V.parent[reindex(V, V.indices, I)...] - r - end -else - function getindex(V::SlowSubQuasiArray{T,N}, I::Vararg{Number,N}) where {T,N} - @_inline_meta - @boundscheck checkbounds(V, I...) - @inbounds r = V.parent[reindex(V.indices, I)...] - r - end +function getindex(V::SlowSubQuasiArray{T,N}, I::Vararg{Number,N}) where {T,N} + @_inline_meta + @boundscheck checkbounds(V, I...) + @inbounds r = V.parent[reindex(V.indices, I)...] + r end FastSubQuasiArray{T,N,P,I} = SubQuasiArray{T,N,P,I,true} @@ -208,21 +243,13 @@ function getindex(V::FastContiguousSubQuasiArray, i::Number) r end -if VERSION < v"1.2-" - function setindex!(V::SlowSubQuasiArray{T,N}, x, I::Vararg{Number,N}) where {T,N} - @_inline_meta - @boundscheck checkbounds(V, I...) - @inbounds V.parent[reindex(V, V.indices, I)...] = x - V - end -else - function setindex!(V::SlowSubQuasiArray{T,N}, x, I::Vararg{Number,N}) where {T,N} - @_inline_meta - @boundscheck checkbounds(V, I...) - @inbounds V.parent[reindex(V.indices, I)...] = x - V - end +function setindex!(V::SlowSubQuasiArray{T,N}, x, I::Vararg{Number,N}) where {T,N} + @_inline_meta + @boundscheck checkbounds(V, I...) + @inbounds V.parent[reindex(V.indices, I)...] = x + V end + function setindex!(V::FastSubQuasiArray, x, i::Number) @_inline_meta @boundscheck checkbounds(V, i) diff --git a/test/test_abstractquasiarray.jl b/test/test_abstractquasiarray.jl index a51816f..05a3403 100644 --- a/test/test_abstractquasiarray.jl +++ b/test/test_abstractquasiarray.jl @@ -1,4 +1,5 @@ using QuasiArrays, Test +import QuasiArrays: QuasiCartesianIndex @testset "AbstractQuasiArray" begin A = QuasiArray(rand(5,4,3), (range(0;stop=1,length=5), Base.OneTo(4), [2,3,6])) @@ -57,7 +58,6 @@ using QuasiArrays, Test @test QuasiArrays.checkbounds(Bool, A, 1:6, 1:4, 1:3) == false @test QuasiArrays.checkbounds(Bool, A, 1:5, 1:5, 1:3) == false @test QuasiArrays.checkbounds(Bool, A, 1:5, 1:4, 1:4) == false - @test QuasiArrays.checkbounds(Bool, A, 1:60) == true @test QuasiArrays.checkbounds(Bool, A, 1:61) == false @test QuasiArrays.checkbounds(Bool, A, 0.25, 2, 2, 1:1) == true # extra indices @test QuasiArrays.checkbounds(Bool, A, 2, 2, 2, 1:2) == false @@ -103,4 +103,30 @@ using QuasiArrays, Test @test axes(v) == axes(v[:]) == axes(v[Inclusion(0:0.5:1)]) == (Inclusion(0:0.5:1),) @test v[0.5] == v[:][0.5] == v[Inclusion(0:0.5:1)][0.5] == 2 end + + @testset "Vec indexing" begin + A = QuasiArray(rand(2), ([[1,2],[3,4]],)) + @test @inferred(indextype(A)) == Tuple{Vector{Int}} + @test QuasiArrays.checkbounds(Bool, A, [1,2]) + @test A[[1,2]] == parent(A)[1] + @test_throws BoundsError A[[1,2,3]] + @test A[[[1,2],[3,4]]] == parent(A) + + @test Base.to_indices(A,(QuasiCartesianIndex([1,2]),)) == ([1,2],) + @test A[QuasiCartesianIndex([1,2])] == A[[1,2]] + + A = QuasiArray(rand(2,2), ([[1,2],[3,4]],[[5,6],[7,8]])) + @test A[[1,2], [5,6]] == parent(A)[1] + @test_throws BoundsError A[[1,2]] + V = view(A,[[1,2],[3,4]],[[5,6],[7,8]]) + @test A[[[1,2],[3,4]],[[5,6],[7,8]]] == V == parent(A) + V = view(A,[1,2], [[5,6],[7,8]]) + @test axes(V) == (Base.OneTo(2),) + @test V == parent(A)[1,:] + + A = QuasiArray(rand(2), ([[1.0,2],[3.0,4]],)) + @test Base.to_indices(A, axes(A), ([1,2],)) isa Tuple{Vector{Float64}} + @test parentindices(view(A,[1,2])) isa Tuple{Vector{Float64}} + @test A[[1,2]] == A[[1.0,2.0]] == parent(A)[1] + end end diff --git a/test/test_quasibroadcast.jl b/test/test_quasibroadcast.jl index aab1ed6..e022860 100644 --- a/test/test_quasibroadcast.jl +++ b/test/test_quasibroadcast.jl @@ -24,7 +24,7 @@ import QuasiArrays: QuasiCartesianIndex, QuasiCartesianIndices, DefaultQuasiArra @test_throws DimensionMismatch check_broadcast_axes(ax, zeros(3,4,2)) @test_throws DimensionMismatch check_broadcast_axes(ax, zeros(3,5), zeros(2)) - ci(x) = QuasiCartesianIndex(x) + ci(x) = QuasiCartesianIndex(x...) @test @inferred(newindex(ci((2,2)), (true, true), (-1,-1))) == ci((2,2)) @test @inferred(newindex(ci((2,2)), (true, false), (-1,-1))) == ci((2,-1)) @test @inferred(newindex(ci((2,2)), (false, true), (-1,-1))) == ci((-1,2))