lunatic_common_api/
lib.rs

1use anyhow::{anyhow, Context, Result};
2use std::{fmt::Display, future::Future, io::Write, pin::Pin};
3use wasmtime::{Caller, Memory, Val};
4
5const ALLOCATOR_FUNCTION_NAME: &str = "lunatic_alloc";
6const FREEING_FUNCTION_NAME: &str = "lunatic_free";
7
8// Get exported memory
9pub fn get_memory<T>(caller: &mut Caller<T>) -> Result<Memory> {
10    caller
11        .get_export("memory")
12        .or_trap("No export `memory` found")?
13        .into_memory()
14        .or_trap("Export `memory` is not a memory")
15}
16
17// Call guest to allocate a Vec of size `size`
18pub fn allocate_guest_memory<'a, T: Send>(
19    caller: &'a mut Caller<T>,
20    size: u32,
21) -> Pin<Box<dyn Future<Output = Result<u32>> + Send + 'a>> {
22    Box::pin(async move {
23        let mut results = [Val::I32(0)];
24        caller
25            .get_export(ALLOCATOR_FUNCTION_NAME)
26            .or_trap(format!("no export named {ALLOCATOR_FUNCTION_NAME} found"))?
27            .into_func()
28            .or_trap("cannot turn export into func")?
29            .call_async(caller, &[Val::I32(size as i32)], &mut results)
30            .await
31            .or_trap(format!("failed to call {ALLOCATOR_FUNCTION_NAME}"))?;
32
33        Ok(results[0]
34            .i32()
35            .or_trap(format!("result of {ALLOCATOR_FUNCTION_NAME} is not i32"))? as u32)
36    })
37}
38
39// Call guest to free a slice of memory at location ptr
40pub fn free_guest_memory<'a, T: Send>(
41    caller: &'a mut Caller<T>,
42    ptr: u32,
43) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
44    Box::pin(async move {
45        let mut results = [];
46        let result = caller
47            .get_export(FREEING_FUNCTION_NAME)
48            .or_trap(format!("no export named {FREEING_FUNCTION_NAME} found"))?
49            .into_func()
50            .or_trap("cannot turn export into func")?
51            .call_async(caller, &[Val::I32(ptr as i32)], &mut results)
52            .await;
53
54        result.or_trap(format!("failed to call {FREEING_FUNCTION_NAME}"))?;
55        Ok(())
56    })
57}
58
59// Allocates and writes data to guest memory, updating the len_ptr and returning the allocated ptr.
60pub async fn write_to_guest_vec<T: Send>(
61    caller: &mut Caller<'_, T>,
62    memory: &Memory,
63    data: &[u8],
64    len_ptr: u32,
65) -> Result<u32> {
66    let alloc_len = data.len();
67    let alloc_ptr = allocate_guest_memory(caller, alloc_len as u32).await?;
68
69    let (memory_slice, _) = memory.data_and_store_mut(&mut (*caller));
70    let mut alloc_vec = memory_slice
71        .get_mut(alloc_ptr as usize..(alloc_ptr as usize + alloc_len))
72        .context("allocated memory does not exist")?;
73
74    alloc_vec.write_all(data)?;
75
76    memory.write(caller, len_ptr as usize, &alloc_len.to_le_bytes())?;
77
78    Ok(alloc_ptr)
79}
80
81pub trait IntoTrap<T> {
82    fn or_trap<S: Display>(self, info: S) -> Result<T>;
83}
84
85impl<T, E: Display> IntoTrap<T> for Result<T, E> {
86    fn or_trap<S: Display>(self, info: S) -> Result<T> {
87        match self {
88            Ok(result) => Ok(result),
89            Err(error) => Err(anyhow!(
90                "Trap raised during host call: {} ({}).",
91                error,
92                info
93            )),
94        }
95    }
96}
97
98impl<T> IntoTrap<T> for Option<T> {
99    fn or_trap<S: Display>(self, info: S) -> Result<T> {
100        match self {
101            Some(result) => Ok(result),
102            None => Err(anyhow!(
103                "Trap raised during host call: Expected `Some({})` got `None` ({}).",
104                std::any::type_name::<T>(),
105                info
106            )),
107        }
108    }
109}