Lazy segmented storage of residue numbers

[AntonOresten]

e.g. store [1,2,3,4,5,7,8,9] as SegmentVector([1:5, 7:9])

module SegmentVectors

export SegmentVector, segmentize

import Base: iterate, length, size, firstindex, lastindex, getindex,
             IndexStyle, @propagate_inbounds, searchsortedfirst

struct SegmentVector{T,V<:AbstractVector{T}} <: AbstractVector{T}
    segments :: Vector{V}
    cumlens  :: Vector{Int}

    function SegmentVector(segments::AbstractVector{V}) where {T,V<:AbstractVector{T}}
        cumlens  = cumsum(length.(segments))
        new{T,V}(segments, cumlens)
    end
end

IndexStyle(::Type{<:SegmentVector}) = IndexLinear()

@inline length(v::SegmentVector)        = v.cumlens[end]
@inline size(v::SegmentVector)          = (length(v),)
@inline firstindex(::SegmentVector)     = 1
@inline lastindex(v::SegmentVector)     = length(v)

@propagate_inbounds function _find_segment(v::SegmentVector, i::Int)
    s = searchsortedfirst(v.cumlens, i)          # O(log k)
    local_i = i - (s == 1 ? 0 : v.cumlens[s-1])
    return s, local_i
end

@propagate_inbounds function getindex(v::SegmentVector, i::Int)
    1 ≤ i ≤ length(v) || throw(BoundsError(v, i))
    s, local_i = _find_segment(v, i)
    @inbounds return v.segments[s][local_i]
end

@propagate_inbounds function getindex(v::SegmentVector, I::AbstractVector{<:Integer})
    similar(I, eltype(v)) .= (v[i] for i in I)   # fast broadcasted lookup
end

function iterate(v::SegmentVector, state=(1, iterate(v.segments[1])))
    seg, seg_state = state
    seg > length(v.segments) && return nothing
    nxt = iterate(v.segments[seg], seg_state)
    if nxt === nothing
        return iterate(v, (seg+1, iterate(v.segments[seg+1])))
    else
        val, newstate = nxt
        return val, (seg, newstate)
    end
end

"""
    segmentize(v::AbstractVector{<:Integer}, m::Integer = 1) -> SegmentVector

Partition `v` into maximal contiguous runs whose consecutive elements differ by
`m`, and return those runs packed into a `SegmentVector`.  Duplicates are
dropped and the input is sorted ascending first.
"""
function segmentize(v::AbstractVector{T}, ::Val{1}) where {T<:Integer}
    isempty(v) && return SegmentVector{T,UnitRange{T}}(UnitRange{T}[], Int[])

    segs = Vector{UnitRange{T}}()
    start = prev = v[1]

    for x in v[2:end]
        if x - prev == 1
            prev = x
        else
            push!(segs, start:prev)
            start = prev = x
        end
    end

    push!(segs, start:prev)
    return SegmentVector(segs)
end

function segmentize(v::AbstractVector{T}, ::Val{M}) where {T<:Integer, M}
    M == 0 && throw(ArgumentError("step size must be non-zero"))
    M == 1 && error("Use Val{1} method to get UnitRanges")

    isempty(v) && return SegmentVector{T,StepRange{T,Int}}(StepRange{T,Int}[], Int[])

    segs  = Vector{StepRange{T,Int}}()
    step  = M
    start = prev = v[1]

    for x in v[2:end]
        if x - prev == step
            prev = x
        else
            push!(segs, start:step:prev)
            start = prev = x
        end
    end

    push!(segs, start:step:prev)
    return SegmentVector(segs)
end

segmentize(v::AbstractVector{T}, m::Integer=1) where {T<:Integer} = segmentize(v, Val(m))

end