Make the imported memory available in functions

[Lohann]

Feature

Hello, I need to access the imported memory inside a functions, then I realized the Caller only show exports, not imports, so I had a lot of trouble to get Wasmtime store working with MaybeUnit, this simple solution segfaults:

pub struct State {
    pub memory: Memory,
}

let state = MaybeUninit::<State>::uninit();
let mut store = Store::new(engine, state);

let memory_type = MemoryType::new(16, None);
store.data_mut().memory.write(Memory::new(&mut store, memory_type)?);

let store = unsafe {
    std::mem::transmute::<Store<MaybeUninit<State>>, Store<State>>(store)
};

// ...
let imports = [memory.into()];

// segfaults below
let instance = Instance::new(&mut store, &module, &imports)?;

Then I realized the issue is that there's no way for me transmute only the State, I need to transmute the Store which is not recommended, once rust doesn't guarantee the same memory layout:

assert_eq!(size_of::<Option<bool>>(), 1);
assert_eq!(size_of::<Option<MaybeUninit<bool>>>(), 2);

Actually I haven't find any way to use MaybeUnit that doesn't look hacky, and I want to avoid the usage of Option and unwraps in the code, once it bloats the binary with panic data.

Alternatives

1. Make the imported memory easily available inside Functions, ex: expose it in the Caller.
2. Use #[repr(C)] on Store, so we can safely transmute it.

[alexcrichton]

Where possible I'd recommend avoiding unsafe. If things are segfaulting it's probably due to that, so for example you could store Option<Memory> instead of using MaybeUninit and then there's no need for transmute and this probably won't segfault.

Otherwise though is there a problem with storing the memory in State?

[Lohann]

Otherwise though is there a problem with storing the memory in State?

There's no problem, but rust encourages the use of Typestate Pattern, where the state of an object guarantees it is valid. In this case I want to guarantee the memory ALWAYS exists, that's why I don't want to use Option.
One example is the NonZeroU32 by knowing the number is never zero, the rust compiler can do some neat optimizations:

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=6778305b1980ec60413183a0a1127a4d

assert_eq!(size_of::<u32>(), 4);
assert_eq!(size_of::<Option<u32>>(), 8);
assert_eq!(size_of::<Option<NonZeroU32>>(), 4);

In my case I want to guarantee the memory always exists, I don't want to handle the Option::None everywhere, neither use unwrap() everywhere, that's ok if there's no other option, but I think this is something that should be supported somehow by wasmtime, once the Store owns the memory, makes sense I be able to store it together with the store.

[alexcrichton]

One thing you could perhaps do is to create a dummy Memory with a throwaway Store which is placed within future Stores as they're created. That would then be overwritten to the "real" memory once the store is created. That way you can store just Memory without having to deal with Option and you won't have to deal with any unsafety either.

[Lohann]

Update: I managed to store the imported memory in the State without Option<Memory> by using MaybeUninit, the code is sound but requires two unsafe blocks.

For me it worthy the price, better than having to handle Option<Memory> in every imported function. But I still would love to have a more ergonomic way to access imported memory without Option or unsafe code.

pub struct State {
    /// Memory that will be IMPORTED to the WASM instance, then it must be
    /// created before instantiate the WASM, must be stored here otherwise
    /// you can't access the memory inside WasmTime imported functions.
    pub memory: Memory,
}

impl State {
    /// Initializes a Store with a memory owned by the `store`.
    pub fn new(engine: &Engine, memory_type: MemoryType) -> anyhow::Result<Store<Self>> {
        // The `Store` needs `State`, the `State` needs `Memory`, but  `Memory` can only be
        // created after `Store`, so we have a ciclc dependency issue. Solved this bellow by first
        // initialize the `State` with a dummy memory using `MaybeUninit::zeroed().assume_init()`.
        // References:
        // - https://github.com/bytecodealliance/wasmtime/issues/4922#issuecomment-1251086171
        // - https://doc.rust-lang.org/std/mem/union.MaybeUninit.html
        #[allow(clippy::uninit_assumed_init)]
        let state = Self {
            // SAFETY: This `memory` will be correctly initialized after the `Store` below.
            memory: unsafe { MaybeUninit::<Memory>::zeroed().assume_init() },
        };

        // Now use the `State` to create the `Store`.
        let mut store = Store::new(engine, state);

        // Use the `Store` to create a valid `Memory` then replaces the dummy memory.
        // Ref: https://doc.rust-lang.org/std/mem/union.MaybeUninit.html#initializing-a-struct-field-by-field
        unsafe {
            // Get the address of state's dummy memory.
            let ptr = std::ptr::addr_of_mut!(store.data_mut().memory);
            // replaces the  dummy memory by a valid memory
            ptr.write(Memory::new(&mut store, memory_type)?);
        }

        Ok(store)
    }
}

Before you ask, no, I can't simply export the memory like everyone else, because the wasm code is a ink smart-contract, and the memory is always imported.

[ZylosLumen]

Just ran into this issue and wanted to avoid unsafe Option::unwrap_unchecked() everywhere but it will have to do for now.
Really curious as importing the memory (or any other store-bound type) and also needing to access it from the host does not seem like a rare use case so having a low cost api would be really nice

[Lohann]

@alexcrichton
I noticed Memory is simple a identifier, it doesn't have the actual buffer, (that's why it implements Copy), so a possible solution is create both the state and memory atomically, or create the Store last, ex:

Proposal 1 - Store::new_with

pub struct State {
    pub memory: Memory,
}

let memory_type = MemoryType::new(16, None);
let mut store = Store::new_with(engine, move |store_ref: &mut Store<_>| -> State  {
    // obs: in method requires `State` to implement Unpin, this prevents State from reference Store.
    State {
        memory: Memory::new(store_ref, memory_type)
    }
});

Proposal 2 - StoreBuilder

pub struct State {
    pub memory: Memory,
}

let memory_type = MemoryType::new(16, None);
let (mut store, mut linker):  = StoreBuilder::new()
    .with_memory("env", "memory", memory_type)
    .with_func("env", "add", |x: i32, y: i32| x + y)
    .init_state(|linker|: State {
        // SAFETY: the memory was imported above
        let memory = linker.get("env", "memory").unwrap();
        State { memory }
    })
    .build()?;
let instance = linker.instantiate(&mut store, &module)?;