How to get a tuple's field using a runtime index in a type stable way?

General Usage
1

Suppose I have a tuple, which contains some data. I want to get one of the fields of the tuple, with an index at runtime, but I know the type of the field the index points to. Is there a way to do this?

My best attempt so far is the follow:

tup = (10, 100, 1., 2., 3.)

function test(tup, idx, type::Type{T}) where T
    ret::T = tup[idx]
    ret
end

@time test(tup, 1, Int)

But this allocates:

0.000008 seconds (1 allocation: 48 bytes)

I want to tell the compiler this is always going to be type T, and if its not it can throw an error BEFORE it gets the field. Is there a way to do this with generated functions or some Core compiler hint or some other nasty trick?

2

This probably doesn’t help your actual use case (though maybe), but I think you get an allocation here becease tup is a global variable. Try @testing this:

function outer_test()
    tup = (10, 100, 1., 2., 3.)
    ret = test(tup, 1, Int)
end

and you’ll get zero allocations. However, in this case, 1 is a constant index, which is why you’re not getting an allocation.

This won’t help you if your actual index is truly generated at runtime though.

3

Thanks for the reply!

I modified your example so the idx is randomly chosen between 1 and 2

function test(tup, idx, type::Type{T}) where T
    # @inline Core.getfield(tup, idx)::T
    # Core.unsafe_convert(T, Core.getfield(tup, idx))
    # Core.get_binding_type
    ret::T = Base.getindex(tup, idx)
end

function outer_test()
    tup = (10, 100, 1., 2., 3.)
    idx = rand((1, 2))
    @noinline ret = test(tup, idx, Int)
end

@time outer_test()

With the @noinline, we still get 1 allocation. Without it, we don’t…

4

Intuitively, I feel like this should be possible. Please, if anyone experienced can tell me why it isn’t, I would appreciate it.

5

As you can tell, it took a few tries :smiley:

function test4(tup,idx,typ::Type{T}) where T
       idx == 1 && (x=first(tup); isa(x,T) && return x)
       test4(Base.tail(tup),idx-1,T)
end

This seems to do the trick and not allocate (throws an error otherwise), it is also type stable. outer_test would not allocate regardless of inlining.

Hope this helps…

Edit:
This seems to do the same and simpler:

function test4(tup,idx,typ::Type{T}) where T
       idx == 1 && (x=first(tup); return x::T) 
       #or if you want an explicit error
       #isa(x,T) ? (return x) : error("Bla"))
       test4(Base.tail(tup),idx-1,T)
end
1 Like
6

Have you considered, alternatively, the use of named tuples?

nt=(;zip([:a,:b,:c,:d,:e],tup)...)
julia> @btime $nt.e
  1.500 ns (0 allocations: 0 bytes)
3.0

julia> @btime $nt.b
  1.500 ns (0 allocations: 0 bytes)
100

or


function test(tup, idx)
    ret = getfield(tup,idx)
    ret
end

@btime test($tup, 4)


function test1(tup, idx)
     r=first(iterate(tup,idx))
     r
end
7

I think this is only possible if you know the returned type.

8

Are you sure that test(...) allocates?

julia> const tup_const = (10, 100, 1., 2., 3.)
(10, 100, 1.0, 2.0, 3.0)

julia> function test(tup, idx, type::Type{T}) where T
           ret::T = tup[idx]
           ret
       end
test (generic function with 1 method)

julia> @time test(tup_const, 3, Float64)
  0.000003 seconds
1.0

julia> @btime test(tup_const, 3, Float64)
  1.041 ns (0 allocations: 0 bytes)
1.0

You may also want to consider the effect of constant propagation.

Note the union type below.

julia> @code_warntype test(tup_const, 3, Float64)
MethodInstance for test(::Tuple{Int64, Int64, Float64, Float64, Float64}, ::Int64, ::Type{Float64})
  from test(tup, idx, type::Type{T}) where T @ Main REPL[27]:1
Static Parameters
  T = Float64
Arguments
  #self#::Core.Const(test)
  tup::Tuple{Int64, Int64, Float64, Float64, Float64}
  idx::Int64
  type::Core.Const(Float64)
Locals
  ret::Float64
  @_6::Union{Float64, Int64}
Body::Float64
1 ─      Core.NewvarNode(:(ret))
│   %2 = Base.getindex(tup, idx)::Union{Float64, Int64}
│        (@_6 = %2)
│   %4 = (@_6 isa $(Expr(:static_parameter, 1)))::Bool
└──      goto #3 if not %4
2 ─      goto #4
3 ─ %7 = Base.convert($(Expr(:static_parameter, 1)), @_6::Int64)::Float64
└──      (@_6 = Core.typeassert(%7, $(Expr(:static_parameter, 1))))
4 ┄      (ret = @_6::Float64)
└──      return ret

If we create a function with the index three hard-coded in, we see the absence of instability.

julia> function test2(tup, idx)
           tup[idx]
       end
test2 (generic function with 2 methods)

julia> f(tup) = @inline test2(tup, 3)
f (generic function with 1 method)

julia> @code_warntype f(tup_const)
MethodInstance for f(::Tuple{Int64, Int64, Float64, Float64, Float64})
  from f(tup) @ Main REPL[43]:1
Arguments
  #self#::Core.Const(f)
  tup::Tuple{Int64, Int64, Float64, Float64, Float64}
Locals
  val::Float64
Body::Float64
1 ─     nothing
│       (val = Main.test2(tup, 3))
│       nothing
└──     return val
1 Like
9

The test call was simple enough to inline, resulting in the getindex call being directly in outer_test. tup’s type only contains Int or Float64, so the compiler can infer that an inlined getindex only returns a Int or Float64. Hypothetically the compiler could’ve leveraged the niche detail that rand((1,2)) only selects the indices with Int, but the compiler evidently doesn’t go that far. On the other hand, if the tuple contained 4+ types, the getindex return type is inferred as Any. You can see the effect directly in the getindex call:

julia> @code_warntype getindex((10, 100, 1., 2., 3.), 2)
MethodInstance for getindex(::Tuple{Int64, Int64, Float64, Float64, Float64}, ::Int64)
  from getindex(t::Tuple, i::Int64) @ Base tuple.jl:31
Arguments
  #self#::Core.Const(getindex)
  t::Tuple{Int64, Int64, Float64, Float64, Float64}
  i::Int64
Body::Union{Float64, Int64}
1 ─      nothing
│   %2 = Base.getfield(t, i, $(Expr(:boundscheck)))::Union{Float64, Int64}
└──      return %2

julia> @code_warntype getindex((10, 100im, 1., 2.0im, 3.), 2)
MethodInstance for getindex(::Tuple{Int64, Complex{Int64}, Float64, ComplexF64, Float64}, ::Int64)
  from getindex(t::Tuple, i::Int64) @ Base tuple.jl:31
Arguments
  #self#::Core.Const(getindex)
  t::Tuple{Int64, Complex{Int64}, Float64, ComplexF64, Float64}
  i::Int64
Body::Any
1 ─      nothing
│   %2 = Base.getfield(t, i, $(Expr(:boundscheck)))::Any
└──      return %2

If your index isn’t a literal or a constant, limiting the number of types in the tuple to 3 (a limit that can change across minor revisions) is the best you can hope for. To @time this properly, you’d need a global variable to hold the index instead of inputting a literal, otherwise the compiler knows the exact index and type, making the type inference report above irrelevant. Note that this specialization over the literal or constant value rather than its type Int is not a general feature, the compiler treats the underlying built-in function Core.getfield specially.

julia> @time getindex((10, 100, 1., 2., 3.), 2);
  0.000001 seconds

julia> @time getindex((10, 100im, 1., 2.0im, 3.), 2);
  0.000004 seconds

julia> i::Int = 2;

julia> @time getindex((10, 100, 1., 2., 3.), i);
  0.000009 seconds

julia> @time getindex((10, 100im, 1., 2.0im, 3.), i);
  0.000014 seconds (1 allocation: 32 bytes)
1 Like
10

That’s awesome @raminammour! I think it can be further simplified a bit:

function test5(tup, idx, type)
    idx == 1 && return tup[1]::type
    test5(Base.tail(tup), idx-1, type)
end

I think we can even drop the type argument while maintaining the same performance:

function test6(tup, idx)
    idx == 1 && return tup[1]
    test6(Base.tail(tup), idx-1)
end

(and call with test6(tup, idx)::type if necessary).

2 Likes
11

Wow! Thank you, this is great :smiley:

12

Yes, I stated that I do know the return type.

13

The problem is explicitly how to solve this when we can’t hard code in the index, but you (somehow) know the return type. If you can hard code the index it is trivial

14

@Larbino1 the way I’d recommend approaching this would be to re-order your Tuple{Int, Int, Float64, Float64, Float64} into a Tuple{Tuple{Int, Int}, Tuple{Float64, Float64, Float64}} because then the compiler can better understand what is going on.

You’d then index this thing like new_tup[1][i] where 1 and 2 are the only valid indices of new_tup[1].

If you want to get kinda fancy though, an alternative would be to make a MaskedTuple{T, Tup} type that behaves like Tup <: Tuple except only lets you grab the entries of type T.

e.g. something like

struct MaskedTuple{T, Tup <: Tuple, N}
    masked_tup::Tup
    inds::NTuple{N, Int}
    function MaskedTuple{T}(tup::Tuple) where {T}
        masked_tup = filter(x -> x isa T, tup)
        counter = Ref(0)
        inds = ntuple(length(tup)) do i
            if tup[i] isa T
                counter[] += 1
            else
                0
            end
        end
        new{T, typeof(masked_tup), length(inds)}(masked_tup, inds)
    end
end
Base.length(mt::MaskedTuple) = length(mt.inds)
Base.getindex(m::MaskedTuple{T}, i::Int) where {T} = let j = m.inds[i]
    if j === 0
        @noinline ((i, m, T) -> error(lazy"Index $i of $m is not of type $T"))(i, m, T)
    end
    m.masked_tup[j]
end

function Base.show(io::IO, mt::MaskedTuple{T}) where {T}
    print(io, "MaskedTuple{$T}(" * join((mt.inds[i] == 0 ? "_" : mt[i] for i ∈ 1:length(mt)), ", ")* ")")
end

and then

julia> function test(tup, idx, ::Type{T}) where {T}
           MaskedTuple{T}(tup)[idx]
       end;

julia> let tup = Ref((1, 2, 3.0, 4.0, 5.0))
           @btime test($tup[], rand(1:2), Int)
       end;
  4.970 ns (0 allocations: 0 bytes)
1

This works because we’ve made it so that we’re only ever accessing a Tuple{Int, Int}, at an unkown index rather than a Tuple{Int, Int, Float64, Float64, Float64}.

1 Like
15

This is a pattern I’ve ended up creating several times for may package. I wonder if it has a name or a better implementation.

Consider this: I have a machine with many buttons, of many (but fewer) different types. I want to store a collection of buttons and press all the buttons, but in a type stable way.

The way I end up solving this is with a Tuple{Vector{ButtonType1}, Vector{ButtonType2}...}, then to add a button, I use a generated function to add it to the right vector. I can keep an index to this button so I can get back to it later which is something like CollectionIndex(Val{tupidx}(), vecidx)

This way, we don’t create huge tuples (only as big as the numbe of types), we can still map/foreach etc. over the collection;
map(press, button_collection)
by implementing map on our collection, and we can still go back inside the collection and get the items in a type stable way (as tupidx is known at compile time due to being a Val)

16

I did something like this for Dicts instead of Vectors about a year ago: ConcreteTupleDicts.jl

An equivalent backing for Arrays or Vectors (or generalization which subsumes both the Array and Dict implementations) would probably be a useful contribution. My construction code is a bit messy because I was trying to get it type stable, but the retrieval worked nicely for the use-case I had

2 Likes
17

Another ref I ran into when working on the concrete tuple-dict stuff: TypeSortedCollections.jl

I think this implements the pattern you described with tuples of vectors of homogeneous types

1 Like
18

Amazing, thank you!

EDIT: This is exactly what I was looking for, shame I didn’t manage to find it. Would have saved me many headaches, but now I can clean things up using it :slight_smile:

1 Like
19

I would argue that we do know the return type from the tuple type. There should not be a need to encode the type at a particular index twice.

julia> begin
           second(x) = x[2]
           third(x) = x[3]
           fourth(x) = x[4]
           fifth(x) = x[5]
       end
fifth (generic function with 1 method)

julia> square(f, tup) = f(tup)^2
square (generic function with 1 method)

julia> @code_warntype square(first, tup)
MethodInstance for square(::typeof(first), ::Tuple{Int64, Int64, Float64, Float64, Float64})
  from square(f, tup) @ Main REPL[13]:1
Arguments
  #self#::Core.Const(square)
  f::Core.Const(first)
  tup::Tuple{Int64, Int64, Float64, Float64, Float64}
Body::Int64
1 ─ %1 = Main.:^::Core.Const(^)
│   %2 = (f)(tup)::Int64
│   %3 = Core.apply_type(Base.Val, 2)::Core.Const(Val{2})
│   %4 = (%3)()::Core.Const(Val{2}())
│   %5 = Base.literal_pow(%1, %2, %4)::Int64
└──      return %5


julia> @code_warntype square(fifth, tup)
MethodInstance for square(::typeof(fifth), ::Tuple{Int64, Int64, Float64, Float64, Float64})
  from square(f, tup) @ Main REPL[13]:1
Arguments
  #self#::Core.Const(square)
  f::Core.Const(fifth)
  tup::Tuple{Int64, Int64, Float64, Float64, Float64}
Body::Float64
1 ─ %1 = Main.:^::Core.Const(^)
│   %2 = (f)(tup)::Float64
│   %3 = Core.apply_type(Base.Val, 2)::Core.Const(Val{2})
│   %4 = (%3)()::Core.Const(Val{2}())
│   %5 = Base.literal_pow(%1, %2, %4)::Float64
└──      return %5
20

See (my) UnionCollections.jl as well (: Appears to be more fully-featured for Arrays, also supports Dictionaries.

2 Likes