portal_pc_koffle/
lib.rs

1use impl_trait_for_tuples::impl_for_tuples;
2use once_map::OnceMap;
3use std::{
4    collections::BTreeMap,
5    iter::once,
6    mem::{replace, take},
7};
8use waffle::{
9    Block, BlockTarget, ExportKind, Func, FuncDecl, FunctionBody, Global, GlobalData, HeapType,
10    ImportKind, Memory, Module, Operator, Table, TableData, Type, Value, WithNullable,
11};
12use waffle_ast::tutils::{talloc, tfree};
13pub mod bulk_memory_lowering;
14#[cfg(feature = "corpack")]
15pub mod corpack;
16pub mod detta;
17pub mod hustle;
18pub mod inline;
19pub mod rand;
20pub mod importify;
21#[impl_for_tuples(12)]
22pub trait FuncCollector {
23    fn add_func(&mut self, f: Func);
24}
25impl<'a, T: FuncCollector> FuncCollector for &'a mut T {
26    fn add_func(&mut self, f: Func) {
27        FuncCollector::add_func(&mut **self, f);
28    }
29}
30pub fn init_with(
31    module: &mut Module,
32    collector: &mut (dyn FuncCollector + '_),
33    init: Func,
34    run_on_imports: bool,
35) {
36    let idg = module.globals.push(GlobalData {
37        ty: Type::I32,
38        mutable: true,
39        value: Some(1),
40    });
41    for fi in module
42        .funcs
43        .iter()
44        .filter(|f| {
45            if run_on_imports {
46                true
47            } else {
48                module
49                    .imports
50                    .iter()
51                    .filter_map(|a| match &a.kind {
52                        ImportKind::Func(f) => Some(f),
53                        _ => None,
54                    })
55                    .all(|g| *g != *f)
56            }
57        })
58        // .filter(|a| funcs.contains(a))
59        .collect::<Vec<_>>()
60    {
61        // let mut f = take(&mut module.funcs[fi]);
62        // if let Some(f) = f.body_mut() {
63        with_swizz(module, fi, collector, |(module, f, fi, collector)| {
64            let new = f.add_block();
65            let params = f.blocks[f.entry]
66                .params
67                .iter()
68                .map(|a| a.0)
69                .collect::<Vec<_>>();
70            let params = params
71                .into_iter()
72                .map(|a| f.add_blockparam(new, a))
73                .collect::<Vec<_>>();
74            let old = f.add_block();
75            let mut id = f.add_op(
76                f.entry,
77                Operator::GlobalGet { global_index: idg },
78                &[],
79                &[Type::I32],
80            );
81            for _ in 0..2{
82                id = f.add_op(f.entry, Operator::I32Eqz, &[id], &[Type::I32]);
83            }
84            let fix = f.add_block();
85            f.set_terminator(
86                f.entry,
87                waffle::Terminator::CondBr {
88                    cond: id,
89                    if_true: BlockTarget {
90                        block: fix,
91                        args: vec![],
92                    },
93                    if_false: BlockTarget {
94                        block: old,
95                        args: params.clone(),
96                    },
97                },
98            );
99            let id = f.add_op(fix, Operator::I32Const { value: 0 }, &[], &[Type::I32]);
100            f.add_op(fix, Operator::GlobalSet { global_index: idg }, &[id], &[]);
101            f.add_op(
102                fix,
103                Operator::Call {
104                    function_index: init,
105                },
106                &[],
107                &[],
108            );
109            // let retvals = f
110            //     .rets
111            //     .clone()
112            //     .into_iter()
113            //     .map(|a| default_val(a, f, fix))
114            //     .collect();
115            f.set_terminator(
116                fix,
117                waffle::Terminator::Br {
118                    target: BlockTarget {
119                        block: f.entry,
120                        args: params,
121                    },
122                },
123            );
124            let params = f.blocks[f.entry]
125                .params
126                .iter()
127                .map(|a| a.0)
128                .collect::<Vec<_>>();
129            let params = params
130                .into_iter()
131                .map(|p| f.add_blockparam(old, p))
132                .collect();
133            f.set_terminator(
134                old,
135                waffle::Terminator::ReturnCall {
136                    func: fi,
137                    args: params,
138                },
139            );
140        });
141        // }
142
143        // module.funcs[fi] = f;
144    }
145    // Ok(())
146}
147#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
148pub struct TableInfo {
149    pub table: Table,
150    pub talloc: Func,
151    pub tfree: Func,
152    pub ty: WithNullable<HeapType>,
153}
154#[derive(Default)]
155pub struct TableMap {
156    tables: OnceMap<WithNullable<HeapType>, Box<TableInfo>>,
157}
158impl TableMap {
159    pub fn table_in(
160        &self,
161        module: &mut Module,
162        collector: &mut (dyn FuncCollector + '_),
163        ty: WithNullable<HeapType>,
164    ) -> TableInfo {
165        *self.tables.insert(ty, |ty| {
166            Box::new({
167                let t = module.tables.push(TableData {
168                    ty: Type::Heap(ty.clone()),
169                    initial: 0,
170                    max: None,
171                    func_elements: None,
172                    table64: false,
173                });
174                let i = TableInfo {
175                    table: t,
176                    talloc: talloc(module, t, &[]).unwrap(),
177                    tfree: tfree(module, t, &[]).unwrap(),
178                    ty: *ty,
179                };
180                for f in [(i.talloc), (i.tfree)] {
181                    collector.add_func(f);
182                }
183
184                i
185            })
186        })
187        // (*t, *alloc, *free)
188    }
189    pub fn into_iter(self) -> impl Iterator<Item = (WithNullable<HeapType>, Box<TableInfo>)> {
190        self.tables.into_iter()
191    }
192}
193
194pub fn ift(module: &Module) -> Option<Table> {
195    module.exports.iter().find_map(|x| {
196        if x.name == "__indirect_function_table" {
197            match &x.kind {
198                ExportKind::Table(t) => Some(*t),
199                _ => None,
200            }
201        } else {
202            None
203        }
204    })
205}
206
207pub fn memory(module: &Module) -> Option<Memory> {
208    module.exports.iter().find_map(|x| {
209        if x.name == "memory" {
210            match &x.kind {
211                ExportKind::Memory(t) => Some(*t),
212                _ => None,
213            }
214        } else {
215            None
216        }
217    })
218}
219pub fn default_val(ty: Type, dst: &mut FunctionBody, k: Block) -> Value {
220    dst.add_op(
221        k,
222        match ty.clone() {
223            Type::I32 => Operator::I32Const { value: 0 },
224            Type::I64 => Operator::I64Const { value: 0 },
225            Type::F32 => Operator::F32Const { value: 0 },
226            Type::F64 => Operator::F64Const { value: 0 },
227            Type::V128 => todo!(),
228            Type::Heap(_) => Operator::RefNull { ty: ty.clone() },
229            _ => todo!(),
230        },
231        &[],
232        &[ty],
233    )
234}
235
236pub fn with_swizz<R>(
237    module: &mut Module,
238    f: Func,
239    collector: &mut (dyn FuncCollector + '_),
240    shim: impl FnOnce(
241        (
242            &mut Module,
243            &mut FunctionBody,
244            Func,
245            &mut (dyn FuncCollector + '_),
246        ),
247    ) -> R,
248) -> R {
249    let sig = module.funcs[f].sig();
250    let name = module.funcs[f].name().to_owned();
251    let g = replace(
252        &mut module.funcs[f],
253        waffle::FuncDecl::Import(sig, name.clone()),
254    );
255    let g = module.funcs.push(g);
256    for i in module.imports.iter_mut() {
257        if let ImportKind::Func(i) = &mut i.kind {
258            if *i == f {
259                *i = g;
260            }
261        }
262    }
263    let mut dst = FunctionBody::new(module, sig);
264    let r = shim((module, &mut dst, g, collector));
265    module.funcs[f] = FuncDecl::Body(sig, name, dst);
266    collector.add_func(f);
267    return r;
268}
269#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
270pub enum Cond {
271    Func { func: Func, pass_args: bool },
272    // ViaTable { table: Table, dispatch: Func },
273}
274pub fn guard_table(
275    module: &mut Module,
276    f: Func,
277    collector: &mut (dyn FuncCollector + '_),
278    table: Table,
279) -> Global {
280    with_swizz(module, f, collector, |(module, b, f, collector)| {
281        let idx = need(module, table, f);
282        let ty = if module.tables[table].table64 {
283            Type::I64
284        } else {
285            Type::I32
286        };
287        let g = module.globals.push(GlobalData {
288            ty: ty.clone(),
289            value: Some(idx as u64),
290            mutable: true,
291        });
292        let gv = b.add_op(
293            b.entry,
294            Operator::GlobalGet { global_index: g },
295            &[],
296            &[ty.clone()],
297        );
298        let mut p = b.blocks[b.entry]
299            .params
300            .iter()
301            .map(|p| p.1)
302            .collect::<Vec<_>>();
303        p.push(gv);
304        b.set_terminator(
305            b.entry,
306            waffle::Terminator::ReturnCallIndirect {
307                sig: module.funcs[f].sig(),
308                table: table,
309                args: p,
310            },
311        );
312        return g;
313    })
314}
315#[derive(Default)]
316pub struct GuardMap {
317    all: OnceMap<(Func, Table), Box<Global>>,
318}
319impl GuardMap {
320    pub fn guard(
321        &self,
322        module: &mut Module,
323        collector: &mut (dyn FuncCollector + '_),
324        f: Func,
325        t: Table,
326    ) -> Global {
327        return *self
328            .all
329            .insert((f, t), |_| Box::new(guard_table(module, f, collector, t)));
330    }
331}
332pub fn swap_fns(
333    module: &mut Module,
334    a: Func,
335    b: Func,
336    collector: &mut (dyn FuncCollector + '_),
337    cond: Option<Cond>,
338) {
339    if a == b {
340        return;
341    }
342    with_swizz(module, a, collector, |(module, ab, a, collector)| {
343        with_swizz(module, b, collector, |(module, bb, b, collector)| {
344            let (a, b) = (b, a);
345            for (x, b, y) in [(a, ab, b), (b, bb, a)] {
346                let p = b.blocks[b.entry].params.iter().map(|a| a.1).collect();
347                match cond {
348                    None => {
349                        b.set_terminator(
350                            b.entry,
351                            waffle::Terminator::ReturnCall { func: x, args: p },
352                        );
353                    }
354                    Some(cond) => match cond {
355                        Cond::Func {
356                            func: cond,
357                            pass_args,
358                        } => {
359                            let [xb, yb] = [x, y].map(|f| {
360                                let k = b.add_block();
361                                b.set_terminator(
362                                    k,
363                                    waffle::Terminator::ReturnCall {
364                                        func: f,
365                                        args: p.clone(),
366                                    },
367                                );
368                                BlockTarget {
369                                    block: k,
370                                    args: vec![],
371                                }
372                            });
373                            let cond = b.add_op(
374                                b.entry,
375                                Operator::Call {
376                                    function_index: cond,
377                                },
378                                if pass_args { &p } else { &[] },
379                                &[Type::I32],
380                            );
381                            b.set_terminator(
382                                b.entry,
383                                waffle::Terminator::CondBr {
384                                    cond: cond,
385                                    if_true: xb,
386                                    if_false: yb,
387                                },
388                            );
389                        }
390                    },
391                }
392            }
393        })
394    })
395}
396pub fn guard_fn(
397    module: &mut Module,
398    x: Func,
399    y: Func,
400    collector: &mut (dyn FuncCollector + '_),
401    cond: Option<Cond>,
402) {
403    with_swizz(module, x, collector, |(module, b, x, collector)| {
404        let p = b.blocks[b.entry].params.iter().map(|a| a.1).collect();
405        match cond {
406            None => {
407                b.set_terminator(b.entry, waffle::Terminator::ReturnCall { func: x, args: p });
408            }
409            Some(cond) => match cond {
410                Cond::Func {
411                    func: cond,
412                    pass_args,
413                } => {
414                    let [xb, yb] = [x, y].map(|f| {
415                        let k = b.add_block();
416                        b.set_terminator(
417                            k,
418                            waffle::Terminator::ReturnCall {
419                                func: f,
420                                args: p.clone(),
421                            },
422                        );
423                        BlockTarget {
424                            block: k,
425                            args: vec![],
426                        }
427                    });
428                    let cond = b.add_op(
429                        b.entry,
430                        Operator::Call {
431                            function_index: cond,
432                        },
433                        if pass_args { &p } else { &[] },
434                        &[Type::I32],
435                    );
436                    b.set_terminator(
437                        b.entry,
438                        waffle::Terminator::CondBr {
439                            cond: cond,
440                            if_true: xb,
441                            if_false: yb,
442                        },
443                    );
444                }
445            },
446        }
447    })
448}
449pub fn replace_fns(
450    module: &mut Module,
451    f: impl Iterator<Item = Func>,
452    mut a: Func,
453    collector: &mut (dyn FuncCollector + '_),
454    cond: Option<Cond>,
455) -> Func {
456    for f in f {
457        a = with_swizz(module, f, collector, |(module, b, x, collector)| {
458            let (x, a) = (a, x);
459
460            let p = b.blocks[b.entry].params.iter().map(|a| a.1).collect();
461            match cond {
462                None => {
463                    b.set_terminator(b.entry, waffle::Terminator::ReturnCall { func: x, args: p });
464                }
465                Some(cond) => match cond {
466                    Cond::Func {
467                        func: cond,
468                        pass_args,
469                    } => {
470                        let [xb, yb] = [x, a].map(|f| {
471                            let k = b.add_block();
472                            b.set_terminator(
473                                k,
474                                waffle::Terminator::ReturnCall {
475                                    func: f,
476                                    args: p.clone(),
477                                },
478                            );
479                            BlockTarget {
480                                block: k,
481                                args: vec![],
482                            }
483                        });
484                        let cond = b.add_op(
485                            b.entry,
486                            Operator::Call {
487                                function_index: cond,
488                            },
489                            if pass_args { &p } else { &[] },
490                            &[Type::I32],
491                        );
492                        b.set_terminator(
493                            b.entry,
494                            waffle::Terminator::CondBr {
495                                cond: cond,
496                                if_true: xb,
497                                if_false: yb,
498                            },
499                        );
500                    }
501                },
502            }
503
504            return a;
505        });
506    }
507
508    return a;
509}
510pub fn loop_fns(
511    module: &mut Module,
512    mut a: impl Iterator<Item = Func>,
513    collector: &mut (dyn FuncCollector + '_),
514    cond: Option<Cond>,
515) {
516    if let Some(b) = a.next() {
517        let c = replace_fns(module, a, b, collector, cond);
518        swap_fns(module, c, b, collector, cond);
519    }
520}
521pub fn need(module: &mut Module, table: Table, x: Func) -> usize {
522    let tab = module.tables[table].func_elements.as_mut().unwrap();
523    for (idx, y) in tab.iter().enumerate() {
524        if *y == x {
525            return idx;
526        }
527    }
528    let idx = tab.len();
529    tab.push(x);
530    return idx;
531}
portal_pc_koffle/lib.rs

portal_pc_koffle/
lib.rs
