pit_rust_host_lib/
lib.rs

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