wasm_bindgen_wasm_conventions/
lib.rs

1//! A tiny crate of utilities for working with implicit Wasm codegen conventions
2//! (often established by LLVM and lld).
3//!
4//! Examples conventions include:
5//!
6//! * The stack pointer
7//! * The canonical linear memory that contains the stack
8
9use std::io::Cursor;
10
11use anyhow::{anyhow, bail, Context, Result};
12use walrus::{
13    ir::Value, ConstExpr, ElementId, ElementItems, FunctionBuilder, FunctionId, FunctionKind,
14    GlobalId, GlobalKind, MemoryId, Module, RawCustomSection, ValType,
15};
16use wasmparser::{BinaryReader, WasmFeatures};
17
18/// Get a Wasm module's canonical linear memory.
19pub fn get_memory(module: &Module) -> Result<MemoryId> {
20    let mut memories = module.memories.iter().map(|m| m.id());
21    let memory = memories.next();
22    if memories.next().is_some() {
23        bail!(
24            "expected a single memory, found multiple; multiple memories \
25             currently not supported"
26        );
27    }
28    memory.ok_or_else(|| {
29        anyhow!(
30            "module does not have a memory; must have a memory \
31             to transform return pointers into Wasm multi-value"
32        )
33    })
34}
35
36/// Get the `__stack_pointer`.
37pub fn get_stack_pointer(module: &Module) -> Option<GlobalId> {
38    if let Some(g) = module
39        .globals
40        .iter()
41        .find(|g| matches!(g.name.as_deref(), Some("__stack_pointer")))
42    {
43        return Some(g.id());
44    }
45
46    let candidates = module
47        .globals
48        .iter()
49        .filter(|g| g.ty == ValType::I32)
50        .filter(|g| g.mutable)
51        // The stack pointer is guaranteed to not be initialized to 0, and it's
52        // guaranteed to have an i32 initializer, so find globals which are
53        // locally defined, are an i32, and have a nonzero initializer
54        .filter(|g| match g.kind {
55            GlobalKind::Local(ConstExpr::Value(Value::I32(n))) => n != 0,
56            _ => false,
57        })
58        .collect::<Vec<_>>();
59
60    match candidates.len() {
61        0 => None,
62        1 => Some(candidates[0].id()),
63        2 => {
64            log::warn!("Unable to accurately determine the location of `__stack_pointer`");
65            Some(candidates[0].id())
66        }
67        _ => None,
68    }
69}
70
71/// Get the `__tls_base`.
72pub fn get_tls_base(module: &Module) -> Option<GlobalId> {
73    let candidates = module
74        .exports
75        .iter()
76        .filter(|ex| ex.name == "__tls_base")
77        .filter_map(|ex| match ex.item {
78            walrus::ExportItem::Global(id) => Some(id),
79            _ => None,
80        })
81        .filter(|id| {
82            let global = module.globals.get(*id);
83
84            global.ty == ValType::I32
85        })
86        .collect::<Vec<_>>();
87
88    match candidates.len() {
89        1 => Some(candidates[0]),
90        _ => None,
91    }
92}
93
94pub struct FunctionTableEntry {
95    pub element: ElementId,
96    pub idx: usize,
97    pub func: Option<FunctionId>,
98}
99
100/// Looks up a function table entry by index in the main function table.
101pub fn get_function_table_entry(module: &Module, idx: u32) -> Result<FunctionTableEntry> {
102    let table = module
103        .tables
104        .main_function_table()?
105        .ok_or_else(|| anyhow!("no function table found in module"))?;
106    let table = module.tables.get(table);
107    for &segment in table.elem_segments.iter() {
108        let segment = module.elements.get(segment);
109        let offset = match &segment.kind {
110            walrus::ElementKind::Active {
111                offset: ConstExpr::Value(Value::I32(n)),
112                ..
113            } => *n as u32,
114            _ => continue,
115        };
116        let idx = (idx - offset) as usize;
117
118        let slot = match &segment.items {
119            ElementItems::Functions(items) => items.get(idx).map(Some),
120            ElementItems::Expressions(_, items) => items.get(idx).map(|item| {
121                if let ConstExpr::RefFunc(target) = item {
122                    Some(target)
123                } else {
124                    None
125                }
126            }),
127        };
128
129        match slot {
130            Some(slot) => {
131                return Ok(FunctionTableEntry {
132                    element: segment.id(),
133                    idx,
134                    func: slot.cloned(),
135                })
136            }
137            None => continue,
138        }
139    }
140    bail!("failed to find `{}` in function table", idx);
141}
142
143pub fn get_start(module: &mut Module) -> Result<FunctionId, Option<FunctionId>> {
144    match module.start {
145        Some(start) => match module.funcs.get_mut(start).kind {
146            FunctionKind::Import(_) => Err(Some(start)),
147            FunctionKind::Local(_) => Ok(start),
148            FunctionKind::Uninitialized(_) => unimplemented!(),
149        },
150        None => Err(None),
151    }
152}
153
154pub fn get_or_insert_start_builder(module: &mut Module) -> &mut FunctionBuilder {
155    let prev_start = get_start(module);
156
157    let id = match prev_start {
158        Ok(id) => id,
159        Err(prev_start) => {
160            let mut builder = FunctionBuilder::new(&mut module.types, &[], &[]);
161
162            if let Some(prev_start) = prev_start {
163                builder.func_body().call(prev_start);
164            }
165
166            let id = builder.finish(Vec::new(), &mut module.funcs);
167            module.start = Some(id);
168            id
169        }
170    };
171
172    module
173        .funcs
174        .get_mut(id)
175        .kind
176        .unwrap_local_mut()
177        .builder_mut()
178}
179
180pub fn target_feature(module: &Module, feature: &str) -> Result<bool> {
181    // Taken from <https://github.com/bytecodealliance/wasm-tools/blob/f1898f46bb9d96f0f09682415cb6ccfd6a4dca79/crates/wasmparser/src/limits.rs#L27>.
182    anyhow::ensure!(feature.len() <= 100_000, "feature name too long");
183
184    // Try to find an existing section.
185    let section = module
186        .customs
187        .iter()
188        .find(|(_, custom)| custom.name() == "target_features");
189
190    if let Some((_, section)) = section {
191        let section: &RawCustomSection = section
192            .as_any()
193            .downcast_ref()
194            .context("failed to read section")?;
195        let mut reader = BinaryReader::new(&section.data, 0, WasmFeatures::default());
196        // The first integer contains the target feature count.
197        let count = reader.read_var_u32()?;
198
199        // Try to find if the target feature is already present.
200        for _ in 0..count {
201            // First byte is the prefix.
202            let prefix = reader.read_u8()?;
203            // Read the feature.
204            let length = reader.read_var_u32()?;
205            let this_feature = reader.read_bytes(length as usize)?;
206
207            // If we found the target feature, we are done here.
208            if this_feature == feature.as_bytes() {
209                // Make sure we set any existing prefix to "enabled".
210                if prefix == b'-' {
211                    return Ok(false);
212                }
213
214                return Ok(true);
215            }
216        }
217
218        Ok(false)
219    } else {
220        Ok(false)
221    }
222}
223
224pub fn insert_target_feature(module: &mut Module, new_feature: &str) -> Result<()> {
225    // Taken from <https://github.com/bytecodealliance/wasm-tools/blob/f1898f46bb9d96f0f09682415cb6ccfd6a4dca79/crates/wasmparser/src/limits.rs#L27>.
226    anyhow::ensure!(new_feature.len() <= 100_000, "feature name too long");
227
228    // Try to find an existing section.
229    let section = module
230        .customs
231        .iter_mut()
232        .find(|(_, custom)| custom.name() == "target_features");
233
234    // If one exists, check if the target feature is already present.
235    let section = if let Some((_, section)) = section {
236        let section: &mut RawCustomSection = section
237            .as_any_mut()
238            .downcast_mut()
239            .context("failed to read section")?;
240        let mut reader = BinaryReader::new(&section.data, 0, WasmFeatures::default());
241        // The first integer contains the target feature count.
242        let count = reader.read_var_u32()?;
243
244        // Try to find if the target feature is already present.
245        for _ in 0..count {
246            // First byte is the prefix.
247            let prefix_index = reader.current_position();
248            let prefix = reader.read_u8()?;
249            // Read the feature.
250            let length = reader.read_var_u32()?;
251            let feature = reader.read_bytes(length as usize)?;
252
253            // If we found the target feature, we are done here.
254            if feature == new_feature.as_bytes() {
255                // Make sure we set any existing prefix to "enabled".
256                if prefix == b'-' {
257                    section.data[prefix_index] = b'+';
258                }
259
260                return Ok(());
261            }
262        }
263
264        section
265    } else {
266        let mut data = Vec::new();
267        leb128::write::unsigned(&mut data, 0).unwrap();
268        let id = module.customs.add(RawCustomSection {
269            name: String::from("target_features"),
270            data,
271        });
272        module.customs.get_mut(id).unwrap()
273    };
274
275    // If we couldn't find the target feature, insert it.
276
277    // The first byte contains an integer describing the target feature count, which we increase by one.
278    let mut data = Cursor::new(&section.data);
279    let count = leb128::read::unsigned(&mut data).unwrap();
280    let mut new_count = Vec::new();
281    leb128::write::unsigned(&mut new_count, count + 1).unwrap();
282    section.data.splice(0..data.position() as usize, new_count);
283    // Then we insert the "enabled" prefix at the end.
284    section.data.push(b'+');
285    // The next byte contains the length of the target feature string.
286    leb128::write::unsigned(&mut section.data, new_feature.len() as u64).unwrap();
287    // Lastly the target feature string is inserted.
288    section.data.extend(new_feature.as_bytes());
289
290    Ok(())
291}