llvm-native-core 0.1.2

//! LLVM DataLayout — encodes target-specific size, alignment, and
//! offset information for types.
//!
//! The DataLayout string specifies how types are laid out in memory
//! on a particular target. It controls:
//! - Endianness (little-endian or big-endian)
//! - Pointer size and alignment for each address space
//! - Integer/float/vector type alignment
//! - Aggregate type alignment (natural vs ABI)
//! - Stack alignment
//! - Mangling scheme
//!
//! Format: "e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128"
//!
//! Clean-room behavioral reconstruction from LLVM documentation.
//! No LLVM source code is consulted.

use std::collections::HashMap;
use std::fmt;

// ═══════════════════════════════════════════════════════════════════════════
// Endianness
// ═══════════════════════════════════════════════════════════════════════════

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Endianness {
    Little,
    Big,
}

impl Endianness {
    pub fn is_little(&self) -> bool {
        *self == Endianness::Little
    }
    pub fn is_big(&self) -> bool {
        *self == Endianness::Big
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Mangling mode
// ═══════════════════════════════════════════════════════════════════════════

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ManglingMode {
    ELF,
    MIPS,
    MachO,
    COFF,
    XCOFF,
    WASM,
    GOFF,
}

// ═══════════════════════════════════════════════════════════════════════════
// DataLayout
// ═══════════════════════════════════════════════════════════════════════════

/// Parsed target data layout describing type sizes and alignments.
#[derive(Debug, Clone)]
pub struct DataLayout {
    /// Endianness.
    pub endianness: Endianness,
    /// Mangling mode for symbol names.
    pub mangling: ManglingMode,
    /// Pointer size in bits for each address space (default space 0).
    pub pointer_sizes: HashMap<u32, u32>,
    /// Pointer ABI alignment in bits for each address space.
    pub pointer_abi_align: HashMap<u32, u32>,
    /// Pointer preferred alignment in bits for each address space.
    pub pointer_pref_align: HashMap<u32, u32>,
    /// Integer alignment overrides: (bit_width, abi_align, pref_align).
    pub integer_aligns: Vec<(u32, u32, u32)>,
    /// Float alignment overrides.
    pub float_aligns: Vec<(u32, u32, u32)>,
    /// Vector alignment overrides.
    pub vector_aligns: Vec<(u32, u32, u32)>,
    /// Aggregate alignment: 0 = natural, >0 = specific.
    pub aggregate_align: u32,
    /// Stack alignment in bits.
    pub stack_align: u32,
    /// Whether non-integral pointers exist in address space 0.
    pub non_integral_address_spaces: Vec<u32>,
}

impl DataLayout {
    /// Parse a DataLayout string (e.g., "e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128").
    pub fn parse(s: &str) -> Result<Self, String> {
        let mut dl = DataLayout {
            endianness: Endianness::Little,
            mangling: ManglingMode::ELF,
            pointer_sizes: HashMap::new(),
            pointer_abi_align: HashMap::new(),
            pointer_pref_align: HashMap::new(),
            integer_aligns: Vec::new(),
            float_aligns: Vec::new(),
            vector_aligns: Vec::new(),
            aggregate_align: 0,
            stack_align: 0,
            non_integral_address_spaces: Vec::new(),
        };
        dl.parse_impl(s)?;
        Ok(dl)
    }

    fn parse_impl(&mut self, s: &str) -> Result<(), String> {
        for part in s.split('-') {
            if part.is_empty() {
                continue;
            }
            // Special handling for single-char keys that embed values
            // like "S128" or "p0:64:64" or "i64:64"
            let first_char = part.chars().next().unwrap_or(' ');
            match first_char {
                'e' => {
                    if part == "e" {
                        self.endianness = Endianness::Little;
                    }
                }
                'E' => {
                    if part == "E" {
                        self.endianness = Endianness::Big;
                    }
                }
                'm' => {
                    let val = &part[2..]; // skip "m:"
                    self.mangling = match val {
                        "e" => ManglingMode::ELF,
                        "m" => ManglingMode::MIPS,
                        "o" => ManglingMode::MachO,
                        "c" => ManglingMode::COFF,
                        "x" => ManglingMode::XCOFF,
                        "w" => ManglingMode::WASM,
                        "g" => ManglingMode::GOFF,
                        _ => return Err(format!("unknown mangling: {}", val)),
                    };
                }
                'p' => {
                    let rest = &part[2..]; // skip "p:"
                    let pv: Vec<&str> = rest.split(':').collect();
                    if pv.len() >= 2 {
                        let addr_space: u32 = if part.len() > 2 && part.as_bytes()[1] != b':' {
                            part[1..part.find(':').unwrap_or(part.len())]
                                .parse()
                                .unwrap_or(0)
                        } else {
                            0
                        };
                        let size: u32 = pv[0].parse().unwrap_or(64);
                        let abi: u32 = pv[1].parse().unwrap_or(size);
                        let pref: u32 = pv.get(2).and_then(|s| s.parse().ok()).unwrap_or(abi);
                        self.pointer_sizes.insert(addr_space, size);
                        self.pointer_abi_align.insert(addr_space, abi);
                        self.pointer_pref_align.insert(addr_space, pref);
                    }
                }
                'i' => {
                    let rest = &part[1..]; // skip "i"
                    let colon = rest.find(':').unwrap_or(rest.len());
                    let size: u32 = rest[..colon].parse().unwrap_or(0);
                    let vals: Vec<&str> =
                        rest[colon..].trim_start_matches(':').split(':').collect();
                    let abi: u32 = vals
                        .first()
                        .and_then(|s| s.parse().ok())
                        .unwrap_or(size / 8);
                    let pref: u32 = vals.get(1).and_then(|s| s.parse().ok()).unwrap_or(abi);
                    self.integer_aligns.push((size, abi * 8, pref * 8));
                }
                'f' => {
                    let rest = &part[1..];
                    let colon = rest.find(':').unwrap_or(rest.len());
                    let size: u32 = rest[..colon].parse().unwrap_or(0);
                    let vals: Vec<&str> =
                        rest[colon..].trim_start_matches(':').split(':').collect();
                    let abi: u32 = vals
                        .first()
                        .and_then(|s| s.parse().ok())
                        .unwrap_or(size / 8);
                    let pref: u32 = vals.get(1).and_then(|s| s.parse().ok()).unwrap_or(abi);
                    self.float_aligns.push((size, abi * 8, pref * 8));
                }
                'v' => {
                    let rest = &part[1..];
                    let colon = rest.find(':').unwrap_or(rest.len());
                    let size: u32 = rest[..colon].parse().unwrap_or(0);
                    let vals: Vec<&str> =
                        rest[colon..].trim_start_matches(':').split(':').collect();
                    let abi: u32 = vals
                        .first()
                        .and_then(|s| s.parse().ok())
                        .unwrap_or(size / 8);
                    let pref: u32 = vals.get(1).and_then(|s| s.parse().ok()).unwrap_or(abi);
                    self.vector_aligns.push((size, abi * 8, pref * 8));
                }
                'a' => {
                    let rest = if part.len() > 1 && part.as_bytes()[1] == b':' {
                        &part[2..]
                    } else {
                        &part[1..]
                    };
                    self.aggregate_align = if rest.is_empty() {
                        0
                    } else {
                        rest.parse().unwrap_or(0)
                    };
                }
                'n' => {
                    let rest = &part[1..]; // skip 'n'
                    for sval in rest.split(':') {
                        if let Ok(sz) = sval.parse::<u32>() {
                            self.non_integral_address_spaces.push(sz);
                        }
                    }
                }
                'S' => {
                    let sz_str = &part[1..];
                    self.stack_align = if sz_str.is_empty() {
                        0
                    } else {
                        sz_str.parse().unwrap_or(0)
                    };
                }
                _ => {} // Ignore unknown specifiers
            }
        }
        Ok(())
    }

    /// Get the default DataLayout for x86_64-linux-gnu.
    pub fn x86_64_linux() -> Self {
        Self::parse("e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128").unwrap()
    }

    /// Get the default DataLayout for aarch64-linux-gnu.
    pub fn aarch64_linux() -> Self {
        Self::parse("e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128").unwrap()
    }

    // ── Queries ─────────────────────────────────────────────────────────

    /// Whether this is a little-endian target.
    pub fn is_little_endian(&self) -> bool {
        self.endianness.is_little()
    }
    /// Whether this is a big-endian target.
    pub fn is_big_endian(&self) -> bool {
        self.endianness.is_big()
    }

    /// Get the pointer size in bits for address space 0.
    pub fn pointer_size(&self) -> u32 {
        self.pointer_sizes.get(&0).copied().unwrap_or(64)
    }

    /// Get pointer size for a specific address space.
    pub fn pointer_size_for_space(&self, addr_space: u32) -> u32 {
        self.pointer_sizes
            .get(&addr_space)
            .copied()
            .unwrap_or_else(|| self.pointer_size())
    }

    /// Get pointer ABI alignment in bits.
    pub fn pointer_abi_alignment(&self, addr_space: u32) -> u32 {
        self.pointer_abi_align
            .get(&addr_space)
            .copied()
            .unwrap_or_else(|| self.pointer_size_for_space(addr_space))
    }

    /// Get pointer preferred alignment in bits.
    pub fn pointer_pref_alignment(&self, addr_space: u32) -> u32 {
        self.pointer_pref_align
            .get(&addr_space)
            .copied()
            .unwrap_or_else(|| self.pointer_abi_alignment(addr_space))
    }

    /// Get the ABI alignment for an integer of the given bit width.
    pub fn integer_abi_alignment(&self, bit_width: u32) -> u32 {
        for &(bw, abi, _) in &self.integer_aligns {
            if bw == bit_width {
                return abi;
            }
        }
        (bit_width / 8).min(4) * 8 // default: natural alignment capped at 32 bits
    }

    /// Get the ABI alignment for a float of the given bit width.
    pub fn float_abi_alignment(&self, bit_width: u32) -> u32 {
        for &(bw, abi, _) in &self.float_aligns {
            if bw == bit_width {
                return abi;
            }
        }
        bit_width // default: natural alignment
    }

    /// Get the ABI alignment for a vector of the given bit width.
    pub fn vector_abi_alignment(&self, bit_width: u32) -> u32 {
        for &(bw, abi, _) in &self.vector_aligns {
            if bw == bit_width {
                return abi;
            }
        }
        if bit_width <= 128 {
            128
        } else {
            256
        }
    }

    /// Get the preferred alignment for an integer of the given bit width.
    pub fn integer_pref_alignment(&self, bit_width: u32) -> u32 {
        for &(bw, _, pref) in &self.integer_aligns {
            if bw == bit_width {
                return pref;
            }
        }
        self.integer_abi_alignment(bit_width)
    }

    /// Get the stack alignment in bits.
    pub fn stack_alignment(&self) -> u32 {
        self.stack_align
    }

    /// Get the size of a type index in bits (for type-based alias analysis).
    pub fn index_size_in_bits(&self) -> u32 {
        self.pointer_size()
    }

    /// Generate the canonical string representation.
    pub fn to_string(&self) -> String {
        let mut s = String::new();
        // Endianness
        s.push(if self.endianness.is_little() {
            'e'
        } else {
            'E'
        });
        // Mangling
        let m = match self.mangling {
            ManglingMode::ELF => 'e',
            ManglingMode::MIPS => 'm',
            ManglingMode::MachO => 'o',
            ManglingMode::COFF => 'c',
            ManglingMode::XCOFF => 'x',
            ManglingMode::WASM => 'w',
            ManglingMode::GOFF => 'g',
        };
        s.push_str(&format!("-m:{}", m));
        // Pointer
        if let Some(&sz) = self.pointer_sizes.get(&0) {
            let abi = self.pointer_abi_align.get(&0).copied().unwrap_or(sz);
            let pref = self.pointer_pref_align.get(&0).copied().unwrap_or(abi);
            s.push_str(&format!("-p:{}:{}:{}", sz, abi, pref));
        }
        // Integers
        for &(bw, abi, pref) in &self.integer_aligns {
            s.push_str(&format!("-i{}:{}:{}", bw, abi / 8, pref / 8));
        }
        // Floats
        for &(bw, abi, pref) in &self.float_aligns {
            s.push_str(&format!("-f{}:{}:{}", bw, abi / 8, pref / 8));
        }
        // Stack
        if self.stack_align > 0 {
            s.push_str(&format!("-S{}", self.stack_align));
        }
        s
    }
}

impl Default for DataLayout {
    fn default() -> Self {
        Self::x86_64_linux()
    }
}

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

// ═══════════════════════════════════════════════════════════════════════════
// Type size and layout computation
// ═══════════════════════════════════════════════════════════════════════════

impl DataLayout {
    /// Get the size of a type in bits.
    pub fn type_size_in_bits(&self, ty: &crate::types::Type) -> u64 {
        use crate::types::TypeKind;
        match &ty.kind {
            TypeKind::Void | TypeKind::Label | TypeKind::Metadata | TypeKind::Token => 0,
            TypeKind::Half | TypeKind::BFloat => 16,
            TypeKind::Float => 32,
            TypeKind::Double => 64,
            TypeKind::FP128 | TypeKind::PPCFP128 => 128,
            TypeKind::X86FP80 => 80,
            TypeKind::X86MMX => 64,
            TypeKind::X86AMX => 8192,
            TypeKind::Integer { bits } => *bits as u64,
            TypeKind::Pointer { .. } => self.pointer_size() as u64,
            TypeKind::Array { len, .. } => *len * 8,
            TypeKind::FixedVector { len, .. } => *len as u64 * 32,
            TypeKind::ScalableVector { min_elems, .. } => *min_elems as u64 * 32,
            TypeKind::Struct { .. } | TypeKind::Function { .. } => 0,
        }
    }

    /// Get store size in bits (rounded to alignment).
    pub fn type_store_size_in_bits(&self, ty: &crate::types::Type) -> u64 {
        let size = self.type_size_in_bits(ty);
        let align = self.abi_alignment_in_bits(ty) as u64;
        if align == 0 {
            size
        } else {
            ((size + align - 1) / align) * align
        }
    }

    /// Get ABI alignment in bits for a type.
    pub fn abi_alignment_in_bits(&self, ty: &crate::types::Type) -> u32 {
        use crate::types::TypeKind;
        match &ty.kind {
            TypeKind::Integer { bits } => self.integer_abi_alignment(*bits),
            TypeKind::Half | TypeKind::BFloat => self.float_abi_alignment(16),
            TypeKind::Float => self.float_abi_alignment(32),
            TypeKind::Double => self.float_abi_alignment(64),
            TypeKind::FP128 | TypeKind::PPCFP128 => self.float_abi_alignment(128),
            TypeKind::X86FP80 => self.float_abi_alignment(80),
            TypeKind::Pointer { .. } => self.pointer_abi_alignment(0),
            TypeKind::FixedVector { len, .. } => self.vector_abi_alignment(*len * 32),
            TypeKind::ScalableVector { min_elems, .. } => {
                self.vector_abi_alignment(*min_elems * 32)
            }
            TypeKind::Array { .. } | TypeKind::Struct { .. } => {
                if self.aggregate_align > 0 {
                    self.aggregate_align
                } else {
                    64
                }
            }
            _ => 8,
        }
    }

    /// Get preferred alignment in bits for a type.
    pub fn pref_alignment_in_bits(&self, ty: &crate::types::Type) -> u32 {
        use crate::types::TypeKind;
        match &ty.kind {
            TypeKind::Integer { bits } => self.integer_pref_alignment(*bits),
            _ => self.abi_alignment_in_bits(ty),
        }
    }

    /// Compute struct field offsets and sizes. Returns Vec<(offset_bits, size_bits)>.
    pub fn struct_layout_offsets(&self, fields: &[(crate::types::Type, bool)]) -> Vec<(u64, u64)> {
        let packed = fields.iter().any(|(_, p)| *p);
        let mut offset = 0u64;
        fields
            .iter()
            .map(|(ty, _)| {
                if !packed {
                    let align = self.abi_alignment_in_bits(ty) as u64;
                    if align > 0 {
                        offset = ((offset + align - 1) / align) * align;
                    }
                }
                let size = self.type_size_in_bits(ty);
                let cur = offset;
                offset += size;
                (cur, size)
            })
            .collect()
    }

    /// Total struct size in bits.
    pub fn struct_size_in_bits(&self, fields: &[(crate::types::Type, bool)]) -> u64 {
        let layout = self.struct_layout_offsets(fields);
        if layout.is_empty() {
            return 0;
        }
        let (last_off, last_sz) = layout.last().unwrap();
        let total = last_off + last_sz;
        let packed = fields.iter().any(|(_, p)| *p);
        if packed {
            total
        } else {
            let align = fields
                .iter()
                .map(|(ty, _)| self.abi_alignment_in_bits(ty) as u64)
                .max()
                .unwrap_or(8);
            ((total + align - 1) / align) * align
        }
    }

    pub fn fits_in_register(&self, ty: &crate::types::Type) -> bool {
        self.type_size_in_bits(ty) <= self.pointer_size() as u64
    }

    pub fn is_32bit(&self) -> bool {
        self.pointer_size() == 32
    }
    pub fn is_64bit(&self) -> bool {
        self.pointer_size() == 64
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════════════

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

    #[test]
    fn test_default_x86_64() {
        let dl = DataLayout::x86_64_linux();
        assert!(dl.is_little_endian());
        assert_eq!(dl.pointer_size(), 64);
        assert_eq!(dl.stack_alignment(), 128);
    }

    #[test]
    fn test_parse_basic() {
        let dl = DataLayout::parse("e-m:e-p:64:64-S128").unwrap();
        assert!(dl.is_little_endian());
        assert_eq!(dl.pointer_size(), 64);
        assert_eq!(dl.stack_alignment(), 128);
    }

    #[test]
    fn test_parse_full() {
        let dl = DataLayout::parse("e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128").unwrap();
        // i64:64 means 64-bit integer has 64-byte (512-bit) ABI alignment on this target
        assert_eq!(dl.integer_abi_alignment(64), 512);
        assert_eq!(dl.stack_alignment(), 128);
        assert_eq!(dl.pointer_size(), 64);
    }

    #[test]
    fn test_parse_big_endian() {
        let dl = DataLayout::parse("E-m:e-p:32:32-S64").unwrap();
        assert!(dl.is_big_endian());
        assert_eq!(dl.pointer_size(), 32);
    }

    #[test]
    fn test_pointer_size_for_space() {
        let dl = DataLayout::x86_64_linux();
        assert_eq!(dl.pointer_size_for_space(0), 64);
    }

    #[test]
    fn test_to_string_roundtrip() {
        let orig = "e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128";
        let dl = DataLayout::parse(orig).unwrap();
        let s = dl.to_string();
        // The string may not be identical but should contain key parts
        assert!(s.contains("e"));
        assert!(s.contains("p:64"));
    }

    #[test]
    fn test_aarch64_default() {
        let dl = DataLayout::aarch64_linux();
        assert!(dl.is_little_endian());
        assert_eq!(dl.pointer_size(), 64); // AArch64 Linux is LP64
    }

    #[test]
    fn test_integer_alignment_default() {
        let dl = DataLayout::x86_64_linux();
        let i8_align = dl.integer_abi_alignment(8);
        let i32_align = dl.integer_abi_alignment(32);
        assert!(i32_align >= i8_align);
    }

    #[test]
    fn test_float_alignment() {
        let dl = DataLayout::x86_64_linux();
        let f32_align = dl.float_abi_alignment(32);
        let f64_align = dl.float_abi_alignment(64);
        assert!(f64_align >= f32_align);
    }
}