pit_rust_host/
lib.rs

1use std::iter::once;
2
3use pit_core::{Arg, Interface, ResTy, Sig};
4use proc_macro2::{Span, TokenStream};
5use quasiquote::quasiquote;
6use quote::{format_ident, quote, ToTokens};
7use syn::{spanned::Spanned, Ident, Index};
8pub struct Opts {
9    pub guest: Option<pit_rust_guest::Opts>,
10}
11pub fn render(root: &TokenStream, i: &Interface, opts: &Opts) -> TokenStream {
12    let id = format_ident!("B{}", i.rid_str());
13    // let internal = format_ident!("{id}_utils");
14    let methods = i.methods.iter().map(|(a, b)| {
15        quasiquote! {
16            fn #{format_ident!("{a}")}#{render_sig(root,b,&quote! {&self},quote!{
17                ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>
18            })}
19        }
20    });
21    let impls = i.methods.iter().enumerate().map(|(c, (a, b))| {
22        let init = b
23            .params
24            .iter()
25            .enumerate()
26            .map(|(pi, a)| render_new_val(root, a, quasiquote! {#{format_ident!("p{pi}")}}));
27        let init =
28            once(quote! {#root::wasm_runtime_layer::Value::I32(self.base as i32)}).chain(init);
29        let fini = b.rets.iter().enumerate().map(|(ri, r)| {
30            quasiquote! {
31                #{render_base_val(root, r, quote! {&rets[#ri]})}
32            }
33        });
34        quasiquote! {
35            fn #{format_ident!("{a}")}#{render_sig(root,b,&quote! {&self},quote!{
36                ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>
37            })}{
38                let a = self.all[#{c+1}].clone();
39                let args = vec![#(#init),*];
40                let mut rets = a(ctx,args);
41                return Ok((#(#fini),*))
42            }
43        }
44    });
45    let injects = i.methods.iter().map(|(a,b)|{
46        let init = b.params.iter().enumerate().map(|(pi,a)|quasiquote!{#{render_base_val(root, a, quote! {&args[#pi + 1]})}});
47        let fini = b.rets.iter().enumerate().map(|(ri,r)|{
48            quasiquote!{
49                #{render_new_val(root, r, quasiquote! {r . #{Index{index: ri as u32,span: root.span()}}})}
50            }
51        });
52        quasiquote!{
53            let r = a.clone();
54            Arc::new(ctx,move|ctx,args|{
55                let r = r.#{format_ident!("{a}")}(ctx,#(#init),*)
56                Ok(vec![#(#fini),*])
57            })
58    }});
59    quasiquote! {
60        pub trait #id<U: 'static,E: #root::wasm_runtime_layer::backend::WasmEngine>{
61            #(#methods)*
62            unsafe fn finalize(&self, ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>) -> #root::anyhow::Result<()>;
63        }
64        const _: () = {
65            impl<U: 'static,E: #root::wasm_runtime_layer::backend::WasmEngine> #id<U,E> for #root::RWrapped<U,E>{
66                #(#impls)*
67                unsafe fn finalize(&self, ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>) -> #root::anyhow::Result<()>{
68                    self.all[0](ctx,vec![])?;
69                    Ok(())
70                }
71            }
72            impl<U: 'static,E: #root::wasm_runtime_layer::backend::WasmEngine> From<Arc<dyn $id<U,E>> for #root::RWrapped<U,E>{
73                fn from(a: Arc<dyn $id<U,E>>) -> Self{
74                    Self{
75                        rid: Arc::new(#root::pit_core::Interface::parse_interface(#{i.to_string()}).ok().unwrap()),
76                        all: vec![#{
77                            let all = once(quasiquote!{
78                                let r = a.clone();
79                                unsafe{
80                                    Arc::new(move|ctx,args|{r.finalize(ctx);Ok(vec![])})
81                                }
82                            }).chain(injects);
83
84                            quote!{
85                                #(#all),*
86                            }
87                        }]
88                    }
89                }
90            }
91            #{match opts.guest.as_ref(){
92                None=>quote!{},
93                Some(g) => proxy(root,i,opts,g), 
94            }}
95        }
96    }
97}
98pub fn proxy(root: &TokenStream, i: &Interface, opts: &Opts, g: &pit_rust_guest::Opts) -> TokenStream{
99    let id = format_ident!("B{}", i.rid_str());
100    let pid = format_ident!("R{}", i.rid_str());
101    let root2 = &g.root;
102    let res = if g.tpit {
103        quote! {
104            #root2::tpit_rt::Tpit
105        }
106    } else {
107        quote! {
108            #root2::externref::Resource
109        }
110    };
111    let impl_guest = i.methods.iter().map(|(a, b)| {
112        quasiquote! {
113            fn #{format_ident!("{a}")}#{pit_rust_guest::render_sig(g,root,i,b,&quote! {&mut self},false)}{
114                let ctx = unsafe{
115                    self.get()
116                };
117                let r = self.r.#{format_ident!("{a}")}(ctx,#{
118                    let params = b.params.iter().enumerate().map(|(a,b)|{
119                        let mut c = format_ident!("p{a}");
120                        let mut c = quote!{
121                            #c
122                        };
123                        if let Arg::Resource { ty, nullable, take, ann } = b{
124                            c = quote!{Box::new(#root::W{
125                                r: #root::RWrapped<U,E>::from(Arc::new(#c)),
126                                store: self.store.clone()
127                            }).into()};
128                            if !take{
129                                c = quote!{&mut #c};
130                            };
131                        }
132                        c
133                });
134                    quote! {
135                        #(#params),*
136                    }
137                }).unwrap();
138                return (#{
139                    let xs = b.rets.iter().enumerate().map(|(a,b)|{
140                        let mut c = Index{
141                            index: a as u32,
142                            span: Span::call_site()
143                        };
144                        let mut c = quote!{
145                            r.#c
146                        };
147                        if let Arg::Resource { ty, nullable, take, ann } = b{
148                            c = quote!{
149                                #root::RWrapped<U,E>::from(Compat(::std::cell::UnsafeCell::new(Some(#c)))) 
150                            }
151                        }
152                        c
153                    });
154                    quote!{
155                        #(#xs),*
156                    }
157                })
158            };
159
160        }
161    });
162    let impl_host = i.methods.iter().enumerate().map(|(c, (a, b))| {
163        let params = b.params.iter().enumerate().map(|(a,b)|{
164            let mut c = format_ident!("p{a}");
165            let mut c = quote!{
166                #c
167            };
168            if let Arg::Resource { ty, nullable, take, ann } = b{
169                    c = quote!{
170                        #root::RWrapped<U,E>::from(Compat(::std::cell::UnsafeCell::new(Some(#c)))) 
171                    };
172                if !take{
173                    c = quote!{&mut #c};
174                };
175            }
176            c
177    });
178    let rets = b.rets.iter().enumerate().map(|(a,b)|{
179        let mut c = Index{
180            index: a as u32,
181            span: Span::call_site()
182        };
183        let mut c = quote!{
184            r.#c
185        };
186        if let Arg::Resource { ty, nullable, take, ann } = b{
187            c = quote!{Box::new(#root::W{
188                r: #root::RWrapped<U,E>::from(Arc::new(#c)),
189                store: Arc::new(#root::StoreCell{
190                    wrapped: ::std::cell::UnsafeCell::new(#root::wasm_runtime_layer::Store::new(ctx.engine(),ctx.data().clone()))
191                })
192            }).into()};
193        }
194        c
195    });
196        quasiquote! {
197            fn #{format_ident!("{a}")}#{render_sig(root,b,&quote! {&self},quote!{
198                ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>
199            })}{
200                let x = unsafe{
201                    &mut *(self.0.get())
202                }.as_mut().unwrap().#{format_ident!("{a}")}(#(#params),*);
203
204                Ok((#(#rets),*))
205            }
206        }
207    });
208    quasiquote!{
209        struct Compat<T>(::std::cell::UnsafeCell<Option<T>>);
210        impl<U: 'static + Clone,E: #root::wasm_runtime_layer::backend::WasmEngine> #pid for #root::W<#root::RWrapped<U,E>,U,E>{
211            #(#impl_guest)*
212        }
213        impl<U: 'static + Clone,E: #root::wasm_runtime_layer::backend::WasmEngine, X: #pid> #id<U,E> for Compat<X>{
214            #(#impl_host)*
215            unsafe fn finalize(&self, ctx: #root::wasm_runtime_layer::StoreContextMut<'_,U,E>) -> #root::anyhow::Result<()>{
216                let x = unsafe{
217                    &mut *(self.0.get())
218                }.take();
219                let Some(x) = x else{
220                    return Err(#root::anyhow::anyhow!("double finalized"))
221                }
222                Ok(())
223            }
224        }
225        impl From<Box<dyn #pid>> for Arc<dyn #id>{
226            fn from(p: Box<dyn #pid>) -> Self{
227                return Arc::new(Compat(::std::cell::UnsafeCell::new(Some(#res::from(p)))))
228            }
229        }
230    }
231}
232pub fn render_sig(
233    root: &TokenStream,
234    s: &Sig,
235    self_: &TokenStream,
236    s2: TokenStream,
237) -> TokenStream {
238    let params = s
239        .params
240        .iter()
241        .map(|a| render_ty(root, a))
242        .enumerate()
243        .map(|(a, b)| quasiquote!(#{format_ident!("p{a}")} : #b));
244    let params = once(self_).cloned().chain(once(s2)).chain(params);
245    let rets = s.rets.iter().map(|a| render_ty(root, a));
246    quote! {
247        (#(#params),*) -> #root::anyhow::Result<(#(#rets),*)>
248    }
249}
250pub fn render_blit(root: &TokenStream, p: &Arg) -> TokenStream {
251    match p {
252        Arg::I32 => quote! {
253            #root::wasm_runtime_layer::ValueType::I32
254        },
255        Arg::I64 => quote! {
256            #root::wasm_runtime_layer::ValueType::I64
257        },
258        Arg::F32 => quote! {
259            #root::wasm_runtime_layer::ValueType::F32
260        },
261        Arg::F64 => quote! {
262            #root::wasm_runtime_layer::ValueType::F64
263        },
264        Arg::Resource { .. } => quote! {
265            #root::wasm_runtime_layer::ValueType::ExternRef
266        },
267        _ => todo!()
268    }
269}
270pub fn render_blit_sig(root: &TokenStream, s: &Sig) -> TokenStream {
271    quasiquote! {
272        #root::wasm_runtime_layer::FuncType::new([#{
273            let p = s.params.iter().map(|p|render_blit(root, p));
274            let p = once(quote! {#root::wasm_runtime_layer::ValueType::ExternRef}).chain(p);
275            quote! {
276                #(#p),*
277            }
278        }].into_iter(),[#{            let p = s.rets.iter().map(|p|render_blit(root, p));
279            quote! {
280                #(#p),*
281            }}].into_iter())
282    }
283}
284pub fn render_ty(root: &TokenStream, p: &Arg) -> TokenStream {
285    match p {
286        Arg::I32 => quote! {
287            u32
288        },
289        Arg::I64 => quote! {
290            u64
291        },
292        Arg::F32 => quote! {
293            f32
294        },
295        Arg::F64 => quote! {
296            f64
297        },
298        Arg::Resource {
299            ty,
300            nullable,
301            take,
302            ann,
303        } => match ty {
304            ResTy::None => quote! {
305                #root::wasm_runtime_layer::ExternRef
306            },
307            _ => {
308                let a = quasiquote! {
309                    ::std::sync::Arc<#root::Wrapped<U,E>>
310                };
311                if *nullable {
312                    quote! {Option<#a>}
313                } else {
314                    a
315                }
316            }
317        },
318        _ => todo!()
319    }
320}
321pub fn render_base_val(root: &TokenStream, p: &Arg, x: TokenStream) -> TokenStream {
322    let v = match p {
323        Arg::I32 => quote! {
324            let #root::wasm_runtime_layer::Value::I32(t) = #x else{
325                #root::anyhow::bail!("invalid param")
326            }
327        },
328        Arg::I64 => quote! {
329            let #root::wasm_runtime_layer::Value::I64(t) = #x else{
330                #root::anyhow::bail!("invalid param")
331            }
332        },
333        Arg::F32 => quote! {
334            let #root::wasm_runtime_layer::Value::F32(t) = #x else{
335                #root::anyhow::bail!("invalid param")
336            }
337        },
338        Arg::F64 => quote! {
339            let #root::wasm_runtime_layer::Value::F64(t) = #x else{
340                #root::anyhow::bail!("invalid param")
341            }
342        },
343        Arg::Resource {
344            ty,
345            nullable,
346            take,
347            ann,
348        } => {
349            let mut a = quote! {
350                let #root::wasm_runtime_layer::Value::ExternRef(t) = #x else{
351                    #root::anyhow::bail!("invalid param")
352                }
353            };
354            if !matches!(ty, ResTy::None) {
355                quasiquote!{
356                    let t = match t.downcast::<'_,'_,                ::std::sync::Arc<#root::Wrapped<U,E>>,U,E>(ctx){
357                            Ok(t) => Arc::new(t.clone()),
358                            Err(_) =>                     #root::anyhow::bail!("invalid param")
359                        }
360
361                }.to_tokens(&mut a);
362            }
363            if !*nullable {
364                quote! {
365                    let t = match t{
366                        Some(a) => a,
367                        None => #root::anyhow::bail!("invalid param")
368                    }
369                }
370                .to_tokens(&mut a)
371            }
372            a
373        },
374        _ => todo!()
375    };
376    quote! {
377        {
378            #v;t
379        }
380    }
381}
382pub fn render_new_val(root: &TokenStream, p: &Arg, t: TokenStream) -> TokenStream {
383    match p {
384        Arg::I32 => quote! {
385            #root::wasm_runtime_layer::Value::I32(#t)
386        },
387        Arg::I64 => quote! {
388            #root::wasm_runtime_layer::Value::I64(#t)
389        },
390        Arg::F32 => quote! {
391            #root::wasm_runtime_layer::Value::F32(#t)
392        },
393        Arg::F64 => quote! {
394            #root::wasm_runtime_layer::Value::F64(#t)
395        },
396        Arg::Resource {
397            ty,
398            nullable,
399            take,
400            ann,
401        } => {
402            let tq = |t: TokenStream| {
403                quasiquote! {
404                    {
405                        let t = #t;
406                        #{match ty{
407                            ResTy::None => quote! {                    #root::wasm_runtime_layer::ExternRef::new(ctx,t)},
408                            _ => quote! {
409                                match t.to_any().downcast_ref::<::std::sync::Arc<#root::Wrapped<U,E>>>(){
410                                    None =>                     #root::wasm_runtime_layer::ExternRef::new(ctx,t),
411                                    Some(t) => #root::wasm_runtime_layer::ExternRef::new(ctx,t.clone()),
412                                }
413                            }
414                        }}
415                    }
416                }
417            };
418            if !*nullable {
419                quasiquote! {
420                    #root::wasm_runtime_layer::Value::ExternRef(Some(#{match ty{
421                        ResTy::None => t,
422                        _ => tq(t)
423                    }}))
424                }
425            } else {
426                quasiquote! {
427                    #root::wasm_runtime_layer::Value::ExternRef(#{match ty{
428                        ResTy::None => t,
429                        _ => quasiquote! {#t.map(|t|#{tq(quote! {t})})}
430                    }})
431                }
432            }
433        },
434        _ => todo!()
435    }
436}