Static `findfirst`, `findall` for `Tuple`

[mtfishman]

Here is a use case that I think this package could help with. I have a situation where I want to split a tuple like this:

t = (1, "X", 1, 2, "Y", 2, "Z", 4)

into a series of tuples:

t1, t2, t3, t4 = ((1,), ("X", 1, 2), ("Y", 2), ("Z", 4))

where the start of each split tuple is the location of an element of a certain type (in this case, AbstractString).

Here is a very type unstable version:

julia> t = (1, "X", 1, 2, "Y", 2, "Z", 4)
(1, "X", 1, 2, "Y", 2, "Z", 4)

julia> function split(f, t::Tuple)
         n = findall(f, t)
         ti = t[1:(first(n) - 1)]
         ts = ntuple(i -> t[n[i]:(n[i + 1] - 1)], length(n) - 1)
         tf = t[last(n):end]
         return ti, ts..., tf
       end
split (generic function with 1 method)

julia> split(x -> x isa AbstractString, t)
((1,), ("X", 1, 2), ("Y", 2), ("Z", 4))

It seems like this is a case that could benefit from static numbers, ranges, and slices (like ones discussed in #4). Can something like this be made type stable with this package? It seems like it would require a version of findall for Tuples that output static index locations.

[perrutquist]

I think you could easily achieve this using generated functions. (It would also be possible to do using recursive functions, although trickier to get it right.)
I don't see how using static numbers will help, though.

[mtfishman]

The part that I was imagining that static numbers could help with is having static versions of functions like findfirst and findall, for example:

t = (1, "X", 1, 2)
# n is a 
n = staticfindfirst(x -> x isa String, t) # Maybe `@stat findfirst(x -> x isa String, t)`?
t1 = @stat t[1:n - 1] # (1,)
t2 = @stat t[n:end] # ("X", 1, 2)

So then the slicing could become type stable. Agreed there are probably alternatives with generated functions or recursion but I think the above syntax is pretty natural and closer to the code you might write for vectors.

[perrutquist]

Sure, you could do that. I think a general "staticfindfirst" function would be of very limited use, but you could write one for your specific case like so:

@generated function staticfirstoftype(::Type{T}, v::Tuple) where {T}
    static(findfirst(==(T), fieldtypes(v)))
end

staticfirstoftype(String, t) # returns static(2)

[mtfishman]

Thanks, it is probably sufficient to be able to define things like staticfindfirst ourselves.

Here is a minimal type stable version of staticfindfirst:

using StaticNumbers

_findfirst(f, i, ::Tuple{}) = nothing
_findfirst(f, i, x) = f(first(x)) ? i : _findfirst(f, i+1, Base.tail(x))
tuple_findfirst(f, x::Tuple) = _findfirst(f, 1, x)

@generated function staticfindfirst(f, v::Tuple) 
  return :(static(tuple_findfirst(f, v)))
end

For example:

julia> staticfindfirst(x -> x isa String, (1, "X", 2))
static(2)

julia> @code_warntype staticfindfirst(x -> x isa String, (1, "X", 2))
MethodInstance for staticfindfirst(::var"#29#30", ::Tuple{Int64, String, Int64})
  from staticfindfirst(f, v::Tuple) in Main at /home/mfishman/Dropbox (Simons Foundation)/workdir/StaticNumbers/staticfindfirst.jl:7
Arguments
  #self#::Core.Const(staticfindfirst)
  f::Core.Const(var"#29#30"())
  v::Tuple{Int64, String, Int64}
Body::StaticInteger{2}
1 ─ %1 = Main.tuple_findfirst(f, v)::Core.Const(2)
│   %2 = Main.static(%1)::Core.Const(static(2))
└──      return %2

Note that I found that using Base.findfirst did not lead to type stable results. This may be related to this discourse discussion: https://discourse.julialang.org/t/why-does-findfirst-t-on-a-tuple-of-typed-only-constant-fold-for-the-first/68893/3

[mtfishman]

Actually it seems like it is unnecessary for it to be an @generated function:

function staticfindfirst2(f, v::Tuple)
  return static(tuple_findfirst(f, v))
end

gives:

julia> staticfindfirst2(x -> x isa String, (1, "X", 2))
static(2)

julia> @code_warntype staticfindfirst2(x -> x isa String, (1, "X", 2))
MethodInstance for staticfindfirst2(::var"#40#41", ::Tuple{Int64, String, Int64})
  from staticfindfirst2(f, v::Tuple) in Main at /home/mfishman/Dropbox (Simons Foundation)/workdir/StaticNumbers/staticfindfirst.jl:11
Arguments
  #self#::Core.Const(staticfindfirst2)
  f::Core.Const(var"#40#41"())
  v::Tuple{Int64, String, Int64}
Body::StaticInteger{2}
1 ─ %1 = Main.tuple_findfirst(f, v)::Core.Const(2)
│   %2 = Main.static(%1)::Core.Const(static(2))
└──      return %2

[perrutquist]

Yeah, it is usually possible to avoid generated functions by using recursive functions instead. Just be aware that the compiler might not specialise on long tuples.