|
| 1 | +function instrument_method(m::Method) |
| 2 | + # Step 1: pick a new name |
| 3 | + name = Symbol(m.name, "#instrumented") |
| 4 | + f = Core.eval(m.module, :(function $name end)) # create the function object (with no methods) |
| 5 | + # Step 2: add an extra argument, `#framedata#` |
| 6 | + # (For simplicity here I'm cheating by assuming no internal locals and no type parameters. |
| 7 | + # In reality, among other things you'll need to renumber the slots) |
| 8 | + # For help, see https://docs.julialang.org/en/latest/devdocs/ast/#Expr-types-1, section on `method` |
| 9 | + sigm, src = m.sig, copy_codeinfo(Base.uncompressed_ast(m)) |
| 10 | + sigt = Core.svec(typeof(f), sigm.parameters[2:end]..., FrameData) # first is #self# |
| 11 | + sigp = Core.svec() |
| 12 | + sigsv = Core.svec(sigt, sigp) |
| 13 | + push!(src.slotnames, Symbol("#framedata#")) |
| 14 | + # Step 3: instrument the body |
| 15 | + # - for each store to a slotnumber (SlotNumber on the LHS of :(=)), add a |
| 16 | + # "real" assignment to `#framedata#.locals` |
| 17 | + # - for each line that is `used`, add a store to `#framedata#.ssavalues` |
| 18 | + # Both of these will require renumbering the ssavalues. Again, for simplicity I'll just cheat: |
| 19 | + # change `x + 1` into `x + 2` for my `inc1` example |
| 20 | + src.code[1].args[end] = 2 |
| 21 | + # Step 4: create the new method |
| 22 | + ccall(:jl_method_def, Cvoid, (Any, Any, Any), sigsv, src, m.module) |
| 23 | + # See below about "wrapper methods" that set up `#framedata#` for this method; |
| 24 | + # you might also create the wrapper here? |
| 25 | + return f |
| 26 | +end |
| 27 | + |
| 28 | +function typeinf_instrumented(linfo::MethodInstance, params::Core.Compiler.Params) |
| 29 | + # This modifies the source code to add extra instrumentation |
| 30 | + # Step 1: call regular inference. Here, a major goal is to perform inlining, |
| 31 | + # so that we don't have to create so many `framedata`s |
| 32 | + src = Core.Compiler.typeinf_ext(linfo, params) |
| 33 | + # Step 2: replace the `:invoke` Exprs in `isrc` with invokes to `#framedata#`-instrumented variants. |
| 34 | + # You will also have to insert statements to create the framedata for that method. |
| 35 | + # Presumably, the best approach would be to just :invoke a wrapper method that creates a framedata |
| 36 | + # for the method we instrumented above, and then calls the instrumented method. That wrapper would have the same args as the |
| 37 | + # original function, so it's really just a case of changing which MethodInstance you call. |
| 38 | + # See usage of `Core.Compiler.code_for_method` below to create that MethodInstance. |
| 39 | + # I'm skipping that part |
| 40 | + return src |
| 41 | +end |
| 42 | + |
| 43 | +function precompile_instrumented(f, atypes) |
| 44 | + # delightfully insane! |
| 45 | + # !!Note!!: before executing this, you need to compile `typeinf_instrumented` |
| 46 | + try |
| 47 | + ccall(:jl_set_typeinf_func, Cvoid, (Any,), typeinf_instrumented) |
| 48 | + precompile(f, atypes) |
| 49 | + finally |
| 50 | + ccall(:jl_set_typeinf_func, Cvoid, (Any,), Core.Compiler.typeinf_ext) |
| 51 | + end |
| 52 | +end |
0 commit comments