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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
//! This module is used for having backtraces in the Wasm runtime.
//! Once the Compiler has compiled the ModuleInfo, and we have a set of
//! compiled functions (addresses and function index) and a module,
//! then we can use this to set a backtrace for that module.
//!
//! # Example
//! ```ignore
//! use wasmer_vm::{FRAME_INFO};
//! use wasmer_types::ModuleInfo;
//!
//! let module: ModuleInfo = ...;
//! FRAME_INFO.register(module, compiled_functions);
//! ```

use crate::types::address_map::{
    ArchivedFunctionAddressMap, ArchivedInstructionAddressMap, FunctionAddressMap,
    InstructionAddressMap,
};
use crate::types::function::{ArchivedCompiledFunctionFrameInfo, CompiledFunctionFrameInfo};
use crate::ArtifactBuildFromArchive;
use rkyv::vec::ArchivedVec;
use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use wasmer_types::lib::std::{cmp, ops::Deref};
use wasmer_types::{
    entity::{BoxedSlice, EntityRef, PrimaryMap},
    FrameInfo, LocalFunctionIndex, ModuleInfo, SourceLoc, TrapInformation,
};
use wasmer_vm::FunctionBodyPtr;

lazy_static::lazy_static! {
    /// This is a global cache of backtrace frame information for all active
    ///
    /// This global cache is used during `Trap` creation to symbolicate frames.
    /// This is populated on module compilation, and it is cleared out whenever
    /// all references to a module are dropped.
    pub static ref FRAME_INFO: RwLock<GlobalFrameInfo> = Default::default();
}

#[derive(Default)]
pub struct GlobalFrameInfo {
    /// An internal map that keeps track of backtrace frame information for
    /// each module.
    ///
    /// This map is morally a map of ranges to a map of information for that
    /// module. Each module is expected to reside in a disjoint section of
    /// contiguous memory. No modules can overlap.
    ///
    /// The key of this map is the highest address in the module and the value
    /// is the module's information, which also contains the start address.
    ranges: BTreeMap<usize, ModuleInfoFrameInfo>,
}

/// An RAII structure used to unregister a module's frame information when the
/// module is destroyed.
#[cfg_attr(feature = "artifact-size", derive(loupe::MemoryUsage))]
pub struct GlobalFrameInfoRegistration {
    /// The key that will be removed from the global `ranges` map when this is
    /// dropped.
    key: usize,
}

#[derive(Debug)]
struct ModuleInfoFrameInfo {
    start: usize,
    functions: BTreeMap<usize, FunctionInfo>,
    module: Arc<ModuleInfo>,
    frame_infos: FrameInfosVariant,
}

impl ModuleInfoFrameInfo {
    fn function_debug_info(
        &self,
        local_index: LocalFunctionIndex,
    ) -> CompiledFunctionFrameInfoVariant {
        self.frame_infos.get(local_index).unwrap()
    }

    /// Gets a function given a pc
    fn function_info(&self, pc: usize) -> Option<&FunctionInfo> {
        let (end, func) = self.functions.range(pc..).next()?;
        if func.start <= pc && pc <= *end {
            Some(func)
        } else {
            None
        }
    }
}

#[derive(Debug)]
struct FunctionInfo {
    start: usize,
    local_index: LocalFunctionIndex,
}

impl GlobalFrameInfo {
    /// Fetches frame information about a program counter in a backtrace.
    ///
    /// Returns an object if this `pc` is known to some previously registered
    /// module, or returns `None` if no information can be found.
    pub fn lookup_frame_info(&self, pc: usize) -> Option<FrameInfo> {
        let module = self.module_info(pc)?;
        let func = module.function_info(pc)?;

        // Use our relative position from the start of the function to find the
        // machine instruction that corresponds to `pc`, which then allows us to
        // map that to a wasm original source location.
        let rel_pos = pc - func.start;
        let debug_info = module.function_debug_info(func.local_index);
        let instr_map = debug_info.address_map();
        let pos = match instr_map.instructions().code_offset_by_key(rel_pos) {
            // Exact hit!
            Ok(pos) => Some(pos),

            // This *would* be at the first slot in the array, so no
            // instructions cover `pc`.
            Err(0) => None,

            // This would be at the `nth` slot, so check `n-1` to see if we're
            // part of that instruction. This happens due to the minus one when
            // this function is called form trap symbolication, where we don't
            // always get called with a `pc` that's an exact instruction
            // boundary.
            Err(n) => {
                let instr = &instr_map.instructions().get(n - 1);
                if instr.code_offset <= rel_pos && rel_pos < instr.code_offset + instr.code_len {
                    Some(n - 1)
                } else {
                    None
                }
            }
        };

        let instr = match pos {
            Some(pos) => instr_map.instructions().get(pos).srcloc,
            // Some compilers don't emit yet the full trap information for each of
            // the instructions (such as LLVM).
            // In case no specific instruction is found, we return by default the
            // start offset of the function.
            None => instr_map.start_srcloc(),
        };
        let func_index = module.module.func_index(func.local_index);
        Some(FrameInfo::new(
            module.module.name(),
            func_index.index() as u32,
            module.module.function_names.get(&func_index).cloned(),
            instr_map.start_srcloc(),
            instr,
        ))
    }

    /// Fetches trap information about a program counter in a backtrace.
    pub fn lookup_trap_info(&self, pc: usize) -> Option<TrapInformation> {
        let module = self.module_info(pc)?;
        let func = module.function_info(pc)?;
        let debug_info = module.function_debug_info(func.local_index);
        let traps = debug_info.traps();
        let idx = traps
            .binary_search_by_key(&((pc - func.start) as u32), |info| info.code_offset)
            .ok()?;
        Some(traps[idx])
    }

    /// Gets a module given a pc
    fn module_info(&self, pc: usize) -> Option<&ModuleInfoFrameInfo> {
        let (end, module_info) = self.ranges.range(pc..).next()?;
        if module_info.start <= pc && pc <= *end {
            Some(module_info)
        } else {
            None
        }
    }
}

impl Drop for GlobalFrameInfoRegistration {
    fn drop(&mut self) {
        if let Ok(mut info) = FRAME_INFO.write() {
            info.ranges.remove(&self.key);
        }
    }
}

/// Represents a continuous region of executable memory starting with a function
/// entry point.
#[derive(Debug)]
#[repr(C)]
pub struct FunctionExtent {
    /// Entry point for normal entry of the function. All addresses in the
    /// function lie after this address.
    pub ptr: FunctionBodyPtr,
    /// Length in bytes.
    pub length: usize,
}

/// The variant of the frame information which can be an owned type
/// or the explicit framed map
#[derive(Debug)]
pub enum FrameInfosVariant {
    /// Owned frame infos
    Owned(PrimaryMap<LocalFunctionIndex, CompiledFunctionFrameInfo>),
    /// Archived frame infos
    Archived(ArtifactBuildFromArchive),
}

impl FrameInfosVariant {
    /// Gets the frame info for a given local function index
    pub fn get(&self, index: LocalFunctionIndex) -> Option<CompiledFunctionFrameInfoVariant> {
        match self {
            Self::Owned(map) => map.get(index).map(CompiledFunctionFrameInfoVariant::Ref),
            Self::Archived(archive) => archive
                .get_frame_info_ref()
                .get(index)
                .map(CompiledFunctionFrameInfoVariant::Archived),
        }
    }
}

/// The variant of the compiled function frame info which can be an owned type
#[derive(Debug)]
pub enum CompiledFunctionFrameInfoVariant<'a> {
    /// A reference to the frame info
    Ref(&'a CompiledFunctionFrameInfo),
    /// An archived frame info
    Archived(&'a ArchivedCompiledFunctionFrameInfo),
}

impl CompiledFunctionFrameInfoVariant<'_> {
    /// Gets the address map for the frame info
    pub fn address_map(&self) -> FunctionAddressMapVariant<'_> {
        match self {
            CompiledFunctionFrameInfoVariant::Ref(info) => {
                FunctionAddressMapVariant::Ref(&info.address_map)
            }
            CompiledFunctionFrameInfoVariant::Archived(info) => {
                FunctionAddressMapVariant::Archived(&info.address_map)
            }
        }
    }

    /// Gets the traps for the frame info
    pub fn traps(&self) -> VecTrapInformationVariant {
        match self {
            CompiledFunctionFrameInfoVariant::Ref(info) => {
                VecTrapInformationVariant::Ref(&info.traps)
            }
            CompiledFunctionFrameInfoVariant::Archived(info) => {
                let traps = rkyv::deserialize::<_, rkyv::rancor::Error>(&info.traps).unwrap();
                VecTrapInformationVariant::Owned(traps)
            }
        }
    }
}

/// The variant of the trap information which can be an owned type
#[derive(Debug)]
pub enum VecTrapInformationVariant<'a> {
    Ref(&'a Vec<TrapInformation>),
    Owned(Vec<TrapInformation>),
}

// We need to implement it for the `Deref` in `wasmer_types` to support both `core` and `std`.
impl Deref for VecTrapInformationVariant<'_> {
    type Target = [TrapInformation];

    fn deref(&self) -> &Self::Target {
        match self {
            VecTrapInformationVariant::Ref(traps) => traps,
            VecTrapInformationVariant::Owned(traps) => traps,
        }
    }
}

#[derive(Debug)]
pub enum FunctionAddressMapVariant<'a> {
    Ref(&'a FunctionAddressMap),
    Archived(&'a ArchivedFunctionAddressMap),
}

impl FunctionAddressMapVariant<'_> {
    pub fn instructions(&self) -> FunctionAddressMapInstructionVariant {
        match self {
            FunctionAddressMapVariant::Ref(map) => {
                FunctionAddressMapInstructionVariant::Owned(&map.instructions)
            }
            FunctionAddressMapVariant::Archived(map) => {
                FunctionAddressMapInstructionVariant::Archived(&map.instructions)
            }
        }
    }

    pub fn start_srcloc(&self) -> SourceLoc {
        match self {
            FunctionAddressMapVariant::Ref(map) => map.start_srcloc,
            FunctionAddressMapVariant::Archived(map) => {
                rkyv::deserialize::<_, rkyv::rancor::Error>(&map.start_srcloc).unwrap()
            }
        }
    }

    pub fn end_srcloc(&self) -> SourceLoc {
        match self {
            FunctionAddressMapVariant::Ref(map) => map.end_srcloc,
            FunctionAddressMapVariant::Archived(map) => {
                rkyv::deserialize::<_, rkyv::rancor::Error>(&map.end_srcloc).unwrap()
            }
        }
    }

    pub fn body_offset(&self) -> usize {
        match self {
            FunctionAddressMapVariant::Ref(map) => map.body_offset,
            FunctionAddressMapVariant::Archived(map) => map.body_offset.to_native() as usize,
        }
    }

    pub fn body_len(&self) -> usize {
        match self {
            FunctionAddressMapVariant::Ref(map) => map.body_len,
            FunctionAddressMapVariant::Archived(map) => map.body_len.to_native() as usize,
        }
    }
}

#[derive(Debug)]
pub enum FunctionAddressMapInstructionVariant<'a> {
    Owned(&'a Vec<InstructionAddressMap>),
    Archived(&'a ArchivedVec<ArchivedInstructionAddressMap>),
}

impl FunctionAddressMapInstructionVariant<'_> {
    pub fn code_offset_by_key(&self, key: usize) -> Result<usize, usize> {
        match self {
            FunctionAddressMapInstructionVariant::Owned(instructions) => {
                instructions.binary_search_by_key(&key, |map| map.code_offset)
            }
            FunctionAddressMapInstructionVariant::Archived(instructions) => {
                instructions.binary_search_by_key(&key, |map| map.code_offset.to_native() as usize)
            }
        }
    }

    pub fn get(&self, index: usize) -> InstructionAddressMap {
        match self {
            FunctionAddressMapInstructionVariant::Owned(instructions) => instructions[index],
            FunctionAddressMapInstructionVariant::Archived(instructions) => instructions
                .get(index)
                .map(|map| InstructionAddressMap {
                    srcloc: rkyv::deserialize::<_, rkyv::rancor::Error>(&map.srcloc).unwrap(),
                    code_offset: map.code_offset.to_native() as usize,
                    code_len: map.code_len.to_native() as usize,
                })
                .unwrap(),
        }
    }
}

/// Registers a new compiled module's frame information.
///
/// This function will register the `names` information for all of the
/// compiled functions within `module`. If the `module` has no functions
/// then `None` will be returned. Otherwise the returned object, when
/// dropped, will be used to unregister all name information from this map.
pub fn register(
    module: Arc<ModuleInfo>,
    finished_functions: &BoxedSlice<LocalFunctionIndex, FunctionExtent>,
    frame_infos: FrameInfosVariant,
) -> Option<GlobalFrameInfoRegistration> {
    let mut min = usize::MAX;
    let mut max = 0;
    let mut functions = BTreeMap::new();
    for (
        i,
        FunctionExtent {
            ptr: start,
            length: len,
        },
    ) in finished_functions.iter()
    {
        let start = **start as usize;
        // end is "last byte" of the function code
        let end = start + len - 1;
        min = cmp::min(min, start);
        max = cmp::max(max, end);
        let func = FunctionInfo {
            start,
            local_index: i,
        };
        assert!(functions.insert(end, func).is_none());
    }
    if functions.is_empty() {
        return None;
    }

    let mut info = FRAME_INFO.write().unwrap();
    // First up assert that our chunk of jit functions doesn't collide with
    // any other known chunks of jit functions...
    if let Some((_, prev)) = info.ranges.range(max..).next() {
        assert!(prev.start > max);
    }
    if let Some((prev_end, _)) = info.ranges.range(..=min).next_back() {
        assert!(*prev_end < min);
    }

    // ... then insert our range and assert nothing was there previously
    let prev = info.ranges.insert(
        max,
        ModuleInfoFrameInfo {
            start: min,
            functions,
            module,
            frame_infos,
        },
    );
    assert!(prev.is_none());
    Some(GlobalFrameInfoRegistration { key: max })
}