Transform CodeInfo of a function to get a new function by purging unused args from its arg list

General Usage
,
2

This code should work.


function assign_purge_argument!(@nospecialize(p::Pair{<:Function,<:Function}), rule::Pair{Int,<:Any})
    (src, dst) = p
    (nth, value) = rule
    @assert length(methods(src)) == 1 "assign_purge_argument!(src=>dst, n=>val) requires src to have only one method."
    mod = getmodule(dst)
    function change_slotnumber_after_nth(s)
        if s isa Core.SlotNumber
            return (s.id==nth+1) ? value : (s.id>nth+1 ? Core.SlotNumber(s.id-1) : s)
        else
            return s
        end
    end
    for m in methods(src).ms
        # to CodeInfo
        ci = uncompressed_ast(m)
        # transform CodeInfo
        for ex in ci.code
            if ex isa Expr
                map!(change_slotnumber_after_nth, ex.args, ex.args)
            end
        end
        # purge nth argument from slots
        ci.slotnames = ci.slotnames[[i for i = 1:length(ci.slotnames) if i!=nth+1]]
        ci.slotflags = ci.slotflags[[i for i = 1:length(ci.slotflags) if i!=nth+1]]
        # purge nth argument from m.sig
        sig1 = Tuple{[m.sig.parameters[i] for i=1:length(m.sig.parameters) if i!=nth+1]...}
        # inject CodeInfo to mod = getmodule(dst)
        jl_method_def(argdata(sig1, dst), ci, mod)
    end
end

After debuging with @code_lowered and @code_typed I just realized that I forgot to decrease the id of each slot number after the n’th one (which is assigned and then purged). Now the updated version of assign_purge_argument! works (at least for all test cases I have at hand).

1 Like