1use crate::entity::{EntityRef, PrimaryMap};
8use crate::{
9    CustomSectionIndex, DataIndex, ElemIndex, ExportIndex, ExportType, ExternType, FunctionIndex,
10    FunctionType, GlobalIndex, GlobalInit, GlobalType, ImportIndex, ImportType, LocalFunctionIndex,
11    LocalGlobalIndex, LocalMemoryIndex, LocalTableIndex, LocalTagIndex, MemoryIndex, MemoryType,
12    ModuleHash, SignatureIndex, TableIndex, TableInitializer, TableType, TagIndex, TagType,
13};
14
15use indexmap::IndexMap;
16use rkyv::rancor::{Fallible, Source, Trace};
17use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
18#[cfg(feature = "enable-serde")]
19use serde::{Deserialize, Serialize};
20use std::collections::BTreeMap;
21use std::collections::HashMap;
22use std::fmt;
23use std::iter::ExactSizeIterator;
24use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
25
26#[derive(Debug, Clone, RkyvSerialize, RkyvDeserialize, Archive)]
27#[cfg_attr(feature = "artifact-size", derive(loupe::MemoryUsage))]
28#[rkyv(derive(Debug))]
29pub struct ModuleId {
30    id: usize,
31}
32
33impl ModuleId {
34    pub fn id(&self) -> String {
35        format!("{}", &self.id)
36    }
37}
38
39impl Default for ModuleId {
40    fn default() -> Self {
41        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
42        Self {
43            id: NEXT_ID.fetch_add(1, SeqCst),
44        }
45    }
46}
47
48#[derive(Debug, Hash, Eq, PartialEq, Clone, Default, RkyvSerialize, RkyvDeserialize, Archive)]
50#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
51#[rkyv(derive(PartialOrd, Ord, PartialEq, Eq, Hash, Debug))]
52pub struct ImportKey {
53    pub module: String,
55    pub field: String,
57    pub import_idx: u32,
59}
60
61impl From<(String, String, u32)> for ImportKey {
62    fn from((module, field, import_idx): (String, String, u32)) -> Self {
63        Self {
64            module,
65            field,
66            import_idx,
67        }
68    }
69}
70
71#[cfg(feature = "enable-serde")]
72mod serde_imports {
73
74    use crate::ImportIndex;
75    use crate::ImportKey;
76    use indexmap::IndexMap;
77    use serde::{Deserialize, Deserializer, Serialize, Serializer};
78
79    type InitialType = IndexMap<ImportKey, ImportIndex>;
80    type SerializedType = Vec<(ImportKey, ImportIndex)>;
81    pub fn serialize<S: Serializer>(s: &InitialType, serializer: S) -> Result<S::Ok, S::Error> {
84        let vec: SerializedType = s
85            .iter()
86            .map(|(a, b)| (a.clone(), b.clone()))
87            .collect::<Vec<_>>();
88        vec.serialize(serializer)
89    }
90
91    pub fn deserialize<'de, D: Deserializer<'de>>(
92        deserializer: D,
93    ) -> Result<InitialType, D::Error> {
94        let serialized = <SerializedType as Deserialize>::deserialize(deserializer)?;
95        Ok(serialized.into_iter().collect())
96    }
97}
98
99#[derive(Debug, Clone, Default)]
106#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
107#[cfg_attr(feature = "artifact-size", derive(loupe::MemoryUsage))]
108pub struct ModuleInfo {
109    #[cfg_attr(feature = "enable-serde", serde(skip_serializing, skip_deserializing))]
116    pub id: ModuleId,
117
118    pub hash: Option<ModuleHash>,
120
121    pub name: Option<String>,
123
124    #[cfg_attr(feature = "enable-serde", serde(with = "serde_imports"))]
130    pub imports: IndexMap<ImportKey, ImportIndex>,
131
132    pub exports: IndexMap<String, ExportIndex>,
134
135    pub start_function: Option<FunctionIndex>,
137
138    pub table_initializers: Vec<TableInitializer>,
140
141    pub passive_elements: HashMap<ElemIndex, Box<[FunctionIndex]>>,
143
144    pub passive_data: HashMap<DataIndex, Box<[u8]>>,
146
147    pub global_initializers: PrimaryMap<LocalGlobalIndex, GlobalInit>,
149
150    pub function_names: HashMap<FunctionIndex, String>,
152
153    pub signatures: PrimaryMap<SignatureIndex, FunctionType>,
155
156    pub functions: PrimaryMap<FunctionIndex, SignatureIndex>,
158
159    pub tables: PrimaryMap<TableIndex, TableType>,
161
162    pub memories: PrimaryMap<MemoryIndex, MemoryType>,
164
165    pub globals: PrimaryMap<GlobalIndex, GlobalType>,
167
168    pub tags: PrimaryMap<TagIndex, SignatureIndex>,
170
171    pub custom_sections: IndexMap<String, CustomSectionIndex>,
173
174    pub custom_sections_data: PrimaryMap<CustomSectionIndex, Box<[u8]>>,
176
177    pub num_imported_functions: usize,
179
180    pub num_imported_tables: usize,
182
183    pub num_imported_memories: usize,
185
186    pub num_imported_tags: usize,
188
189    pub num_imported_globals: usize,
191}
192
193#[derive(Debug, RkyvSerialize, RkyvDeserialize, Archive)]
195#[rkyv(derive(Debug))]
196pub struct ArchivableModuleInfo {
197    name: Option<String>,
198    hash: Option<ModuleHash>,
199    imports: IndexMap<ImportKey, ImportIndex>,
200    exports: IndexMap<String, ExportIndex>,
201    start_function: Option<FunctionIndex>,
202    table_initializers: Vec<TableInitializer>,
203    passive_elements: BTreeMap<ElemIndex, Box<[FunctionIndex]>>,
204    passive_data: BTreeMap<DataIndex, Box<[u8]>>,
205    global_initializers: PrimaryMap<LocalGlobalIndex, GlobalInit>,
206    function_names: BTreeMap<FunctionIndex, String>,
207    signatures: PrimaryMap<SignatureIndex, FunctionType>,
208    functions: PrimaryMap<FunctionIndex, SignatureIndex>,
209    tables: PrimaryMap<TableIndex, TableType>,
210    memories: PrimaryMap<MemoryIndex, MemoryType>,
211    globals: PrimaryMap<GlobalIndex, GlobalType>,
212    tags: PrimaryMap<TagIndex, SignatureIndex>,
213    custom_sections: IndexMap<String, CustomSectionIndex>,
214    custom_sections_data: PrimaryMap<CustomSectionIndex, Box<[u8]>>,
215    num_imported_functions: usize,
216    num_imported_tables: usize,
217    num_imported_tags: usize,
218    num_imported_memories: usize,
219    num_imported_globals: usize,
220}
221
222impl From<ModuleInfo> for ArchivableModuleInfo {
223    fn from(it: ModuleInfo) -> Self {
224        Self {
225            name: it.name,
226            hash: it.hash,
227            imports: it.imports,
228            exports: it.exports,
229            start_function: it.start_function,
230            table_initializers: it.table_initializers,
231            passive_elements: it.passive_elements.into_iter().collect(),
232            passive_data: it.passive_data.into_iter().collect(),
233            global_initializers: it.global_initializers,
234            function_names: it.function_names.into_iter().collect(),
235            signatures: it.signatures,
236            functions: it.functions,
237            tables: it.tables,
238            memories: it.memories,
239            globals: it.globals,
240            tags: it.tags,
241            custom_sections: it.custom_sections,
242            custom_sections_data: it.custom_sections_data,
243            num_imported_functions: it.num_imported_functions,
244            num_imported_tables: it.num_imported_tables,
245            num_imported_tags: it.num_imported_tags,
246            num_imported_memories: it.num_imported_memories,
247            num_imported_globals: it.num_imported_globals,
248        }
249    }
250}
251
252impl From<ArchivableModuleInfo> for ModuleInfo {
253    fn from(it: ArchivableModuleInfo) -> Self {
254        Self {
255            id: Default::default(),
256            name: it.name,
257            hash: it.hash,
258            imports: it.imports,
259            exports: it.exports,
260            start_function: it.start_function,
261            table_initializers: it.table_initializers,
262            passive_elements: it.passive_elements.into_iter().collect(),
263            passive_data: it.passive_data.into_iter().collect(),
264            global_initializers: it.global_initializers,
265            function_names: it.function_names.into_iter().collect(),
266            signatures: it.signatures,
267            functions: it.functions,
268            tables: it.tables,
269            memories: it.memories,
270            globals: it.globals,
271            tags: it.tags,
272            custom_sections: it.custom_sections,
273            custom_sections_data: it.custom_sections_data,
274            num_imported_functions: it.num_imported_functions,
275            num_imported_tables: it.num_imported_tables,
276            num_imported_tags: it.num_imported_tags,
277            num_imported_memories: it.num_imported_memories,
278            num_imported_globals: it.num_imported_globals,
279        }
280    }
281}
282
283impl From<&ModuleInfo> for ArchivableModuleInfo {
284    fn from(it: &ModuleInfo) -> Self {
285        Self::from(it.clone())
286    }
287}
288
289impl Archive for ModuleInfo {
290    type Archived = <ArchivableModuleInfo as Archive>::Archived;
291    type Resolver = <ArchivableModuleInfo as Archive>::Resolver;
292
293    fn resolve(&self, resolver: Self::Resolver, out: rkyv::Place<Self::Archived>) {
294        ArchivableModuleInfo::from(self).resolve(resolver, out)
295    }
296}
297
298impl<S: rkyv::ser::Allocator + rkyv::ser::Writer + Fallible + ?Sized> RkyvSerialize<S>
299    for ModuleInfo
300where
301    <S as Fallible>::Error: rkyv::rancor::Source + rkyv::rancor::Trace,
302{
303    fn serialize(&self, serializer: &mut S) -> Result<Self::Resolver, S::Error> {
304        ArchivableModuleInfo::from(self).serialize(serializer)
305    }
306}
307
308impl<D: Fallible + ?Sized> RkyvDeserialize<ModuleInfo, D> for ArchivedArchivableModuleInfo
309where
310    D::Error: Source + Trace,
311{
312    fn deserialize(&self, deserializer: &mut D) -> Result<ModuleInfo, D::Error> {
313        let archived = RkyvDeserialize::<ArchivableModuleInfo, D>::deserialize(self, deserializer)?;
314        Ok(ModuleInfo::from(archived))
315    }
316}
317
318impl PartialEq for ModuleInfo {
320    fn eq(&self, other: &Self) -> bool {
321        self.name == other.name
322            && self.imports == other.imports
323            && self.exports == other.exports
324            && self.start_function == other.start_function
325            && self.table_initializers == other.table_initializers
326            && self.passive_elements == other.passive_elements
327            && self.passive_data == other.passive_data
328            && self.global_initializers == other.global_initializers
329            && self.function_names == other.function_names
330            && self.signatures == other.signatures
331            && self.functions == other.functions
332            && self.tables == other.tables
333            && self.memories == other.memories
334            && self.globals == other.globals
335            && self.tags == other.tags
336            && self.custom_sections == other.custom_sections
337            && self.custom_sections_data == other.custom_sections_data
338            && self.num_imported_functions == other.num_imported_functions
339            && self.num_imported_tables == other.num_imported_tables
340            && self.num_imported_tags == other.num_imported_tags
341            && self.num_imported_memories == other.num_imported_memories
342            && self.num_imported_globals == other.num_imported_globals
343    }
344}
345
346impl Eq for ModuleInfo {}
347
348impl ModuleInfo {
349    pub fn new() -> Self {
351        Default::default()
352    }
353
354    pub fn hash(&self) -> Option<ModuleHash> {
356        self.hash
357    }
358
359    pub fn get_passive_element(&self, index: ElemIndex) -> Option<&[FunctionIndex]> {
361        self.passive_elements.get(&index).map(|es| &**es)
362    }
363
364    pub fn exported_signatures(&self) -> Vec<FunctionType> {
366        self.exports
367            .iter()
368            .filter_map(|(_name, export_index)| match export_index {
369                ExportIndex::Function(i) => {
370                    let signature = self.functions.get(*i).unwrap();
371                    let func_type = self.signatures.get(*signature).unwrap();
372                    Some(func_type.clone())
373                }
374                _ => None,
375            })
376            .collect::<Vec<FunctionType>>()
377    }
378
379    pub fn exports(&'_ self) -> ExportsIterator<Box<dyn Iterator<Item = ExportType> + '_>> {
381        let iter = self.exports.iter().map(move |(name, export_index)| {
382            let extern_type = match export_index {
383                ExportIndex::Function(i) => {
384                    let signature = self.functions.get(*i).unwrap();
385                    let func_type = self.signatures.get(*signature).unwrap();
386                    ExternType::Function(func_type.clone())
387                }
388                ExportIndex::Table(i) => {
389                    let table_type = self.tables.get(*i).unwrap();
390                    ExternType::Table(*table_type)
391                }
392                ExportIndex::Memory(i) => {
393                    let memory_type = self.memories.get(*i).unwrap();
394                    ExternType::Memory(*memory_type)
395                }
396                ExportIndex::Global(i) => {
397                    let global_type = self.globals.get(*i).unwrap();
398                    ExternType::Global(*global_type)
399                }
400                ExportIndex::Tag(i) => {
401                    let signature = self.tags.get(*i).unwrap();
402                    let tag_type = self.signatures.get(*signature).unwrap();
403
404                    ExternType::Tag(TagType {
405                        kind: crate::types::TagKind::Exception,
406                        params: tag_type.params().into(),
407                    })
408                }
409            };
410            ExportType::new(name, extern_type)
411        });
412        ExportsIterator::new(Box::new(iter), self.exports.len())
413    }
414
415    pub fn imports(&'_ self) -> ImportsIterator<Box<dyn Iterator<Item = ImportType> + '_>> {
417        let iter =
418            self.imports
419                .iter()
420                .map(move |(ImportKey { module, field, .. }, import_index)| {
421                    let extern_type = match import_index {
422                        ImportIndex::Function(i) => {
423                            let signature = self.functions.get(*i).unwrap();
424                            let func_type = self.signatures.get(*signature).unwrap();
425                            ExternType::Function(func_type.clone())
426                        }
427                        ImportIndex::Table(i) => {
428                            let table_type = self.tables.get(*i).unwrap();
429                            ExternType::Table(*table_type)
430                        }
431                        ImportIndex::Memory(i) => {
432                            let memory_type = self.memories.get(*i).unwrap();
433                            ExternType::Memory(*memory_type)
434                        }
435                        ImportIndex::Global(i) => {
436                            let global_type = self.globals.get(*i).unwrap();
437                            ExternType::Global(*global_type)
438                        }
439                        ImportIndex::Tag(i) => {
440                            let tag_type = self.tags.get(*i).unwrap();
441                            let func_type = self.signatures.get(*tag_type).unwrap();
442                            ExternType::Tag(TagType::from_fn_type(
443                                crate::TagKind::Exception,
444                                func_type.clone(),
445                            ))
446                        }
447                    };
448                    ImportType::new(module, field, extern_type)
449                });
450        ImportsIterator::new(Box::new(iter), self.imports.len())
451    }
452
453    pub fn custom_sections<'a>(
455        &'a self,
456        name: &'a str,
457    ) -> Box<impl Iterator<Item = Box<[u8]>> + 'a> {
458        Box::new(
459            self.custom_sections
460                .iter()
461                .filter_map(move |(section_name, section_index)| {
462                    if name != section_name {
463                        return None;
464                    }
465                    Some(self.custom_sections_data[*section_index].clone())
466                }),
467        )
468    }
469
470    pub fn func_index(&self, local_func: LocalFunctionIndex) -> FunctionIndex {
472        FunctionIndex::new(self.num_imported_functions + local_func.index())
473    }
474
475    pub fn local_func_index(&self, func: FunctionIndex) -> Option<LocalFunctionIndex> {
478        func.index()
479            .checked_sub(self.num_imported_functions)
480            .map(LocalFunctionIndex::new)
481    }
482
483    pub fn is_imported_function(&self, index: FunctionIndex) -> bool {
485        index.index() < self.num_imported_functions
486    }
487
488    pub fn table_index(&self, local_table: LocalTableIndex) -> TableIndex {
490        TableIndex::new(self.num_imported_tables + local_table.index())
491    }
492
493    pub fn local_table_index(&self, table: TableIndex) -> Option<LocalTableIndex> {
496        table
497            .index()
498            .checked_sub(self.num_imported_tables)
499            .map(LocalTableIndex::new)
500    }
501
502    pub fn is_imported_table(&self, index: TableIndex) -> bool {
504        index.index() < self.num_imported_tables
505    }
506
507    pub fn memory_index(&self, local_memory: LocalMemoryIndex) -> MemoryIndex {
509        MemoryIndex::new(self.num_imported_memories + local_memory.index())
510    }
511
512    pub fn local_memory_index(&self, memory: MemoryIndex) -> Option<LocalMemoryIndex> {
515        memory
516            .index()
517            .checked_sub(self.num_imported_memories)
518            .map(LocalMemoryIndex::new)
519    }
520
521    pub fn is_imported_memory(&self, index: MemoryIndex) -> bool {
523        index.index() < self.num_imported_memories
524    }
525
526    pub fn global_index(&self, local_global: LocalGlobalIndex) -> GlobalIndex {
528        GlobalIndex::new(self.num_imported_globals + local_global.index())
529    }
530
531    pub fn local_global_index(&self, global: GlobalIndex) -> Option<LocalGlobalIndex> {
534        global
535            .index()
536            .checked_sub(self.num_imported_globals)
537            .map(LocalGlobalIndex::new)
538    }
539
540    pub fn is_imported_global(&self, index: GlobalIndex) -> bool {
542        index.index() < self.num_imported_globals
543    }
544
545    pub fn tag_index(&self, local_tag: LocalTagIndex) -> TagIndex {
547        TagIndex::new(self.num_imported_tags + local_tag.index())
548    }
549
550    pub fn local_tag_index(&self, tag: TagIndex) -> Option<LocalTagIndex> {
553        tag.index()
554            .checked_sub(self.num_imported_tags)
555            .map(LocalTagIndex::new)
556    }
557
558    pub fn is_imported_tag(&self, index: TagIndex) -> bool {
560        index.index() < self.num_imported_tags
561    }
562
563    pub fn name(&self) -> String {
565        match self.name {
566            Some(ref name) => name.to_string(),
567            None => "<module>".to_string(),
568        }
569    }
570
571    pub fn imported_function_types(&'_ self) -> impl Iterator<Item = FunctionType> + '_ {
573        self.functions
574            .values()
575            .take(self.num_imported_functions)
576            .map(move |sig_index| self.signatures[*sig_index].clone())
577    }
578}
579
580impl fmt::Display for ModuleInfo {
581    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
582        write!(f, "{}", self.name())
583    }
584}
585
586pub struct ExportsIterator<I: Iterator<Item = ExportType> + Sized> {
592    iter: I,
593    size: usize,
594}
595
596impl<I: Iterator<Item = ExportType> + Sized> ExportsIterator<I> {
597    pub fn new(iter: I, size: usize) -> Self {
599        Self { iter, size }
600    }
601}
602
603impl<I: Iterator<Item = ExportType> + Sized> ExactSizeIterator for ExportsIterator<I> {
604    fn len(&self) -> usize {
606        self.size
607    }
608}
609
610impl<I: Iterator<Item = ExportType> + Sized> ExportsIterator<I> {
611    pub fn functions(self) -> impl Iterator<Item = ExportType<FunctionType>> + Sized {
613        self.iter.filter_map(|extern_| match extern_.ty() {
614            ExternType::Function(ty) => Some(ExportType::new(extern_.name(), ty.clone())),
615            _ => None,
616        })
617    }
618    pub fn memories(self) -> impl Iterator<Item = ExportType<MemoryType>> + Sized {
620        self.iter.filter_map(|extern_| match extern_.ty() {
621            ExternType::Memory(ty) => Some(ExportType::new(extern_.name(), *ty)),
622            _ => None,
623        })
624    }
625    pub fn tables(self) -> impl Iterator<Item = ExportType<TableType>> + Sized {
627        self.iter.filter_map(|extern_| match extern_.ty() {
628            ExternType::Table(ty) => Some(ExportType::new(extern_.name(), *ty)),
629            _ => None,
630        })
631    }
632    pub fn globals(self) -> impl Iterator<Item = ExportType<GlobalType>> + Sized {
634        self.iter.filter_map(|extern_| match extern_.ty() {
635            ExternType::Global(ty) => Some(ExportType::new(extern_.name(), *ty)),
636            _ => None,
637        })
638    }
639}
640
641impl<I: Iterator<Item = ExportType> + Sized> Iterator for ExportsIterator<I> {
642    type Item = ExportType;
643    fn next(&mut self) -> Option<Self::Item> {
644        self.iter.next()
645    }
646}
647
648pub struct ImportsIterator<I: Iterator<Item = ImportType> + Sized> {
651    iter: I,
652    size: usize,
653}
654
655impl<I: Iterator<Item = ImportType> + Sized> ImportsIterator<I> {
656    pub fn new(iter: I, size: usize) -> Self {
658        Self { iter, size }
659    }
660}
661
662impl<I: Iterator<Item = ImportType> + Sized> ExactSizeIterator for ImportsIterator<I> {
663    fn len(&self) -> usize {
665        self.size
666    }
667}
668
669impl<I: Iterator<Item = ImportType> + Sized> ImportsIterator<I> {
670    pub fn functions(self) -> impl Iterator<Item = ImportType<FunctionType>> + Sized {
672        self.iter.filter_map(|extern_| match extern_.ty() {
673            ExternType::Function(ty) => Some(ImportType::new(
674                extern_.module(),
675                extern_.name(),
676                ty.clone(),
677            )),
678            _ => None,
679        })
680    }
681    pub fn memories(self) -> impl Iterator<Item = ImportType<MemoryType>> + Sized {
683        self.iter.filter_map(|extern_| match extern_.ty() {
684            ExternType::Memory(ty) => Some(ImportType::new(extern_.module(), extern_.name(), *ty)),
685            _ => None,
686        })
687    }
688    pub fn tables(self) -> impl Iterator<Item = ImportType<TableType>> + Sized {
690        self.iter.filter_map(|extern_| match extern_.ty() {
691            ExternType::Table(ty) => Some(ImportType::new(extern_.module(), extern_.name(), *ty)),
692            _ => None,
693        })
694    }
695    pub fn globals(self) -> impl Iterator<Item = ImportType<GlobalType>> + Sized {
697        self.iter.filter_map(|extern_| match extern_.ty() {
698            ExternType::Global(ty) => Some(ImportType::new(extern_.module(), extern_.name(), *ty)),
699            _ => None,
700        })
701    }
702}
703
704impl<I: Iterator<Item = ImportType> + Sized> Iterator for ImportsIterator<I> {
705    type Item = ImportType;
706    fn next(&mut self) -> Option<Self::Item> {
707        self.iter.next()
708    }
709}