1 use anyhow::Result;
2 use arbitrary::Arbitrary;
3 use std::mem::MaybeUninit;
4 use wasmtime::component::__internal::{
5     CanonicalAbiInfo, InstanceType, InterfaceType, LiftContext, LowerContext,
6 };
7 use wasmtime::component::{ComponentNamedList, ComponentType, Func, Lift, Lower, TypedFunc, Val};
8 use wasmtime::{AsContextMut, Config, Engine};
9 
10 pub trait TypedFuncExt<P, R> {
11     fn call_and_post_return(&self, store: impl AsContextMut, params: P) -> Result<R>;
12 }
13 
14 impl<P, R> TypedFuncExt<P, R> for TypedFunc<P, R>
15 where
16     P: ComponentNamedList + Lower,
17     R: ComponentNamedList + Lift,
18 {
19     fn call_and_post_return(&self, mut store: impl AsContextMut, params: P) -> Result<R> {
20         let result = self.call(&mut store, params)?;
21         self.post_return(&mut store)?;
22         Ok(result)
23     }
24 }
25 
26 pub trait FuncExt {
27     fn call_and_post_return(
28         &self,
29         store: impl AsContextMut,
30         params: &[Val],
31         results: &mut [Val],
32     ) -> Result<()>;
33 }
34 
35 impl FuncExt for Func {
36     fn call_and_post_return(
37         &self,
38         mut store: impl AsContextMut,
39         params: &[Val],
40         results: &mut [Val],
41     ) -> Result<()> {
42         self.call(&mut store, params, results)?;
43         self.post_return(&mut store)?;
44         Ok(())
45     }
46 }
47 
48 pub fn config() -> Config {
49     drop(env_logger::try_init());
50 
51     let mut config = Config::new();
52     config.wasm_component_model(true);
53 
54     // When `WASMTIME_TEST_NO_HOG_MEMORY` is set it means we're in qemu. The
55     // component model tests create a disproportionate number of instances so
56     // try to cut down on virtual memory usage by avoiding 4G reservations.
57     if std::env::var("WASMTIME_TEST_NO_HOG_MEMORY").is_ok() {
58         config.memory_reservation(0);
59         config.memory_guard_size(0);
60     }
61     config
62 }
63 
64 pub fn engine() -> Engine {
65     Engine::new(&config()).unwrap()
66 }
67 
68 pub fn async_engine() -> Engine {
69     let mut config = config();
70     config.async_support(true);
71     Engine::new(&config).unwrap()
72 }
73 
74 /// Newtype wrapper for `f32` whose `PartialEq` impl considers NaNs equal to each other.
75 #[derive(Copy, Clone, Debug, Arbitrary)]
76 pub struct Float32(pub f32);
77 
78 /// Newtype wrapper for `f64` whose `PartialEq` impl considers NaNs equal to each other.
79 #[derive(Copy, Clone, Debug, Arbitrary)]
80 pub struct Float64(pub f64);
81 
82 macro_rules! forward_impls {
83     ($($a:ty => $b:ty,)*) => ($(
84         unsafe impl ComponentType for $a {
85             type Lower = <$b as ComponentType>::Lower;
86 
87             const ABI: CanonicalAbiInfo = <$b as ComponentType>::ABI;
88 
89             #[inline]
90             fn typecheck(ty: &InterfaceType, types: &InstanceType<'_>) -> Result<()> {
91                 <$b as ComponentType>::typecheck(ty, types)
92             }
93         }
94 
95         unsafe impl Lower for $a {
96             fn lower<U>(
97                 &self,
98                 cx: &mut LowerContext<'_, U>,
99                 ty: InterfaceType,
100                 dst: &mut MaybeUninit<Self::Lower>,
101             ) -> Result<()> {
102                 <$b as Lower>::lower(&self.0, cx, ty, dst)
103             }
104 
105             fn store<U>(&self, cx: &mut LowerContext<'_, U>, ty: InterfaceType, offset: usize) -> Result<()> {
106                 <$b as Lower>::store(&self.0, cx, ty, offset)
107             }
108         }
109 
110         unsafe impl Lift for $a {
111             fn lift(cx: &mut LiftContext<'_>, ty: InterfaceType, src: &Self::Lower) -> Result<Self> {
112                 Ok(Self(<$b as Lift>::lift(cx, ty, src)?))
113             }
114 
115             fn load(cx: &mut LiftContext<'_>, ty: InterfaceType, bytes: &[u8]) -> Result<Self> {
116                 Ok(Self(<$b as Lift>::load(cx, ty, bytes)?))
117             }
118         }
119 
120         impl PartialEq for $a {
121             fn eq(&self, other: &Self) -> bool {
122                 self.0 == other.0 || (self.0.is_nan() && other.0.is_nan())
123             }
124         }
125     )*)
126 }
127 
128 forward_impls! {
129     Float32 => f32,
130     Float64 => f64,
131 }
132