vexil-lang 0.5.0

use crate::ast::{PrimitiveType, SemanticType, SubByteType};
use crate::span::Span;
use smol_str::SmolStr;
use std::collections::HashMap;

// Forward-declare TypeDef so TypeRegistry can reference it.
use super::{ImplDef, TraitDef, TypeDef};

// ---------------------------------------------------------------------------
// TypeId + TypeRegistry
// ---------------------------------------------------------------------------

/// Opaque handle to a type definition in the registry.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TypeId(pub(crate) u32);

impl TypeId {
    /// Returns the underlying registry index.
    pub fn index(self) -> u32 {
        self.0
    }
}

/// Sentinel for unresolvable types (poison value).
pub const POISON_TYPE_ID: TypeId = TypeId(u32::MAX);

/// Central type store. All cross-references use TypeId.
#[derive(Debug, Clone)]
pub struct TypeRegistry {
    types: Vec<Option<TypeDef>>,
    by_name: HashMap<SmolStr, TypeId>,
    /// Transparent aliases, including concrete container targets.
    aliases: HashMap<SmolStr, ResolvedType>,
    /// Stable declaration identity for exported aliases.
    alias_origins: HashMap<SmolStr, (SmolStr, SmolStr)>,
    /// Source-faithful return expressions for trait functions.
    ///
    /// Trait function IR keeps its stable resolved return field. This private
    /// side table retains generic return expressions for impl substitution.
    trait_fn_return_types: HashMap<(TypeId, SmolStr), crate::ast::TypeExpr>,
    /// Stable declaration identity for imported and local definitions.
    origins: HashMap<TypeId, (SmolStr, SmolStr)>,
    /// Resolved trait identity for each local impl record.
    impl_trait_ids: HashMap<TypeId, TypeId>,
    /// Source span of the trait reference for each local impl record.
    impl_trait_spans: HashMap<TypeId, Span>,
}

impl Default for TypeRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl TypeRegistry {
    /// Create an empty type registry.
    pub fn new() -> Self {
        Self {
            types: Vec::new(),
            by_name: HashMap::new(),
            aliases: HashMap::new(),
            alias_origins: HashMap::new(),
            trait_fn_return_types: HashMap::new(),
            origins: HashMap::new(),
            impl_trait_ids: HashMap::new(),
            impl_trait_spans: HashMap::new(),
        }
    }

    /// Register a complete type definition, returning its [`TypeId`].
    pub fn register(&mut self, name: SmolStr, def: TypeDef) -> TypeId {
        let id = TypeId(self.types.len() as u32);
        self.types.push(Some(def));
        self.by_name.insert(name, id);
        id
    }

    /// Register a stub (forward declaration) that will be filled later via [`fill_stub`](Self::fill_stub).
    pub fn register_stub(&mut self, name: SmolStr) -> TypeId {
        let id = TypeId(self.types.len() as u32);
        self.types.push(None);
        self.by_name.insert(name, id);
        id
    }

    pub(crate) fn register_unbound_stub(&mut self) -> TypeId {
        let id = TypeId(self.types.len() as u32);
        self.types.push(None);
        id
    }

    /// Look up a type by name, returning its [`TypeId`] if registered.
    /// Also checks aliases whose transparent target is a named type.
    pub fn lookup(&self, name: &str) -> Option<TypeId> {
        self.by_name
            .get(name)
            .copied()
            .or_else(|| match self.aliases.get(name) {
                Some(ResolvedType::Named(id)) => Some(*id),
                _ => None,
            })
    }

    /// Look up a primitive type alias by name.
    pub fn lookup_primitive_alias(&self, name: &str) -> Option<PrimitiveType> {
        match self.aliases.get(name) {
            Some(ResolvedType::Primitive(primitive)) => Some(*primitive),
            _ => None,
        }
    }

    /// Look up the fully resolved transparent target of an alias.
    pub fn lookup_alias(&self, name: &str) -> Option<&ResolvedType> {
        self.aliases.get(name)
    }

    /// Register a type alias (secondary name mapping).
    /// The target TypeId must already exist in the registry.
    pub fn register_alias(&mut self, alias: SmolStr, target: TypeId) {
        self.register_resolved_alias(alias, ResolvedType::Named(target));
    }

    pub(crate) fn bind_name(&mut self, name: SmolStr, target: TypeId) {
        self.by_name.insert(name, target);
    }

    /// Register a primitive type alias.
    pub fn register_primitive_alias(&mut self, alias: SmolStr, primitive: PrimitiveType) {
        self.register_resolved_alias(alias, ResolvedType::Primitive(primitive));
    }

    /// Register a transparent alias to any fully resolved concrete type.
    pub fn register_resolved_alias(&mut self, alias: SmolStr, target: ResolvedType) {
        self.aliases.insert(alias, target);
    }

    pub(crate) fn set_alias_origin(
        &mut self,
        alias: &str,
        namespace: SmolStr,
        declaration: SmolStr,
    ) {
        self.alias_origins
            .insert(SmolStr::new(alias), (namespace, declaration));
    }

    pub(crate) fn find_alias_origin(
        &self,
        namespace: &str,
        declaration: &str,
    ) -> Option<&ResolvedType> {
        self.alias_origins.iter().find_map(|(binding, (ns, name))| {
            (ns == namespace && name == declaration)
                .then(|| self.aliases.get(binding))
                .flatten()
        })
    }

    pub(crate) fn aliases_from_origin<'a>(
        &'a self,
        namespace: &'a str,
    ) -> impl Iterator<Item = (&'a str, &'a ResolvedType)> + 'a {
        self.alias_origins
            .iter()
            .filter_map(move |(binding, (ns, name))| {
                (ns == namespace)
                    .then(|| {
                        self.aliases
                            .get(binding)
                            .map(|target| (name.as_str(), target))
                    })
                    .flatten()
            })
    }

    /// Get a reference to the type definition for `id`, if it exists and is not a stub.
    pub fn get(&self, id: TypeId) -> Option<&TypeDef> {
        self.types.get(id.0 as usize).and_then(|opt| opt.as_ref())
    }

    /// Get a mutable reference to the type definition for `id`.
    pub fn get_mut(&mut self, id: TypeId) -> Option<&mut TypeDef> {
        self.types
            .get_mut(id.0 as usize)
            .and_then(|opt| opt.as_mut())
    }

    /// Returns `true` if `id` is a registered stub that has not yet been filled.
    pub fn is_stub(&self, id: TypeId) -> bool {
        self.types
            .get(id.0 as usize)
            .is_some_and(|opt| opt.is_none())
    }

    /// Returns the total number of slots (filled + stubs) in the registry.
    pub fn len(&self) -> usize {
        self.types.len()
    }

    /// Returns `true` if the registry contains no types.
    pub fn is_empty(&self) -> bool {
        self.types.is_empty()
    }

    /// Rename a type in the by-name index (used for aliased import qualification).
    pub fn rename(&mut self, id: TypeId, old_name: &str, new_name: SmolStr) {
        self.by_name.remove(old_name);
        self.by_name.insert(new_name, id);
    }

    /// Fill a stub slot with a real type definition.
    pub fn fill_stub(&mut self, id: TypeId, def: TypeDef) {
        let idx = id.0 as usize;
        if idx < self.types.len() {
            self.types[idx] = Some(def);
        }
    }

    /// Iterate over all filled type definitions and their IDs.
    pub fn iter(&self) -> impl Iterator<Item = (TypeId, &TypeDef)> {
        self.types
            .iter()
            .enumerate()
            .filter_map(|(i, opt)| opt.as_ref().map(|def| (TypeId(i as u32), def)))
    }

    /// Iterate over all registered type names (excluding stubs).
    pub fn iter_names(&self) -> impl Iterator<Item = &str> {
        self.by_name.keys().map(|k| k.as_str())
    }

    pub(crate) fn set_trait_fn_return_type(
        &mut self,
        trait_id: TypeId,
        function: SmolStr,
        return_type: crate::ast::TypeExpr,
    ) {
        self.trait_fn_return_types
            .insert((trait_id, function), return_type);
    }

    pub(crate) fn trait_fn_return_type(
        &self,
        trait_id: TypeId,
        function: &str,
    ) -> Option<&crate::ast::TypeExpr> {
        self.trait_fn_return_types
            .get(&(trait_id, SmolStr::new(function)))
    }

    pub(crate) fn clone_trait_fn_return_types(
        &mut self,
        source: &TypeRegistry,
        source_id: TypeId,
        target_id: TypeId,
    ) {
        let entries: Vec<_> = source
            .trait_fn_return_types
            .iter()
            .filter(|((id, _), _)| *id == source_id)
            .map(|((_, name), ty)| (name.clone(), ty.clone()))
            .collect();
        for (name, ty) in entries {
            self.trait_fn_return_types.insert((target_id, name), ty);
        }
    }

    pub(crate) fn set_origin(&mut self, id: TypeId, namespace: SmolStr, declaration: SmolStr) {
        self.origins.insert(id, (namespace, declaration));
    }

    pub(crate) fn clone_origin(
        &mut self,
        source: &TypeRegistry,
        source_id: TypeId,
        target_id: TypeId,
    ) {
        if let Some((namespace, declaration)) = source.origins.get(&source_id) {
            self.origins
                .insert(target_id, (namespace.clone(), declaration.clone()));
        }
    }

    pub(crate) fn find_origin(&self, namespace: &str, declaration: &str) -> Option<TypeId> {
        self.origins
            .iter()
            .find_map(|(id, (ns, name))| (ns == namespace && name == declaration).then_some(*id))
    }

    /// Return the defining namespace and bare declaration name for a type.
    pub fn origin(&self, id: TypeId) -> Option<(&str, &str)> {
        self.origins
            .get(&id)
            .map(|(namespace, declaration)| (namespace.as_str(), declaration.as_str()))
    }

    pub(crate) fn set_impl_trait_id(&mut self, impl_id: TypeId, trait_id: TypeId) {
        self.impl_trait_ids.insert(impl_id, trait_id);
    }

    pub(crate) fn set_impl_trait_span(&mut self, impl_id: TypeId, span: Span) {
        self.impl_trait_spans.insert(impl_id, span);
    }

    pub(crate) fn impl_trait_span(&self, impl_id: TypeId) -> Option<Span> {
        self.impl_trait_spans.get(&impl_id).copied()
    }

    /// Return the resolved trait identity associated with an impl record.
    pub fn impl_trait_id(&self, impl_id: TypeId) -> Option<TypeId> {
        self.impl_trait_ids.get(&impl_id).copied()
    }

    /// Resolve an impl reference to its persisted trait identity and definition.
    pub fn trait_for_impl(&self, implementation: &ImplDef) -> Option<(TypeId, &TraitDef)> {
        let impl_id = self.iter().find_map(|(id, definition)| match definition {
            TypeDef::Impl(candidate) if std::ptr::eq(candidate, implementation) => Some(id),
            _ => None,
        })?;
        let trait_id = self.impl_trait_id(impl_id)?;
        match self.get(trait_id) {
            Some(TypeDef::Trait(trait_def)) => Some((trait_id, trait_def)),
            _ => None,
        }
    }
}

// ---------------------------------------------------------------------------
// ResolvedType
// ---------------------------------------------------------------------------

/// A fully resolved type reference in the IR.
///
/// All named types have been resolved to [`TypeId`] handles. Container types
/// (optional, array, map, result) wrap their inner types recursively.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum ResolvedType {
    Primitive(PrimitiveType),
    SubByte(SubByteType),
    Semantic(SemanticType),
    Named(TypeId),
    Optional(Box<ResolvedType>),
    Array(Box<ResolvedType>),
    /// Fixed-size array with compile-time known length: `array<T, N>`
    FixedArray(Box<ResolvedType>, u64),
    Set(Box<ResolvedType>),
    Map(Box<ResolvedType>, Box<ResolvedType>),
    Result(Box<ResolvedType>, Box<ResolvedType>),
    /// Geometric types parameterized by element type
    Vec2(Box<ResolvedType>),
    Vec3(Box<ResolvedType>),
    Vec4(Box<ResolvedType>),
    Quat(Box<ResolvedType>),
    Mat3(Box<ResolvedType>),
    Mat4(Box<ResolvedType>),
    /// Inline bitfield: bits { name1, name2, ... }
    /// Wire size = number of bits, LSB-first packing
    BitsInline(Vec<SmolStr>),
}

// ---------------------------------------------------------------------------
// Encoding
// ---------------------------------------------------------------------------

/// Wire encoding strategy for a field or type.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum Encoding {
    Default,
    Varint,
    ZigZag,
    Delta(Box<Encoding>),
}

/// Per-field encoding configuration (encoding strategy and optional element limit).
#[derive(Debug, Clone, PartialEq)]
pub struct FieldEncoding {
    pub encoding: Encoding,
    pub limit: Option<u64>,
}

impl FieldEncoding {
    /// Create a default field encoding (no varint, no limit).
    pub fn default_encoding() -> Self {
        Self {
            encoding: Encoding::Default,
            limit: None,
        }
    }
}

// ---------------------------------------------------------------------------
// WireSize
// ---------------------------------------------------------------------------

/// The computed wire size of a type, in bits.
///
/// Fixed-size types have a known bit count. Variable-size types (containing
/// arrays, optionals, or varints) have a minimum and optional maximum.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum WireSize {
    Fixed(u64),
    Variable {
        min_bits: u64,
        max_bits: Option<u64>,
    },
}

// ---------------------------------------------------------------------------
// ResolvedAnnotations
// ---------------------------------------------------------------------------

/// Information attached to a `@deprecated` annotation.
#[derive(Debug, Clone, PartialEq)]
pub struct DeprecatedInfo {
    pub reason: SmolStr,
    pub since: Option<SmolStr>,
}

/// A user-defined annotation preserved from source through to IR.
///
/// Unknown annotations (not `doc`, `deprecated`, `since`, `revision`,
/// `non_exhaustive`, `version`, or encoding annotations) are collected
/// here so SDK consumers can access custom metadata without codegen
/// needing to know about them.
#[derive(Debug, Clone, PartialEq)]
pub struct CustomAnnotation {
    pub name: SmolStr,
    pub args: Vec<CustomAnnotationArg>,
}

/// A single argument to a custom annotation.
#[derive(Debug, Clone, PartialEq)]
pub struct CustomAnnotationArg {
    pub key: Option<SmolStr>,
    pub value: CustomAnnotationValue,
}

/// Value of a custom annotation argument.
#[derive(Debug, Clone, PartialEq)]
pub enum CustomAnnotationValue {
    Int(u64),
    Hex(u64),
    Str(SmolStr),
    Bool(bool),
    Ident(SmolStr),
}

/// Annotations resolved from source and available on any IR node.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ResolvedAnnotations {
    pub deprecated: Option<DeprecatedInfo>,
    pub since: Option<SmolStr>,
    pub doc: Vec<SmolStr>,
    pub revision: Option<u64>,
    pub non_exhaustive: bool,
    pub version: Option<SmolStr>,
    pub custom: Vec<CustomAnnotation>,
}

// ---------------------------------------------------------------------------
// TombstoneDef
// ---------------------------------------------------------------------------

/// A tombstoned (removed) field or variant ordinal.
#[derive(Debug, Clone, PartialEq)]
pub struct TombstoneDef {
    pub span: Span,
    pub ordinal: u32,
    pub reason: SmolStr,
    pub since: Option<SmolStr>,
    /// Resolved original field type, retained as wire-inert history metadata.
    pub original_type: Option<ResolvedType>,
}