zingen/codegen/
dispatcher.rs1use crate::{
4 wasm::{self, Env, Functions},
5 JumpTable, MacroAssembler, Result,
6};
7use std::collections::BTreeMap;
8use wasmparser::FuncType;
9use zabi::Abi;
10
11pub struct Dispatcher {
13 pub abi: Vec<Abi>,
15 pub asm: MacroAssembler,
17 pub env: Env,
19 pub funcs: BTreeMap<u32, FuncType>,
21 pub table: JumpTable,
23}
24
25impl Dispatcher {
26 pub fn new(env: Env, funcs: &Functions<'_>) -> Result<Self> {
28 let funcs = funcs
29 .values()
30 .map(|func| Ok((func.index(), func.sig()?)))
31 .collect::<Result<_>>()?;
32
33 Ok(Self {
34 abi: Default::default(),
35 asm: Default::default(),
36 env,
37 funcs,
38 table: Default::default(),
39 })
40 }
41
42 pub fn finish(&mut self, selectors: Functions<'_>, table: &mut JumpTable) -> Result<Vec<u8>> {
44 if selectors.is_empty() {
45 return Ok(Default::default());
46 }
47
48 self.asm._push0()?;
49 self.asm._calldataload()?;
50 self.asm.push(&[0xe0])?;
51 self.asm._shr()?;
52 let mut len = selectors.len();
53 for (_, func) in selectors.iter() {
54 self.emit_selector(func, len == 1)?;
55 len -= 1;
56 }
57
58 table.merge(self.table.clone(), 0)?;
59 Ok(self.asm.buffer().into())
60 }
61
62 fn emit_selector(&mut self, selector: &wasm::Function<'_>, last: bool) -> Result<()> {
64 let abi = self.env.load_abi(selector)?;
65 self.abi.push(abi.clone());
66
67 let selector_bytes = abi.selector();
68 tracing::debug!(
69 "Emitting selector {:?} for function: {}",
70 selector_bytes,
71 abi.signature(),
72 );
73
74 let func = self.env.query_func(&abi.name)?;
75 self.asm.increment_sp(1)?;
76
77 self.table.call(self.asm.pc(), func);
79
80 if last {
81 self.asm._swap1()?;
82 } else {
83 self.asm._dup2()?;
84 }
85
86 self.asm.push(&selector_bytes)?;
87 self.asm._eq()?;
88 self.asm._swap1()?;
89 self.asm._jumpi()?;
90
91 Ok(())
92 }
93}