Merge pull request #534 from JuliaArrays/ajf/SOneTo

Add `SOneTo` as `AbstractUnitRange` to represent `axes(::StaticArray)`

Author: andyferris

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

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}

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)