Last active
January 13, 2021 17:37
-
-
Save fcard/2756d9895ba450a4f1880812b4fa47b6 to your computer and use it in GitHub Desktop.
Applications for IteratorIndexable
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Base: iterate | |
struct Reverse{T} | |
itr::T | |
end | |
reverse(itr) = Reverse(itr) | |
reverse(r::AbstractRange) = Base.reverse(r) | |
Base.eltype(::Type{Reverse{T}}) where T = eltype(T) | |
Base.IteratorEltype(::Type{Reverse{T}}) where T = Base.IteratorEltype(T) | |
Base.IteratorSize(::Type{Reverse{T}}) where T = Base.IteratorSize(T) | |
Base.length(it::Reverse) = length(it.itr) | |
Base.size(it::Reverse) = size(it.itr) | |
function iterate(rev::Reverse, state...) | |
_iterate_reverse(rev.itr, state..., IteratorIndexable(rev.itr), EachIndexSupport(rev.itr)) | |
end | |
_iterate_reverse(it, ::IteratorIndexable, ::EachIndexSupport) = | |
throw(MethodError(iterate, (Reverse(it),))) | |
_iterate_reverse(it, state, ::IteratorIndexable, ::EachIndexSupport) = | |
throw(MethodError(iterate, (Reverse(it), state))) | |
function _iterate_reverse(it, ::IsIndexable, ::HasEachIndex) | |
ei = reverse(eachindex(it)) | |
y = iterate(ei) | |
y === nothing && return nothing | |
it[y[1]], (ei, y[2]) | |
end | |
function _iterate_reverse(it, state, ::IsIndexable, ::HasEachIndex) | |
ei, state = state | |
y = iterate(ei, state) | |
y === nothing && return nothing | |
it[y[1]], (ei, y[2]) | |
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import .Base: iterate | |
struct Drop{T} | |
xs::T | |
n::Int | |
end | |
Base.eltype(::Type{Drop{I}}) where {I} = eltype(I) | |
Base.IteratorEltype(::Type{Drop{I}}) where {I} = Base.IteratorEltype(I) | |
Base.IteratorSize(::Type{Drop{I}}) where {I} = | |
Iterators.drop_iteratorsize(Base.IteratorSize(I)) | |
Base.length(it::Drop) = | |
Iterators._diff_length(it.xs, 1:it.n, Base.IteratorSize(it.xs), Base.HasLength()) | |
function iterate(it::Drop, state...) | |
_iterate_drop(it, state..., eachindex_iteratortype(it.xs)) | |
end | |
@inline function _drop1(idx::AbstractRange, n) | |
(n+1) > length(idx) ? nothing : @inbounds iterate(idx[(n+1):end]) | |
end | |
@inline function _drop1(idx, n) | |
y = iterate(idx) | |
for i in 1:n | |
y === nothing && return nothing | |
y = iterate(idx, y[2]) | |
end | |
y | |
end | |
function _iterate_drop(it, ::Type{<:AbstractRange}) | |
ei = eachindex(it.xs) | |
y = _drop1(ei, it.n) | |
y === nothing && return nothing | |
@inbounds it.xs[y[1]], (ei, y[2]) | |
end | |
function _iterate_drop(it, state, ::Type{<:AbstractRange}) | |
ei, state = state | |
y = iterate(ei, state) | |
y === nothing && return nothing | |
@inbounds it.xs[y[1]], (ei, y[2]) | |
end | |
function _iterate_drop(it, _) | |
_drop1(it.xs, it.n) | |
end | |
function _iterate_drop(it, state, _) | |
iterate(it.xs, state) | |
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Base: getindex, Generator | |
import .Iterators: take, Take, drop, Drop | |
function getindex_rethrowing_boundserror(A, B, idx) | |
try | |
getindex(B, idx...) | |
catch e | |
if length(idx) < 2 && (e == BoundsError(B, idx...) || e == BoundsError(B, idx)) | |
rethrow(BoundsError(A, idx...)) | |
elseif e == BoundsError(B, idx) | |
rethrow(BoundsError(A, idx)) | |
else | |
rethrow(e) | |
end | |
end | |
end | |
function getindex_rethrowing_boundserror(A, i, B, j) | |
try | |
getindex(B, j) | |
catch e | |
if e == BoundsError(B, j) || e == BoundsError(B, (j,)) | |
rethrow(BoundsError(A, i)) | |
else | |
rethrow(e) | |
end | |
end | |
end | |
### Generator | |
function getindex(it::Generator{T,F}, idx...) where {T,F} | |
_getindex_map(it, IteratorIndexable(T), idx...) | |
end | |
function _getindex_map(it, ::IteratorIndexable, idx...) | |
throw(MethodError(getindex, (it, idx...))) | |
end | |
function _getindex_map(it, ::IsIndexable, idx...) | |
value = getindex_rethrowing_boundserror(it, it.iter, idx) | |
it.f(value) | |
end | |
### Take | |
function getindex(it::Take{T}, idx) where T | |
_getindex_take(it, idx, IteratorIndexable(T), EachIndexSupport(T)) | |
end | |
function _getindex_take(it, idx, ::IteratorIndexable, ::EachIndexSupport) where T | |
throw(MethodError(getindex, (it, idx))) | |
end | |
function _checkbounds_take(it, idx) | |
if !_checkindex_take(eachindex(it.xs), idx, it.n) | |
throw(BoundsError(it, idx)) | |
end | |
end | |
for T in [Real, AbstractArray, AbstractUnitRange] | |
@eval _checkindex_take(ei::AbstractUnitRange, idx::$T, n) = checkindex(Bool, ei[1:n], idx) | |
end | |
_checkindex_take(ei, idx::Real, n) = idx in take(ei, n) | |
_checkindex_take(ei, idx::AbstractArray, n) = all(in(take(ei, n)), idx) | |
_checkindex_take(ei, idx::AbstractUnitRange, n) = | |
all(in(take(ei, n)), (first(idx), last(idx))) | |
function _getindex_take(it, idx, ::IsIndexable, ::HasEachIndex) where T | |
@boundscheck _checkbounds_take(it, idx) | |
getindex_rethrowing_boundserror(it, it.xs, idx) | |
end | |
### Drop | |
function getindex(it::Drop{T}, idx) where T | |
_getindex_drop(it, idx, IteratorIndexable(T)) | |
end | |
function _getindex_drop(it, idx, ::IteratorIndexable) where T | |
throw(MethodError(getindex, (it, idx))) | |
end | |
_drop_idx(itr, idx::Real, n) = foldl((i,_)->nextind(itr, i), 1:n, init=idx) | |
_drop_idx(itr, idx::AbstractArray, n) = map(i->_drop_idx(itr, i, n), idx) | |
_drop_idx(itr, idx::AbstractUnitRange, n) = _drop_idx(itr, first(idx), n):_drop_idx(itr, last(idx), n) | |
function _getindex_drop(it, idx, ::IsIndexable) where T | |
j = _drop_idx(it.xs, idx, it.n) | |
getindex_rethrowing_boundserror(it, idx, it.xs, j) | |
end | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment