pit_rust_host_lib/
lib.rs

1#![no_std]
2#[doc(hidden)]
3pub extern crate alloc;
4use alloc::borrow::ToOwned;
5use alloc::string::String;
6use alloc::sync::Arc;
7use alloc::vec;
8use alloc::format;
9use alloc::vec::Vec;
10use core::{
11    cell::UnsafeCell,
12    iter::{empty, once},
13};
14
15use pit_core::{Arg, Interface};
16use wasm_runtime_layer::{
17    backend::WasmEngine, AsContext, AsContextMut, Extern, ExternRef, Func, FuncType, Imports,
18    Instance, Module, Store, StoreContext, StoreContextMut, Value, ValueType,
19};
20pub fn init<U: AsRef<Instance> + 'static, E: WasmEngine>(
21    l: &mut Imports,
22    ctx: &mut impl AsContextMut<UserState = U, Engine = E>,
23) {
24    l.define(
25        "pit",
26        "drop",
27        Extern::Func(Func::new(
28            &mut *ctx,
29            FuncType::new(once(ValueType::ExternRef), empty()),
30            move |mut ctx, args, rets| {
31                let Value::ExternRef(Some(a)) = args[0].clone() else {
32                    anyhow::bail!("invalid type")
33                };
34                let Ok(x): Result<&Wrapped<U, E>, anyhow::Error> =
35                    a.downcast::<'_, '_, Wrapped<U, E>, U, E>(ctx.as_context())
36                else {
37                    return Ok(());
38                };
39                let f = x.all[0].clone();
40                f(ctx.as_context_mut(), vec![])?;
41                Ok(())
42            },
43        )),
44    );
45}
46pub fn emit_ty(a: &Arg) -> ValueType {
47    match a {
48        Arg::I32 => ValueType::I32,
49        Arg::I64 => ValueType::I64,
50        Arg::F32 => ValueType::F32,
51        Arg::F64 => ValueType::F64,
52        Arg::Resource {
53            ty,
54            nullable,
55            take,
56            ann,
57        } => ValueType::ExternRef,
58        _ => todo!(), // Arg::Func(f) => ValueType::FuncRef,
59    }
60}
61pub fn emit<U: AsRef<Instance> + 'static, E: WasmEngine>(
62    l: &mut Imports,
63    rid: Arc<Interface>,
64    m: &Module,
65    ctx: &mut impl AsContextMut<UserState = U, Engine = E>,
66) {
67    let n = format!("pit/{}", rid.rid_str());
68    for (j, (i, m)) in rid.methods.iter().enumerate() {
69        l.define(
70            &n,
71            i.as_str(),
72            Extern::Func(Func::new(
73                &mut *ctx,
74                FuncType::new(
75                    once(ValueType::ExternRef).chain(m.params.iter().map(emit_ty)),
76                    m.rets.iter().map(emit_ty),
77                ),
78                move |mut ctx, args, rets| {
79                    let Value::ExternRef(Some(a)) = args[0].clone() else {
80                        anyhow::bail!("invalid type")
81                    };
82                    let x: &Wrapped<U, E> = a.downcast(ctx.as_context())?;
83                    let t = x.all[j + 1].clone();
84                    let rets2 = t(ctx.as_context_mut(), args[1..].iter().cloned().collect())?;
85                    for (r, s) in rets2.into_iter().zip(rets.iter_mut()) {
86                        *s = r;
87                    }
88                    Ok(())
89                },
90            )),
91        )
92    }
93    let i = m
94        .imports(ctx.as_context().engine())
95        .map(|a| (a.module.to_owned(), a.name.to_owned()))
96        .collect::<Vec<_>>();
97    for i in i {
98        if i.0 == n {
99            if let Some(t) = i.1.strip_prefix("~") {
100                let t = t.to_owned();
101                let rid = rid.clone();
102                l.define(
103                    &n,
104                    &i.1,
105                    Extern::Func(Func::new(
106                        &mut *ctx,
107                        FuncType::new(once(ValueType::I32), once(ValueType::ExternRef)),
108                        move |mut ctx, args, rets| {
109                            let i = ctx.data().as_ref().clone();
110                            let object = Wrapped::new(
111                                args.to_owned(),
112                                rid.clone(),
113                                t.to_owned(),
114                                i,
115                                ctx.as_context_mut(),
116                            );
117                            rets[0] = Value::ExternRef(Some(ExternRef::new(ctx, object)));
118                            Ok(())
119                        },
120                    )),
121                )
122            };
123        };
124    }
125}
126pub struct Wrapped<U: 'static, E: wasm_runtime_layer::backend::WasmEngine> {
127    pub rid: Arc<Interface>,
128    pub all: Vec<
129        Arc<
130            dyn Fn(StoreContextMut<'_, U, E>, Vec<Value>) -> anyhow::Result<Vec<Value>>
131                + Send
132                + Sync,
133        >,
134    >,
135    // },
136}
137impl<U: 'static, E: wasm_runtime_layer::backend::WasmEngine> Wrapped<U, E> {
138    pub fn new(
139        base: Vec<wasm_runtime_layer::Value>,
140        rid: Arc<Interface>,
141        rs: String,
142        instance: ::wasm_runtime_layer::Instance,
143        store: ::wasm_runtime_layer::StoreContextMut<'_, U, E>,
144    ) -> Self {
145        // impl<U: 'static, E: WasmEngine> Copy for X<U, E> {}
146        // let ctx: X<U, E> = X {
147        //     base: Arc::new(Mutex::new(unsafe { std::mem::transmute(store) })),
148        // };
149        let rid2 = rid.clone();
150        let instance2 = instance.clone();
151        // let ctx2 = ctx.clone();
152        let rs2 = rs.clone();
153        let base2 = base.clone();
154        return Self {
155            rid: rid.clone(),
156            all: once(Arc::new(
157                move |mut ctx: StoreContextMut<'_, U, E>,
158                      vals: Vec<Value>|
159                      -> anyhow::Result<Vec<Value>> {
160                    // let _ = &ctx2;
161                    // let mut b = ctx2.base.lock().unwrap();
162                    let vals: Vec<_> = base.clone().into_iter().chain(vals.into_iter()).collect();
163                    let mut rets = vec![];
164                    let f = instance2.get_export(
165                        ctx.as_context_mut(),
166                        &format!("pit/{}/~{rs2}.drop", rid2.rid_str()),
167                    );
168                    let Some(Extern::Func(f)) = f else {
169                        anyhow::bail!("invalid func")
170                    };
171                    f.call(ctx.as_context_mut(), &vals, &mut rets)?;
172                    Ok(rets)
173                },
174            )
175                as Arc<
176                    dyn Fn(StoreContextMut<'_, U, E>, Vec<Value>) -> anyhow::Result<Vec<Value>>
177                        + Send
178                        + Sync,
179                >)
180            .chain(rid.methods.iter().map(|(a, b)| {
181                let rid = rid.clone();
182                let a = a.clone();
183                let instance = instance.clone();
184                let b = Arc::new(b.clone());
185                // let ctx = ctx.clone();
186                let rs = rs.clone();
187                let base = base2.clone();
188                (Arc::new(
189                    move |mut ctx: StoreContextMut<'_, U, E>,
190                          vals: Vec<Value>|
191                          -> anyhow::Result<Vec<Value>> {
192                        // let _ = &ctx;
193                        // let mut bi = ctx.base.lock().unwrap();
194                        let vals: Vec<_> =
195                            base.clone().into_iter().chain(vals.into_iter()).collect();
196                        let mut rets = vec![Value::I32(0); b.rets.len()];
197                        let f = instance.get_export(
198                            ctx.as_context_mut(),
199                            &format!("pit/{}/~{rs}/{a}", rid.rid_str()),
200                        );
201                        let Some(Extern::Func(f)) = f else {
202                            anyhow::bail!("invalid func")
203                        };
204                        f.call(ctx.as_context_mut(), &vals, &mut rets)?;
205                        Ok(rets)
206                    },
207                )
208                    as Arc<
209                        dyn Fn(StoreContextMut<'_, U, E>, Vec<Value>) -> anyhow::Result<Vec<Value>>
210                            + Send
211                            + Sync,
212                    >)
213            }))
214            .collect(),
215        };
216    }
217}
218// impl Drop for Wrapped {
219//     fn drop(&mut self) {
220//         self.all[0](vec![]).unwrap();
221//     }
222// }
223pub type RWrapped<U, E> = ::alloc::sync::Arc<Wrapped<U, E>>;
224pub extern crate anyhow;
225pub extern crate wasm_runtime_layer;
226
227pub struct W<X, U: 'static, E: WasmEngine> {
228    pub r: X,
229    pub store: Arc<StoreCell<U, E>>,
230}
231impl<U: 'static, E: WasmEngine, X: Clone> Clone for W<X, U, E> {
232    fn clone(&self) -> Self {
233        Self {
234            r: self.r.clone(),
235            store: self.store.clone(),
236        }
237    }
238}
239pub struct StoreCell<U, E: WasmEngine> {
240    pub wrapped: UnsafeCell<Store<U, E>>,
241}
242unsafe impl<
243        U: Send,
244        E: Send + wasm_runtime_layer::backend::WasmEngine + wasm_runtime_layer::backend::WasmEngine,
245    > Send for StoreCell<U, E>
246{
247}
248unsafe impl<
249        U: Sync,
250        E: Sync + wasm_runtime_layer::backend::WasmEngine + wasm_runtime_layer::backend::WasmEngine,
251    > Sync for StoreCell<U, E>
252{
253}
254impl<U, E: WasmEngine> StoreCell<U, E> {
255    pub unsafe fn get(&self) -> StoreContextMut<'_, U, E> {
256        unsafe { &mut *self.wrapped.get() }.as_context_mut()
257    }
258}
pit_rust_host_lib/lib.rs

pit_rust_host_lib/
lib.rs
