r2fas 0.2.1

//! Assembly dump rows (FAS.TXT table 4).
//!
//! Each row says: this preprocessed line was assembled at file offset X and
//! virtual address `$`. A trailing `u32` after the row array is the output
//! file offset where assembly ended.

use crate::error::{FasError, Result};
use crate::fas::bytes;
use crate::fas::header::Header;
use std::collections::BTreeSet;

/// Size of one dump row in bytes.
pub const DUMP_ROW_SIZE: usize = 28;

/// Dump flags bit 0: assembled inside a `virtual` block (file offset meaningless).
pub const DUMP_VIRTUAL: u8 = 1;
/// Dump flags bit 1: not present in the output file (reserved tail, etc.).
pub const DUMP_NOT_IN_OUTPUT: u8 = 2;

/// One 28-byte assembly-dump row.
#[derive(Debug, Clone)]
pub struct DumpRow {
    /// Stable zero-based row index.
    pub id: usize,
    /// Offset in the output file.
    pub file_offset: u32,
    /// Offset of the preprocessed line that produced these bytes.
    pub line_off: u32,
    /// Low 64 bits of the `$` address.
    pub dollar: u64,
    /// High 8 bits of `$` (almost always 0 on 64-bit targets).
    pub dollar_hi: u8,
    /// Extended SIB associated with the `$` value.
    pub extended_sib: u32,
    /// Section/external relocation information for `$`.
    pub reloc: u32,
    /// Value type of `$` (table 2.2).
    pub addr_type: u8,
    /// Code size at this point: 16, 32, or 64.
    pub code_type: u8,
    /// Dump flags (virtual / not-in-output).
    pub flags: u8,
}

impl DumpRow {
    /// Parse one row from a 28-byte slice.
    pub fn parse(rec: &[u8]) -> Result<Self> {
        Self::parse_with_id(0, rec)
    }

    /// Parse one row and assign its stable table index.
    pub fn parse_with_id(id: usize, rec: &[u8]) -> Result<Self> {
        if rec.len() < DUMP_ROW_SIZE {
            return Err(FasError::Truncated("dump row"));
        }
        Ok(Self {
            id,
            file_offset: bytes::u32_at(rec, 0)?,
            line_off: bytes::u32_at(rec, 4)?,
            dollar: bytes::u64_at(rec, 8)?,
            dollar_hi: bytes::u8_at(rec, 27)?,
            extended_sib: bytes::u32_at(rec, 16)?,
            reloc: bytes::u32_at(rec, 20)?,
            addr_type: bytes::u8_at(rec, 24)?,
            code_type: bytes::u8_at(rec, 25)?,
            flags: bytes::u8_at(rec, 26)?,
        })
    }

    /// `$` as a 64-bit VA. High-byte extensions beyond 64 bits are ignored.
    pub fn address(&self) -> u64 {
        self.dollar
    }

    /// Assembled inside a `virtual` block.
    pub fn is_virtual(&self) -> bool {
        self.flags & DUMP_VIRTUAL != 0
    }

    /// Bytes were not written to the output file.
    pub fn not_in_output(&self) -> bool {
        self.flags & DUMP_NOT_IN_OUTPUT != 0
    }

    /// Useful for mapping a real load address back to source.
    pub fn is_mappable(&self) -> bool {
        !self.is_virtual() && !self.not_in_output() && self.address() != 0
    }
}

/// Whether and how many output bytes an assembly row owns.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmittedSpan {
    /// The row does not own bytes in the generated output.
    None,
    /// The row owns this many output bytes.
    Bytes(u32),
    /// Output coordinates did not allow an unambiguous extent.
    Ambiguous,
}

/// Parsed dump: rows plus the trailing end-of-assembly file offset.
#[derive(Debug, Clone, Default)]
pub struct AssemblyDump {
    /// Dump rows in assembly order.
    pub rows: Vec<DumpRow>,
    /// Output file offset at which assembly ended.
    pub end_file_offset: u32,
    /// Derived emitted span for each row, indexed by [`DumpRow::id`].
    pub spans: Vec<EmittedSpan>,
}

impl AssemblyDump {
    /// Parse table 4. `raw` includes the trailing end-offset dword.
    pub fn parse(raw: &[u8]) -> Result<Self> {
        if raw.len() < 4 {
            return Err(FasError::Truncated("dump"));
        }
        let body_len = raw.len() - 4;
        if body_len % DUMP_ROW_SIZE != 0 {
            return Err(FasError::Geometry("dump rows not a multiple of 28"));
        }
        let rows = raw[..body_len]
            .chunks_exact(DUMP_ROW_SIZE)
            .enumerate()
            .map(|(id, record)| DumpRow::parse_with_id(id, record))
            .collect::<Result<Vec<_>>>()?;
        let end_file_offset = bytes::u32_at(raw, body_len)?;
        let spans = derive_spans(&rows, end_file_offset);
        Ok(Self {
            rows,
            end_file_offset,
            spans,
        })
    }

    /// Derived output span for a row.
    pub fn span(&self, row: &DumpRow) -> EmittedSpan {
        self.spans
            .get(row.id)
            .copied()
            .unwrap_or(EmittedSpan::Ambiguous)
    }

    /// Whether a row emitted at least one byte into the output.
    pub fn emitted(&self, row: &DumpRow) -> bool {
        matches!(self.span(row), EmittedSpan::Bytes(length) if length > 0)
    }

    /// Set of `$` addresses at which FASM emitted output bytes.
    pub fn assembled_addresses(&self) -> BTreeSet<u64> {
        self.rows
            .iter()
            .filter(|row| row.is_mappable() && self.emitted(row))
            .map(DumpRow::address)
            .collect()
    }

    /// Most common `address - file_offset` among mappable rows.
    ///
    /// For `format ELF64 executable` this is the original image base
    /// (`0x400000` on a typical FASM Linux binary).
    pub fn inferred_baddr(&self) -> Option<u64> {
        use std::collections::BTreeMap;
        let mut counts: BTreeMap<u64, usize> = BTreeMap::new();
        for row in &self.rows {
            if !row.is_mappable() {
                continue;
            }
            let base = row.address().wrapping_sub(u64::from(row.file_offset));
            *counts.entry(base).or_insert(0) += 1;
        }
        counts.into_iter().max_by_key(|(_, n)| *n).map(|(b, _)| b)
    }
}

fn derive_spans(rows: &[DumpRow], end_file_offset: u32) -> Vec<EmittedSpan> {
    let mut spans = vec![EmittedSpan::None; rows.len()];
    let mut owners = Vec::<(usize, u32)>::new();
    for (row_index, row) in rows.iter().enumerate().filter(|(_, row)| row.is_mappable()) {
        if let Some((owner_index, offset)) = owners.last_mut() {
            if *offset == row.file_offset {
                *owner_index = row_index;
                continue;
            }
        }
        owners.push((row_index, row.file_offset));
    }
    for (position, (row_index, offset)) in owners.iter().enumerate() {
        let next_offset = owners[position + 1..]
            .iter()
            .find_map(|(_, next)| (*next > *offset).then_some(*next))
            .unwrap_or(end_file_offset);
        spans[*row_index] = if next_offset > *offset {
            EmittedSpan::Bytes(next_offset - *offset)
        } else if next_offset == *offset {
            EmittedSpan::None
        } else {
            EmittedSpan::Ambiguous
        };
    }
    spans
}

/// Parse the dump table when the header says it exists.
pub fn parse_dump(header: &Header, data: &[u8]) -> Result<AssemblyDump> {
    match header.dump(data)? {
        Some(raw) => AssemblyDump::parse(raw),
        None => Ok(AssemblyDump::default()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn row(file_offset: u32, address: u64, flags: u8) -> DumpRow {
        DumpRow {
            id: 0,
            file_offset,
            line_off: 0,
            dollar: address,
            dollar_hi: 0,
            extended_sib: 0,
            reloc: 0,
            addr_type: 0,
            code_type: 64,
            flags,
        }
    }

    #[test]
    fn repeated_offsets_leave_only_the_last_row_owning_bytes() {
        let mut rows = vec![row(10, 0x100a, 0), row(10, 0x100a, 0), row(14, 0x100e, 0)];
        for (id, row) in rows.iter_mut().enumerate() {
            row.id = id;
        }
        assert_eq!(
            derive_spans(&rows, 18),
            [
                EmittedSpan::None,
                EmittedSpan::Bytes(4),
                EmittedSpan::Bytes(4)
            ]
        );
        let dump = AssemblyDump {
            rows,
            end_file_offset: 18,
            spans: vec![
                EmittedSpan::None,
                EmittedSpan::Bytes(4),
                EmittedSpan::Bytes(4),
            ],
        };
        assert_eq!(dump.assembled_addresses(), [0x100a, 0x100e].into());
    }

    #[test]
    fn virtual_and_reserved_rows_never_emit_output_bytes() {
        let mut rows = vec![
            row(10, 0x100a, DUMP_VIRTUAL),
            row(10, 0x100a, DUMP_NOT_IN_OUTPUT),
            row(10, 0x100a, 0),
        ];
        for (id, row) in rows.iter_mut().enumerate() {
            row.id = id;
        }
        assert_eq!(
            derive_spans(&rows, 14),
            [EmittedSpan::None, EmittedSpan::None, EmittedSpan::Bytes(4)]
        );
    }

    #[test]
    fn decreasing_offsets_mark_the_late_row_ambiguous() {
        let rows = vec![row(20, 0x1014, 0), row(10, 0x100a, 0)];
        assert_eq!(
            derive_spans(&rows, 16),
            [EmittedSpan::Ambiguous, EmittedSpan::Bytes(6)]
        );
    }

    #[test]
    fn final_row_at_output_end_is_zero_byte() {
        let rows = vec![row(10, 0x100a, 0), row(14, 0x100e, 0)];
        assert_eq!(
            derive_spans(&rows, 14),
            [EmittedSpan::Bytes(4), EmittedSpan::None]
        );
    }
}