Modify callee MIR in optimization pass?

compiler
1

I'm working on a research project where I modify the Rust compiler to (1) look for any object definitions of a certain type, (2) trace their uses, then (3) modify some related statements before/after the use.

Right now I can successfully use a depth-first search to follow a def to all its uses using rustc_mir::util::DefUseAnalysis in the MIR, including through function calls and closures. However, now I'm stuck on how to modify statements in multiple MIR Body structs (step 3).

I have defined a struct that implements the MirPass trait and added it as a new MIR optimization pass. The run_pass function takes &mut Body<'tcx> as an argument, so obviously you can modify the MIR of the current function, but can I get a mutable reference to the called functions' MIR? I can't see a way to do this without calling something like tcx.mir_built(fn_call_def_id).steal(), which will later cause a compiler error because the value can only be stolen once and it needs to be stolen later.

For example, suppose we have the following code:

fn foo() {
    let my_object = MyType::new();
    bar(my_object);
}

fn bar(my_object: MyType) {
    my_object.baz(1);
}

I can recognize the definition of my_object in foo and its use in bar. Now I want to modify the MIR in bar, for example, by changing the argument from 1 to 2. However, we're only running an MIR pass on foo, not bar, so we don't have a mutable reference to the Body struct of bar. Is it even possible to get a mutable reference to bar inside foo's run_pass?

Please let me know if I can better clarify my problem. Since this is a research project, I'm open to "hacky" solutions that work for a proof of concept.

1 Like
2

You could always modify Steal to expose a borrow_mut accessor. If you're lucky, that shouldn't actually break anything in the compiler.

A more principled design (though I have no idea if it works in the querified world) would be to build up a queue of edits to apply later when you do have access to the body in question.

3 Likes
3

@CAD97 Thank you for your suggestion! I tried it out this week and unfortunately hit a dead end :frowning:

To call any mutator method on Steal I would need a mutable reference to the Steal. The current. queries return &Steal. Query return values have to implement Copy and I don't think a mutable reference can be copied, though, so I don't think I can return &mut Steal from a query.

I would love to be able to enqueue my edits to some sort of "global" data structure (accessible from TyCtxt?) that, say, can be mutated by a query on foo and later accessed by a query on bar. Do you (or anyone else) have other ideas? I've been looking into arena allocation and type interning.

EDIT: Copy, not Clone.

4

This topic was automatically closed 90 days after the last reply. New replies are no longer allowed.