wamex_cli/analysis/
mod.rs

1//!
2//! Module with external info usefull to build dep graph, and request information about function and data entries.
3//!
4
5use std::{cmp::Ordering, collections::HashMap, fmt::Debug, ops::Range};
6
7use anyhow::{Context, Result, anyhow, bail, ensure};
8pub use symbols::{StaticModuleInfo, SymbolMap};
9use wasmparser::{ElementItems, ElementKind, TypeRef};
10
11use crate::{
12    index::{
13        AnySymbolId, DefinedFuncId, ElementId, ExportId, FuncTypeId, IdMap, ImportId, InputFuncId,
14        InputGlobalId, TableId,
15    },
16    read,
17};
18mod debug;
19pub mod dep_graph;
20pub mod split_point;
21pub mod symbols;
22#[cfg(test)]
23mod testing;
24
25#[derive(Debug, PartialEq, Eq, Clone)]
26pub struct ImportInfo {
27    // List of imported functions
28    pub imported_funcs: Vec<ImportId>,
29    pub imported_func_map: IdMap<ImportId, InputFuncId>,
30
31    pub imported_globals: Vec<ImportId>,
32    pub imported_global_map: IdMap<ImportId, InputGlobalId>,
33}
34
35/// Provides a additional info about module.
36/// Like ordered_data_symbols - ordered by offsets where symbol is defined (relative to module start)
37/// and info about imported functions
38pub struct ModuleInfo<'src> {
39    pub import_info: ImportInfo,
40
41    pub wasm: read::InputModule<'src>,
42    pub export_map: HashMap<(isize, AnySymbolId), (ExportId, &'src str)>,
43    pub symbols: SymbolMap<'src>,
44
45    pub indirect_function_table_id: (TableId, ElementId),
46    pub indirect_function_list: Vec<InputFuncId>,
47}
48
49impl<'src> ModuleInfo<'src> {
50    pub fn from_wasm_bytes(wasm_bytes: &'src [u8]) -> Result<Self> {
51        let module = read::InputModule::parse(&wasm_bytes)?;
52        Self::from_raw_module(module)
53    }
54    pub fn from_raw_module(module: read::InputModule<'src>) -> Result<Self> {
55        //TODO: Maybe we should use `IdMap` here?
56        let mut imported_funcs: Vec<ImportId> = Vec::new();
57        let mut imported_globals: Vec<ImportId> = Vec::new();
58
59        for (import_id, import) in module.imports.iter() {
60            match import.ty {
61                TypeRef::Global(_) => {
62                    imported_globals.push(import_id);
63                    continue;
64                }
65                TypeRef::Func(_) => {
66                    imported_funcs.push(import_id);
67                }
68                _ => {}
69            }
70        }
71        let imported_func_map = imported_funcs
72            .iter()
73            .enumerate()
74            .map(|(func_id, &import_id)| (import_id, InputFuncId::from_index(func_id)))
75            .collect();
76        let imported_global_map = imported_globals
77            .iter()
78            .enumerate()
79            .map(|(global_id, &import_id)| (import_id, InputGlobalId::from_index(global_id)))
80            .collect();
81
82        let import_funcs_info = ImportInfo {
83            imported_funcs,
84            imported_func_map,
85            imported_globals,
86            imported_global_map,
87        };
88        let export_map = module
89            .exports
90            .iter()
91            .map(|(i, export)| {
92                (
93                    (export.kind as isize, export.index as AnySymbolId),
94                    (i, export.name),
95                )
96            })
97            .collect();
98
99        let (_table_name, table_id) = module
100            .tables
101            .iter()
102            .filter_map(|(id, _)| module.names.tables.get(id).map(|name| (*name, id)))
103            .find(|(name, _)| *name == "__indirect_function_table")
104            .unwrap_or_else(|| {
105                assert!(
106                    module.tables.len() == 1,
107                    "No named __indirect_function_table was found, and there is not one table in the module."
108                );
109                (
110                    "__indirect_function_table",
111                    module.tables.iter().next().unwrap().0,
112                )
113            });
114
115        let mut indirect_element = None;
116        for (id, element) in module.elements.iter() {
117            let ElementKind::Active {
118                table_index,
119                offset_expr,
120            } = &element.kind
121            else {
122                continue;
123            };
124
125            if !table_index.is_none()  // None for first index.
126               && table_index.unwrap() == table_id.as_raw_index() as u32
127            {
128                continue;
129            }
130
131            let offset = Self::read_const_expr(offset_expr)
132                .with_context(|| format!("Failed to read offset expression for element {id:?}"))?;
133
134            ensure!(
135                offset == 1,
136                "Element segment {id:?} should be inited with 1 offset, but got {offset}, which is not supported"
137            );
138
139            let ElementItems::Functions(functions) = &element.items else {
140                bail!("Only function elements are supported, but got constant instead");
141            };
142
143            let mut function_list = Vec::with_capacity(functions.count() as usize);
144            for function_id in functions.clone().into_iter() {
145                let raw_function_id = function_id
146                    .with_context(|| format!("Failed to read function ID from element {id:?}"))?;
147                function_list.push(InputFuncId::from_index(raw_function_id));
148            }
149            indirect_element = Some((id, function_list));
150            break;
151        }
152        let (indirect_element_id, indirect_function_list) = indirect_element
153            .ok_or_else(|| anyhow!("No element segment with __indirect_function_table found"))?;
154
155        let symbols_map = symbols::SymbolMap::new(&module, import_funcs_info.imported_funcs.len())?;
156
157        Ok(ModuleInfo {
158            import_info: import_funcs_info,
159            symbols: symbols_map,
160            wasm: module,
161            export_map,
162            indirect_function_list,
163            indirect_function_table_id: (table_id, indirect_element_id),
164        })
165    }
166
167    pub(crate) fn read_const_expr(offset_expr: &wasmparser::ConstExpr<'_>) -> Result<i32> {
168        let mut reader = offset_expr.get_operators_reader();
169
170        let val = match reader.read()? {
171            wasmparser::Operator::I32Const { value } => Ok(value),
172            op => bail!("Expected only I32.const operator, found: {:?}", op),
173        };
174        match reader.read()? {
175            wasmparser::Operator::End => {}
176            op => bail!("Expected End after I32.const: {:?}", op),
177        }
178        val
179    }
180    pub fn function_id_iter<'any>(
181        &'any self,
182    ) -> impl Iterator<Item = InputFuncId> + use<'any, 'src> {
183        (0..self.import_info.imported_funcs.len())
184            .map(InputFuncId::from_index)
185            .chain(
186                self.wasm
187                    .code
188                    .section_payload
189                    .defined_funcs
190                    .iter()
191                    .enumerate()
192                    .map(|(index, _)| {
193                        InputFuncId::from_index(index + self.import_info.imported_funcs.len())
194                    }),
195            )
196    }
197
198    pub fn is_imported_function(&self, func_id: InputFuncId) -> bool {
199        func_id.as_raw_index() < self.import_info.imported_funcs.len()
200    }
201
202    pub fn as_defined_function_id(&self, func_id: InputFuncId) -> Option<DefinedFuncId> {
203        if self.is_imported_function(func_id) {
204            None
205        } else {
206            Some(DefinedFuncId::from_index(
207                func_id
208                    .as_raw_index()
209                    .checked_sub(self.import_info.imported_funcs.len())
210                    .expect("Function ID is out of bounds") as u32,
211            ))
212        }
213    }
214
215    pub fn get_function_type_id(&self, func_id: InputFuncId) -> FuncTypeId {
216        let Some(defined_index) = self.as_defined_function_id(func_id) else {
217            // It's import function - recover from import id.
218            let import_id = self.import_info.imported_funcs[func_id.as_raw_index()];
219            let TypeRef::Func(ty) = self.wasm.imports[import_id].ty else {
220                panic!("Expected function type")
221            };
222            return FuncTypeId::from_index(ty);
223        };
224        // It's a defined function.
225        self.wasm.defined_func_type_id(defined_index)
226    }
227
228    pub fn get_function_import_id(&self, func_id: InputFuncId) -> Option<ImportId> {
229        self.import_info
230            .imported_funcs
231            .get(func_id.as_raw_index())
232            .copied()
233    }
234    pub fn get_global_import_id(&self, global_id: InputGlobalId) -> Option<ImportId> {
235        self.import_info
236            .imported_globals
237            .get(global_id.as_raw_index())
238            .copied()
239    }
240
241    pub fn find_function_id_by_name(&self, name: &str) -> Option<InputFuncId> {
242        let func = self.wasm.names.functions.iter().find(|f| *f.1 == name)?;
243        Some(func.0)
244    }
245
246    pub fn find_global_id_by_name(&self, name: &str) -> Option<InputGlobalId> {
247        let global = self.wasm.names.globals.iter().find(|f| *f.1 == name)?;
248        Some(global.0)
249    }
250
251    pub fn find_function_id_containing_range(&self, range: Range<usize>) -> Result<InputFuncId> {
252        let func_index = Self::find_by_range(
253            self.wasm.code.section_payload.defined_funcs.as_slice(),
254            &range,
255            |defined_func| defined_func.body.range(),
256        )
257        .with_context(|| format!("No match for function relocation range {range:?}"))?;
258        Ok(InputFuncId::from_index(
259            func_index + self.import_info.imported_funcs.len(),
260        ))
261    }
262
263    fn find_by_range<T: Debug, U: Debug + Ord, F: Fn(&T) -> Range<U>>(
264        items: &[T],
265        range: &Range<U>,
266        get_range: F,
267    ) -> anyhow::Result<usize> {
268        let index = items
269            .binary_search_by(|item| {
270                let item_range = get_range(item);
271                if item_range.end <= range.start {
272                    Ordering::Less
273                } else if item_range.start <= range.start {
274                    Ordering::Equal
275                } else {
276                    Ordering::Greater
277                }
278            })
279            .or_else(|index| {
280                bail!(
281                    "Prev range is: {:?}, next range is: {:?}",
282                    index
283                        .checked_sub(1)
284                        .and_then(|i| items.get(i).map(|item| (item, get_range(item)))),
285                    items.get(index).map(|item| (item, get_range(item)))
286                )
287            })?;
288        if range.end > get_range(&items[index]).end {
289            bail!(
290                "Item {:?} has incompatible range {:?}",
291                items[index],
292                get_range(&items[index])
293            )
294        }
295        Ok(index)
296    }
297}
298
299impl<'src> Debug for ModuleInfo<'src> {
300    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
301        f.debug_struct("ModuleInfo")
302            .field("import_funcs_info", &self.import_info)
303            .finish()
304    }
305}
wamex_cli/analysis/mod.rs

wamex_cli/analysis/
mod.rs
