From dc00652285ac7f4d68f439462580c43d8e9989e1 Mon Sep 17 00:00:00 2001 From: Andy Ferris Date: Mon, 29 Oct 2018 14:52:49 +1000 Subject: [PATCH] Add SOneTo as AbstractUnitRange to represent axes(::StaticArray) --- src/SOneTo.jl | 60 ++++++++++++++++++++++++++++++++++++++++++++ src/StaticArrays.jl | 3 +++ src/abstractarray.jl | 32 +++++++++++++++++++++++ 3 files changed, 95 insertions(+) create mode 100644 src/SOneTo.jl diff --git a/src/SOneTo.jl b/src/SOneTo.jl new file mode 100644 index 00000000..84501f92 --- /dev/null +++ b/src/SOneTo.jl @@ -0,0 +1,60 @@ +""" + SOneTo(n) + +Return a statically-sized `AbstractUnitRange` starting at `1`, functioning as the `axes` of +a `StaticArray`. +""" +struct SOneTo{n} <: AbstractUnitRange{Int} +end + +SOneTo(n::Int) = SOneTo{n}() + +Base.axes(s::SOneTo) = (s,) +Base.size(s::SOneTo) = (length(s),) +Base.length(s::SOneTo{n}) where {n} = n + +function Base.getindex(s::SOneTo, i::Int) + @boundscheck checkbounds(s, i) + return i +end +function Base.getindex(s::SOneTo, s2::SOneTo) + @boundscheck checkbounds(s, s2) + return s2 +end + +Base.first(::SOneTo) = 1 +Base.last(::SOneTo{n}) where {n} = n::Int + +@pure function Base.iterate(::SOneTo{n}) where {n} + if n::Int < 1 + return nothing + else + (1, 1) + end +end +function Base.iterate(::SOneTo{n}, s::Int) where {n} + if s < n::Int + s2 = s + 1 + return (s2, s2) + else + return nothing + end +end + +function Base.getproperty(::SOneTo{n}, s::Symbol) where {n} + if s === :start + return 1 + elseif s === :stop + return n::Int + else + error("type SOneTo has no property $s") + end +end + +function Base.show(io::IO, ::SOneTo{n}) where {n} + print(io, "SOneTo(", n::Int, ")") +end + +Base.@pure function Base.checkindex(::Type{Bool}, ::SOneTo{n1}, ::SOneTo{n2}) where {n1, n2} + return n1::Int >= n2::Int +end diff --git a/src/StaticArrays.jl b/src/StaticArrays.jl index 99218c0b..ad852d1c 100644 --- a/src/StaticArrays.jl +++ b/src/StaticArrays.jl @@ -25,6 +25,7 @@ import LinearAlgebra: transpose, adjoint, dot, eigvals, eigen, lyap, tr, import LinearAlgebra: eye end +export SOneTo export StaticScalar, StaticArray, StaticVector, StaticMatrix export Scalar, SArray, SVector, SMatrix export MArray, MVector, MMatrix @@ -40,6 +41,8 @@ export @MVector, @MMatrix, @MArray export similar_type export push, pop, pushfirst, popfirst, insert, deleteat, setindex +include("SOneTo.jl") + """ abstract type StaticArray{S, T, N} <: AbstractArray{T, N} end StaticScalar{T} = StaticArray{Tuple{}, T, 0} diff --git a/src/abstractarray.jl b/src/abstractarray.jl index b000b831..bb2b4d5d 100644 --- a/src/abstractarray.jl +++ b/src/abstractarray.jl @@ -9,6 +9,16 @@ end @inline size(a::StaticArray) = size(typeof(a)) @inline size(a::StaticArray, d::Int) = size(typeof(a), d) +Base.axes(s::StaticArray) = _axes(Size(s)) +@pure function _axes(::Size{sizes}) where {sizes} + map(SOneTo, sizes) +end + +function Base.summary(io::IO, a, inds::Tuple{SOneTo, Vararg{SOneTo}}) + print(io, Base.dims2string(length.(inds)), " ") + Base.showarg(io, a, true) +end + # This seems to confuse Julia a bit in certain circumstances (specifically for trailing 1's) @inline function Base.isassigned(a::StaticArray, i::Int...) ii = LinearIndices(size(a))[i...] @@ -51,6 +61,17 @@ similar_type(::Type{A},s::Size{S}) where {A<:AbstractArray,S} = similar_type(A,e similar_type(::A,::Type{T},s::Size{S}) where {A<:AbstractArray,T,S} = similar_type(A,T,s) +# We should be able to deal with SOneTo axes +similar_type(s::SOneTo) = similar_type(typeof(s)) +similar_type(::Type{SOneTo{n}}) where {n} = similar_type(SOneTo{n}, Int, Size(n)) + +similar_type(::A, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar_type(A, eltype(A), shape) +similar_type(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar_type(A, eltype(A), shape) + +similar_type(::A,::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar_type(A, T, Size(last.(shape))) +similar_type(::Type{A},::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar_type(A, T, Size(last.(shape))) + + # Default types # Generally, use SArray similar_type(::Type{A},::Type{T},s::Size{S}) where {A<:AbstractArray,T,S} = default_similar_type(T,s,length_val(s)) @@ -98,6 +119,17 @@ similar(::Type{A},::Type{T},s::Size{S}) where {A<:AbstractArray,T,S} = mutable_s similar(::Type{SA},::Type{T},s::Size{S}) where {SA<:SizedArray,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef) similar(::Type{A},::Type{T},s::Size{S}) where {A<:Array,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef) +# We should be able to deal with SOneTo axes +similar(::A, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape) +similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape) + +similar(::A,::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape))) +similar(::Type{A},::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape))) + +# Handle mixtures of SOneTo and other ranges (probably should make Base more robust here) +similar(::Type{A}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray} = similar(A, length.(shape)) # Jumps back to 2-argument form in Base +similar(::Type{A},::Type{T}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray,T} = similar(A, length.(shape)) + @inline reshape(a::StaticArray, s::Size) = similar_type(a, s)(Tuple(a)) @inline reshape(a::AbstractArray, s::Size) = _reshape(a, IndexStyle(a), s)