Search code examples
functionjuliatype-stability

Making closures type-stable dependent on the captured variable


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?


Solution

  • 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.