Skip to content

Commit

Permalink
Simplify selectindices (#856)
Browse files Browse the repository at this point in the history 
* simplify selectindices dispatch

* reorgaise selectindices

* simplify
  • Loading branch information
@rafaqz
rafaqz authored Nov 12, 2024
1 parent 02b730f commit fa1f52d
Showing 1 changed file with 46 additions and 69 deletions.
115 changes: 46 additions & 69 deletions src/Lookups/selector.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,16 +79,6 @@ const SelectorOrInterval = Union{Selector,Interval,Not}

const SelTuple = Tuple{SelectorOrInterval,Vararg{SelectorOrInterval}}

# `Not` form InvertedIndices.jl
@inline function selectindices(l::Lookup, sel::Not; kw...)
indices = selectindices(l, sel.skip; kw...)
return first(to_indices(l, (Not(indices),)))
end
@inline function selectindices(l::Lookup, sel; kw...)
selstr = sprint(show, sel)
throw(ArgumentError("Invalid index `$selstr`. Did you mean `At($selstr)`? Use stardard indices, `Selector`s, or `Val` for compile-time `At`."))
end

"""
At <: IntSelector
Expand Down Expand Up @@ -136,24 +126,6 @@ Base.show(io::IO, x::At) = print(io, "At(", val(x), ", ", atol(x), ", ", rtol(x)
struct _True end
struct _False end

@inline selectindices(l::Lookup, sel::At; kw...) = at(l, sel; kw...)
@inline selectindices(l::Lookup, sel::At{<:AbstractVector}; kw...) =
_selectvec(l, sel; kw...)
@inline selectindices(l::Lookup, sel::At{<:Tuple{<:Any,<:Any}}; kw...) =
_selecttuple(l, sel; kw...)
# Handle lookups of Tuple
@inline selectindices(l::Lookup{<:Tuple}, sel::At{<:Tuple}; kw...) = at(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::At{<:Tuple{<:Any,<:Any}}; kw...) =
at(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::At{<:Tuple{<:Tuple,<:Tuple}}; kw...) =
_selecttuple(l, sel; kw...)

@inline _selectvec(l, sel; kw...) = [selectindices(l, rebuild(sel, v); kw...) for v in val(sel)]
@inline function _selecttuple(l, sel; kw...)
v1, v2 = _maybeflipbounds(l, val(sel))
selectindices(l, rebuild(sel, v1); kw...):selectindices(l, rebuild(sel, v2); kw...)
end

function at(lookup::AbstractCyclic{Cycling}, sel::At; kw...)
cycled_sel = rebuild(sel, cycle_val(lookup, val(sel)))
return at(no_cycling(lookup), cycled_sel; kw...)
Expand Down Expand Up @@ -288,13 +260,6 @@ end
Near() = Near(nothing)
Near(a, b) = Near((a, b))

@inline selectindices(l::Lookup, sel::Near; kw...) = near(l, sel; kw...)
@inline selectindices(l::Lookup, sel::Near{<:AbstractVector}; kw...) = _selectvec(l, sel; kw...)
@inline selectindices(l::Lookup, sel::Near{<:Tuple}; kw...) = _selecttuple(l, sel; kw...)
# Handle lookups of Tuple
@inline selectindices(l::Lookup{<:Tuple}, sel::Near{<:Tuple}; kw...) = near(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::Near{<:Tuple{<:Tuple,<:Tuple}}; kw...) = _selecttuple(l, sel; kw...)

Base.show(io::IO, x::Near) = print(io, "Near(", val(x), ")")

function near(lookup::AbstractCyclic{Cycling}, sel::Near; kw...)
Expand Down Expand Up @@ -392,14 +357,11 @@ struct Contains{T} <: IntSelector{T}
end
Contains() = Contains(nothing)
Contains(a, b) = Contains((a, b))

# Filter based on sampling and selector -----------------
@inline selectindices(l::Lookup, sel::Contains; kw...) = contains(l, sel; kw...)
@inline selectindices(l::Lookup, sel::Contains{<:AbstractVector}; kw...) = _selectvec(l, sel; kw...)
@inline selectindices(l::Lookup, sel::Contains{<:Tuple}; kw...) = _selecttuple(l, sel; kw...)
# Handle lookups of Tuple
@inline selectindices(l::Lookup{<:Tuple}, sel::Contains{<:Tuple}; kw...) = contains(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::Contains{<:Tuple{<:Tuple,<:Tuple}}; kw...) = _selecttuple(l, sel; kw...)
@inline _selectvec(l, sel; kw...) = [selectindices(l, rebuild(sel, v); kw...) for v in val(sel)]
@inline function _selecttuple(l, sel; kw...)
v1, v2 = _maybeflipbounds(l, val(sel))
selectindices(l, rebuild(sel, v1); kw...):selectindices(l, rebuild(sel, v2); kw...)
end

Base.show(io::IO, x::Contains) = print(io, "Contains(", val(x), ")")

Expand Down Expand Up @@ -612,14 +574,6 @@ abstract type _Side end
struct _Upper <: _Side end
struct _Lower <: _Side end

@inline selectindices(l::Lookup, sel::Union{Between{<:Tuple},Interval}) = between(l, sel)
@inline function selectindices(lookup::Lookup, sel::Between{<:AbstractVector})
inds = Int[]
for v in val(sel)
append!(inds, selectindices(lookup, rebuild(sel, v)))
end
end

# between
# returns a UnitRange from an Interval
function between(l::Lookup, sel::Between)
Expand Down Expand Up @@ -871,14 +825,6 @@ Touches(a, b) = Touches((a, b))
Base.first(sel::Touches) = first(val(sel))
Base.last(sel::Touches) = last(val(sel))

@inline selectindices(l::Lookup, sel::Touches) = touches(l, sel)
@inline function selectindices(lookup::Lookup, sel::Touches{<:AbstractVector})
inds = Int[]
for v in val(sel)
append!(inds, selectindices(lookup, rebuild(sel, v)))
end
end

# touches for tuple intervals
# returns a UnitRange like Touches/Interval but for cells contained
# NoIndex behaves like `Sampled` `ForwardOrdered` `Points` of 1:N Int
Expand Down Expand Up @@ -1048,7 +994,7 @@ val(sel::Where) = sel.f
Base.show(io::IO, x::Where) = print(io, "Where(", repr(val(x)), ")")

# Yes this is everything. `Where` doesn't need lookup specialisation
@inline function selectindices(lookup::Lookup, sel::Where)
@inline function _selectindices(lookup::Lookup, sel::Where)
[i for (i, v) in enumerate(parent(lookup)) if sel.f(v)]
end

Expand Down Expand Up @@ -1089,12 +1035,11 @@ All(args::SelectorOrInterval...) = All(args)

Base.show(io::IO, x::All) = print(io, "All(", x.selectors, ")")

@inline function selectindices(lookup::Lookup, sel::All)
@inline function _selectindices(lookup::Lookup, sel::All)
results = map(s -> selectindices(lookup, s), sel.selectors)
sort!(union(results...))
end


# selectindices ==========================================================================


Expand All @@ -1104,17 +1049,49 @@ end
Converts [`Selector`](@ref) to regular indices.
"""
function selectindices end
@inline selectindices(lookups::LookupTuple, s1, ss...) = selectindices(lookups, (s1, ss...))
@inline selectindices(lookups::LookupTuple, selectors::Tuple) =
map((l, s) -> selectindices(l, s), lookups, selectors)
@inline selectindices(lookups::LookupTuple, selectors::Tuple{}) = ()
@inline selectindices(lookups::LookupTuple, s1, ss...; kw...) =
selectindices(lookups, (s1, ss...); kw...)
@inline selectindices(lookups::LookupTuple, selectors::Tuple; kw...) =
map((l, s) -> selectindices(l, s), lookups, selectors; kw...)
@inline selectindices(lookups::LookupTuple, selectors::Tuple{}; kw...) = ()
# @inline selectindices(dim::Lookup, sel::Val) = selectindices(val(dim), At(sel))
# Standard indices are just returned.
@inline selectindices(::Lookup, sel::StandardIndices) = sel
@inline selectindices(::Lookup, sel::StandardIndices; kw...) = sel

# Vectors are mapped
@inline selectindices(lookup::Lookup, sel::Selector{<:AbstractVector}) =
[selectindices(lookup, rebuild(sel; val=v)) for v in val(sel)]
@inline selectindices(lookup::Lookup, sel::Selector{<:AbstractVector}; kw...) =
Int[selectindices(lookup, rebuild(sel; val=v); kw...) for v in val(sel)]
# Tuples may become ranges, unless the lookup is a tuple
@inline selectindices(l::Lookup, sel::IntSelector{<:Tuple}; kw...) =
_selecttuple(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::IntSelector{<:Tuple}; kw...) =
_selectindices(l, sel; kw...)
@inline selectindices(l::Lookup{<:Tuple}, sel::IntSelector{<:Tuple{<:Tuple,<:Tuple}}; kw...) =
_selecttuple(l, sel; kw...)
selectindices(l::Lookup, sel::SelectorOrInterval; kw...) = _selectindices(l, sel; kw...)
# @inline function selectindices(l::Lookup, sel; kw...)
# selstr = sprint(show, sel)
# throw(ArgumentError("Invalid index `$selstr`. Did you mean `At($selstr)`? Use stardard indices, `Selector`s, or `Val` for compile-time `At`."))
# end

# Separated _selectindices methods reduce ambiguity issues

# `Not` form InvertedIndices.jl
@inline function _selectindices(l::Lookup, sel::Not; kw...)
indices = selectindices(l, sel.skip; kw...)
return first(to_indices(l, (Not(indices),)))
end
@inline _selectindices(l::Lookup, sel::Touches) = touches(l, sel)
@inline _selectindices(l::Lookup, sel::Union{Between{<:Tuple},Interval}) = between(l, sel)
@inline function _selectindices(lookup::Lookup, sel::Between{<:AbstractVector})
inds = Int[]
for v in val(sel)
append!(inds, selectindices(lookup, rebuild(sel, v)))
end
end
@inline _selectindices(l::Lookup, sel::Contains; kw...) = contains(l, sel; kw...)
@inline _selectindices(l::Lookup, sel::Near; kw...) = near(l, sel; kw...)
@inline _selectindices(l::Lookup, sel::At; kw...) = at(l, sel; kw...)

# Unaligned Lookup ------------------------------------------

Expand Down

0 comments on commit fa1f52d

Please sign in to comment.