cranelift_jit: the JITModule::free_memory method should take &mut self

[tertsdiepraam]

Currently, the free_memory method on JITModule takes self. This makes some sense, because the whole thing is then invalidated. However, it is preventing me from (easily) making a wrapper type that frees the memory on drop:

struct MyJIT(JITModule);

impl Drop for MyJit {
    fn drop(&mut self) {
        unsafe { self.0.free_memory() } // error! Can't move out of &mut self
    }
}

The above is simplified from what I'm actually trying to do because the above is obviously unsafe. I was trying to do this:

struct MyJIT(Arc<JITModule>);

impl Drop for MyJit {
    fn drop(&mut self) {
        if let Some(module) = self.0.get_mut() {
            unsafe { module.free_memory() }
        }
    }
}

where all functions that I give out would have a MyJIT, so that the memory automatically gets dropped once the last function that uses it would drop the memory.

I can work around this by swapping in a new JITModule, but that doesn't seem ideal. Since free_memory is already unsafe, I think it's alright to let it take &mut self and allow people to build safe abstractions over it.

[tertsdiepraam]

To illustrate why the workaround is annoying: it requires creating a JITModule, which requires creating a JITBuilder, which requires a function for libcalls and can fail. So the final function becomes something like this:

#[derive(Clone)]
pub struct ModuleData(Arc<JITModule>);

impl Drop for ModuleData {
    fn drop(&mut self) {
        // get_mut returns None if we are not the last Arc, so the JITModule
        // shouldn't be dropped yet.
        let Some(module) = Arc::get_mut(&mut self.0) else {
            return;
        };

        // If this fails we'd rather just return and leak some memory than
        // panic.
        let Ok(builder) =
            JITBuilder::new(cranelift_module::default_libcall_names())
        else {
            return;
        };

        // Swap in an empty JITModule so that we can take ownership
        // so that we can call `free_memory` on the old one.
        // The new one hasn't allocated anything so it will just get dropped
        // normally.
        let mut other = JITModule::new(builder);
        std::mem::swap(&mut other, module);

        // SAFETY: We only give out functions that hold a ModuleData and
        // therefore an Arc to this module. By `get_mut`, we know that we are
        // the last Arc to this memory and hence it is safe to free its
        // memory. New Arcs cannot have been created in the meantime because
        // that requires access to the last Arc, which we know that we have.
        unsafe { other.free_memory() };
    }
}

[bjorn3]

The canonical solution for this is wrapping JITModule in ManuallyDrop and then doing something like

let Some(module) = Arc::get_mut(&mut self.0) else {
    return;
};

unsafe {
    let module = ManuallyDrop::take(module);
    module.free_memory();
}

In this case given that free_memory is unsafe anyway and it is already possible for pointers to outlive the jit module, I think accepting &mut self in free_memory is fine.

[tertsdiepraam]

Oh that sounds reasonable! Thanks!

[tertsdiepraam]

I've got it working. Feel free to close this, though it might still be useful. I'll leave that decision up to you. If you want this I'd be happy to attempt a PR.