waffle_trie/
lib.rs

1use std::mem::take;
2
3use anyhow::Context;
4use trie_rs::map::{Trie, TrieBuilder};
5use waffle::{
6    util::new_sig, Block, BlockTarget, ExportKind, Func, FuncDecl, FunctionBody, ImportKind,
7    MemoryArg, Module, Operator, Signature, SignatureData, Type, Value,
8};
9use waffle_ast::{add_op, Builder};
10
11pub fn internal_trie(m: &Module) -> Trie<u8, Func> {
12    let mut b = TrieBuilder::new();
13    for f in m.funcs.iter() {
14        b.push(m.funcs[f].name().as_bytes(), f);
15    }
16    return b.build();
17}
18pub fn export_trie(m: &Module) -> Trie<u8, Func> {
19    let mut b = TrieBuilder::new();
20    for f in m.exports.iter() {
21        let ExportKind::Func(func) = &f.kind else {
22            continue;
23        };
24        b.push(f.name.as_bytes(), *func);
25    }
26    return b.build();
27}
28pub fn tx<T>(
29    m: &mut Module,
30    t: &Trie<u8, T>,
31    x: &[u8],
32    f: &mut FunctionBody,
33    k: Block,
34    v: Value,
35    s: &mut impl FnMut(&T, &mut Module) -> bool,
36    go: &mut impl FnMut(&T, &mut FunctionBody, &mut Module, Block) -> anyhow::Result<()>,
37) -> anyhow::Result<()> {
38    if let Some(0) = x.last() {
39        if let Some(x) = t.exact_match(&x[..(x.len() - 2)]) {
40            if s(x, m) {
41                return go(x, f, m, k);
42            }
43            return Ok(());
44        }
45    }
46    let p = t
47        .postfix_search::<Vec<_>, _>(x)
48        .filter(|a| s(a.1, m))
49        .map(|_| ())
50        .collect::<Vec<_>>()
51        .len();
52    if p == 0 {
53        return Ok(());
54    }
55    let ty = f.values[v]
56        .ty(&f.type_pool)
57        .context("in getting the type")?;
58    let mem = m
59        .exports
60        .iter()
61        .find_map(|x| {
62            if x.name == "memory" {
63                match &x.kind {
64                    waffle::ExportKind::Memory(m) => Some(*m),
65                    _ => None,
66                }
67            } else {
68                None
69            }
70        })
71        .context("in getting the main memory")?;
72    let a = add_op(
73        f,
74        &[v],
75        &[ty.clone()],
76        Operator::I32Load8U {
77            memory: MemoryArg {
78                memory: mem,
79                align: 0,
80                offset: x.len() as u64,
81            },
82        },
83    );
84    f.append_to_block(k, a);
85    let xs = 0..=255;
86    let xs = xs.map(|i| {
87        let mut x = x.to_vec();
88        x.push(i);
89        let b = f.add_block();
90        tx(m, t, &x, f, b, v, s, go)?;
91        anyhow::Ok(BlockTarget {
92            block: b,
93            args: vec![],
94        })
95    });
96    let xs = xs.collect::<anyhow::Result<Vec<_>>>()?;
97    let xa = xs[0].clone();
98    f.set_terminator(
99        k,
100        waffle::Terminator::Select {
101            value: a,
102            targets: xs,
103            default: xa,
104        },
105    );
106    Ok(())
107}
108pub fn emit(m: &mut Module) -> anyhow::Result<()> {
109    let t_internal = internal_trie(&m);
110    let t_export = export_trie(&m);
111    for i in take(&mut m.imports) {
112        if let Some(a) = i.module.strip_prefix("dyn_call") {
113            let t = match a {
114                "func" => &t_internal,
115                "export" => &t_export,
116                _ => {
117                    m.imports.push(i);
118                    continue;
119                }
120            };
121            if let ImportKind::Func(f) = i.kind {
122                let s = m.signatures[m.funcs[f].sig()].clone();
123                let SignatureData::Func { params, returns } = s else{
124                    anyhow::bail!("not a func")
125                };
126                let s = new_sig(
127                    m,
128                    SignatureData::Func {
129                        returns: returns,
130                        params: params[1..].iter().cloned().collect(),
131                    },
132                );
133                let os = m.funcs[f].sig();
134                let n = m.funcs[f].name().to_owned();
135                let mut body = FunctionBody::new(&m, os);
136                let e = body.entry;
137                let mut p = body.blocks[e].params.iter();
138                let pv = p.next().context("in getting the first param")?.1;
139                let ps = p.map(|a| a.1).collect::<Vec<_>>();
140                tx(
141                    m,
142                    &*t,
143                    &[],
144                    &mut body,
145                    e,
146                    pv,
147                    &mut |a,m| m.funcs[*a].sig() == s,
148                    &mut |f, b, m, k| {
149                        b.set_terminator(
150                            k,
151                            waffle::Terminator::ReturnCall {
152                                func: *f,
153                                args: ps.clone(),
154                            },
155                        );
156                        Ok(())
157                    },
158                )?;
159                m.funcs[f] = FuncDecl::Body(os, n, body);
160                continue;
161            }
162        }
163        m.imports.push(i)
164    }
165    Ok(())
166}
waffle_trie/lib.rs

waffle_trie/
lib.rs
