zingen/wasm/
mod.rs

1//! WASM related primitives.
2
3mod abi;
4mod data;
5mod func;
6mod host;
7
8pub use self::{
9    abi::{ToLSBytes, Type},
10    data::Data,
11    func::{Function, Functions},
12    host::HostFunc,
13};
14use crate::{Error, Result};
15use host::CompilerLabel;
16use smallvec::SmallVec;
17use std::collections::BTreeMap;
18use wasmparser::Operator;
19use zabi::Abi;
20
21macro_rules! impl_deref {
22    ($doc:literal, $name:ident, $target:ty) => {
23        #[derive(Clone, Debug, Default)]
24        #[doc = concat!(" ", $doc)]
25        pub struct $name($target);
26
27        impl core::ops::Deref for $name {
28            type Target = $target;
29
30            fn deref(&self) -> &Self::Target {
31                &self.0
32            }
33        }
34
35        impl core::ops::DerefMut for $name {
36            fn deref_mut(&mut self) -> &mut Self::Target {
37                &mut self.0
38            }
39        }
40    };
41    ($(($doc:literal, $name:ident, $target:ty)),*) => {
42        $( impl_deref!($doc, $name, $target); )*
43    };
44}
45
46impl_deref! {
47    ("WASM import section", Imports, BTreeMap<u32, HostFunc>),
48    ("WASM export section", Exports, BTreeMap<u32, String>),
49    ("WASM slot registry", Slots, BTreeMap<u32, u32>),
50    ("WASM function registry", Funcs, BTreeMap<u32, (u32, u32)>)
51}
52
53/// A struct that holds the environment wasm module.
54#[derive(Clone, Debug, Default)]
55pub struct Env {
56    /// WASM imports
57    pub imports: Imports,
58    /// WASM exports
59    pub exports: Exports,
60    /// Function memory slots
61    pub slots: Slots,
62    /// Function params count
63    pub funcs: Funcs,
64    /// WASM data slots
65    pub data: Data,
66    /// Current function index
67    pub index: Option<u32>,
68}
69
70impl Env {
71    /// Load abis from functions
72    pub fn load_abis(&self, funs: &Functions<'_>) -> Result<Vec<Abi>> {
73        let mut abis: Vec<_> = Default::default();
74        for (_, fun) in funs.iter() {
75            abis.push(self.load_abi(fun)?);
76        }
77
78        Ok(abis)
79    }
80
81    /// Load abi from function
82    pub fn load_abi(&self, fun: &Function<'_>) -> Result<Abi> {
83        let mut reader = fun.body.get_operators_reader()?;
84
85        let Operator::I32Const { value: offset } = reader.read()? else {
86            return Err(Error::InvalidSelector);
87        };
88        let Operator::I32Const { value: length } = reader.read()? else {
89            return Err(Error::InvalidSelector);
90        };
91
92        // Validate zinkc helper `emit_abi`
93        let Operator::Call {
94            function_index: index,
95        } = reader.read()?
96        else {
97            return Err(Error::InvalidSelector);
98        };
99
100        if !self.imports.is_emit_abi(index) {
101            return Err(Error::FuncNotImported("emit_abi".into()));
102        }
103
104        let abi = self.data.load(offset, length as usize)?;
105        Abi::from_hex(String::from_utf8_lossy(&abi)).map_err(Into::into)
106    }
107
108    /// Query exported function from selector.
109    pub fn query_func(&self, name: &str) -> Result<u32> {
110        for (index, export) in self.exports.iter() {
111            if export == name {
112                return Ok(*index);
113            }
114        }
115
116        Err(Error::FuncNotImported(name.into()))
117    }
118
119    /// Check if the input function is external function
120    pub fn is_external(&self, index: u32) -> bool {
121        // self.exports.get(&index).is_some()
122        let Some(name) = self.exports.get(&index) else {
123            return false;
124        };
125
126        let selector = name.to_owned() + "_selector";
127        self.exports.iter().any(|(_, n)| **n == selector)
128    }
129
130    /// If the present function index is the main function
131    ///
132    /// NOTE: in wasm the indexes of the imports will be ordered
133    /// before the functions
134    pub fn is_main(&self, index: u32) -> bool {
135        self.imports.len() as u32 == index
136    }
137
138    /// Clone a new environment with function index provided
139    pub fn with_index(&self, index: u32) -> Self {
140        let mut this = self.clone();
141        this.index = Some(index);
142        this
143    }
144
145    /// Get reserved slots
146    pub fn reserved(&self) -> u32 {
147        let Some(index) = self.index else {
148            return 0;
149        };
150
151        *self.slots.get(&index).unwrap_or(&0)
152    }
153
154    /// Allocate memory slots from local index
155    pub fn alloc(&self, index: u32) -> SmallVec<[u8; 4]> {
156        let slots = index + self.reserved();
157        tracing::debug!(
158            "allocating memory for local {index} of function {:?}, slot: {slots}",
159            self.index
160        );
161        (slots * 0x20).to_ls_bytes()
162    }
163}
164
165impl Imports {
166    /// If the function is `emit_abi`.
167    pub fn is_emit_abi(&self, index: u32) -> bool {
168        self.get(&index) == Some(&HostFunc::EmitABI)
169    }
170
171    /// Get reserved slots in memory for storage calculations
172    pub fn reserved(&self) -> u32 {
173        let mut reserved = 0;
174        for host_fn in self.0.values() {
175            match *host_fn {
176                HostFunc::Label(CompilerLabel::ReserveMemory32) => reserved = 1,
177                HostFunc::Label(CompilerLabel::ReserveMemory64) => {
178                    return 2;
179                }
180                _ => {}
181            }
182        }
183
184        reserved
185    }
186}
187
188impl Exports {
189    /// Get all function selectors
190    pub fn selectors(&self) -> Vec<u32> {
191        self.iter()
192            .filter_map(|(index, export)| {
193                if export.ends_with("_selector") {
194                    Some(*index)
195                } else {
196                    None
197                }
198            })
199            .collect::<Vec<_>>()
200    }
201}