fix update_slots

Author: Roger-luo

src/CompilerPluginTools.jl

@@ -3,12 +3,15 @@ module CompilerPluginTools
 using MLStyle
 using Expronicon
 
-export NewCodeInfo,
+export
+    # interp
+    NewCodeInfo,
     NewSlotNumber,
     JuliaLikeInterpreter,
     AbstractInterpreter,
     isintrinsic,
     # reexport SSA nodes from Core
+    MethodInstance,
     CodeInfo,
     SSAValue,
     Const,
@@ -50,6 +53,8 @@ export NewCodeInfo,
     const_invoke!,
     finish,
     IntrinsicError,
+    # utils
+    method_instance,
     @make_codeinfo,
     @make_ircode,
     @intrinsic_stub

src/codeinfo.jl

@@ -249,6 +249,8 @@ function update_slots(e, newslotmap, changemap)
     @match e begin
         SlotNumber(id) => SlotNumber(id + changemap[id])
         NewSlotNumber(id) => SlotNumber(newslotmap[id])
+        ReturnNode(x) => ReturnNode(update_slots(x, newslotmap, changemap))
+        GotoIfNot(cond, dest) => GotoIfNot(update_slots(cond, newslotmap, changemap), dest)
         NewvarNode(SlotNumber(id)) => NewvarNode(SlotNumber(id+slotmap[id]))
         Expr(head, args...) => Expr(head, map(x->update_slots(x, newslotmap, changemap), e.args)...)
         _ => e

src/utils.jl

@@ -59,3 +59,18 @@ function make_codeinfo_m(ex)
     end
 end
 
+"""
+    method_instance(f, tt, world=Base.get_world_counter())
+
+Return the `MethodInstance`, unlike `Base.method_instances`, `tt` must be specified type.
+"""
+function method_instance(@nospecialize(f), @nospecialize(tt), world=Base.get_world_counter())
+    # get the method instance
+    meth = which(f, tt)
+    sig = Base.signature_type(f, tt)::Type
+    (ti, env) = ccall(:jl_type_intersection_with_env, Any, (Any, Any), sig,
+        meth.sig)::Core.SimpleVector
+    meth = Base.func_for_method_checked(meth, ti, env)
+    return ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
+        (Any, Any, Any, UInt), meth, ti, env, world)
+end