How to write a macro like #[global_allocator]

help
1

The #[global_allocator] macro can define how Box, Vec, String etc. find the static variable with GlobalAlloc trait. I want to define something similar to #[global_allocator] (e.g.: global_object attribute) so that I can find a static variable with such attribute.
It seems I should use attribute macros, but I have no idea how to use it in my case. I failed to find a relevant example in the docs.

1 Like
2

#[global_allocator] is internal compiler magic.

The closest thing you can get in user-space is something like the inventory crate.

3

If you know that the relevant item will be defined exactly once in the final binary, you can use an extern ABI for this:

mod a {
    extern "C" { static MY_GLOBAL: i32; }
    pub fn f() {
        unsafe { dbg!(MY_GLOBAL); }
    }
}

mod b {
    #[unsafe(export_name = "MY_GLOBAL")]
    static DEFINED_IN_B: i32 = 42;
}

fn main() {
    a::f()
}
4 Likes
4

Well, it does sound like a fair workaround, but what if I want a global object in generics that implements certain traits? I am asking "how to write a macro like #[global_allocator]" in that such object can be arbitrary generics that implements GlobalAlloc trait.
I tried to declare MY_GLOBAL with &'static dyn XxxTrait type, but the program crashed. The debugger shows that the vtable is null.

5

This sounds very similar to something I did a while ago. Here's my original comment from wit-bindgen-rust: remove "standalone" mode in the future · Issue #243 · bytecodealliance/wit-bindgen · GitHub

You might find our story interesting.

For some background, we've currently got 20+ implementations for a particular set of WIT files, and not being able to put wit_bindgen_export!() in a shared crate was becoming a real pain (can't share code for trait implementations, can't use ?, no convenience methods, code is super verbose, etc.). Wasmer also haven't updated their fork of wit-bindgen in quite a while so we couldn't use #193 without switching our WebAssembly VM from Wasmer to Wasmtime.

To deal with this, we took inspiration from how Rust's #[global_allocator] works...

First, we call wit_bindgen_rust::export!() in some shared crate.

// support/src/guest/bindings.rs

wit_bindgen_rust::import!("../wit-files/rune/runtime-v2.wit");
wit_bindgen_rust::export!("../wit-files/rune/proc-block-v2.wit");

Then we created a trait which represents our resource and would be implemented downstream.

// support/src/guest/proc_block.rs

pub trait ProcBlock {
    fn tensor_constraints(&self) -> TensorConstraints;
    fn run(&self, inputs: Vec<Tensor>) -> Result<Vec<Tensor>, RunError>;
}

Next comes extern "Rust" definitions for our ProcBlock's constructor and a free function the host uses to find out more about its functionality (name, description, input tensors, etc.).

// support/src/guest/proc_block.rs

extern "Rust" {
    fn __proc_block_metadata() -> Metadata;
    fn __proc_block_new(
        args: Vec<Argument>,
    ) -> Result<Box<dyn ProcBlock>, CreateError>;
}

From there, we can write a shim implementation in our support crate for the ProcBlockV2 trait that wit-bindgen generates.

// support/src/guest/proc_block.rs

struct ProcBlockV2;

impl proc_block_v2::ProcBlockV2 for ProcBlockV2 {
    fn metadata() -> Metadata {
        logging::initialize_logger();
        unsafe { __proc_block_metadata() }
    }

    fn create_node(
        args: Vec<Argument>,
    ) -> Result<wit_bindgen_rust::Handle<self::Node>, CreateError> {
        logging::initialize_logger();
        let proc_block = unsafe { __proc_block_new(args)? };
        Ok(Handle::new(Node(Box::new(proc_block))))
    }
}

pub struct Node(Box<dyn ProcBlock>);

impl proc_block_v2::Node for Node {
    fn tensor_constraints(&self) -> TensorConstraints {
        self.0.tensor_constraints()
    }

    fn run(&self, inputs: Vec<Tensor>) -> Result<Vec<Tensor>, RunError> {
        self.0.run(inputs)
    }
}

And to wrap it all up, we have a macro which end users can use to generate __proc_block_metadata() and friends.

// support/src/guest/proc_block.rs

/// Tell the runtime that a WebAssembly module contains a proc-block.
#[macro_export]
macro_rules! export_proc_block {
    (metadata: $metadata_func:expr, proc_block: $proc_block:ty $(,)?) => {
        #[doc(hidden)]
        #[no_mangle]
        pub fn __proc_block_metadata() -> $crate::guest::Metadata { $metadata_func() }

        #[doc(hidden)]
        #[no_mangle]
        pub fn __proc_block_new(
            args: Vec<$crate::guest::Argument>,
        ) -> Result<Box<dyn $crate::guest::ProcBlock>, $crate::guest::CreateError> {
            fn assert_impl_proc_block(_: &impl $crate::guest::ProcBlock) {}

            let proc_block = <$proc_block>::try_from(args)?;
            assert_impl_proc_block(&proc_block);

            Ok(Box::new(proc_block) as Box<dyn $crate::guest::ProcBlock>)
        }
    };
}

I really like this approach because, besides the export_proc_block!() macro, it all looks like normal Rust code and you don't need to know about wit-bindgen types (see here for an example). It also solves the awkward super issue where downstream crates need to somehow inject their types and free function implementations into a crate higher up the dependency tree.

1 Like