wars_pit_plugin/
lib.rs

1use std::{
2    borrow::Cow,
3    collections::{BTreeMap, BTreeSet},
4    convert::Infallible,
5    f32::consts::E,
6    iter::once,
7    sync::{Arc, Mutex, OnceLock},
8};
9use wars::*;
10
11use pit_core::{Arg, Interface};
12use proc_macro2::{Span, TokenStream};
13use quasiquote::quasiquote;
14use quote::{format_ident, quote, ToTokens};
15use relooper::{reloop, BranchMode, ShapedBlock};
16use sha3::Digest;
17use syn::{Ident, Lifetime};
18use portal_pc_waffle::{
19    cfg::CFGInfo, entity::EntityRef, Block, BlockTarget, Export, ExportKind, Func, ImportKind,
20    Memory, Module, Operator, Signature, SignatureData, Type, Value,
21};
22#[derive(Default)]
23pub struct PitPlugin {
24    pub tpit: OnceLock<BTreeSet<pit_core::Interface>>,
25    pub extra: Vec<Arc<dyn PitPluginPlugin>>,
26}
27
28pub trait PitPluginPlugin {
29    fn pre(&self, pit: &PitPlugin, module: &mut Opts<Module<'static>>) -> anyhow::Result<()> {
30        return Ok(());
31    }
32    fn choose_type(&self, opts: &Opts<Module<'static>>) -> anyhow::Result<Option<TokenStream>>;
33    fn emit_method(
34        &self,
35        opts: &Opts<Module<'static>>,
36        idx: usize,
37        i: &Interface,
38        s: &str,
39        value: TokenStream,
40        params: &[TokenStream],
41    ) -> anyhow::Result<TokenStream>;
42    fn emit_drop(
43        &self,
44        opts: &Opts<Module<'static>>,
45        idx: usize,
46        value: TokenStream,
47    ) -> anyhow::Result<TokenStream> {
48        return Ok(quasiquote!(#{opts.fp()}::ret(Ok(()))));
49    }
50    fn post(
51        &self,
52        parent: &PitPlugin,
53        idx: usize,
54        opts: &Opts<Module<'static>>,
55    ) -> anyhow::Result<TokenStream>;
56}
57
58impl PitPlugin {
59    pub fn tpit(&self, opts: &Opts<Module<'static>>) -> &BTreeSet<Interface> {
60        return self.tpit.get_or_init(|| {
61            pit_patch::get_interfaces(&opts.module)
62                .unwrap()
63                .into_iter()
64                .collect()
65        });
66    }
67    pub fn host_tpit(&self, opts: &Opts<Module<'static>>) -> anyhow::Result<TokenStream> {
68        let mut a;
69        let root = &opts.crate_path;
70        a = quote! {#root::Infallible};
71        for e in self.extra.iter() {
72            if let Some(c) = e.choose_type(opts)? {
73                a = quote! {
74                    #root::Either<#c,#a>
75                };
76            }
77        }
78        return Ok(a);
79    }
80    pub fn apply_host(
81        &self,
82        mut x: TokenStream,
83        opts: &Opts<Module<'static>>,
84        i: &Interface,
85        s: &str,
86        params: &[TokenStream],
87    ) -> anyhow::Result<TokenStream> {
88        let root = &opts.crate_path;
89        for (idx, e) in self.extra.iter().enumerate() {
90            if let Some(c) = e.choose_type(opts)? {
91                x = quasiquote! {
92                    match host{
93                        #root::Either::Right(host) => #x,
94                        #root::Either::Left(host) => #{e.emit_method(opts,idx, i, s, quote! {host}, params)?},
95                    }
96                };
97            }
98        }
99        return Ok(x);
100    }
101    pub fn apply_drop(
102        &self,
103        mut x: TokenStream,
104        opts: &Opts<Module<'static>>,
105    ) -> anyhow::Result<TokenStream> {
106        let root = &opts.crate_path;
107        for (idx, e) in self.extra.iter().enumerate() {
108            if let Some(c) = e.choose_type(opts)? {
109                x = quasiquote! {
110                    match host{
111                        #root::Either::Right(host) => #x,
112                        #root::Either::Left(host) => #{e.emit_drop(opts,idx, quote! {host})?},
113                    }
114                };
115            }
116        }
117        return Ok(x);
118    }
119    pub fn wrap(
120        &self,
121        opts: &Opts<Module<'static>>,
122        mut x: TokenStream,
123    ) -> anyhow::Result<TokenStream> {
124        let root = &opts.crate_path;
125        for e in self.extra.iter() {
126            if let Some(c) = e.choose_type(opts)? {
127                x = quasiquote! {
128                    #root::Either::Right(#x)
129                };
130            }
131        }
132        return Ok(x);
133    }
134}
135impl Plugin for PitPlugin {
136    fn exref_bounds(&self, opts: &Opts<Module<'static>>) -> anyhow::Result<Option<TokenStream>> {
137        let root = &opts.crate_path;
138        Ok(Some(
139            quasiquote! {From<#root::Pit<Vec<#{opts.fp()}::Value<Self>>,#{self.host_tpit(opts)?}>> + TryInto<#root::Pit<Vec<#{opts.fp()}::Value<Self>>,#{self.host_tpit(opts)?}>>},
140        ))
141    }
142    fn pre(&self, module: &mut Opts<Module<'static>>) -> anyhow::Result<()> {
143        for p in self.extra.iter() {
144            p.pre(self, module)?;
145        }
146        Ok(())
147    }
148
149    fn import(
150        &self,
151        opts: &Opts<Module<'static>>,
152        module: &str,
153        name: &str,
154        params: Vec<TokenStream>,
155    ) -> anyhow::Result<Option<TokenStream>> {
156        let root = &opts.crate_path;
157        let mut params = params.into_iter();
158        if let Some(i) = module.strip_prefix("pit/") {
159            let x: [u8; 32] = hex::decode(i).unwrap().try_into().unwrap();
160            if let Some(s) = name.strip_prefix("~") {
161                let s = {
162                    let mut h = sha3::Sha3_256::default();
163                    h.update(s);
164                    h.finalize()
165                };
166                return Ok(Some(quasiquote! {
167                    #{opts.fp()}::ret(Ok(#{opts.fp()}::Value::<C>::ExternRef(#root::Pit::Guest{
168                        id: [#(#x),*],
169                        x: #{opts.fp()}::CoeVec::coe(#root::tuple_list::tuple_list!(#(#params),*)),
170                        s: [#(#s),*],
171                    }.into())))
172                }));
173            }
174
175            let mut f = params.next().unwrap();
176            // let id = format_ident!("{}", bindname(&format!("pit/{i}/~{PIT_NS}/{name}")));
177            // ctx.#id(x.x,#(#params),*)
178            let params = params.collect::<Vec<_>>();
179            let cases = opts
180                .module
181                .exports
182                .iter()
183                .filter_map(|x| x.name.strip_prefix(&format!("pit/{i}/~")))
184                .filter_map(|x| x.strip_suffix(&format!("/{name}")))
185                .map(|s| {
186                    let id = format_ident!("{}", bindname(&format!("pit/{i}/~{s}/{name}")));
187                    let s = {
188                        let mut h = sha3::Sha3_256::default();
189                        h.update(s);
190                        h.finalize()
191                    };
192                    quasiquote! {
193                        [#(#s),*] => {
194                            let mut y = #{opts.fp()}::CoeVec::coe(#root::tuple_list::tuple_list!(#(#params),*));
195                            y.extend(&mut x.clone());
196                            ctx.#id(#{opts.fp()}::CoeVec::uncoe(y))
197                        }
198                    }
199                });
200            let interface = self.tpit(opts).iter().find(|a| a.rid() == x);
201            let meth = interface.and_then(|a| a.methods.get(name));
202            return Ok(Some(quasiquote! {
203                'a: {
204                    let x = #f;
205                    let #{opts.fp()}::Value::<C>::ExternRef(x) = x else{
206                        break 'a #{opts.fp()}::ret(Err(#root::_rexport::anyhow::anyhow!("not an externref")))
207                    };
208                    let Ok(x) = x.try_into() else{
209                        break 'a #{opts.fp()}::ret(Err(#root::_rexport::anyhow::anyhow!("not a pit externref")))
210                    };
211                    match x{
212                        #root::Pit::Guest{s,x,id} => match s{
213                            #(#cases),*,
214                            _ => break 'a #{opts.fp()}::ret(Err(#root::_rexport::anyhow::anyhow!("invalid target")))
215                        },
216                        #root::Pit::Host{host} => #{let t = quote!{match host{}};self.apply_host(t,opts,interface.unwrap(),name,&params)?}
217            _ => todo!()
218                    }
219                }
220            }));
221        }
222        if module == "pit" && name == "drop" {
223            let mut f = params.next().unwrap();
224            let cases = opts
225                .module
226                .exports
227                .iter()
228                .filter_map(|x| {
229                    let x = x.name.as_str();
230                    let x = x.strip_prefix("pit/")?;
231                    let (a, x) = x.split_once("/~")?;
232                    let s = x.strip_suffix(".drop")?;
233                    return Some((a, s));
234                })
235                .map(|(a, s)| {
236                    let x = hex::decode(a).unwrap();
237                    let id = format_ident!("{}", bindname(&format!("pit/{a}/~{s}.drop")));
238                    let s = {
239                        let mut h = sha3::Sha3_256::default();
240                        h.update(s);
241                        h.finalize()
242                    };
243                    // let id = format_ident!(
244                    //     "{}",
245                    //     bindname(&format!("pit/{}/~{PIT_NS}.drop", i.rid_str()))
246                    // ); ctx.#id(x.x)
247                    quasiquote!(
248                        ([#(#x),*],[#(#s),*]) => ctx.#id(#{opts.fp()}::CoeVec::uncoe(x))
249                    )
250                });
251            return Ok(Some(quasiquote! {
252                'a: {
253                    let x = #f;
254                    let #{opts.fp()}::Value::<C>::ExternRef(x) = x else{
255                        break 'a #{opts.fp()}::ret(Ok(()));
256                    };
257                    if let Ok(x) = x.try_into(){
258                        match x{
259                            #root::Pit::Guest{s,x,id} => => break 'a match (id,s){
260                                #(#cases),*,
261                                _ => #{opts.fp()}::ret(Ok(()))
262                            },
263                            #root::Pit::Host{host} => break 'a #{self.apply_drop(quasiquote!(#{opts.fp()}::ret(Ok(()))),opts)?}
264                        }
265                    }else{
266                        break 'a #{opts.fp()}::ret(Ok(()))
267                    }
268                }
269            }));
270        };
271        return Ok(None);
272    }
273
274    fn post(&self, opts: &Opts<Module<'static>>) -> anyhow::Result<TokenStream> {
275        let root = &opts.crate_path;
276        let name = opts.name.clone();
277
278        let bs = self
279            .extra
280            .iter()
281            .enumerate()
282            .map(|(i, x)| x.post(self, i, opts))
283            .collect::<anyhow::Result<Vec<_>>>()?;
284        return Ok(quote! {
285            #(#bs)*
286        });
287    }
288}
289pub fn indexed_lookup(root: &TokenStream, i: usize, a: TokenStream) -> TokenStream {
290    if i == 0 {
291        return quote! {#root::Either::Right(#a)};
292    }
293    return quasiquote!(#root::Either::Left(#{indexed_lookup(root, i - 1, a)}));
294}
295pub struct Passthrough {}
296impl PitPluginPlugin for Passthrough {
297    fn choose_type(&self, opts: &Opts<Module<'static>>) -> anyhow::Result<Option<TokenStream>> {
298        Ok(Some(quasiquote!(
299            #{opts.host_tpit()}
300        )))
301    }
302
303    fn emit_method(
304        &self,
305        opts: &Opts<Module<'static>>,
306        idx: usize,
307        i: &Interface,
308        s: &str,
309        value: TokenStream,
310        params: &[TokenStream],
311    ) -> anyhow::Result<TokenStream> {
312        let meth = i.methods.get(s);
313        let i2 = i.rid_str();
314        let root = &opts.crate_path;
315        Ok(match opts.roots.get("tpit_rt") {
316            None => quote! {
317                match #value{
318
319                }
320            },
321            Some(r) => quasiquote! {
322                    match #value{
323                        host => {
324                    let casted = unsafe{
325                        host.cast::<Box<dyn #{format_ident!("R{}",i2)}>>()
326                    };
327                    let a = casted.#{format_ident!("{s}")}(#{
328                        let p = params.iter().zip(meth.unwrap().params.iter()).map(|(x,y)|match y{
329                            Arg::Resource { ty, nullable, take, ann } => quasiquote!{
330                                Box::new(Shim{wrapped: ctx, x: #x}).into()
331                            },
332                            _ => quote!{
333                                #x
334                            }
335                        });
336
337                        quote!{
338                            #(#p),*
339                        }
340                    });
341                    break 'a #{opts.fp()}::ret(Ok(#root::tuple_list::tuple_list!(#{
342                        let r = meth.unwrap().rets.iter().enumerate().map(|(i,r)|{
343                            let i = syn::Index{index: i as u32, span: Span::call_site()};
344                            let i = quote!{
345                                a.#i
346                            };
347                            match r{
348                                Arg::Resource { ty, nullable, take, ann } => quote!{
349                                    #{opts.fp()}::Value::<C>::ExternRef(#root::Pit::Host{host: #{indexed_lookup(root,idx,quasiquote!{
350                                        unsafe{i.cast()}
351                                    })}})
352                                },
353                                _ => i
354                            }
355                        });
356
357                        quote!{
358                            #(#r),*
359                        }
360                    })));
361                }
362            }
363                },
364        })
365    }
366
367    fn post(
368        &self,
369        parent: &PitPlugin,
370        idx: usize,
371        opts: &Opts<Module<'static>>,
372    ) -> anyhow::Result<TokenStream> {
373        let root = &opts.crate_path;
374        let name = opts.name.clone();
375        let a = match opts.roots.get("tpit_rt") {
376            None => quote! {},
377            Some(tpit_rt) => quasiquote! {
378                impl<T: #name + ?Sized> Into<#tpit_rt::Tpit<()>> for Box<Shim<T>>{
379                    fn into(self) -> #tpit_rt::Tpit<()>{
380                        if let #{opts.fp()}::Value::<T>::ExternRef(e) = *self{
381                            if let Ok(a) = e.try_into(){
382                                if let #root::Pit::Host{host} = a{
383                                    return host;
384                                }
385                            }
386                        }
387                        Default::default()
388                    }
389                }
390                impl<T: #name + ?Sized> Drop for Shim<T>{
391                    fn drop(&mut self){
392                        let ctx = unsafe{
393                            &mut *self.wrapped
394                        };
395                        #root::rexport::tramp::tramp(#{opts.import("pit","drop",once(quote!{
396                            self.x.clone()
397                        }))?})
398                    }
399                }
400                #{
401                    let a = parent.tpit(opts).iter().map(|i|{
402                        let tname = format_ident!("R{}",i.rid_str());
403                        let meths = i.methods.iter().map(|(a,b)|
404                            Ok(quasiquote!{
405                                fn #{format_ident!("{a}")}#{pit_rust_guest::render_sig(&pit_rust_guest::Opts { root: tpit_rt.clone(), salt: vec![], tpit: true },&tpit_rt.clone(),i,b,&quote! {&mut self},false)}{
406                                    let ctx = unsafe{
407                                        &mut *self.wrapped
408                                    };
409                                    let res = #{opts.import(&format!("pit/{}",i.rid_str()),&format!("{a}"),once(Ok(quote!{self.x.clone()})).chain(b.params.iter().enumerate().map(|(i,p)|{
410                                        let i = format_ident!("p{i}");
411                                        Ok(match p{
412                                            Arg::Resource{ty,nullable,take,ann} => {
413                                                quote!{
414                                                    #{opts.fp()}::Value::<C>::ExternRef(Pit::Host{host:#{indexed_lookup(opts,i,dxquasiquote!{
415                                                        unsafe{
416                                                            #i.cast()
417                                                        }
418                                                    })?}}.into())
419                                                }
420                                            }
421                                            _ => quote!{
422                                                #i
423                                            }
424                                        })
425                                    })).collect::<anyhow::Result<Vec<_>>>()?.into_iter())?};
426                                    let res = #root::rexport::tramp::tramp(res).unwrap().into_tuple()
427                                    ;
428                                    #{                                        let r = b.rets.iter().enumerate().map(|(i,r)|{
429                                        let i = syn::Index{index: i as u32, span: Span::call_site()};
430                                        let i = quote!{
431                                            res.#i
432                                        };
433                                        match r{
434                                            Arg::Resource { ty, nullable, take, ann } => quote!{
435                                                Box::new(Shim{wrapped:self.wrapped,x: #i}).into()
436                                            },
437                                            _ => i
438                                        }
439                                    });
440
441                                    quote!{
442                                        #(#r),*
443                                    }}
444                                }
445                        })).collect::<anyhow::Result<Vec<_>>>()?;
446                        Ok(quote!{
447                            impl<C: #name + ?Sized> #tname for Shim<C>{
448                                #(#meths),*
449                            }
450                        })
451                    }).collect::<anyhow::Result<Vec<_>>>()?;
452                    quote!{
453                        #(#a)*
454                    }
455                }
456            },
457        };
458        Ok(a)
459    }
460}
wars_pit_plugin/lib.rs

wars_pit_plugin/
lib.rs
