1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use crate::{core::reader::types::export::ExportDesc, execution::execution_info::ExecutionInfo};
use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};

pub struct Lut {
    /// function_lut\[local_module_idx\]\[function_local_idx\] = (foreign_module_idx, function_foreign_idx)
    ///
    /// - Module A imports a function "foo". Inside module A, the function has the index "function_local_idx". Module A
    ///   is assigned the index "local_module_idx".
    /// - Module B exports a function "foo". Inside module B, the function has the index "function_foreign_idx". Module
    ///   B is assigned the index "foreign_module_idx".
    function_lut: Vec<Vec<(usize, usize)>>,
}

impl Lut {
    /// Create a new linker lookup-table.
    ///
    /// # Arguments
    /// - `modules`: The modules to link together.
    /// - `module_map`: A map from module name to module index within the `modules` array.
    ///
    /// # Returns
    /// A new linker lookup-table. Can return `None` if there are import directives that cannot be resolved.
    pub fn new(modules: &[ExecutionInfo], module_map: &BTreeMap<String, usize>) -> Option<Self> {
        let mut function_lut = Vec::new();
        for module in modules {
            let module_lut = module
                .store
                .funcs
                .iter()
                .filter_map(|f| f.try_into_imported())
                .map(|import| {
                    Self::manual_lookup(
                        modules,
                        module_map,
                        &import.module_name,
                        &import.function_name,
                    )
                })
                .collect::<Option<Vec<_>>>()?;

            // If there is a missing import/export pair,  we fail the entire
            // operation. Better safe than sorry...

            function_lut.push(module_lut);
        }

        Some(Self { function_lut })
    }

    /// Lookup a function by its module and function index.
    ///
    /// # Arguments
    /// - `module_idx`: The index of the module within the `modules` array passed in [Lut::new].
    /// - `function_idx`: The index of the function within the module. This index is considered in-bounds only if it is
    ///   an index of an imported function.
    ///
    /// # Returns
    /// - `None`, if the indicies are out of bound
    /// - `Some(export_module_idx, export_function_idx)`, where the new indicies are the indicies of the module which
    ///   contains the implementation of the imported function, and the implementation has the returned index within.
    pub fn lookup(&self, module_idx: usize, function_idx: usize) -> Option<(usize, usize)> {
        self.function_lut
            .get(module_idx)?
            .get(function_idx)
            .copied()
    }

    /// Manually lookup a function by its module and function name.
    ///
    /// This function is used to resolve import directives before the [Lut] is created, and can be used to resolve
    /// imports even after the [Lut] is created at the cost of speed.
    ///
    /// # Arguments
    /// - `modules`: The modules to link together.
    /// - `module_map`: A map from module name to module index within the `modules` array.
    /// - `module_name`: The name of the module which imports the function.
    /// - `function_name`: The name of the function to import.
    ///
    /// # Returns
    /// - `None`, if the module or function is not found.
    /// - `Some(export_module_idx, export_function_idx)`, where the new indicies are the indicies of the module which
    ///    contains the implementation of the imported function, and the implementation has the returned index within.
    ///    Note that this function returns the first matching function, if there are multiple functions with the same
    ///    name.
    pub fn manual_lookup(
        modules: &[ExecutionInfo],
        module_map: &BTreeMap<String, usize>,
        module_name: &str,
        function_name: &str,
    ) -> Option<(usize, usize)> {
        let module_idx = module_map.get(module_name)?;
        let module = &modules[*module_idx];

        module
            .store
            .exports
            .iter()
            .filter_map(|export| {
                if export.name == function_name {
                    Some(&export.desc)
                } else {
                    None
                }
            })
            .find_map(|desc| {
                if let ExportDesc::FuncIdx(func_idx) = desc {
                    Some((*module_idx, *func_idx))
                } else {
                    None
                }
            })
    }
}