Modifying AST

[seccompgeek]

Hi, I am working on a project that requires me to somehow modify AST (or I think that's the best option I have). Specifically, I would like to wrap/replace some specific function calls with a call to a closure. I found that closures cannot be defined at MIR level as they require DefIds. My plan so far is, running the compiler the first time and find out which calls I want to modify and saving their NodeIds. Then, I run the compiler again, and during AST expansion, I try to replace the calls I'm interested in. The challenge is, this is breaking the compiler, either because I'm not doing it right, or I'm supposed to tackle it a different way. Any advice on how to proceed would be most appreciated.

This is what I have so far, but it fails to pass the typechecker, saying DefId{0:0} does not exist. For some reason I'm messing up the original crate.

`pub fn wrap_extern_calls(
krate: &mut ast::Crate,
sess: &Session,
resolver: &mut dyn ResolverExpand,
crate_name: String,
features: &Features
) {
struct ExternCallVisitor<'a> {
extern_calls: FxHashSet,
ext_cx: ExtCtxt<'a>
}

impl<'a> MutVisitor for ExternCallVisitor<'a> {
    fn visit_expr(&mut self, expr: &mut P<Expr>){
        let id = expr.id;
        if self.extern_calls.contains(&id){
            let local_expr = expr.clone();
            let closure = self.ext_cx.expr(expr.span, ExprKind::Closure(
                Box::new(Closure {
                   binder: ast::ClosureBinder::NotPresent,
                   capture_clause: ast::CaptureBy::Ref,
                   constness: ast::Const::No,
                   asyncness: ast::Async::No,
                   movability: ast::Movability::Movable,
                   fn_decl: P(
                    FnDecl {
                        inputs: ThinVec::new(),
                        output: ast::FnRetTy::Default(DUMMY_SP)
                    }
                   ),
                   body: local_expr,
                   fn_decl_span: expr.span,
                   fn_arg_span: expr.span
                })
            ));
            let wrapper_fn = self.ext_cx.path(expr.span,  vec![Ident::from_str("std"), Ident::from_str("metasafe"), Ident::with_dummy_span(sym::metasafe_extern_stack_run)]);
            let expr_path = self.ext_cx.expr_path(wrapper_fn);
            let wrapper_call = self.ext_cx.expr_call(expr.span, expr_path, thin_vec![closure]);
            let wrapper_call_frag = AstFragment::Expr(wrapper_call);
            let wrapper_call = self.ext_cx.expander().fully_expand_fragment(wrapper_call_frag).make_expr();
            let _ = std::mem::replace(expr, wrapper_call);
        }
    }
}

let extern_calls = load_extern_calls(crate_name.clone());
if extern_calls.is_empty() {
    return;
}

let econfig = ExpansionConfig::default(crate_name, features);
let ext_cx = ExtCtxt::new(&sess, econfig, resolver, None);
let mut extern_call_visitor = ExternCallVisitor {
    extern_calls,
    ext_cx
};
extern_call_visitor.visit_crate(krate);`

}

Thanks in advance.