item_tree.rs

//! Position-independent item storage.
//!
//! The `ItemTree` contains minimal representations of all items in a container.
//! It acts as an "invalidation barrier" - only changes to item signatures
//! cause the `ItemTree` to change, not edits to whitespace, comments, or bodies.

use std::ops::Index;

use baml_base::{FieldAttr, Name, TyAttr};
use indexmap::IndexMap;
use rowan::TextRange;
use rustc_hash::FxHashMap;

use crate::{
    ids::{ItemKind, LocalItemId, allocate_local_id},
    loc::{
        ClassMarker, ClientMarker, EnumMarker, FunctionMarker, GeneratorMarker, RetryPolicyMarker,
        TemplateStringMarker, TestMarker, TypeAliasMarker,
    },
    type_ref::TypeRef,
};

//
// ──────────────────────────────────────────────────────── ATTRIBUTE TYPE ─────
//

/// Represents an attribute that may or may not be explicitly set in source.
///
/// This generic type captures whether an attribute was present in the BAML source.
/// `T` is the value type: `String` for attributes like `@alias("name")`,
/// or `()` for presence-only attributes like `@@dynamic`.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub enum Attribute<T> {
    /// Attribute was not present in source.
    #[default]
    Unset,
    /// Attribute was explicitly set with the given value.
    Explicit(T),
}

impl<T> Attribute<T> {
    /// Returns the value if explicitly set, None otherwise.
    pub fn value(&self) -> Option<&T> {
        match self {
            Attribute::Unset => None,
            Attribute::Explicit(v) => Some(v),
        }
    }

    /// Returns true if the attribute was explicitly set.
    pub fn is_explicit(&self) -> bool {
        matches!(self, Attribute::Explicit(_))
    }
}

//
// ──────────────────────────────────────────────────────── ITEM TREE ─────
//

/// Position-independent item storage for a container.
///
/// This is the core HIR data structure. Items are stored in hash maps
/// keyed by name-based IDs, following rust-analyzer's architecture.
///
/// **Key property:** Items are indexed by name, not source position.
/// Adding an item in the middle of the file doesn't change the `LocalItemIds`
/// of other items because `LocalItemIds` are derived from names.
///
/// **Collision handling:** IDs pack a 16-bit hash and 16-bit collision index.
/// The `next_index` map tracks the next available index per `(ItemKind, hash)`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ItemTree {
    pub(crate) functions: FxHashMap<LocalItemId<FunctionMarker>, Function>,
    pub(crate) classes: FxHashMap<LocalItemId<ClassMarker>, Class>,
    pub(crate) enums: FxHashMap<LocalItemId<EnumMarker>, Enum>,
    pub(crate) type_aliases: FxHashMap<LocalItemId<TypeAliasMarker>, TypeAlias>,
    pub(crate) clients: FxHashMap<LocalItemId<ClientMarker>, Client>,
    pub(crate) generators: FxHashMap<LocalItemId<GeneratorMarker>, Generator>,
    pub(crate) tests: FxHashMap<LocalItemId<TestMarker>, Test>,
    pub(crate) template_strings: FxHashMap<LocalItemId<TemplateStringMarker>, TemplateString>,
    pub(crate) retry_policies: FxHashMap<LocalItemId<RetryPolicyMarker>, RetryPolicy>,

    /// Collision tracker: (`ItemKind`, hash) -> next available index.
    /// Single map for all item types, following rust-analyzer's pattern.
    next_index: FxHashMap<(ItemKind, u16), u16>,
}

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

impl ItemTree {
    /// Create a new empty `ItemTree`.
    pub fn new() -> Self {
        Self {
            functions: FxHashMap::default(),
            classes: FxHashMap::default(),
            enums: FxHashMap::default(),
            type_aliases: FxHashMap::default(),
            clients: FxHashMap::default(),
            generators: FxHashMap::default(),
            tests: FxHashMap::default(),
            template_strings: FxHashMap::default(),
            retry_policies: FxHashMap::default(),
            next_index: FxHashMap::default(),
        }
    }

    /// Allocate a collision-resistant ID for an item.
    /// Returns a `LocalItemId` with the name's hash and a unique collision index.
    fn alloc_id<T>(&mut self, kind: ItemKind, name: &Name) -> LocalItemId<T> {
        allocate_local_id(&mut self.next_index, kind, name)
    }

    /// Add a function and return its local ID.
    pub fn alloc_function(&mut self, func: Function) -> LocalItemId<FunctionMarker> {
        let id = self.alloc_id(ItemKind::Function, &func.name);
        self.functions.insert(id, func);
        id
    }

    /// Add a class and return its local ID.
    pub fn alloc_class(&mut self, class: Class) -> LocalItemId<ClassMarker> {
        let id = self.alloc_id(ItemKind::Class, &class.name);
        self.classes.insert(id, class);
        id
    }

    /// Add an enum and return its local ID.
    pub fn alloc_enum(&mut self, enum_def: Enum) -> LocalItemId<EnumMarker> {
        let id = self.alloc_id(ItemKind::Enum, &enum_def.name);
        self.enums.insert(id, enum_def);
        id
    }

    /// Add a type alias and return its local ID.
    pub fn alloc_type_alias(&mut self, alias: TypeAlias) -> LocalItemId<TypeAliasMarker> {
        let id = self.alloc_id(ItemKind::TypeAlias, &alias.name);
        self.type_aliases.insert(id, alias);
        id
    }

    /// Add a client and return its local ID.
    pub fn alloc_client(&mut self, client: Client) -> LocalItemId<ClientMarker> {
        let id = self.alloc_id(ItemKind::Client, &client.name);
        self.clients.insert(id, client);
        id
    }

    /// Add a test and return its local ID.
    pub fn alloc_test(&mut self, test: Test) -> LocalItemId<TestMarker> {
        let id = self.alloc_id(ItemKind::Test, &test.name);
        self.tests.insert(id, test);
        id
    }

    /// Add a generator and return its local ID.
    pub fn alloc_generator(&mut self, generator: Generator) -> LocalItemId<GeneratorMarker> {
        let id = self.alloc_id(ItemKind::Generator, &generator.name);
        self.generators.insert(id, generator);
        id
    }

    /// Add a template string and return its local ID.
    pub fn alloc_template_string(
        &mut self,
        template_string: TemplateString,
    ) -> LocalItemId<TemplateStringMarker> {
        let id = self.alloc_id(ItemKind::TemplateString, &template_string.name);
        self.template_strings.insert(id, template_string);
        id
    }

    /// Add a retry policy and return its local ID.
    pub fn alloc_retry_policy(
        &mut self,
        retry_policy: RetryPolicy,
    ) -> LocalItemId<RetryPolicyMarker> {
        let id = self.alloc_id(ItemKind::RetryPolicy, &retry_policy.name);
        self.retry_policies.insert(id, retry_policy);
        id
    }

    /// Iterate over all classes in the item tree.
    pub fn iter_classes(&self) -> impl Iterator<Item = (&LocalItemId<ClassMarker>, &Class)> {
        self.classes.iter()
    }
}

/// Metadata for compiler-generated functions.
///
/// These functions are created by the compiler during HIR lowering,
/// not by the user. They skip type inference and have special handling.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CompilerGenerated {
    /// Client resolve function - evaluates options and returns `PrimitiveClient`.
    /// Contains the client name (e.g., "GPT4" for "GPT4.resolve").
    ClientResolve { client_name: Name },
    /// LLM main call - function named `base_name` (e.g. "Foo") with body
    /// `baml.llm.call_llm_function(base_name, args)`.
    LlmCall { base_name: Name },
    /// LLM `render_prompt` - function named `base_name.render_prompt` with body
    /// `baml.llm.render_prompt(base_name, args)`.
    LlmRenderPrompt { base_name: Name },
    /// LLM `build_request` - function named `base_name.build_request` with body
    /// `baml.llm.build_request(base_name, args)`.
    LlmBuildRequest { base_name: Name },
}

/// A function definition in the `ItemTree`.
///
/// This is the MINIMAL representation - ONLY the name.
/// Everything else (params, return type, body) is in separate queries for incrementality.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Function {
    pub name: Name,
    /// If this function is compiler-generated, contains the metadata.
    /// `None` for user-defined functions.
    pub compiler_generated: Option<CompilerGenerated>,
}

/// A class definition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Class {
    pub name: Name,
    pub fields: Vec<Field>,

    // Block attributes (@@dynamic, @@alias, @@description)
    /// @@dynamic - marks class as dynamically extensible
    pub is_dynamic: Attribute<()>,
    /// @@alias("name") - alternative name for serialization
    pub alias: Attribute<String>,
    /// @@description("text") - documentation for the class
    pub description: Attribute<String>,
    /// Class-level type attribute (e.g., from @@stream.done).
    pub ty_attr: TyAttr,
    // Note: Generic parameters are queried separately via generic_params()
    // for incrementality - changes to generics don't invalidate ItemTree
}

/// Class field.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Field {
    pub name: Name,
    pub type_ref: TypeRef,

    // Field attributes (@alias, @description, @skip)
    /// @alias("name") - alternative name for serialization
    pub alias: Attribute<String>,
    /// @description("text") - documentation for the field
    pub description: Attribute<String>,
    /// @skip - exclude field from serialization
    pub skip: Attribute<()>,
    /// Field attributes for streaming (e.g., from @sap.*)
    pub field_attr: FieldAttr,
}

/// An enum definition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Enum {
    pub name: Name,
    pub variants: Vec<EnumVariant>,

    // Block attributes (@@alias)
    /// @@alias("name") - alternative name for serialization
    pub alias: Attribute<String>,
    /// Enum-level type attribute.
    pub ty_attr: TyAttr,
    // Note: Generic parameters are queried separately via generic_params()
}

/// Enum variant.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EnumVariant {
    pub name: Name,

    // Variant attributes (@alias, @description, @skip)
    /// @alias("name") - alternative name for serialization
    pub alias: Attribute<String>,
    /// @description("text") - documentation for the variant
    pub description: Attribute<String>,
    /// @skip - exclude variant from serialization
    pub skip: Attribute<()>,
}

/// Type alias definition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TypeAlias {
    pub name: Name,
    pub type_ref: TypeRef,
    // Note: Generic parameters are queried separately via generic_params()
}

/// Client configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Client {
    pub name: Name,
    pub provider: Name,
    /// Default role for chat messages (e.g., "user").
    pub default_role: Option<String>,
    /// Allowed roles for chat messages.
    pub allowed_roles: Vec<String>,
    /// Name of the retry policy (references a top-level `retry_policy` definition).
    pub retry_policy_name: Option<Name>,
    /// Span of the retry policy reference (for diagnostics).
    pub retry_policy_span: Option<TextRange>,
    /// Sub-client names for composite clients (fallback/round-robin).
    /// Empty for primitive clients.
    pub sub_client_names: Vec<Name>,
    /// Optional round-robin start index (`options { start ... }`).
    /// Only used by `round-robin` providers.
    pub round_robin_start: Option<i32>,
}

/// Retry policy configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetryPolicy {
    pub name: Name,
    /// Maximum number of retries.
    pub max_retries: Option<String>,
    /// Initial delay before first retry (milliseconds).
    pub initial_delay_ms: Option<String>,
    /// Delay multiplier for exponential backoff.
    pub multiplier: Option<String>,
    /// Maximum delay between retries (milliseconds).
    pub max_delay_ms: Option<String>,
}

/// A test argument value, extracted from the CST during HIR lowering.
///
/// These are untyped constant values — type checking against function
/// signatures happens at a later stage (during emission or at runtime).
#[derive(Debug, Clone)]
pub enum TestArgValue {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(String),
    Null,
    Array(Vec<TestArgValue>),
    Map(IndexMap<String, TestArgValue>),
}

// Manual PartialEq/Eq implementation to handle f64 comparison via bit pattern.
// This satisfies ItemTree's Eq requirement (needed for salsa early cutoff).
impl PartialEq for TestArgValue {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Int(a), Self::Int(b)) => a == b,
            (Self::Float(a), Self::Float(b)) => a.to_bits() == b.to_bits(),
            (Self::Bool(a), Self::Bool(b)) => a == b,
            (Self::String(a), Self::String(b)) => a == b,
            (Self::Null, Self::Null) => true,
            (Self::Array(a), Self::Array(b)) => a == b,
            (Self::Map(a), Self::Map(b)) => a == b,
            _ => false,
        }
    }
}

impl Eq for TestArgValue {}

/// Test definition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Test {
    pub name: Name,

    /// Unresolved function references.
    pub function_refs: Vec<Name>,

    /// Test arguments, keyed by parameter name.
    pub args: IndexMap<String, TestArgValue>,

    /// Type builder block containing dynamic type definitions.
    pub type_builder: Option<TypeBuilderBlock>,
}

/// A `type_builder` block inside a test definition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TypeBuilderBlock {
    pub entries: Vec<TypeBuilderEntry>,
}

/// An entry in a `type_builder` block.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TypeBuilderEntry {
    /// A class definition (non-dynamic).
    Class(Class),
    /// An enum definition (non-dynamic).
    Enum(Enum),
    /// A dynamic class definition.
    DynamicClass(Class),
    /// A dynamic enum definition.
    DynamicEnum(Enum),
    /// A type alias.
    TypeAlias(TypeAlias),
}

/// Generator configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Generator {
    pub name: Name,
    /// The output type (e.g., "python/pydantic", "typescript").
    pub output_type: Option<String>,
    /// The output directory (relative to `baml_src`).
    pub output_dir: Option<String>,
    /// The version string.
    pub version: Option<String>,
    /// Default client mode: "sync" or "async".
    pub default_client_mode: Option<String>,
    /// Command to run after code generation.
    pub on_generate: Option<String>,
    /// Project identifier for boundary-cloud.
    pub project: Option<String>,
    /// Go package name (required for Go generator).
    pub client_package_name: Option<String>,
    /// Module format for TypeScript: "cjs" or "esm".
    pub module_format: Option<String>,
}

/// Template string definition.
///
/// Template strings are reusable prompt fragments that can be called
/// from within function prompts. They have parameters and a body template.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TemplateString {
    pub name: Name,
}

//
// ──────────────────────────────────────────────────────── INDEX IMPLS ─────
//

/// Index `ItemTree` by `FunctionMarker` to get Function data.
impl Index<LocalItemId<FunctionMarker>> for ItemTree {
    type Output = Function;
    fn index(&self, index: LocalItemId<FunctionMarker>) -> &Self::Output {
        self.functions
            .get(&index)
            .expect("Function not found in ItemTree")
    }
}

/// Index `ItemTree` by `ClassMarker` to get Class data.
impl Index<LocalItemId<ClassMarker>> for ItemTree {
    type Output = Class;
    fn index(&self, index: LocalItemId<ClassMarker>) -> &Self::Output {
        self.classes
            .get(&index)
            .expect("Class not found in ItemTree")
    }
}

/// Index `ItemTree` by `EnumMarker` to get Enum data.
impl Index<LocalItemId<EnumMarker>> for ItemTree {
    type Output = Enum;
    fn index(&self, index: LocalItemId<EnumMarker>) -> &Self::Output {
        self.enums.get(&index).expect("Enum not found in ItemTree")
    }
}

/// Index `ItemTree` by `TypeAliasMarker` to get `TypeAlias` data.
impl Index<LocalItemId<TypeAliasMarker>> for ItemTree {
    type Output = TypeAlias;
    fn index(&self, index: LocalItemId<TypeAliasMarker>) -> &Self::Output {
        self.type_aliases
            .get(&index)
            .expect("TypeAlias not found in ItemTree")
    }
}

/// Index `ItemTree` by `ClientMarker` to get Client data.
impl Index<LocalItemId<ClientMarker>> for ItemTree {
    type Output = Client;
    fn index(&self, index: LocalItemId<ClientMarker>) -> &Self::Output {
        self.clients
            .get(&index)
            .expect("Client not found in ItemTree")
    }
}

/// Index `ItemTree` by `TestMarker` to get Test data.
impl Index<LocalItemId<TestMarker>> for ItemTree {
    type Output = Test;
    fn index(&self, index: LocalItemId<TestMarker>) -> &Self::Output {
        self.tests.get(&index).expect("Test not found in ItemTree")
    }
}

/// Index `ItemTree` by `GeneratorMarker` to get Generator data.
impl Index<LocalItemId<GeneratorMarker>> for ItemTree {
    type Output = Generator;
    fn index(&self, index: LocalItemId<GeneratorMarker>) -> &Self::Output {
        self.generators
            .get(&index)
            .expect("Generator not found in ItemTree")
    }
}

/// Index `ItemTree` by `TemplateStringMarker` to get `TemplateString` data.
impl Index<LocalItemId<TemplateStringMarker>> for ItemTree {
    type Output = TemplateString;
    fn index(&self, index: LocalItemId<TemplateStringMarker>) -> &Self::Output {
        self.template_strings
            .get(&index)
            .expect("TemplateString not found in ItemTree")
    }
}

/// Index `ItemTree` by `RetryPolicyMarker` to get `RetryPolicy` data.
impl Index<LocalItemId<RetryPolicyMarker>> for ItemTree {
    type Output = RetryPolicy;
    fn index(&self, index: LocalItemId<RetryPolicyMarker>) -> &Self::Output {
        self.retry_policies
            .get(&index)
            .expect("RetryPolicy not found in ItemTree")
    }
}