pit_rust_guest/
lib.rs

1use std::iter::once;
2
3use pit_core::{Arg, Interface, ResTy, Sig};
4use proc_macro2::TokenStream;
5use quasiquote::quasiquote;
6use quote::{format_ident, quote};
7use sha3::Digest;
8use std::io::Write;
9pub struct Opts {
10    pub root: TokenStream,
11    pub salt: Vec<u8>,
12    pub tpit: bool,
13}
14pub fn render(opts: &Opts, i: &Interface) -> TokenStream {
15    let root = &opts.root;
16    let id = i.rid_str();
17    let mut ha = sha3::Sha3_256::default();
18    write!(ha, "~{}", id);
19    ha.write(&opts.salt);
20    let ha = hex::encode(ha.finalize());
21    let id2 = format_ident!("R{}", i.rid_str());
22    let methods = i.methods.iter().map(|(a, b)| {
23        quasiquote! {
24            fn #{format_ident!("{a}")}#{render_sig(opts,root,i,b,&quote! {&mut self},false)};
25        }
26    });
27    let xref = if opts.tpit {
28        quote! {}
29    } else {
30        quasiquote!(
31            #[#root::externref::externref(crate = #{quote! {#root::externref}.to_string()})]
32        )
33    };
34    let res = if opts.tpit {
35        quote! {
36            #root::tpit_rt::Tpit
37        }
38    } else {
39        quote! {
40            #root::externref::Resource
41        }
42    };
43    let rx = if opts.tpit {
44        quote! {
45            u32
46        }
47    } else {
48        quote! {
49            &mut #res<Box<dyn #id2>>
50        }
51    };
52    let t = if opts.tpit { "t" } else { "" };
53    let impl_dyns = i.methods.iter().map(|(a, b)| {
54        quasiquote! {
55            fn #{format_ident!("{a}")}#{render_sig(opts,root,i,b,&quote! {&mut self},false)}{
56                #xref
57                #[link(wasm_import_module = #{format!("{t}pit/{}",i.rid_str())})]
58                extern "C"{
59                    #[link_name = #a]
60                    fn go #{render_sig(opts,root, i,b, &quote! {this: #rx},true)};
61                }
62                return unsafe{go(#{
63                    if opts.tpit{
64                        quote!{self.ptr()}
65                    }else{
66                        quote!{self}
67                    }
68                },#{
69                    let params = b.params.iter().enumerate().map(|(a,b)|{
70                        let mut c = format_ident!("p{a}");
71                        let mut c = quote!{
72                            #c
73                        };
74                        if let Arg::Resource { ty, nullable, take, ann } = b{
75                            if opts.tpit{
76                                if !*take{
77                                    c = quote!{
78                                        #root::tpit_rt::Tpit::summon(&mut #c)
79                                    }
80                                }
81                            }
82                        }
83                        c
84                });
85                    quote! {
86                        #(#params),*
87                    }
88                })};
89            }
90        }
91    });
92    let chains2 = i.methods.iter().map(|(a,b)|quasiquote! {
93       #xref
94        #[export_name = #{format!("{t}pit/{id}/~{ha}/{a}")}]
95        extern "C" fn #{format_ident!("{a}")}#{render_sig(opts,root,i,b,&quote! {id: u32},true)}{
96            return unsafe{&mut *(TABLE.all.get())}.get_mut(&id).unwrap().#{format_ident!("{a}")}(#{
97                let params = b.params.iter().enumerate().map(|(a,b)|{
98                    let mut c = format_ident!("p{a}");
99                    let mut c = quote!{
100                        #c
101                    };
102                    if let Arg::Resource { ty, nullable, take, ann } = b{
103                        if opts.tpit{
104                            if !*take{
105                                c = quote!{
106                                    #c.ptr()
107                                }
108                            }
109                        }
110                    }
111                    c
112                });
113                quote! {
114                    #(#params),*
115                }
116            });
117        }
118    });
119    let sb = i.to_string();
120    let sc = sb
121        .as_bytes()
122        .iter()
123        .cloned()
124        .chain(once(0u8))
125        .collect::<Vec<_>>();
126    quasiquote! {
127        pub trait #id2{
128            #(#methods)*
129        }
130            const _: () = {
131                #[link_section = ".pit-types"]
132                static SECTION_CONTENT: [u8; #{sc.len()}] = #{
133                    quote!{
134                        [#(#sc),*]
135                    }
136                };
137                fn alloc<T>(m: &mut ::std::collections::BTreeMap<u32,T>, x: T) -> u32{
138                    let mut u = 0;
139                    while m.contains_key(&u){
140                        u += 1;
141                    };
142                    m.insert(u,x);
143                    return u;
144                }
145                #[derive(Default)]
146                struct TableCell{
147                    all: std::cell::UnsafeCell<::std::collections::BTreeMap<u32,Box<dyn #id2>>>
148                };
149                unsafe impl Send for TableCell{}
150                unsafe impl Sync for TableCell{}
151                static TABLE: ::std::sync::LazyLock<TableCell> = ::std::sync::LazyLock::new(||TableCell::default());
152                impl #id2 for #res<Box<dyn #id2>>{
153                    #(#impl_dyns)*
154                }
155                #xref
156                #[export_name = #{format!("{t}pit/{id}/~{ha}.drop")}]
157                extern "C" fn _drop(a: u32){
158                    unsafe{
159                        (&mut *(TABLE.all.get())).remove(&a)
160                    };
161                }
162                #(#chains2)*
163                impl From<Box<dyn #id2>> for #res<Box<dyn #id2>>{
164                    fn from(a: Box<dyn #id2>) -> Self{
165                        #xref
166                        #[link(wasm_import_module = #{format!("pit/{}",i.rid_str())})]
167                        extern "C"{
168                            #[link_name = #{format!("~{ha}")}]
169                            fn _push(a: u32) -> #res<Box<dyn #id2>>;
170                        }
171                        return unsafe{
172                            _push(alloc(&mut *(TABLE.all.get()),a))
173                        }
174                    }
175                }
176            };
177            // mod #internal{
178            //     use super::#id2;
179            //     pub fn push(a: Box<dyn #id2>) -> #root::externref::Resource<Box<dyn #id2>>{
180            //         return unsafe{
181            //             _push(Box::into_raw(Box::new(a)))
182            //         }
183            //     }
184            // }
185    }
186}
187pub fn render_sig(
188    opts: &Opts,
189    root: &TokenStream,
190    base: &Interface,
191    s: &Sig,
192    self_: &TokenStream,
193    ffi: bool,
194) -> TokenStream {
195    let params = s
196        .params
197        .iter()
198        .map(|a| render_ty(opts, root, base, a, ffi))
199        .enumerate()
200        .map(|(a, b)| quasiquote!(#{format_ident!("p{a}")} : #b));
201    let params = once(self_).cloned().chain(params);
202    let rets = s.rets.iter().map(|a| render_ty(opts, root, base, a, ffi));
203    quote! {
204        (#(#params),*) -> (#(#rets),*)
205    }
206}
207pub fn render_ty(
208    opts: &Opts,
209    root: &TokenStream,
210    base: &Interface,
211    p: &Arg,
212    ffi: bool,
213) -> TokenStream {
214    match p {
215        Arg::I32 => quote! {
216            u32
217        },
218        Arg::I64 => quote! {
219            u64
220        },
221        Arg::F32 => quote! {
222            f32
223        },
224        Arg::F64 => quote! {
225            f64
226        },
227        Arg::Resource {
228            ty,
229            nullable,
230            take,
231            ann,
232        } => {
233            if !opts.tpit {
234                let ty = match ty {
235                    ResTy::Of(a) => quasiquote! {
236                        #root::externref::Resource<Box<dyn #{format_ident!("R{}",hex::encode(a))}>>
237                    },
238                    ResTy::None => quote! {
239                        #root::externref::Resource<()>
240                    },
241                    ResTy::This => quasiquote! {
242                        #root::externref::Resource<Box<dyn #{format_ident!("R{}",base.rid_str())}>>
243                    },
244                    _ => todo!()
245                };
246                let ty = if *nullable {
247                    quote! {Option<#ty>}
248                } else {
249                    ty
250                };
251                let ty = if *take {
252                    ty
253                } else {
254                    quote! {&mut #ty}
255                };
256                ty
257            } else {
258                let ty = match ty {
259                    ResTy::Of(a) => quasiquote! {
260                        #root::tpit_rt::Tpit<Box<dyn #{format_ident!("R{}",hex::encode(a))}>>
261                    },
262                    ResTy::None => quote! {
263                        #root::tpit_rt::Tpit<()>
264                    },
265                    ResTy::This => quasiquote! {
266                        #root::tpit_rt::Tpit<Box<dyn #{format_ident!("R{}",base.rid_str())}>>
267                    },
268                    _ => todo!()
269                };
270                let mut ty = if *take {
271                    ty
272                } else {
273                    quote! {&mut #ty}
274                };
275                if *take && ffi {
276                    ty = quote! {u32}
277                }
278                ty
279            }
280        }
281        _ => todo!()
282        // Arg::Func(_) => todo!()
283    }
284}