pit_patch/
tpit.rs

1use alloc::borrow::ToOwned;
2use alloc::collections::{BTreeMap, BTreeSet};
3use alloc::format;
4use alloc::vec;
5use alloc::vec::Vec;
6use core::{
7    iter::once,
8    mem::{replace, take},
9};
10
11use anyhow::Context;
12use pit_core::{Arg, ResTy};
13use portal_pc_waffle::{
14    util::*, Block, BlockTarget, Export, ExportKind, Func, FuncDecl, FunctionBody, ImportKind,
15    Module, Operator, SignatureData, Table, TableData, Type, Value, WithNullable,
16};
17
18use crate::get_interfaces;
19use crate::tutils::{talloc, tfree};
20use crate::util::{add_op, to_waffle_sig, waffle_funcs};
21
22// use waffle_ast::{add_op, results_ref_2};
23
24// use crate::{
25//     get_interfaces,
26//     util::{talloc, tfree, to_waffle_sig, waffle_funcs},
27// };
28
29pub fn shim(
30    retref: bool,
31    f: &mut FunctionBody,
32    mut k: Block,
33    r: &Arg,
34    mut v: Value,
35    talloc: Func,
36    tfree: Func,
37    table: Table,
38) -> anyhow::Result<(Value, Block)> {
39    let Arg::Resource {
40        ty,
41        nullable,
42        take,
43        ann,
44    } = r
45    else {
46        return Ok((v, k));
47    };
48    let end = f.add_block();
49    let ep = f.add_blockparam(
50        end,
51        if retref {
52            portal_pc_waffle::Type::Heap(WithNullable {
53                nullable: true,
54                value: portal_pc_waffle::HeapType::ExternRef,
55            })
56        } else {
57            Type::I32
58        },
59    );
60    // if *nullable {
61    let s = if retref {
62        let a = add_op(f, &[], &[Type::I32], Operator::I32Const { value: 0 });
63        f.append_to_block(k, a);
64        let a = add_op(f, &[a, v], &[Type::I32], Operator::I32Eq);
65        f.append_to_block(k, a);
66        a
67    } else {
68        let a = add_op(f, &[v], &[Type::I32], Operator::RefIsNull);
69        f.append_to_block(k, a);
70        a
71    };
72    let n = f.add_block();
73    if retref {
74        let a = add_op(
75            f,
76            &[],
77            &[portal_pc_waffle::Type::Heap(WithNullable {
78                nullable: true,
79                value: portal_pc_waffle::HeapType::ExternRef,
80            })],
81            Operator::RefNull {
82                ty: portal_pc_waffle::Type::Heap(WithNullable {
83                    nullable: true,
84                    value: portal_pc_waffle::HeapType::ExternRef,
85                }),
86            },
87        );
88        f.append_to_block(n, a);
89        f.set_terminator(
90            n,
91            portal_pc_waffle::Terminator::Br {
92                target: BlockTarget {
93                    block: end,
94                    args: vec![a],
95                },
96            },
97        )
98    } else {
99        let a = add_op(f, &[], &[Type::I32], Operator::I32Const { value: 0 });
100        f.append_to_block(n, a);
101        f.set_terminator(
102            n,
103            portal_pc_waffle::Terminator::Br {
104                target: BlockTarget {
105                    block: end,
106                    args: vec![a],
107                },
108            },
109        )
110    }
111    let kk = f.add_block();
112    f.set_terminator(
113        k,
114        portal_pc_waffle::Terminator::CondBr {
115            cond: s,
116            if_true: BlockTarget {
117                block: n,
118                args: vec![],
119            },
120            if_false: BlockTarget {
121                block: kk,
122                args: vec![],
123            },
124        },
125    );
126    k = kk;
127    if retref {
128        let a = add_op(f, &[], &[Type::I32], Operator::I32Const { value: 1 });
129        f.append_to_block(k, a);
130        let a = add_op(f, &[v, a], &[Type::I32], Operator::I32Sub);
131        f.append_to_block(k, a);
132        v = a;
133    }
134    // }
135
136    match (take, retref) {
137        (true, true) => {
138            v = add_op(
139                f,
140                &[v],
141                &[portal_pc_waffle::Type::Heap(WithNullable {
142                    nullable: true,
143                    value: portal_pc_waffle::HeapType::ExternRef,
144                })],
145                Operator::Call {
146                    function_index: tfree,
147                },
148            );
149            f.append_to_block(k, v);
150        }
151        (_, false) => {
152            v = add_op(
153                f,
154                &[v],
155                &[Type::I32],
156                Operator::Call {
157                    function_index: talloc,
158                },
159            );
160            f.append_to_block(k, v);
161        }
162        (false, true) => {
163            v = add_op(
164                f,
165                &[v],
166                &[portal_pc_waffle::Type::Heap(WithNullable {
167                    nullable: true,
168                    value: portal_pc_waffle::HeapType::ExternRef,
169                })],
170                Operator::TableGet { table_index: table },
171            );
172            f.append_to_block(k, v);
173        }
174    }
175    if !retref {
176        let a = add_op(f, &[], &[Type::I32], Operator::I32Const { value: 1 });
177        f.append_to_block(k, a);
178        let a = add_op(f, &[v, a], &[Type::I32], Operator::I32Add);
179        f.append_to_block(k, a);
180        v = a;
181    }
182    f.set_terminator(
183        k,
184        portal_pc_waffle::Terminator::Br {
185            target: BlockTarget {
186                block: end,
187                args: vec![v],
188            },
189        },
190    );
191    Ok((ep, end))
192}
193pub fn wrap(m: &mut Module) -> anyhow::Result<()> {
194    let t = m.tables.push(TableData {
195        ty: portal_pc_waffle::Type::Heap(WithNullable {
196            nullable: true,
197            value: portal_pc_waffle::HeapType::ExternRef,
198        }),
199        initial: 0,
200        max: None,
201        func_elements: None,
202        table64: false,
203    });
204    let talloc = talloc(m, t, &[])?;
205    let tfree = tfree(m, t, &[])?;
206    let is = get_interfaces(m)?.into_iter().collect::<BTreeSet<_>>();
207    for mut import in take(&mut m.imports) {
208        if import.module == "tpit" && import.name == "void" {
209            if let ImportKind::Func(f) = &mut import.kind {
210                let s = m.funcs[*f].sig();
211                let o = *f;
212                let mut b = FunctionBody::new(&m, s);
213                let e = b.entry;
214                let arg = b.blocks[b.entry].params[0].1;
215                let arg = add_op(
216                    &mut b,
217                    &[arg],
218                    &[portal_pc_waffle::Type::Heap(WithNullable {
219                        nullable: true,
220                        value: portal_pc_waffle::HeapType::ExternRef,
221                    })],
222                    Operator::Call {
223                        function_index: tfree,
224                    },
225                );
226                b.append_to_block(e, arg);
227                b.set_terminator(e, portal_pc_waffle::Terminator::Return { values: vec![] });
228                m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
229            }
230            continue;
231        }
232        if import.module == "tpit" && import.name == "clone" {
233            if let ImportKind::Func(f) = &mut import.kind {
234                let s = m.funcs[*f].sig();
235                let o = *f;
236                let mut b = FunctionBody::new(&m, s);
237                let e = b.entry;
238                let arg = b.blocks[b.entry].params[0].1;
239                let arg = add_op(
240                    &mut b,
241                    &[arg],
242                    &[portal_pc_waffle::Type::Heap(WithNullable {
243                        nullable: true,
244                        value: portal_pc_waffle::HeapType::ExternRef,
245                    })],
246                    Operator::TableGet { table_index: t },
247                );
248                b.append_to_block(e, arg);
249                b.set_terminator(
250                    e,
251                    portal_pc_waffle::Terminator::ReturnCall {
252                        func: talloc,
253                        args: vec![arg],
254                    },
255                );
256                m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
257            }
258            continue;
259        }
260        if import.module == "tpit" {
261            import.module = format!("pit");
262            let p = new_sig(
263                m,
264                SignatureData::Func {
265                    params: vec![Type::I32],
266                    returns: vec![],
267                    shared: true,
268                },
269            );
270            let p = m
271                .funcs
272                .push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
273            if let ImportKind::Func(f) = &mut import.kind {
274                let s = m.funcs[*f].sig();
275                let o = replace(f, p);
276                let mut b = FunctionBody::new(&m, s);
277                let e = b.entry;
278                let arg = b.blocks[b.entry].params[0].1;
279                let arg = add_op(
280                    &mut b,
281                    &[arg],
282                    &[portal_pc_waffle::Type::Heap(WithNullable {
283                        nullable: true,
284                        value: portal_pc_waffle::HeapType::ExternRef,
285                    })],
286                    if import.name == "drop" {
287                        Operator::Call {
288                            function_index: tfree,
289                        }
290                    } else {
291                        Operator::TableGet { table_index: t }
292                    },
293                );
294                b.append_to_block(e, arg);
295                b.set_terminator(
296                    e,
297                    portal_pc_waffle::Terminator::ReturnCall {
298                        func: p,
299                        args: vec![arg],
300                    },
301                );
302                m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
303            }
304        }
305        if let Some(a) = import.name.strip_suffix(".tpit-res").map(|a| a.to_owned()) {
306            import.name = a;
307            if let ImportKind::Func(f) = &mut import.kind {
308                if let SignatureData::Func { params, returns,.. } =
309                    m.signatures[m.funcs[*f].sig()].clone()
310                {
311                    let p = params;
312                    let p = new_sig(
313                        m,
314                        SignatureData::Func {
315                            params: p,
316                            returns: vec![Type::I32],
317                            shared: true,
318                        },
319                    );
320                    let p = m
321                        .funcs
322                        .push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
323                    let s = m.funcs[*f].sig();
324                    let o = replace(f, p);
325                    let mut b = FunctionBody::new(&m, s);
326                    let e = b.entry;
327                    let arg = b.blocks[b.entry]
328                        .params
329                        .iter()
330                        .map(|a| a.1)
331                        .collect::<Vec<_>>();
332                    let arg = add_op(
333                        &mut b,
334                        &arg,
335                        &[portal_pc_waffle::Type::Heap(WithNullable {
336                            nullable: true,
337                            value: portal_pc_waffle::HeapType::ExternRef,
338                        })],
339                        Operator::Call { function_index: p },
340                    );
341                    b.append_to_block(e, arg);
342                    b.set_terminator(
343                        e,
344                        portal_pc_waffle::Terminator::ReturnCall {
345                            func: talloc,
346                            args: vec![arg],
347                        },
348                    );
349                    m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
350                }
351            }
352        }
353        m.imports.push(import);
354    }
355    for mut export in take(&mut m.exports) {
356        if let Some(a) = export.name.strip_suffix(".tpit-res").map(|a| a.to_owned()) {
357            export.name = a;
358            if let ExportKind::Func(f) = &mut export.kind {
359                if let SignatureData::Func { params, returns,.. } =
360                    m.signatures[m.funcs[*f].sig()].clone()
361                {
362                    let p = params;
363                    let p = new_sig(
364                        m,
365                        SignatureData::Func {
366                            params: p,
367                            returns: vec![portal_pc_waffle::Type::Heap(WithNullable {
368                                nullable: true,
369                                value: portal_pc_waffle::HeapType::ExternRef,
370                            })],
371                            shared: true,
372                        },
373                    );
374                    let p = m
375                        .funcs
376                        .push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
377                    let s = m.funcs[*f].sig();
378                    let o = replace(f, p);
379                    let mut b = FunctionBody::new(&m, s);
380                    let e = b.entry;
381                    let arg = b.blocks[b.entry]
382                        .params
383                        .iter()
384                        .map(|a| a.1)
385                        .collect::<Vec<_>>();
386                    let arg = add_op(
387                        &mut b,
388                        &arg,
389                        &[Type::I32],
390                        Operator::Call { function_index: p },
391                    );
392                    b.append_to_block(e, arg);
393                    b.set_terminator(
394                        e,
395                        portal_pc_waffle::Terminator::ReturnCall {
396                            func: tfree,
397                            args: vec![arg],
398                        },
399                    );
400                    m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
401                }
402            }
403        }
404        m.exports.push(export);
405    }
406    for i in is {
407        let f = waffle_funcs(m, &i, false);
408        let mut ss = BTreeMap::new();
409        for mut import in take(&mut m.imports) {
410            if import.module == format!("tpit/{}", i.rid_str()) {
411                match import.name.strip_prefix("~") {
412                    Some(a) => {
413                        if let ImportKind::Func(f) = &mut import.kind {
414                            if let SignatureData::Func { params, returns,.. } =
415                                m.signatures[m.funcs[*f].sig()].clone()
416                            {
417                                let p = params;
418                                ss.insert(a.to_owned(), p.clone());
419                                let p = new_sig(
420                                    m,
421                                    SignatureData::Func {
422                                        params: p,
423                                        returns: vec![portal_pc_waffle::Type::Heap(WithNullable {
424                                            nullable: true,
425                                            value: portal_pc_waffle::HeapType::ExternRef,
426                                        })],
427                                        shared: true,
428                                    },
429                                );
430                                let p = m
431                                    .funcs
432                                    .push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
433                                let s = m.funcs[*f].sig();
434                                let o = replace(f, p);
435                                let mut b = FunctionBody::new(&m, s);
436                                let e = b.entry;
437                                let arg = b.blocks[b.entry].params[0].1;
438                                let arg = add_op(
439                                    &mut b,
440                                    &[arg],
441                                    &[portal_pc_waffle::Type::Heap(WithNullable {
442                                        nullable: true,
443                                        value: portal_pc_waffle::HeapType::ExternRef,
444                                    })],
445                                    Operator::Call { function_index: p },
446                                );
447                                b.append_to_block(e, arg);
448                                b.set_terminator(
449                                    e,
450                                    portal_pc_waffle::Terminator::ReturnCall {
451                                        func: talloc,
452                                        args: vec![arg],
453                                    },
454                                );
455                                m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
456                            }
457                        }
458                    }
459                    None => {
460                        let x = i
461                            .methods
462                            .get(&import.name)
463                            .context("in getting the method")?;
464                        let p = to_waffle_sig(m, x, false);
465                        let p = m.signatures[p].clone();
466                        let SignatureData::Func { params, returns,.. } = p else {
467                            continue;
468                        };
469                        let p = new_sig(
470                            m,
471                            SignatureData::Func {
472                                params: vec![Type::I32]
473                                    .into_iter()
474                                    .chain(params.into_iter().map(|a| {
475                                        if a == portal_pc_waffle::Type::Heap(WithNullable {
476                                            nullable: true,
477                                            value: portal_pc_waffle::HeapType::ExternRef,
478                                        }) {
479                                            Type::I32
480                                        } else {
481                                            a
482                                        }
483                                    }))
484                                    .collect(),
485                                returns: returns
486                                    .into_iter()
487                                    .map(|a| {
488                                        if a == portal_pc_waffle::Type::Heap(WithNullable {
489                                            nullable: true,
490                                            value: portal_pc_waffle::HeapType::ExternRef,
491                                        }) {
492                                            Type::I32
493                                        } else {
494                                            a
495                                        }
496                                    })
497                                    .collect(),
498                                    shared: true,
499                            },
500                        );
501                        let p = m
502                            .funcs
503                            .push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
504                        if let ImportKind::Func(f) = &mut import.kind {
505                            let s = m.funcs[*f].sig();
506                            let o = replace(f, p);
507                            let mut b = FunctionBody::new(&m, s);
508                            let mut k = b.entry;
509                            let args = b.blocks[b.entry]
510                                .params
511                                .iter()
512                                .map(|a| a.1)
513                                .collect::<Vec<_>>()
514                                .into_iter()
515                                .zip(
516                                    once(Arg::Resource {
517                                        ty: ResTy::This,
518                                        nullable: false,
519                                        take: false,
520                                        ann: vec![],
521                                    })
522                                    .chain(x.params.iter().cloned()),
523                                );
524                            let mut v2 = vec![];
525                            for (v, r) in args {
526                                let a;
527                                (a, k) = shim(false, &mut b, k, &r, v, talloc, tfree, t)?;
528                                v2.push(a);
529                            }
530                            let rets = b.rets.clone();
531                            let rets =
532                                add_op(&mut b, &v2, &rets, Operator::Call { function_index: p });
533                            b.append_to_block(k, rets);
534                            let rets = results_ref_2(&mut b, rets);
535                            let mut r2 = vec![];
536                            for (v, r) in rets.iter().cloned().zip(x.rets.iter()) {
537                                let a;
538                                (a, k) = shim(true, &mut b, k, r, v, talloc, tfree, t)?;
539                                r2.push(a);
540                            }
541                            b.set_terminator(
542                                k,
543                                portal_pc_waffle::Terminator::Return { values: r2 },
544                            );
545                            m.funcs[o] = FuncDecl::Body(s, format!("_pit"), b);
546                        }
547                    }
548                }
549                import.module = format!("pit/{}", i.rid_str());
550            }
551            m.imports.push(import)
552        }
553        for mut export in take(&mut m.exports) {
554            if let Some(a) = export.name.strip_prefix(&format!("tpit/{}/~", i.rid_str())) {
555                let a = a.to_owned();
556                export.name = format!("pit/{}/~{a}", i.rid_str());
557                match a.strip_prefix(".") {
558                    Some(a) => {
559                        if let Some((a, b)) = a.split_once("@") {
560                            let Ok((_, x)) = pit_core::parse_sig(b) else {
561                                anyhow::bail!("invalid sig")
562                            };
563                            export.name = format!("pit/{}/~{a}", i.rid_str());
564                            let p = to_waffle_sig(m, &x, false);
565                            let p = m.signatures[p].clone();
566                            let SignatureData::Func { params, returns,.. } = p else {
567                                continue;
568                            };
569                            let p = new_sig(
570                                m,
571                                SignatureData::Func {
572                                    params: ss
573                                        .get(a)
574                                        .cloned()
575                                        .into_iter()
576                                        .flatten()
577                                        .chain(params.into_iter())
578                                        .collect(),
579                                    returns: returns.iter().cloned().collect(),
580                                    shared: true,
581                                },
582                            );
583                            // let p = m.funcs.push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
584                            if let ExportKind::Func(f) = &mut export.kind {
585                                let s = m.funcs[*f].sig();
586                                let p = *f;
587                                let mut b = FunctionBody::new(&m, s);
588                                let mut k = b.entry;
589                                let args = b.blocks[b.entry]
590                                    .params
591                                    .iter()
592                                    .map(|a| a.1)
593                                    .collect::<Vec<_>>()
594                                    .into_iter()
595                                    .zip(once(Arg::I32).chain(x.params.iter().cloned()));
596                                let mut v2 = vec![];
597                                for (v, r) in args {
598                                    let a;
599                                    (a, k) = shim(true, &mut b, k, &r, v, talloc, tfree, t)?;
600                                    v2.push(a);
601                                }
602                                let rets = add_op(
603                                    &mut b,
604                                    &v2,
605                                    &returns,
606                                    Operator::Call { function_index: p },
607                                );
608                                b.append_to_block(k, rets);
609                                let rets = results_ref_2(&mut b, rets);
610                                let mut r2 = vec![];
611                                for (v, r) in rets.iter().cloned().zip(x.rets.iter()) {
612                                    let a;
613                                    (a, k) = shim(false, &mut b, k, r, v, talloc, tfree, t)?;
614                                    r2.push(a);
615                                }
616                                b.set_terminator(
617                                    k,
618                                    portal_pc_waffle::Terminator::Return { values: r2 },
619                                );
620                                *f = m.funcs.push(FuncDecl::Body(s, format!("_pit"), b));
621                            }
622                        }
623                    }
624                    None => {
625                        let (a, b) = a.split_once("/").context("in getting the stuff")?;
626                        let x = i.methods.get(b).context("in getting the method")?;
627                        let p = to_waffle_sig(m, x, false);
628                        let p = m.signatures[p].clone();
629                        let SignatureData::Func { params, returns,.. } = p else {
630                            continue;
631                        };
632                        let p = new_sig(
633                            m,
634                            SignatureData::Func {
635                                params: ss
636                                    .get(a)
637                                    .cloned()
638                                    .into_iter()
639                                    .flatten()
640                                    .chain(params.into_iter())
641                                    .collect(),
642                                returns: returns.iter().cloned().collect(),
643                                shared: true,
644                            },
645                        );
646                        // let p = m.funcs.push(portal_pc_waffle::FuncDecl::Import(p, format!("_pit")));
647                        if let ExportKind::Func(f) = &mut export.kind {
648                            let s = m.funcs[*f].sig();
649                            let p = *f;
650                            let mut b = FunctionBody::new(&m, s);
651                            let mut k = b.entry;
652                            let args = b.blocks[b.entry]
653                                .params
654                                .iter()
655                                .map(|a| a.1)
656                                .collect::<Vec<_>>()
657                                .into_iter()
658                                .zip(once(Arg::I32).chain(x.params.iter().cloned()));
659                            let mut v2 = vec![];
660                            for (v, r) in args {
661                                let a;
662                                (a, k) = shim(true, &mut b, k, &r, v, talloc, tfree, t)?;
663                                v2.push(a);
664                            }
665                            let rets =
666                                add_op(&mut b, &v2, &returns, Operator::Call { function_index: p });
667                            b.append_to_block(k, rets);
668                            let rets = results_ref_2(&mut b, rets);
669                            let mut r2 = vec![];
670                            for (v, r) in rets.iter().cloned().zip(x.rets.iter()) {
671                                let a;
672                                (a, k) = shim(false, &mut b, k, r, v, talloc, tfree, t)?;
673                                r2.push(a);
674                            }
675                            b.set_terminator(
676                                k,
677                                portal_pc_waffle::Terminator::Return { values: r2 },
678                            );
679                            *f = m.funcs.push(FuncDecl::Body(s, format!("_pit"), b));
680                        }
681                    }
682                }
683            }
684            m.exports.push(export);
685        }
686    }
687    m.exports.push(Export {
688        name: format!("tpit_alloc"),
689        kind: ExportKind::Func(talloc),
690    });
691    m.exports.push(Export {
692        name: format!("tpit_free"),
693        kind: ExportKind::Func(tfree),
694    });
695    m.exports.push(Export {
696        name: format!("tpit_table"),
697        kind: ExportKind::Table(t),
698    });
699    Ok(())
700}
pit_patch/tpit.rs

pit_patch/
tpit.rs
