Force recompiling functions in conditional compilation

I want to do conditional compilation. Besides, when developing, I’d like to be able to exam all branches. Can I do that?

An explanation by case:
My code will have two different modes when used in production, let’s call them on and off. I may write a macro to do that:

const IS_ON = true
macro if_on(ex1, ex2)
    if IS_ON
        return :($(esc(ex1)))
    else
        return :($(esc(ex2)))
    end
end

function inc_depend_on_mode!(a)
    @if_on begin
        a .+= 1
    end begin
        a .+= 2
    end
end

I can then change the value of IS_ON to get code for different modes.
This works fine for production. However, when developing, I want to be able to examine code for both modes without restarting my REPL.
Then first, I need to change IS_ON from a global constant, to something like

mutable struct Control
    is_on::Bool
end
const control = Control(true)

so that I may dynamically set its value.
There remains one question: how can I recompile function inc_depend_on_mode! after I change the mode?

Very Important: I want to do real conditional precompilation, which means, in my point of view, that the generated lowered code should be free of branches. In the previous example, this can be confirmed:

julia> a = zeros(Int, 5);
julia> @code_lowered inc_depend_on_mode!(a)
CodeInfo(
1 ─ %1 = Base.broadcasted(Main.:+, a, 1)
│   %2 = Base.materialize!(a, %1)
└──      return %2
)

If not, I can simply write something like

mutable struct Control
    is_on::Bool
end
const control = Control(true)
function inc_depend_on_mode!(a)
    if control.is_on
        a .+= 1
    else
        a .+= 2
    end
end

It works fine, but it is not conditional “compilation”, because branches exist in lowered code:

julia> a = zeros(Int, 5);

julia> @code_lowered inc_depend_on_mode!(a)
CodeInfo(
1 ─ %1 = Base.getproperty(Main.control, :is_on)
└──      goto #3 if not %1
2 ─ %3 = Base.broadcasted(Main.:+, a, 1)
│   %4 = Base.materialize!(a, %3)
└──      return %4
3 ─ %6 = Base.broadcasted(Main.:+, a, 2)
│   %7 = Base.materialize!(a, %6)
└──      return %7
)

This thread seems to be relevant

In it, you’ll find the following technique, which also forms the core of the implementation of ToggleableAsserts. Re-writing your example, it would look like this:

julia> mode_on() = true
mode_on (generic function with 1 method)

julia> function inc!(a)
           if mode_on()
               a .+= 1
           else
               a .+= 2
           end
       end
inc! (generic function with 1 method)

You’ll notice that lowered code is not entirely free of branches:

julia> @code_lowered inc!([1,2,3])
CodeInfo(
1 ─ %1 = Main.mode_on()
└──      goto #3 if not %1
2 ─ %3 = Base.broadcasted(Main.:+, a, 1)
│   %4 = Base.materialize!(a, %3)
└──      return %4
3 ─ %6 = Base.broadcasted(Main.:+, a, 2)
│   %7 = Base.materialize!(a, %6)
└──      return %7
)

But later stages of compilation (starting with typed code) do eliminate branches (note how the second block below is considered unreachable):

julia> @code_warntype inc!([1,2,3])
MethodInstance for inc!(::Vector{Int64})
  from inc!(a) in Main at REPL[3]:1
Arguments
  #self#::Core.Const(inc!)
  a::Vector{Int64}
Body::Vector{Int64}
1 ─ %1 = Main.mode_on()::Core.Const(true)
│        Core.typeassert(%1, Core.Bool)
│   %3 = Base.broadcasted(Main.:+, a, 1)::Core.PartialStruct(Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(+), Tuple{Vector{Int64}, Int64}}, Any[Core.Const(+), Core.PartialStruct(Tuple{Vector{Int64}, Int64}, Any[Vector{Int64}, Core.Const(1)]), Core.Const(nothing)])
│   %4 = Base.materialize!(a, %3)::Vector{Int64}
└──      return %4
2 ─      Core.Const(:(Base.broadcasted(Main.:+, a, 2)))
│        Core.Const(:(Base.materialize!(a, %6)))
└──      Core.Const(:(return %7))

So I’d say this qualifies as conditional compilation. Furthermore, if the mode_on() method is redefined, then it will invalidate inc! and force its re-compilation next time it is called (note how the else clause is now the only reachable part of the code):

julia> mode_on() = false
mode_on (generic function with 1 method)

julia> @code_warntype inc!([1,2,3])
MethodInstance for inc!(::Vector{Int64})
  from inc!(a) in Main at REPL[3]:1
Arguments
  #self#::Core.Const(inc!)
  a::Vector{Int64}
Body::Vector{Int64}
1 ─ %1 = Main.mode_on()::Core.Const(false)
└──      goto #3 if not %1
2 ─      Core.Const(:(Base.broadcasted(Main.:+, a, 1)))
│        Core.Const(:(Base.materialize!(a, %3)))
└──      Core.Const(:(return %4))
3 ┄ %6 = Base.broadcasted(Main.:+, a, 2)::Core.PartialStruct(Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(+), Tuple{Vector{Int64}, Int64}}, Any[Core.Const(+), Core.PartialStruct(Tuple{Vector{Int64}, Int64}, Any[Vector{Int64}, Core.Const(2)]), Core.Const(nothing)])
│   %7 = Base.materialize!(a, %6)::Vector{Int64}
└──      return %7
1 Like