# Advanced iteration interface in Julia

**URL:** https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618
**Category:** Performance
**Tags:** iterators, inline, box
**Created:** [September 23, 2021, 7:07am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618 "2021-09-23T07:07:30Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![ettersi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ettersi/32/6829_2.png) [@ettersi](https://discourse.julialang.org/u/ettersi)
#### Post date: [September 23, 2021, 7:07am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/1 "2021-09-23T07:07:30Z")

</div>

I’m working on a function to simplify iterating over binary trees. Ideally, this function should have the following signature:

```julia
walk_down(
    tree, starting_vertex; 
    two_children, # Function vertex -> Union{Val(:left), Val(:right), Val(:break)}
    one_child, # Function vertex -> Union{Val(:continue), Val(:break)}
    leaf # Function vertex -> nothing
)

```

The values returned by the callback functions indicate in which direction the traversal should continue, if at all.

Making the above work is not difficult, but this will be a very low-level function and so I want to ensure that it is optimally efficient. I therefore cooked up the following simple toy problem to beta-test the interface:

```julia
# Alternative to `for vi in v`, mimicking the above interface
function foreach(f,v)
    for vi in v
        if f(vi) == Val(:break)
            break
        end
    end
end

# Benchmarking code
function foo(n)
    s = 0
    foreach(1:n) do i
        s += i
        if i == n÷2
            return Val(:break)
        else
            return Val(:continue)
        end
    end
    return s
end
function foo_ref(n)
    s = 0
    for i in 1:n
        s += i
        if i == n÷2
            break
        end
    end
    return s
end

```

It turns out that the proposed interface incurs a ~44x slowdown:

```julia
julia> @btime foo($n)
       @btime foo_ref($n)
  11.275 μs (470 allocations: 7.34 KiB)
  248.822 ns (0 allocations: 0 bytes)

```

Any suggestions on how to avoid this slowdown?

I’m aware that I can get much more fine-grained control by turning `foreach()` into a macro, but that also means more work and probably a more brittle abstraction, so ideally I’d like to avoid that.

---

<div class="post-metadata">

### Author: ![Sukera](https://avatars.discourse-cdn.com/v4/letter/s/ce7236/32.png) [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [September 23, 2021, 8:03am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/2 "2021-09-23T08:03:35Z")

</div>

Do you know about [AbstractTrees.jl](https://juliahub.com/ui/Packages/AbstractTrees/tEX7J/0.3.4)?

* * *

In case that’s not helpful/you don’t want to use it - the `foreach` you’re using is slow because it is type unstable and you’re constructing `Val` instances at runtime, which leads to a very bad kind of type instability. You’ll need dynamic dispatch everywhere - why not simply return `true` or `false`, making it type stable?

---

<div class="post-metadata">

### Author: ![ettersi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ettersi/32/6829_2.png) [@ettersi](https://discourse.julialang.org/u/ettersi)
#### Post date: [September 23, 2021, 8:37am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/3 "2021-09-23T08:37:28Z")

</div>

I’m aware of AbstractTrees.jl. As far as I am aware, that package does not provide an interface for iterating over only part of a tree (e.g. a single branch from the root), so it doesn’t help for my purposes.

I’m also aware of the type instability, but it turns out this instability is fully handled by Julia’s [union-splitting optimisation](https://julialang.org/blog/2018/08/union-splitting/). Rewriting `foreach()` so it operates on `Bool` instead of `Union{Val(:break), Val(:continue)}` leads to exactly the same runtime and number of allocations:

```julia
function foreach_Bool(f,v)
    for vi in v
        if f(vi)
            break
        end
    end
end

function foo_Bool(n)
    s = 0
    foreach_Bool(1:n) do i
        s += i
        if i == n÷2
            return true
        else
            return false
        end
    end
    return s
end

julia> @btime foo_Bool(1000)
  11.023 μs (470 allocations: 7.34 KiB)

```

My gut feeling is that the performance penalty comes from the fact that the `f` passed to `foreach(f,v)` is a closure over `s` and hence `s` is boxed and moved to the heap. I was hoping that the compiler would inline both `foreach` and `f` so that `s` can be unboxed again, but looking at `@code_native foo(1000)`, it seems that this is not the case.

---

<div class="post-metadata">

### Author: ![ettersi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ettersi/32/6829_2.png) [@ettersi](https://discourse.julialang.org/u/ettersi)
#### Post date: [September 23, 2021, 9:14am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/4 "2021-09-23T09:14:43Z")

</div>

Ah, the issue that I’m facing is [performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub](https://github.com/JuliaLang/julia/issues/15276). Making `s` a `Ref{Int}` seems to fully solve the problem:

```julia
function foo_Ref(n)
    s = Ref(0)
    foreach(1:n) do i
        s[] += i
        if i == n÷2
            return Val(:break)
        else
            return Val(:continue)
        end
    end
    return s[]
end

julia> @btime foo_Ref(1000)
  257.358 ns (0 allocations: 0 bytes)

```

I don’t think I entirely understand why this works, though. Any reference to why `Ref()` jells well with closures would be appreciated.

---

<div class="post-metadata">

### Author: ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)
#### Post date: [September 23, 2021, 9:21am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/5 "2021-09-23T09:21:27Z")

</div>

According to an answer to my earlier performance question:

> [@Cannot achieve type-stability and get rid of allocations](https://discourse.julialang.org/t/cannot-achieve-type-stability-and-get-rid-of-allocations/44548/6):
>
> You can wrap those variables in Ref s that you then mutate instead of rebinding.

, the issue is that `s += 1` assigns a new value to the same variable, and this isn’t treated well by the compiler.

---

<div class="post-metadata">

### Author: ![ettersi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ettersi/32/6829_2.png) [@ettersi](https://discourse.julialang.org/u/ettersi)
#### Post date: [September 23, 2021, 9:39am UTC](https://discourse.julialang.org/t/advanced-iteration-interface-in-julia/68618/6 "2021-09-23T09:39:15Z")

</div>

> [@aplavin](#):
>
> `s += 1` assigns a new value to the same variable, and this isn’t treated well by the compiler.

Which begs the question, why not? To answer this question myself (or at least trying to, after reading a bit more about [15276](https://github.com/JuliaLang/julia/issues/15276)): it seems that the Julia compiler has to decide whether or not to box `s` _before_ type inference is run, thus if `s` is assigned it must defensively assume that it may also change type and hence it must be boxed.
