type_metadata.rs

//! Compile-time type knowledge for code generation.
//!
//! Maps program values (registers or MIR locals) to statically-known types
//! so that backends (JIT, AOT) can emit unboxed native code instead of
//! NaN-tag-aware runtime operations.
//!
//! Shared between the JIT (`src/jit/typed_compiler.rs`) and the AOT
//! compiler (`src/aot/`).

/// The static type of a value known at compile time.
///
/// - `Int`: NaN-tagged integer → strip tag, use direct i64 ops.
/// - `Float`: Raw f64 bits → use direct f64 ops.
/// - `Bool`: NaN-tagged boolean → compare directly against tagged constants.
/// - `Unknown`: Fall back to runtime helpers / boxed representation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum KnownType {
    #[default]
    Unknown,
    Int,
    Float,
    Bool,
}

/// Number of registers in the VM frame.
pub const REG_COUNT: usize = 256;

/// Static type information for the 256 VM registers.
///
/// A flat `[KnownType; 256]` array replaces the previous `HashMap<usize, KnownType>`
/// for deterministic O(1) access with no hashing overhead. Unknown slots are
/// represented by `KnownType::Unknown` (the default).
///
/// # Example
/// ```
/// use nulang::type_metadata::{TypeMetadata, KnownType};
///
/// let mut meta = TypeMetadata::new();
/// meta.set_type(0, KnownType::Int);   // R0 is known Int
/// meta.set_type(1, KnownType::Float); // R1 is known Float
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TypeMetadata {
    pub regs: [KnownType; REG_COUNT],
}

impl Default for TypeMetadata {
    fn default() -> Self {
        Self {
            regs: [KnownType::Unknown; REG_COUNT],
        }
    }
}

/// Convert a language-level `Type` to a `KnownType` for code generation.
///
/// Only primitive types are statically known; polymorphic, compound, and
/// effectful types all map to `Unknown`.
pub fn type_to_known_type(ty: &crate::types::Type) -> KnownType {
    match ty {
        crate::types::Type::Primitive(p) => match p {
            crate::types::PrimitiveType::Int => KnownType::Int,
            crate::types::PrimitiveType::Float => KnownType::Float,
            crate::types::PrimitiveType::Bool => KnownType::Bool,
            _ => KnownType::Unknown,
        },
        _ => KnownType::Unknown,
    }
}

impl TypeMetadata {
    /// Create a new TypeMetadata with all registers Unknown.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the known type for a register index.
    /// Silently ignores indices beyond REG_COUNT (defense against MIR overflow).
    pub fn set_type(&mut self, reg: usize, ty: KnownType) {
        if reg < REG_COUNT {
            self.regs[reg] = ty;
        }
    }

    /// Get the known type for a register index.
    pub fn get_type(&self, reg: usize) -> KnownType {
        if reg < REG_COUNT {
            self.regs[reg]
        } else {
            KnownType::Unknown
        }
    }

    /// Check whether both operands have the same known type.
    pub fn both_known(&self, r1: usize, r2: usize, expected: KnownType) -> bool {
        r1 < REG_COUNT && r2 < REG_COUNT && self.regs[r1] == expected && self.regs[r2] == expected
    }

    /// Check whether a single value has the expected known type.
    pub fn is_known(&self, reg: usize, expected: KnownType) -> bool {
        reg < REG_COUNT && self.regs[reg] == expected
    }

    /// Mark the destination as having a known type after an operation.
    ///
    /// For arithmetic: the result type is the same as the operand type.
    /// For comparisons: the result is always Bool.
    pub fn propagate_result(&mut self, dst: usize, operand_reg: usize) {
        if operand_reg < REG_COUNT && dst < REG_COUNT {
            self.regs[dst] = self.regs[operand_reg];
        }
    }

    /// Mark the destination as Bool (used after comparisons).
    pub fn set_bool_result(&mut self, dst: usize) {
        if dst < REG_COUNT {
            self.regs[dst] = KnownType::Bool;
        }
    }

    /// Returns true if no register has a known type (all Unknown).
    pub fn is_empty(&self) -> bool {
        self.regs.iter().all(|&t| t == KnownType::Unknown)
    }

    /// Build TypeMetadata from an iterator of (register_index, Type) pairs.
    ///
    /// Converts language-level `Type` values to `KnownType` by stripping
    /// away polymorphic wrappers: only primitive `Int`, `Float`, and `Bool`
    /// are statically known; everything else becomes `Unknown`.
    pub fn from_mir_locals<'a>(
        locals: impl Iterator<Item = (usize, &'a crate::types::Type)>,
    ) -> Self {
        let mut meta = TypeMetadata::new();
        for (reg, ty) in locals {
            if reg >= REG_COUNT {
                continue; // MIR locals beyond frame register capacity — skip type tracking
            }
            let known = type_to_known_type(ty);
            if known != KnownType::Unknown {
                meta.set_type(reg, known);
            }
        }
        meta
    }
}

// ---------------------------------------------------------------------------
// Capability Metadata
// ---------------------------------------------------------------------------

pub struct CapabilityMetadata {
    pub caps: [crate::types::Capability; REG_COUNT],
}

impl Default for CapabilityMetadata {
    fn default() -> Self {
        Self {
            caps: [crate::types::Capability::Tag; REG_COUNT],
        }
    }
}

impl CapabilityMetadata {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn set_cap(&mut self, reg: usize, cap: crate::types::Capability) {
        if reg < REG_COUNT {
            self.caps[reg] = cap;
        }
    }
    pub fn get_cap(&self, reg: usize) -> crate::types::Capability {
        self.caps[reg]
    }
    pub fn from_mir_function(func: &crate::mir::Function) -> Self {
        let mut meta = Self::new();
        let local_base = crate::mir::FunctionBuilder::LOCAL_BASE;
        for local in &func.locals {
            let reg = local_base as usize + local.id.0 as usize;
            if reg < REG_COUNT {
                meta.caps[reg] = local.cap;
            }
        }
        meta
    }
}