r2fas 0.2.1

//! Preprocessed source lines (FAS.TXT table 3).

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

/// One preprocessed line, keyed by its offset inside the preprocessed blob.
#[derive(Debug, Clone)]
pub struct PrepLine {
    /// Offset of this record inside the preprocessed source.
    pub offset: u32,
    /// Zero = main input file; otherwise offset of an ASCIIZ file name inside
    /// the preprocessed blob. Ignored when [`Self::generated_by_macro`].
    pub file_or_macro: u32,
    /// 1-based line number in the originating source file (high bit stripped).
    pub line_number: u32,
    /// Token stream was produced by a macro rather than loaded from a file.
    pub generated_by_macro: bool,
    /// Byte position of this line in the source file, or the invoking line
    /// offset when this line is macro-generated.
    pub source_pos_or_invoke: u32,
    /// Offset of the macro-definition line when macro-generated.
    pub macro_def_offset: u32,
    /// Offset of the first token byte in the preprocessed blob.
    pub tokens_start: u32,
    /// Offset immediately after the terminating zero token.
    pub tokens_end: u32,
}

impl PrepLine {
    /// Parse a single table-3 header at `offset` inside `prep`.
    ///
    /// Token bytes after offset 16 are skipped until a 0 terminator so the
    /// caller can find the next line. Quoted tokens (`22h`) carry a 32-bit
    /// length; symbol tokens (`1Ah` / `3Bh`) carry a 8-bit length.
    pub fn parse_at(prep: &[u8], offset: u32) -> Result<(Self, u32)> {
        let o = offset as usize;
        if o + 16 > prep.len() {
            return Err(FasError::Truncated("preprocessed line header"));
        }
        let file_or_macro = bytes::u32_at(prep, o)?;
        let line_field = bytes::u32_at(prep, o + 4)?;
        let source_pos_or_invoke = bytes::u32_at(prep, o + 8)?;
        let macro_def_offset = bytes::u32_at(prep, o + 12)?;
        let generated_by_macro = (line_field & 0x8000_0000) != 0;
        let line_number = line_field & 0x7FFF_FFFF;

        let mut i = o + 16;
        while i < prep.len() {
            match prep[i] {
                0 => {
                    let next = (i + 1) as u32;
                    return Ok((
                        Self {
                            offset,
                            file_or_macro,
                            line_number,
                            generated_by_macro,
                            source_pos_or_invoke,
                            macro_def_offset,
                            tokens_start: (o + 16) as u32,
                            tokens_end: next,
                        },
                        next,
                    ));
                }
                0x1A | 0x3B => {
                    let n = *prep.get(i + 1).ok_or(FasError::Truncated("sym token"))? as usize;
                    i = i
                        .checked_add(2 + n)
                        .ok_or(FasError::Truncated("sym token"))?;
                }
                0x22 => {
                    let n = bytes::u32_at(prep, i + 1)? as usize;
                    i = i
                        .checked_add(5 + n)
                        .ok_or(FasError::Truncated("quote token"))?;
                }
                _ => i += 1,
            }
        }
        Err(FasError::Truncated("unterminated preprocessed line"))
    }

    /// Token bytes including the terminating zero byte.
    pub fn token_bytes<'a>(&self, prep: &'a [u8]) -> Result<&'a [u8]> {
        prep.get(self.tokens_start as usize..self.tokens_end as usize)
            .ok_or(FasError::Truncated("preprocessed tokens"))
    }

    /// Decode the tokenized line into structured tokens.
    pub fn tokens(&self, prep: &[u8]) -> Result<Vec<SourceToken>> {
        decode_tokens(self.token_bytes(prep)?)
    }

    /// Deterministic textual rendering used when original source is absent.
    pub fn detokenize(&self, prep: &[u8]) -> Result<String> {
        Ok(render_tokens(&self.tokens(prep)?))
    }
}

/// One token from a preprocessed source line.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SourceToken {
    /// Normal symbol token (`1Ah`).
    Symbol(Vec<u8>),
    /// Symbol token already interpreted by the preprocessor (`3Bh`).
    InterpretedSymbol(Vec<u8>),
    /// Quoted byte sequence (`22h`).
    Quoted(Vec<u8>),
    /// One special-character token.
    Character(u8),
}

/// One macro provenance hop.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MacroFrame {
    /// Generated line offset.
    pub generated_line: u32,
    /// Macro name, when its Pascal string is valid.
    pub macro_name: Option<String>,
    /// Invoking line offset.
    pub invocation_line: u32,
    /// Macro-definition line offset.
    pub definition_line: u32,
}

/// Complete macro chain and its physical-source origin.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Provenance {
    /// Generated macro frames, outermost traversal order.
    pub frames: Vec<MacroFrame>,
    /// Physical-source line reached by following invocation links.
    pub origin_offset: u32,
    /// Why traversal stopped early, when the chain was incomplete.
    pub diagnostic: Option<ProvenanceDiagnostic>,
}

/// Macro-chain traversal issue.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProvenanceDiagnostic {
    /// An invocation offset did not resolve to a parsed source row.
    MissingLine(u32),
    /// An invocation link repeated a row already visited.
    Cycle(u32),
}

/// Index of every preprocessed line keyed by blob offset.
#[derive(Debug, Clone, Default)]
pub struct PrepSource {
    /// Offset → line.
    pub lines: BTreeMap<u32, PrepLine>,
}

impl PrepSource {
    /// Walk the whole preprocessed blob.
    pub fn parse(prep: &[u8]) -> Result<Self> {
        let mut lines = BTreeMap::new();
        let mut off = 0u32;
        while (off as usize) < prep.len() {
            // File names stored *inside* the blob as ASCIIZ are not line
            // records. A line always starts with a 16-byte header; if we
            // cannot parse, stop. FASM packs names at offsets referenced by
            // lines, typically overlapping the stream. In practice the blob
            // is a sequence of lines only — names live in the same blob at
            // offsets that point *into* a previous line's token area or at
            // dedicated spots. We parse strictly as a line sequence.
            match PrepLine::parse_at(prep, off) {
                Ok((line, next)) => {
                    if next <= off {
                        return Err(FasError::Geometry("preprocessed line did not advance"));
                    }
                    lines.insert(off, line);
                    off = next;
                }
                Err(_) => break,
            }
        }
        Ok(Self { lines })
    }

    /// Look up a line by preprocessed-blob offset.
    pub fn get(&self, offset: u32) -> Option<&PrepLine> {
        self.lines.get(&offset)
    }

    /// File name for a non-macro line. `input_name` is used when the line
    /// belongs to the main source file (`file_or_macro == 0`).
    pub fn file_name<'a>(
        &'a self,
        line: &PrepLine,
        prep: &'a [u8],
        input_name: &'a str,
    ) -> Result<&'a str> {
        if line.generated_by_macro {
            return Ok(input_name);
        }
        if line.file_or_macro == 0 {
            return Ok(input_name);
        }
        bytes::cstring_at(prep, line.file_or_macro as usize)
    }

    /// Walk macro invocation links until a source-file line is found.
    pub fn origin<'a>(&'a self, mut line: &'a PrepLine) -> &'a PrepLine {
        let mut guard = 0u32;
        while line.generated_by_macro && guard < 64 {
            guard += 1;
            match self.get(line.source_pos_or_invoke) {
                Some(invoker) => line = invoker,
                None => break,
            }
        }
        line
    }

    /// Pascal-style macro name stored at `file_or_macro` when the line is
    /// macro-generated. `None` if the pointer is invalid.
    pub fn macro_name<'a>(&self, line: &PrepLine, prep: &'a [u8]) -> Option<&'a str> {
        if !line.generated_by_macro {
            return None;
        }
        bytes::pascal_at(prep, line.file_or_macro as usize).ok()
    }

    /// Preserve every macro invocation/definition hop and report broken chains.
    pub fn provenance(&self, line: &PrepLine, prep: &[u8]) -> Provenance {
        let mut current = line;
        let mut frames = Vec::new();
        let mut seen = BTreeSet::new();
        while current.generated_by_macro {
            if !seen.insert(current.offset) {
                return Provenance {
                    frames,
                    origin_offset: current.offset,
                    diagnostic: Some(ProvenanceDiagnostic::Cycle(current.offset)),
                };
            }
            frames.push(MacroFrame {
                generated_line: current.offset,
                macro_name: self.macro_name(current, prep).map(str::to_owned),
                invocation_line: current.source_pos_or_invoke,
                definition_line: current.macro_def_offset,
            });
            let Some(invoker) = self.get(current.source_pos_or_invoke) else {
                return Provenance {
                    frames,
                    origin_offset: current.offset,
                    diagnostic: Some(ProvenanceDiagnostic::MissingLine(
                        current.source_pos_or_invoke,
                    )),
                };
            };
            current = invoker;
        }
        Provenance {
            frames,
            origin_offset: current.offset,
            diagnostic: None,
        }
    }
}

fn decode_tokens(raw: &[u8]) -> Result<Vec<SourceToken>> {
    let mut tokens = Vec::new();
    let mut index = 0usize;
    while index < raw.len() {
        match raw[index] {
            0 => return Ok(tokens),
            kind @ (0x1A | 0x3B) => {
                let length =
                    *raw.get(index + 1)
                        .ok_or(FasError::Truncated("symbol token"))? as usize;
                let bytes = raw
                    .get(index + 2..index + 2 + length)
                    .ok_or(FasError::Truncated("symbol token"))?
                    .to_vec();
                tokens.push(if kind == 0x1A {
                    SourceToken::Symbol(bytes)
                } else {
                    SourceToken::InterpretedSymbol(bytes)
                });
                index += 2 + length;
            }
            0x22 => {
                let length = bytes::u32_at(raw, index + 1)? as usize;
                let quoted = raw
                    .get(index + 5..index + 5 + length)
                    .ok_or(FasError::Truncated("quoted token"))?
                    .to_vec();
                tokens.push(SourceToken::Quoted(quoted));
                index += 5 + length;
            }
            character => {
                tokens.push(SourceToken::Character(character));
                index += 1;
            }
        }
    }
    Err(FasError::Truncated("unterminated token stream"))
}

fn render_tokens(tokens: &[SourceToken]) -> String {
    let mut output = String::new();
    let mut previous_word = false;
    for token in tokens {
        let (text, word) = match token {
            SourceToken::Symbol(bytes) | SourceToken::InterpretedSymbol(bytes) => {
                (String::from_utf8_lossy(bytes).into_owned(), true)
            }
            SourceToken::Quoted(bytes) => {
                let mut quoted = String::from("\"");
                for byte in bytes {
                    match byte {
                        b'\\' => quoted.push_str("\\\\"),
                        b'\"' => quoted.push_str("\\\""),
                        0x20..=0x7e => quoted.push(char::from(*byte)),
                        _ => quoted.push_str(&format!("\\x{byte:02x}")),
                    }
                }
                quoted.push('\"');
                (quoted, true)
            }
            SourceToken::Character(character) => (char::from(*character).to_string(), false),
        };
        if word && previous_word {
            output.push(' ');
        }
        output.push_str(&text);
        previous_word = word;
    }
    output
}

/// Convenience wrapper used by [`crate::fas::FasFile::parse`].
pub fn parse_preprocessed(header: &Header, data: &[u8]) -> Result<PrepSource> {
    PrepSource::parse(header.preprocessed(data)?)
}

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

    fn line(offset: u32, generated: bool, invocation: u32) -> PrepLine {
        PrepLine {
            offset,
            file_or_macro: 0,
            line_number: 1,
            generated_by_macro: generated,
            source_pos_or_invoke: invocation,
            macro_def_offset: 0,
            tokens_start: 0,
            tokens_end: 1,
        }
    }

    #[test]
    fn decodes_and_renders_token_streams() {
        let raw = [
            0x1a, 3, b'm', b'o', b'v', 0x1a, 3, b'e', b'a', b'x', b',', 0x22, 4, 0, 0, 0, b'A',
            b'B', b'C', b'D', 0,
        ];
        let tokens = decode_tokens(&raw).unwrap();
        assert_eq!(
            tokens,
            [
                SourceToken::Symbol(b"mov".to_vec()),
                SourceToken::Symbol(b"eax".to_vec()),
                SourceToken::Character(b','),
                SourceToken::Quoted(b"ABCD".to_vec()),
            ]
        );
        assert_eq!(render_tokens(&tokens), "mov eax,\"ABCD\"");
    }

    #[test]
    fn provenance_reports_cycles_and_missing_invocations() {
        let mut source = PrepSource::default();
        source.lines.insert(0, line(0, true, 10));
        source.lines.insert(10, line(10, true, 0));
        let cycle = source.provenance(source.get(0).unwrap(), b"");
        assert_eq!(cycle.diagnostic, Some(ProvenanceDiagnostic::Cycle(0)));
        assert_eq!(cycle.frames.len(), 2);

        let missing_line = line(20, true, 99);
        let missing = source.provenance(&missing_line, b"");
        assert_eq!(
            missing.diagnostic,
            Some(ProvenanceDiagnostic::MissingLine(99))
        );
    }
}