Generalize the element type of BlockedUnitRange (#337)

* Allow more general BlockUnitRange element types

* Restrict element type

* Get tests passing

* Fix some tests

* Fix some doctests

* Skip broken test in Julia v1.6

* Better support for unitful numbers

* Fix tests

* Stricter types in _BlockedUnitRange

* Improve tests coverage

Author: mtfishman

src/blockaxis.jl

@@ -16,7 +16,7 @@ function findblockindex(b::AbstractVector, k::Integer)
     bl = blocklasts(b)
     blockidx = _searchsortedfirst(bl, k)
     @assert blockindex != lastindex(bl) + 1 # guaranteed by the @boundscheck above
-    prevblocklast = blockidx == firstindex(bl) ? first(b)-1 : bl[blockidx-1]
+    prevblocklast = blockidx == firstindex(bl) ? first(b)-oneunit(eltype(b)) : bl[blockidx-1]
     local_index = k - prevblocklast
     return BlockIndex(blockidx, local_index)
 end
@@ -36,7 +36,7 @@ a vector of block lengths to a `BlockedUnitRange`.
 # Examples
 ```jldoctest
 julia> blockedrange(2, [2,2,3]) # first value and block lengths
-3-blocked 7-element BlockedUnitRange{Vector{Int64}}:
+3-blocked 7-element BlockedUnitRange{Int64, Vector{Int64}}:
  2
  3
  ─
@@ -50,31 +50,59 @@ julia> blockedrange(2, [2,2,3]) # first value and block lengths
 
 See also [`BlockedOneTo`](@ref).
 """
-struct BlockedUnitRange{CS} <: AbstractBlockedUnitRange{Int,CS}
-    first::Int
+struct BlockedUnitRange{T<:Integer,CS} <: AbstractBlockedUnitRange{T,CS}
+    first::T
     lasts::CS
     # assume that lasts is sorted, no checks carried out here
-    global function _BlockedUnitRange(f, cs::CS) where CS
+    global function _BlockedUnitRange(f::T, cs::CS) where {T,CS<:AbstractVector{T}}
         Base.require_one_based_indexing(cs)
-        new{CS}(f, cs)
+        return new{T,CS}(f, cs)
+    end
+    global function _BlockedUnitRange(f::T, cs::CS) where {T,CS<:Tuple{T,Vararg{T}}}
+        return new{T,CS}(f, cs)
+    end
+    global function _BlockedUnitRange(f::T, cs::Tuple{}) where {T}
+        return new{T,Tuple{}}(f, cs)
     end
 end
 
-@inline _BlockedUnitRange(cs) = _BlockedUnitRange(1,cs)
+@inline function _BlockedUnitRange(f::T, cs::AbstractVector{S}) where {T,S}
+  U = promote_type(T, S)
+  return _BlockedUnitRange(convert(U, f), convert.(U, cs))
+end
+@inline function _BlockedUnitRange(f::T, cs::Tuple{S,Vararg{S}}) where {T,S}
+  U = promote_type(T, S)
+  return _BlockedUnitRange(convert(U, f), convert.(U, cs))
+end
+@inline function _BlockedUnitRange(f, cs::Tuple)
+  return _BlockedUnitRange(f, promote(cs...))
+end
+@inline _BlockedUnitRange(cs::AbstractVector) = _BlockedUnitRange(oneunit(eltype(cs)), cs)
+@inline _BlockedUnitRange(cs::NTuple) = _BlockedUnitRange(oneunit(eltype(cs)), cs)
+_BlockedUnitRange(cs::Tuple) = _BlockedUnitRange(promote(cs...))
 
 first(b::BlockedUnitRange) = b.first
 @inline blocklasts(a::BlockedUnitRange) = a.lasts
 
 BlockedUnitRange(::BlockedUnitRange) = throw(ArgumentError("Forbidden due to ambiguity"))
+# Use `accumulate` instead of `cumsum` because it preserves the element type of the block lengths
+_blocklengths2blocklasts(blocks) = accumulate(+, blocks) # extra level to allow changing default accumulate behaviour
+# Use `cumsum` for fill arrays to output lazy `StepRangeLen` representation
+_blocklengths2blocklasts(blocks::Fill) = cumsum(blocks)
+_blocklengths2blocklasts(blocks::Ones) = cumsum(blocks)
+
+@inline blockfirsts(a::AbstractBlockedUnitRange) = [first(a); @views(blocklasts(a)[1:end-1]) .+ oneunit(eltype(a))]
 
-@inline blockfirsts(a::AbstractBlockedUnitRange) = [first(a); @views(blocklasts(a)[1:end-1]) .+ 1]
 # optimize common cases
 @inline function blockfirsts(a::AbstractBlockedUnitRange{<:Any,<:Union{Vector, RangeCumsum{<:Any, <:UnitRange}}})
     v = Vector{eltype(a)}(undef, length(blocklasts(a)))
     v[1] = first(a)
-    v[2:end] .= @views(blocklasts(a)[oneto(end-1)]) .+ 1
+    v[2:end] .= @views(blocklasts(a)[oneto(end-1)]) .+ oneunit(eltype(a))
     return v
 end
+@inline function blockfirsts(a::AbstractBlockedUnitRange{<:Any,<:Tuple})
+    return (first(a), (blocklasts(a)[oneto(end-1)] .+ oneunit(eltype(a)))...)
+end
 
 """
     BlockedOneTo
@@ -124,8 +152,6 @@ BlockedOneTo(::BlockedOneTo) = throw(ArgumentError("Forbidden due to ambiguity")
 
 axes(b::BlockedOneTo) = (b,)
 
-_blocklengths2blocklasts(blocks) = cumsum(blocks) # extra level to allow changing default cumsum behaviour
-
 """
     blockedrange(blocklengths::Union{Tuple, AbstractVector})
     blockedrange(first::Int, blocklengths::Union{Tuple, AbstractVector})
@@ -144,27 +170,27 @@ julia> blockedrange([1,2])
  3
 
 julia> blockedrange(2, (1,2))
-2-blocked 3-element BlockedUnitRange{Tuple{Int64, Int64}}:
+2-blocked 3-element BlockedUnitRange{Int64, Tuple{Int64, Int64}}:
  2
  ─
  3
  4
 ```
 """
 @inline blockedrange(blocks::Union{Tuple,AbstractVector}) = BlockedOneTo(_blocklengths2blocklasts(blocks))
-@inline blockedrange(f::Int, blocks::Union{Tuple,AbstractVector}) = _BlockedUnitRange(f, f-1 .+ _blocklengths2blocklasts(blocks))
+@inline blockedrange(f::Integer, blocks::Union{Tuple,AbstractVector}) = _BlockedUnitRange(f, f-oneunit(f) .+ _blocklengths2blocklasts(blocks))
 
 _diff(a::AbstractVector) = diff(a)
 _diff(a::Tuple) = diff(collect(a))
-@inline _blocklengths(a, bl, dbl) = isempty(bl) ? [dbl;] : [first(bl)-first(a)+1; dbl]
+@inline _blocklengths(a, bl, dbl) = isempty(bl) ? [dbl;] : [first(bl)-first(a)+oneunit(eltype(a)); dbl]
 @inline function _blocklengths(a::BlockedOneTo, bl::RangeCumsum, ::OrdinalRange)
     # the 1:0 is hardcoded here to enable conversions to a Base.OneTo
     isempty(bl) ? oftype(bl.range, 1:0) : bl.range
 end
 @inline _blocklengths(a, bl) = _blocklengths(a, bl, _diff(bl))
 @inline blocklengths(a::AbstractBlockedUnitRange) = _blocklengths(a, blocklasts(a))
 
-length(a::AbstractBlockedUnitRange) = isempty(blocklasts(a)) ? 0 : Integer(last(blocklasts(a))-first(a)+1)
+length(a::AbstractBlockedUnitRange) = isempty(blocklasts(a)) ? zero(eltype(a)) : Integer(last(blocklasts(a))-first(a)+oneunit(eltype(a)))
 
 """
     blockisequal(a::AbstractUnitRange{Int}, b::AbstractUnitRange{Int})
@@ -217,15 +243,15 @@ function Base.convert(::Type{BlockedUnitRange}, axis::AbstractUnitRange{Int})
     f = first(axis)
     _BlockedUnitRange(f, _shift_blocklengths(axis, bl, f))
 end
-function Base.convert(::Type{BlockedUnitRange{CS}}, axis::AbstractUnitRange{Int}) where CS
+function Base.convert(::Type{BlockedUnitRange{T,CS}}, axis::AbstractUnitRange{Int}) where {T,CS}
     bl = blocklasts(axis)
     f = first(axis)
-    _BlockedUnitRange(f, convert(CS, _shift_blocklengths(axis, bl, f)))
+    _BlockedUnitRange(convert(T, f), convert(CS, _shift_blocklengths(axis, bl, f)))
 end
 
 Base.unitrange(b::AbstractBlockedUnitRange) = first(b):last(b)
 
-Base.promote_rule(::Type{<:AbstractBlockedUnitRange}, ::Type{Base.OneTo{Int}}) = UnitRange{Int}
+Base.promote_rule(::Type{<:AbstractBlockedUnitRange{T}}, ::Type{Base.OneTo{Int}}) where {T} = UnitRange{promote_type(T, Int)}
 
 function Base.convert(::Type{BlockedOneTo}, axis::AbstractUnitRange{Int})
     first(axis) == 1 || throw(ArgumentError("first element of range is not 1"))
@@ -352,10 +378,10 @@ julia> blocksizes(A,2)
 blocksizes(A) = map(blocklengths, axes(A))
 blocksizes(A,i) = blocklengths(axes(A,i))
 
-axes(b::AbstractBlockedUnitRange) = (BlockedOneTo(blocklasts(b) .- (first(b)-1)),)
+axes(b::AbstractBlockedUnitRange) = (BlockedOneTo(blocklasts(b) .- (first(b)-oneunit(eltype(b)))),)
 unsafe_indices(b::AbstractBlockedUnitRange) = axes(b)
 # ::Integer works around case where blocklasts might return different type
-last(b::AbstractBlockedUnitRange)::Integer = isempty(blocklasts(b)) ? first(b)-1 : last(blocklasts(b))
+last(b::AbstractBlockedUnitRange)::Integer = isempty(blocklasts(b)) ? first(b)-oneunit(eltype(b)) : last(blocklasts(b))
 
 # view and indexing are identical for a unitrange
 view(b::AbstractBlockedUnitRange, K::Block{1}) = b[K]
@@ -367,19 +393,19 @@ view(b::AbstractBlockedUnitRange, K::Block{1}) = b[K]
     @boundscheck K in bax || throw(BlockBoundsError(b, k))
     S = first(bax)
     K == S && return first(b):first(cs)
-    return cs[k-1]+1:cs[k]
+    return cs[k-1]+oneunit(eltype(b)):cs[k]
 end
 
 @propagate_inbounds function getindex(b::AbstractBlockedUnitRange, KR::BlockRange{1})
     cs = blocklasts(b)
-    isempty(KR) && return _BlockedUnitRange(1,cs[1:0])
+    isempty(KR) && return _BlockedUnitRange(oneunit(eltype(b)),cs[1:0])
     K,J = first(KR),last(KR)
     k,j = Integer(K),Integer(J)
     bax = blockaxes(b,1)
     @boundscheck K in bax || throw(BlockBoundsError(b,K))
     @boundscheck J in bax || throw(BlockBoundsError(b,J))
     K == first(bax) && return _BlockedUnitRange(first(b),cs[k:j])
-    _BlockedUnitRange(cs[k-1]+1,cs[k:j])
+    _BlockedUnitRange(cs[k-1]+oneunit(eltype(b)),cs[k:j])
 end
 
 @propagate_inbounds function getindex(b::AbstractBlockedUnitRange, KR::BlockRange{1,Tuple{Base.OneTo{Int}}})
@@ -442,7 +468,7 @@ function findblock(b::AbstractUnitRange{Int}, k::Integer)
 end
 
 """
-    blockfirsts(a::AbstractUnitRange{Int})
+    blockfirsts(a::AbstractUnitRange{<:Integer})
 
 Return the first index of each block of `a`.
 
@@ -466,9 +492,9 @@ julia> blockfirsts(b)
  4
 ```
 """
-blockfirsts(a::AbstractUnitRange{Int}) = Ones{Int}(1)
+blockfirsts(a::AbstractUnitRange{<:Integer}) = Ones{eltype(a)}(1)
 """
-    blocklasts(a::AbstractUnitRange{Int})
+    blocklasts(a::AbstractUnitRange{<:Integer})
 
 Return the last index of each block of `a`.
 
@@ -492,9 +518,9 @@ julia> blocklasts(b)
  6
 ```
 """
-blocklasts(a::AbstractUnitRange{Int}) = Fill(length(a),1)
+blocklasts(a::AbstractUnitRange{<:Integer}) = Fill(eltype(a)(length(a)),1)
 """
-    blocklengths(a::AbstractUnitRange{Int})
+    blocklengths(a::AbstractUnitRange{<:Integer})
 
 Return the length of each block of `a`.
 
@@ -518,7 +544,7 @@ julia> blocklengths(b)
  3
 ```
 """
-blocklengths(a::AbstractUnitRange) = blocklasts(a) .- blockfirsts(a) .+ 1
+blocklengths(a::AbstractUnitRange{<:Integer}) = blocklasts(a) .- blockfirsts(a) .+ oneunit(eltype(a))
 
 Base.summary(io::IO, a::AbstractBlockedUnitRange) =  _block_summary(io, a)
 
@@ -538,7 +564,6 @@ blockaxes(S::Base.Slice) = blockaxes(S.indices)
 _broadcaststyle(_) = Broadcast.DefaultArrayStyle{1}()
 Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = _broadcaststyle(Base.BroadcastStyle(R))
 
-
 ###
 # Special Fill/Range cases
 #
@@ -548,22 +573,22 @@ Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = _broa
 _blocklengths2blocklasts(blocks::AbstractRange) = RangeCumsum(blocks)
 function blockfirsts(a::AbstractBlockedUnitRange{<:Any,Base.OneTo{Int}})
     first(a) == 1 || error("Offset axes not supported")
-    Base.OneTo{Int}(length(blocklasts(a)))
+    Base.OneTo{eltype(a)}(length(blocklasts(a)))
 end
 function blocklengths(a::AbstractBlockedUnitRange{<:Any,Base.OneTo{Int}})
     first(a) == 1 || error("Offset axes not supported")
-    Ones{Int}(length(blocklasts(a)))
+    Ones{eltype(a)}(length(blocklasts(a)))
 end
 function blockfirsts(a::AbstractBlockedUnitRange{<:Any,<:AbstractRange})
     st = step(blocklasts(a))
     first(a) == 1 || error("Offset axes not supported")
-    @assert first(blocklasts(a))-first(a)+1 == st
-    range(1; step=st, length=length(blocklasts(a)))
+    @assert first(blocklasts(a))-first(a)+oneunit(eltype(a)) == st
+    range(oneunit(eltype(a)); step=st, length=eltype(a)(length(blocklasts(a))))
 end
 function blocklengths(a::AbstractBlockedUnitRange{<:Any,<:AbstractRange})
     st = step(blocklasts(a))
     first(a) == 1 || error("Offset axes not supported")
-    @assert first(blocklasts(a))-first(a)+1 == st
+    @assert first(blocklasts(a))-first(a)+oneunit(eltype(a)) == st
     Fill(st,length(blocklasts(a)))
 end