# Confusion about type stability of function within a struct

**URL:** <https://discourse.julialang.org/t/confusion-about-type-stability-of-function-within-a-struct/49003>\
**Category:** New to Julia\
**Created:** [October 25, 2020, 9:00pm UTC](https://discourse.julialang.org/t/confusion-about-type-stability-of-function-within-a-struct/49003 "2020-10-25T21:00:14Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 25, 2020, 9:00pm UTC](https://discourse.julialang.org/t/confusion-about-type-stability-of-function-within-a-struct/49003/1 "2020-10-25T21:00:15Z")

</div>

I have a struct called `Operators` that contains a bunch of functions. However, I can’t get it to be type stable. I have read [Type Unstable: Function within a Struct](https://discourse.julialang.org/t/type-unstable-function-within-a-struct/31110) and implemented the solution but it has not helped for this specific problem, although it did in other situations.

Currently each step of the loop causes 24 allocations (only 9 in my production code, can’t figure out why but the idea is the same).

MWE:

```julia
using Memoization
using FFTW
using BenchmarkTools

# This structure is type unstable
struct Operators{F1, F2}
    K̂::F1
    T̂::F2
end # Simulation

function Operators(α, x_order, ω)
    # Create this function with memoization since it is called repeatedly
    @memoize function K̂(dx)
        println("Computing and caching K(dx = $dx)")
        if α== 0
            ifftshift(cis.(dx*ω.^2/2))
        elseif α > 0
            ifftshift(cis.(dx*α*ω.^3))
        #elseif ...
            #more functions to select from
        end
    end

    # Select function 2
    if x_order == 2 && α == 0
        T̂ = T₂ˢ 
    elseif x_order == 2 && α > 0
        T̂ = T₂ˢʰ
    # elseif ...
        # many more functions to select from
    end

    ops = Operators(K̂, T̂)

    return ops
end

function T₂ˢ(ψ, dx, ops)
    # This line is 18 allocs (only 3 in my production code)
    ψ .= ops.K̂(dx) .* ψ 
end #T₂ˢ

function run(N)
    M = 512
    x_order = 2
    α = 0
    dx = 1e-4
    ψ = Array{Complex{Float64}}(undef, M, N)
    ψ[:, 1] = rand(M) + im*rand(M)
    ω = rand(M)

    # Decide which functions to use
    ops = Operators(α, x_order, ω)

    # 24 allocs per loop iteration (only 9 in production code)
    for i = 1:N-1
        @views ψ[:, i+1] .= ops.T̂(ψ[:, i], dx, ops) # 6 allocs from calling ops.T̂
    end

end

```

`@code_warntype` warns as follows:

```julia
@code_warntype run(4)

  ops::Operators{_A,_B} where _B where _A <-- RED
  @_10::Union{Nothing, Tuple{Int64,Int64}} <-- YELLOW

...
│ %31 = (%28)(%29, %30, ops)::Any <-- RED
│ %32 = Base.broadcasted(Base.identity, %31)::Base.Broadcast.Broadcasted{_A,Nothing,typeof(identity),_B} where _B<:Tuple where _A<:Union{Nothing, Base.Broadcast.BroadcastStyle} <-- RED
│ Base.materialize!(%27, %32)
│ (@_10 = Base.iterate(%17, %24))
...

```

I would appreciate any thoughts, this is quite confusing for me.

---

<div class="post-metadata">

**Author:** ![ash](https://avatars.discourse-cdn.com/v4/letter/a/f19dbf/32.png) [@ash](https://discourse.julialang.org/u/ash)\
**Post date:** [October 26, 2020, 1:03am UTC](https://discourse.julialang.org/t/confusion-about-type-stability-of-function-within-a-struct/49003/2 "2020-10-26T01:03:48Z")

</div>

A friend helped me optimize this. The end result is the following:

```julia
using Memoization
using FFTW
using BenchmarkTools

# This structure is type unstable
struct Operators{F1<:Function, F2<:Function}
    K̂::F1
    T̂::F2
end # Simulation

function Operators(α, x_order, ω)
    # Select function 
    function K̂(α)
        fun = if α == 0 
            @memoize function K_cubic(dx::Real)
                println("Computing and caching K(dx = $dx) for cubic NLSE")
                ifftshift(cis.(dx*ω.^2/2))
            end
            K_cubic
        elseif α > 0
            @memoize function K_hirota(dx::Real)
                println("Computing and caching K(dx = $dx) for Hirota Equation")
                fftshift(cis.(dx*α*ω.^3))
            end
            K_hirota
        end
        return fun
    end

    # Select function 2
    T̂ = T₂ˢ 
    if x_order == 2 && α == 0
        T̂ = T₂ˢ 
    elseif x_order == 2 && α > 0
        T̂ = T₂ˢʰ
    #elseif ...
       # many more functions to select from
    end

    ops = Operators(K̂(α), T̂)

    return ops
end

function T₂ˢ(ψ, dx, ops)
    # 3 allocs
    ψ .= ops.K̂(dx) .* ψ 
end #T₂ˢ

function run(N)
    M = 512
    x_order = 2
    α = 0
    dx = 1e-4
    ψ = Array{Complex{Float64}}(undef, M, N)
    ψ[:, 1] = rand(M) + im*rand(M)
    ω = rand(M)

    # Decide which functions to use
    ops = Operators(α, x_order, ω)

    # 3 allocs per loop iteration
    myloop(N, ψ, dx, ops)

end

function myloop(N, ψ, dx, ops)
    for i = 1:N-1
        @time @views ψ[:, i+1] .= ops.T̂(ψ[:, i], dx, ops) # 0 allocs from calling ops.T̂
    end
end

run(10)

```

Using `let` to define `K` significantly improves the allocations since the compiler doesn’t have to deal with it. We also use a function barrier in the form of `myloop` which decreases allocations since `myloop` already knows the type of `ops`. Currently, there are 3 allocations per loop iteration due to the lookups done by `@memoize`, I think.

Any other suggestions are still welcome however.