//! Data structures to provide transformation of the source

use core::fmt;
use object::{Bytes, LittleEndian, U32};
use serde_derive::{Deserialize, Serialize};

/// Single source location to generated address mapping.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct InstructionAddressMap {
    /// Where in the source wasm binary this instruction comes from, specified
    /// in an offset of bytes from the front of the file.
    pub srcloc: FilePos,

    /// Offset from the start of the function's compiled code to where this
    /// instruction is located, or the region where it starts.
    pub code_offset: u32,
}

/// A position within an original source file,
///
/// This structure is used as a newtype wrapper around a 32-bit integer which
/// represents an offset within a file where a wasm instruction or function is
/// to be originally found.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FilePos(u32);

impl FilePos {
    /// Create a new file position with the given offset.
    pub fn new(pos: u32) -> FilePos {
        assert!(pos != u32::MAX);
        FilePos(pos)
    }

    /// Get the null file position.
    pub fn none() -> FilePos {
        FilePos(u32::MAX)
    }

    /// Is this the null file position?
    #[inline]
    pub fn is_none(&self) -> bool {
        *self == FilePos::none()
    }

    /// Returns the offset that this offset was created with.
    ///
    /// Note that positions created with `FilePos::none` and the `Default`
    /// implementation will return `None` here, whereas positions created with
    /// `FilePos::new` will return `Some`.
    pub fn file_offset(self) -> Option<u32> {
        if self.0 == u32::MAX {
            None
        } else {
            Some(self.0)
        }
    }
}

impl Default for FilePos {
    fn default() -> FilePos {
        FilePos::none()
    }
}

/// A Wasm bytecode offset relative to the start of a component (or
/// top-level module) binary.
///
/// When compiling a component, the Wasm parser returns source
/// positions relative to the entire component binary. This type
/// captures that convention. Use
/// [`ComponentPC::to_module_pc`] to convert to a
/// [`ModulePC`] given the byte offset of the module within the
/// component.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ComponentPC(u32);

impl ComponentPC {
    /// Create a new component-relative PC from a raw offset.
    pub fn new(offset: u32) -> Self {
        Self(offset)
    }

    /// Get the raw u32 offset.
    pub fn raw(self) -> u32 {
        self.0
    }

    /// Convert to a module-relative PC by subtracting the byte offset
    /// of the module within the component binary.
    pub fn to_module_pc(self, wasm_module_offset: u64) -> ModulePC {
        let offset = u32::try_from(wasm_module_offset).unwrap();
        ModulePC(self.0 - offset)
    }
}

impl fmt::Debug for ComponentPC {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ComponentPC({:#x})", self.0)
    }
}

impl fmt::Display for ComponentPC {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

/// A Wasm bytecode offset relative to the start of a core Wasm
/// module binary.
///
/// In the guest-debug system, PCs are always module-relative because
/// the debugger presents a core-Wasm view of the world where
/// components are deconstructed into individual core Wasm modules.
///
/// For standalone (non-component) modules, `ModulePC` and
/// [`ComponentPC`] values are numerically identical.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ModulePC(u32);

impl ModulePC {
    /// Create a new module-relative PC from a raw offset.
    pub fn new(offset: u32) -> Self {
        Self(offset)
    }

    /// Get the raw u32 offset.
    pub fn raw(self) -> u32 {
        self.0
    }
}

impl fmt::Debug for ModulePC {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ModulePC({:#x})", self.0)
    }
}

impl fmt::Display for ModulePC {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

/// Parse an `ELF_WASMTIME_ADDRMAP` section, returning the slice of code offsets
/// and the slice of associated file positions for each offset.
fn parse_address_map(section: &[u8]) -> Option<(&[U32<LittleEndian>], &[U32<LittleEndian>])> {
    let mut section = Bytes(section);
    // NB: this matches the encoding written by `append_to` in the
    // `compile::address_map` module.
    let count = section.read::<U32<LittleEndian>>().ok()?;
    let count = usize::try_from(count.get(LittleEndian)).ok()?;
    let (offsets, section) =
        object::slice_from_bytes::<U32<LittleEndian>>(section.0, count).ok()?;
    let (positions, section) =
        object::slice_from_bytes::<U32<LittleEndian>>(section, count).ok()?;
    debug_assert!(section.is_empty());
    Some((offsets, positions))
}

/// Lookup an `offset` within an encoded address map section, returning the
/// original `FilePos` that corresponds to the offset, if found.
///
/// This function takes a `section` as its first argument which must have been
/// created with `AddressMapSection` above. This is intended to be the raw
/// `ELF_WASMTIME_ADDRMAP` section from the compilation artifact.
///
/// The `offset` provided is a relative offset from the start of the text
/// section of the pc that is being looked up. If `offset` is out of range or
/// doesn't correspond to anything in this file then `None` is returned.
pub fn lookup_file_pos(section: &[u8], offset: usize) -> Option<FilePos> {
    let (offsets, positions) = parse_address_map(section)?;

    // First perform a binary search on the `offsets` array. This is a sorted
    // array of offsets within the text section, which is conveniently what our
    // `offset` also is. Note that we are somewhat unlikely to find a precise
    // match on the element in the array, so we're largely interested in which
    // "bucket" the `offset` falls into.
    let offset = u32::try_from(offset).ok()?;
    let index = match offsets.binary_search_by_key(&offset, |v| v.get(LittleEndian)) {
        // Exact hit!
        Ok(i) => i,

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

        // This would be at the `nth` slot, so we're at the `n-1`th slot.
        Err(n) => n - 1,
    };

    // Using the `index` we found of which bucket `offset` corresponds to we can
    // lookup the actual `FilePos` value in the `positions` array.
    let pos = positions.get(index)?;
    Some(FilePos(pos.get(LittleEndian)))
}

/// Iterate over the address map contained in the given address map section.
///
/// This function takes a `section` as its first argument which must have been
/// created with `AddressMapSection` above. This is intended to be the raw
/// `ELF_WASMTIME_ADDRMAP` section from the compilation artifact.
///
/// The yielded offsets are relative to the start of the text section for this
/// map's code object.
pub fn iterate_address_map<'a>(
    section: &'a [u8],
) -> Option<impl Iterator<Item = (u32, FilePos)> + 'a> {
    let (offsets, positions) = parse_address_map(section)?;

    Some(
        offsets
            .iter()
            .map(|o| o.get(LittleEndian))
            .zip(positions.iter().map(|pos| FilePos(pos.get(LittleEndian)))),
    )
}