Making closures type-stable dependent on the captured variable

53 Views Asked by At

For the function

function function_maker(N)
    if N == 1
        x = 1.0
    else
        x = 1
    end
    f(y) = x+y
end

I want the output of this to not be type-stable, but I want it to generate an f that is type-stable, i.e. uses the type of x determined by the value of N to generate a function dependent on N. Basically, I want the functions that come out of this to be performant, but the function_maker itself doesn't need to be performant because it's only used in the global scope or above a function barrier.

f = function_maker(1)
@code_warntype f(1)

Variables:
  #self#::#f#9
  y::Int64

Body:
  begin 
      return ((Core.getfield)((Core.getfield)(#self#::#f#9, :x)::ANY, :contents)::ANY + y::Int64)::ANY
  end::ANY

This doesn't happen by default. I tried f(y) = x::typeof(x)+y but that didn't work either. Is there a simple way to do this?

1

There are 1 best solutions below

7
On BEST ANSWER

There's:

julia> function function_maker2(N)
           if N == 1
               let x = 1.0
                   return f(y) = x + y
               end
           else
               let x = 1
                   return f(y) = x + y
               end
           end
       end
function_maker2 (generic function with 1 method)

julia> f2 = function_maker2(1)
(::f) (generic function with 1 method)

julia> @code_warntype f2(1)
Variables:
  #self#::#f#5{Float64}
  y::Int64

Body:
  begin 
      return (Base.add_float)((Core.getfield)(#self#::#f#5{Float64}, :x)::Float64, (Base.sitofp)(Float64, y::Int64)::Float64)::Float64
  end::Float64

This version separates the x in each branch inside let blocks. Otherwise the compiler seems to get confused.