Stream: helpdesk (published)

Topic: Make a struct act like a Tuple under iteration

Mason Protter (Jan 21 2021 at 01:31):

Hey, so suppose I have a struct

struct Foo
    a::Int
    b::Float64
end

and I basically want this thing to behave like a Tuple under iteration. e.g. I'd like this to 'just work' and carry no runtime performance overhead.

f((a, b)) = a + b

What's a clean way to do this right?

Mason Protter (Jan 21 2021 at 01:37):

If I do the naive thing, it's very slow:

Base.iterate(f::Foo, s=1) = s <= length(fieldnames(Foo)) ? (getfield(f, s), s+1) : nothing

let foo = Ref(Foo(1,2))
    @btime f($foo[])
end

#+RESULTS:
   802.830 ns (4 allocations: 112 bytes)
 3.0

I feel there must be some clean way to basically say "iterate this thing as if it were a Tuple"

Mason Protter (Jan 21 2021 at 01:55):

Okay, so if I don't define the above iterate method and instead define

Base.indexed_iterate(f::Foo, i::Int, state=1) = (Base.@_inline_meta; (getfield(f, i), i+1))

everything is fast. Good to know I suppose.

Simeon Schaub (Jan 21 2021 at 06:53):

You should probably still define iterate as well though, since indexed_iterate is just an implementation detail and is always assumed to return the same results as regular iteration

Pablo Zubieta (Jan 21 2021 at 16:52):

Still slower that indexed_iterate, but better that the OP

@inline Base.iterate(f::Foo, s = 1) = s ≤ fieldcount(Foo) ? (getfield(f, s), s + 1) : nothing

Fredrik Bagge Carlson (Jan 21 2021 at 17:18):

Some LinearAlgebra types does this, I thing QR does, you can perhaps see how they do it?

Mason Protter (Jan 21 2021 at 17:20):

Hm, it seems that @Pablo Zubieta's method is just as fast as indexed_iterate for me on 1.6.

Mason Protter (Jan 21 2021 at 17:21):

julia> struct Foo
           a::Int
           b::Float64
       end

julia> f((a, b)) = a + b
f (generic function with 1 method)

julia> @inline Base.iterate(f::Foo, s = 1) = s ≤ fieldcount(Foo) ? (getfield(f, s), s + 1) : nothing

julia> let foo = Ref(Foo(1,2))
           @btime f($foo[])
       end
  1.859 ns (0 allocations: 0 bytes)
3.0

and

julia> struct Foo
           a::Int
           b::Float64
       end

julia> f((a, b)) = a + b
f (generic function with 1 method)

julia> Base.indexed_iterate(f::Foo, i::Int, state=1) = (Base.@_inline_meta; (getfield(f, i), i+1))

julia> let foo = Ref(Foo(1,2))
           @btime f($foo[])
       end
  1.849 ns (0 allocations: 0 bytes)
3.0

Pablo Zubieta (Jan 21 2021 at 18:07):

Oh. I was trying on 1.5.3, great to see that the compiler is also better here in 1.6.

Eric Hanson (Jan 23 2021 at 01:10):

I think with https://github.com/JuliaLang/julia/pull/39285 you wouldn’t even need an iterate method, just a semicolon: f((;a, b)) = a + b

Mason Protter (Jan 23 2021 at 01:48):

Wow, that’s very neat. Kinda niche though, since you need to bind the variables to the same name as the property names.

Simeon Schaub (Jan 23 2021 at 09:51):

There was some discussion in https://github.com/JuliaLang/julia/issues/28579 about having some more advanced matching as well. Wouldn't be too hard to add, but would need to come to a decission on how exactly that should work.

Last updated: Nov 22 2024 at 04:41 UTC