zingen/codegen/
dispatcher.rs

1//! Code generator for EVM dispatcher.
2
3use crate::{
4    wasm::{self, Env, Functions},
5    JumpTable, MacroAssembler, Result,
6};
7use std::collections::BTreeMap;
8use wasmparser::FuncType;
9use zabi::Abi;
10
11/// Code generator for EVM dispatcher.
12pub struct Dispatcher {
13    /// ABI for the current function
14    pub abi: Vec<Abi>,
15    /// Code buffer
16    pub asm: MacroAssembler,
17    /// WASM environment
18    pub env: Env,
19    /// Module functions
20    pub funcs: BTreeMap<u32, FuncType>,
21    /// Jump table
22    pub table: JumpTable,
23}
24
25impl Dispatcher {
26    /// Create dispatcher with functions.
27    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    /// Emit compiled code to the given buffer.
43    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    /// Emit selector to buffer.
63    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        // Prepare the `PC` of the callee function.
78        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}