Return a function defined in a function

Is anyone able to explain this behaviour?

function fx(n)
    if n == 1
        f(x) = 2x
    else
        f(x) = 3x
    end
    f
end

Seems simple enough, but…

julia> fx(1)(1)   # expecting 2
3

julia> fx(2)(1)   # expecting 3
ERROR: UndefVarError: f not defined
Stacktrace:
 [1] fx(::Int64) at .\REPL[1]:7
 [2] top-level scope at none:0

I tried a few things like declaring f as global, but I’m still stumped. A pointer to the relevant docs would be really appreciated!

It’s an unfortunate and long-standing bug. See: Confused by behaviour of nested functions for more discussion.

Fortunately, it’s easy to avoid by defining an anonymous function instead:

function fx(n)
  if n == 1
    f = x -> 2x
  else
    f = x -> 3x
  end
  f
end
2 Likes

if the function if not complicated, you can do the conditional first and define a general function:

function fx(n)
    if n == 1
        a = 2
    else
        a = 3
    end
    f(x) = a*x
    return f
end
1 Like

This particular way of doing it is not a good idea since it generates type unstable code:

julia> @code_typed fx(1)(1)
CodeInfo(
1 ─ %1 = (Core.getfield)(#self#, :a)::Box
│   %2 = (Core.isdefined)(%1, :contents)::Bool
└──      goto #3 if not %2
2 ─      goto #4
3 ─      $(Expr(:throw_undef_if_not, :a, false))::Any
4 ┄ %6 = (Core.getfield)(%1, :contents)::Any
│   %7 = (%6 * x)::Any
└──      return %7
) => Any

Compare to @rdeits suggestion:

julia> @code_typed fx(1)(1)
CodeInfo(
1 ─ %1 = (Base.mul_int)(2, x)::Int64
└──      return %1
) => Int64

However, that code may suffer from similar issues if either the input n is not compile time known:

julia> fxwrapped(x) = fx(x)(x);

julia> fxwrapped(2)
6

julia> @code_typed fxwrapped(3)
CodeInfo(
1 ─ %1 = (x === 1)::Bool
└──      goto #3 if not %1
2 ─ %3 = %new(Main.:(##133#135))::##133#135
└──      goto #4
3 ─ %5 = %new(Main.:(##134#136))::##134#136
4 ┄ %6 = φ (#2 => %3, #3 => %5)::Union{##133#135, ##134#136}
└──      goto #5
5 ─ %8 = (%6)(x)::Any
└──      return %8
) => Any

Or if f is subsequently used:

function fx2(n)
  if n == 1
    f = x -> 2x
  else
    f = x -> 3x
  end
  x -> f(2x)
end

Problem:

julia> fx2(1)(1)
4

julia> @code_typed fx2(1)(1)
CodeInfo(
1 ─ %1 = (Core.getfield)(#self#, :f)::Box
│   %2 = (Core.isdefined)(%1, :contents)::Bool
└──      goto #3 if not %2
2 ─      goto #4
3 ─      $(Expr(:throw_undef_if_not, :f, false))::Any
4 ┄ %6 = (Core.getfield)(%1, :contents)::Any
│   %7 = (Base.mul_int)(2, x)::Int64
│   %8 = (%6)(%7)::Any
└──      return %8
) => Any

A workaround for both issues is to have a single definition of f, and also ensure that there’s no uncertainty of existence of variables:

function fx3(n)
    a = ifelse(n == 1, 2, 3)
    f(x) = a*x
    return x -> f(2x)
end

fx3wrapped(x) = fx3(x)(x)

Type stability ensues:

julia> @code_typed fx3(2)(1)
CodeInfo(
1 ─ %1 = (Core.getfield)(#self#, :f)::#f#160{Int64}
│   %2 = (Base.mul_int)(2, x)::Int64
│   %3 = (Core.getfield)(%1, :a)::Int64
│   %4 = (Base.mul_int)(%3, %2)::Int64
└──      return %4
) => Int64

julia> @code_typed fx3wrapped(2)
CodeInfo(
1 ─ %1 = (x === 1)::Bool
│   %2 = (Main.ifelse)(%1, 2, 3)::Int64
│   %3 = (Base.mul_int)(2, x)::Int64
│   %4 = (Base.mul_int)(%2, %3)::Int64
└──      return %4
) => Int64
4 Likes