From e7885dc578965e1be899ac55234341ca46beeec3 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 14:21:11 -0700 Subject: [PATCH 01/20] Refactor HIR item tree members to submodules --- .../baml_compiler2_hir/src/item_tree.rs | 385 ++---------------- .../src/item_tree/classes.rs | 47 +++ .../src/item_tree/clients.rs | 14 + .../src/item_tree/common.rs | 35 ++ .../baml_compiler2_hir/src/item_tree/enums.rs | 27 ++ .../src/item_tree/functions.rs | 56 +++ .../src/item_tree/interfaces.rs | 127 ++++++ .../baml_compiler2_hir/src/item_tree/lets.rs | 14 + .../src/item_tree/retry_policies.rs | 14 + .../src/item_tree/source_map.rs | 25 ++ .../src/item_tree/template_strings.rs | 15 + .../src/item_tree/test_items.rs | 11 + .../src/item_tree/type_aliases.rs | 12 + 13 files changed, 422 insertions(+), 360 deletions(-) create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/classes.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/clients.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/enums.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/interfaces.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/lets.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/retry_policies.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/source_map.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/template_strings.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/test_items.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/type_aliases.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs index cd52ada06ec..a52510a5c15 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs @@ -4,12 +4,37 @@ //! following the same scheme as `baml_compiler_hir::item_tree`. //! Items are indexed by name (not source position) for position-independence. +mod classes; +mod clients; +mod common; +mod enums; +mod functions; +mod interfaces; +mod lets; +mod retry_policies; +mod source_map; +mod template_strings; +mod test_items; +mod type_aliases; + use std::ops::Index; use baml_base::Name; use baml_compiler2_ast as ast; +pub use classes::*; +pub use clients::*; +pub use common::*; +pub use enums::*; +pub use functions::*; +pub use interfaces::*; +pub use lets::*; +pub use retry_policies::*; use rustc_hash::FxHashMap; +pub use source_map::*; +pub use template_strings::*; +pub use test_items::*; use text_size::TextRange; +pub use type_aliases::*; use crate::ids::{ ClassMarker, ClientMarker, EnumMarker, FunctionMarker, ImplMarker, InterfaceMarker, ItemKind, @@ -17,366 +42,6 @@ use crate::ids::{ hash_impl_key, hash_name, }; -// ── Span-free attribute representation ─────────────────────────────────────── - -/// A span-free attribute for position-independent storage in the `ItemTree`. -/// Derived from `ast::RawAttribute` with all `TextRange`s stripped. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Attribute { - pub name: Name, - pub args: Vec, -} - -/// A span-free attribute argument. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AttributeArg { - pub key: Option, - pub value: String, -} - -impl From<&ast::RawAttribute> for Attribute { - fn from(raw: &ast::RawAttribute) -> Self { - Self { - name: raw.name.clone(), - args: raw.args.iter().map(AttributeArg::from).collect(), - } - } -} - -impl From<&ast::RawAttributeArg> for AttributeArg { - fn from(raw: &ast::RawAttributeArg) -> Self { - Self { - key: raw.key.clone(), - value: raw.value.clone(), - } - } -} - -// ── Minimal item data structs ──────────────────────────────────────────────── - -/// Full function data stored in the `ItemTree`. -/// Params and return type are stored for signature queries. -/// Body is stored for body queries (no CST re-parsing needed). -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Function { - pub name: Name, - /// Generic type parameters (e.g., `["T", "U"]`). - /// Empty for non-generic functions. - pub generic_params: Vec, - /// BEP-044 generic bounds parallel to `generic_params`. `Some(te)` - /// means the parameter at the matching index was declared with - /// `T extends `; `None` means unbounded. - pub generic_param_bounds: Vec>, - /// Function parameters with optional type annotations and spans. - pub params: Vec, - /// Function parameter default expression arena. - pub defaults: ast::FunctionDefaults, - /// Return type with its source span. - pub return_type: Option, - /// Throws contract type with its source span. - pub throws: Option, - /// Function body — either an expression or a builtin. - pub body: Option, - /// Declarative metadata, if this function was declared with declarative syntax. - pub declarative_meta: Option, - pub origin: ast::FunctionOrigin, - /// Joined `///` doc-comment lines preceding this declaration. - pub docstring: Option, - /// BEP-049 §10: set when the fn def had a `//baml:tagged_string` marker. - /// Mirrors `ast::FunctionDef::is_tagged_template_tag` so TIR can validate - /// tagged-template tags without re-reading the CST. - pub is_tagged_template_tag: bool, - /// Full source span of the function. - pub span: TextRange, -} - -/// A function parameter entry in the `ItemTree`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FunctionParam { - pub name: Name, - pub type_expr: Option, - pub default: Option, - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DefaultExprRef { - pub function: LocalItemId, - pub expr: ast::DefaultExprId, -} - -/// A class field stored in the `ItemTree`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ClassField { - pub name: Name, - pub type_expr: Option, - pub attributes: Vec, - /// Joined `///` doc-comment lines preceding this declaration. - pub docstring: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Class { - pub name: Name, - /// Generic type parameters (e.g., `["T"]` for `Array`). - /// Empty for non-generic classes. - pub generic_params: Vec, - /// Generic bounds parallel to `generic_params`. `Some(te)` means - /// `T extends `; `None` means unbounded. - pub generic_param_bounds: Vec>, - /// Fields of the class, in declaration order. - pub fields: Vec, - /// Methods defined inside this class, referencing their `Function` entries - /// in the same `ItemTree`. Includes both class-level methods and methods - /// declared inside `implements I { ... }` blocks (BEP-044) — flattened so - /// downstream code (e.g. signature queries, method dispatch) can iterate - /// uniformly. - pub methods: Vec>, - /// `implements I { ... }` blocks, in declaration order. Each block keeps - /// the raw target `TypeExpr` so generic parameters like `Container` - /// survive name resolution, plus field redeclarations from the block. - pub implements: Vec, - /// Block-level attributes (@@description, @@alias, etc.). - pub attributes: Vec, - /// Joined `///` doc-comment lines preceding this declaration. - pub docstring: Option, - /// Full source span of the class declaration. - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ImplementsBlock { - pub target: ast::TypeExpr, - pub field_links: Vec, - pub associated_type_bindings: Vec, - pub is_out_of_body: bool, - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct InterfaceFieldLink { - pub interface_field: Name, - pub class_field: Name, - pub span: TextRange, - pub interface_field_span: TextRange, - pub class_field_span: TextRange, -} - -impl InterfaceFieldLink { - pub(crate) fn from_ast(link: &ast::InterfaceFieldLinkDef) -> Self { - Self { - interface_field: link.interface_field.clone(), - class_field: link.class_field.clone(), - span: link.span, - interface_field_span: link.interface_field_span, - class_field_span: link.class_field_span, - } - } -} - -/// A generic parameter on an out-of-body `implements` block, paired with its set -/// of `&`-separated interface bounds (`` → `bounds = []`; `` -/// → `bounds = [A, B]`). Pairing name and bounds makes a length mismatch between -/// the two unrepresentable. Bounds are unresolved `TypeExpr`s here; they resolve -/// to `baml_type::Interface` constraints downstream. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct GenericParam { - pub name: Name, - pub bounds: Vec, -} - -/// What an `implements` block applies to. Unifying the owner with the for-target -/// makes "in-body with an explicit for-target" and "out-of-body without one" -/// both unrepresentable. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ImplSubject { - /// `implements I { … }` in a class body (or a simple `implement I for C` - /// merged onto `C`). The for-type is the class itself; the generics are the - /// class's. `out_of_body` records the syntactic origin for diagnostics only - /// — it must NOT influence resolution, dispatch, or coherence. - InClass { - class: LocalItemId, - out_of_body: bool, - }, - /// `implement<…> I for { … }`: an explicit for-type plus the - /// block's own generic parameters. - Free { - for_target: ast::TypeExpr, - generics: Vec, - }, -} - -/// A unified `implements` block (both kinds) stored in the `ItemTree`, keyed by -/// a stable `LocalItemId`. The interface target is kept as a raw -/// `TypeExpr` here; resolution to an `InterfaceLoc` + `Ty` happens lazily in the -/// `impl_data` query so HIR construction stays independent of name resolution. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ImplBlock { - pub subject: ImplSubject, - pub interface_target: ast::TypeExpr, - pub field_links: Vec, - pub associated_type_bindings: Vec, - pub methods: Vec>, - pub span: TextRange, -} - -/// An enum variant stored in the `ItemTree`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct EnumVariant { - pub name: Name, - /// Field-level attributes (@description, @alias, @skip, etc.). - pub attributes: Vec, - /// Joined `///` doc-comment lines preceding this declaration. - pub docstring: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Enum { - pub name: Name, - /// Variants of the enum, in declaration order. - pub variants: Vec, - /// Block-level attributes (@@description, @@alias, etc.). - pub attributes: Vec, - /// Joined `///` doc-comment lines preceding this declaration. - pub docstring: Option, - /// Full source span of the enum declaration. - pub span: TextRange, -} - -/// A required (no-body) method signature on an interface (BEP-044). -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct InterfaceMethodSig { - pub name: Name, - /// Generic type parameters local to this method. - pub generic_params: Vec, - /// BEP-044 generic bounds parallel to `generic_params`. - pub generic_param_bounds: Vec>, - pub params: Vec, - pub return_type: Option, - pub throws: Option, - pub attributes: Vec, - pub docstring: Option, - pub span: TextRange, -} - -/// An interface (BEP-044) stored in the `ItemTree`. -/// -/// Default methods are stored as full `FunctionMarker` entries in `methods` -/// (same as `Class::methods`). Required signatures live in `required_methods`. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Interface { - pub name: Name, - /// Generic type parameters declared on the interface. - pub generic_params: Vec, - /// BEP-044 generic bounds parallel to `generic_params`. - pub generic_param_bounds: Vec>, - /// Required interfaces from `requires I1, I2, …`. - pub requires: Vec, - /// Field signatures declared on the interface. Interface fields cannot - /// have default values. - pub fields: Vec, - /// Associated type declarations on the interface (BEP-057). - pub associated_types: Vec, - /// Default methods (with bodies). Implementing classes inherit them. - pub default_methods: Vec>, - /// Required methods (no body). Implementing classes must provide a body. - pub required_methods: Vec, - pub attributes: Vec, - pub docstring: Option, - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TypeAlias { - pub name: Name, - /// The type expression on the RHS of the alias, if present. - pub type_expr: Option, - /// Full source span of the type alias declaration. - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Client { - pub name: Name, - /// Provider name (e.g., "openai", "anthropic", "fallback", "round-robin"). - pub provider: Option, - /// Sub-client names for fallback/round-robin clients. - pub sub_client_names: Vec, - /// Retry policy name, if configured. - pub retry_policy_name: Option, - /// Starting index for round-robin clients. - pub round_robin_start: Option, -} - -pub use baml_compiler2_ast::ast::TestArgValue; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Test { - pub name: Name, - /// The function(s) this test exercises. - pub function_refs: Vec, - /// Test arguments as key-value pairs. - pub args: Vec<(Name, TestArgValue)>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TemplateString { - pub name: Name, - /// Template parameters with optional type annotations and spans. - pub params: Vec, - /// Template body text (Jinja template). - pub body: Option, - /// Full source span of the template string declaration. - pub span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct RetryPolicy { - pub name: Name, - /// Raw string value of `max_retries` (parsed at emit time). - pub max_retries: Option, - /// Raw string value of `initial_delay_ms`. - pub initial_delay_ms: Option, - /// Raw string value of multiplier. - pub multiplier: Option, - /// Raw string value of `max_delay_ms`. - pub max_delay_ms: Option, -} - -/// A top-level let binding stored in the `ItemTree`. -/// Carries the optional initializer `ExprBody` for body queries. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct Let { - pub name: Name, - pub initializer: Option<(ast::ExprBody, ast::AstSourceMap)>, - pub origin: ast::LetOrigin, - pub span: TextRange, - pub name_span: TextRange, -} - -// ── ItemTreeSourceMap ───────────────────────────────────────────────────────── - -/// Parallel source map for `ItemTree` — stores name spans that are -/// deliberately excluded from the semantic `ItemTree` to avoid polluting -/// Salsa's early-cutoff comparisons with position data. -/// -/// Follows the same body/signature source-map pattern used by -/// `function_body` / `function_body_source_map`. -#[derive(Debug, Clone, Default)] -pub struct ItemTreeSourceMap { - /// `name_span` for each class's fields, parallel to `Class::fields`. - pub class_field_spans: FxHashMap, Vec>, - /// `name_span` for each enum's variants, parallel to `Enum::variants`. - pub enum_variant_spans: FxHashMap, Vec>, - /// `name_span` for each function. - pub function_name_spans: FxHashMap, TextRange>, - /// `name_span` for each interface's fields, parallel to `Interface::fields`. - pub interface_field_spans: FxHashMap, Vec>, - /// `name_span` for each interface's required methods, parallel to - /// `Interface::required_methods`. - pub interface_method_spans: FxHashMap, Vec>, -} - // ── ItemTree ───────────────────────────────────────────────────────────────── /// Position-independent item storage for a single file. diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/classes.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/classes.rs new file mode 100644 index 00000000000..829f58a2d11 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/classes.rs @@ -0,0 +1,47 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; +use text_size::TextRange; + +use crate::{ + ids::{FunctionMarker, LocalItemId}, + item_tree::{Attribute, ImplementsBlock}, +}; + +/// A class field stored in the `ItemTree`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClassField { + pub name: Name, + pub type_expr: Option, + pub attributes: Vec, + /// Joined `///` doc-comment lines preceding this declaration. + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Class { + pub name: Name, + /// Generic type parameters (e.g., `["T"]` for `Array`). + /// Empty for non-generic classes. + pub generic_params: Vec, + /// Generic bounds parallel to `generic_params`. `Some(te)` means + /// `T extends `; `None` means unbounded. + pub generic_param_bounds: Vec>, + /// Fields of the class, in declaration order. + pub fields: Vec, + /// Methods defined inside this class, referencing their `Function` entries + /// in the same `ItemTree`. Includes both class-level methods and methods + /// declared inside `implements I { ... }` blocks (BEP-044) — flattened so + /// downstream code (e.g. signature queries, method dispatch) can iterate + /// uniformly. + pub methods: Vec>, + /// `implements I { ... }` blocks, in declaration order. Each block keeps + /// the raw target `TypeExpr` so generic parameters like `Container` + /// survive name resolution, plus field redeclarations from the block. + pub implements: Vec, + /// Block-level attributes (@@description, @@alias, etc.). + pub attributes: Vec, + /// Joined `///` doc-comment lines preceding this declaration. + pub docstring: Option, + /// Full source span of the class declaration. + pub span: TextRange, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/clients.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/clients.rs new file mode 100644 index 00000000000..f1613b6073b --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/clients.rs @@ -0,0 +1,14 @@ +use baml_base::Name; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Client { + pub name: Name, + /// Provider name (e.g., "openai", "anthropic", "fallback", "round-robin"). + pub provider: Option, + /// Sub-client names for fallback/round-robin clients. + pub sub_client_names: Vec, + /// Retry policy name, if configured. + pub retry_policy_name: Option, + /// Starting index for round-robin clients. + pub round_robin_start: Option, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs new file mode 100644 index 00000000000..b9003ce80bb --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs @@ -0,0 +1,35 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; + +/// A span-free attribute for position-independent storage in the `ItemTree`. +/// Derived from `ast::RawAttribute` with all `TextRange`s stripped. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Attribute { + pub args: Vec, + pub name: Name, +} + +/// A span-free attribute argument. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AttributeArg { + pub key: Option, + pub value: String, +} + +impl From<&ast::RawAttribute> for Attribute { + fn from(raw: &ast::RawAttribute) -> Self { + Self { + name: raw.name.clone(), + args: raw.args.iter().map(AttributeArg::from).collect(), + } + } +} + +impl From<&ast::RawAttributeArg> for AttributeArg { + fn from(raw: &ast::RawAttributeArg) -> Self { + Self { + key: raw.key.clone(), + value: raw.value.clone(), + } + } +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/enums.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/enums.rs new file mode 100644 index 00000000000..9e5dcd72648 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/enums.rs @@ -0,0 +1,27 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::item_tree::Attribute; + +/// An enum variant stored in the `ItemTree`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EnumVariant { + pub name: Name, + /// Field-level attributes (@description, @alias, @skip, etc.). + pub attributes: Vec, + /// Joined `///` doc-comment lines preceding this declaration. + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Enum { + pub name: Name, + /// Variants of the enum, in declaration order. + pub variants: Vec, + /// Block-level attributes (@@description, @@alias, etc.). + pub attributes: Vec, + /// Joined `///` doc-comment lines preceding this declaration. + pub docstring: Option, + /// Full source span of the enum declaration. + pub span: TextRange, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs new file mode 100644 index 00000000000..559ccc9ec49 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs @@ -0,0 +1,56 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; +use text_size::TextRange; + +use crate::ids::{FunctionMarker, LocalItemId}; + +/// Full function data stored in the `ItemTree`. +/// Params and return type are stored for signature queries. +/// Body is stored for body queries (no CST re-parsing needed). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Function { + pub name: Name, + /// Generic type parameters (e.g., `["T", "U"]`). + /// Empty for non-generic functions. + pub generic_params: Vec, + /// BEP-044 generic bounds parallel to `generic_params`. `Some(te)` + /// means the parameter at the matching index was declared with + /// `T extends `; `None` means unbounded. + pub generic_param_bounds: Vec>, + /// Function parameters with optional type annotations and spans. + pub params: Vec, + /// Function parameter default expression arena. + pub defaults: ast::FunctionDefaults, + /// Return type with its source span. + pub return_type: Option, + /// Throws contract type with its source span. + pub throws: Option, + /// Function body — either an expression or a builtin. + pub body: Option, + /// Declarative metadata, if this function was declared with declarative syntax. + pub declarative_meta: Option, + pub origin: ast::FunctionOrigin, + /// Joined `///` doc-comment lines preceding this declaration. + pub docstring: Option, + /// BEP-049 §10: set when the fn def had a `//baml:tagged_string` marker. + /// Mirrors `ast::FunctionDef::is_tagged_template_tag` so TIR can validate + /// tagged-template tags without re-reading the CST. + pub is_tagged_template_tag: bool, + /// Full source span of the function. + pub span: TextRange, +} + +/// A function parameter entry in the `ItemTree`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionParam { + pub name: Name, + pub type_expr: Option, + pub default: Option, + pub span: TextRange, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DefaultExprRef { + pub function: LocalItemId, + pub expr: ast::DefaultExprId, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/interfaces.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/interfaces.rs new file mode 100644 index 00000000000..9633e654018 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/interfaces.rs @@ -0,0 +1,127 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; +use text_size::TextRange; + +use crate::{ + ids::{ClassMarker, FunctionMarker, LocalItemId}, + item_tree::{Attribute, ClassField, FunctionParam}, +}; + +/// An interface (BEP-044) stored in the `ItemTree`. +/// +/// Default methods are stored as full `FunctionMarker` entries in `methods` +/// (same as `Class::methods`). Required signatures live in `required_methods`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Interface { + pub name: Name, + /// Generic type parameters declared on the interface. + pub generic_params: Vec, + /// BEP-044 generic bounds parallel to `generic_params`. + pub generic_param_bounds: Vec>, + /// Required interfaces from `requires I1, I2, …`. + pub requires: Vec, + /// Field signatures declared on the interface. Interface fields cannot + /// have default values. + pub fields: Vec, + /// Associated type declarations on the interface (BEP-057). + pub associated_types: Vec, + /// Default methods (with bodies). Implementing classes inherit them. + pub default_methods: Vec>, + /// Required methods (no body). Implementing classes must provide a body. + pub required_methods: Vec, + pub attributes: Vec, + pub docstring: Option, + pub span: TextRange, +} + +/// A required (no-body) method signature on an interface (BEP-044). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceMethodSig { + pub name: Name, + /// Generic type parameters local to this method. + pub generic_params: Vec, + /// BEP-044 generic bounds parallel to `generic_params`. + pub generic_param_bounds: Vec>, + pub params: Vec, + pub return_type: Option, + pub throws: Option, + pub attributes: Vec, + pub docstring: Option, + pub span: TextRange, +} + +/// What an `implements` block applies to. Unifying the owner with the for-target +/// makes "in-body with an explicit for-target" and "out-of-body without one" +/// both unrepresentable. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ImplSubject { + /// `implements I { … }` in a class body (or a simple `implement I for C` + /// merged onto `C`). The for-type is the class itself; the generics are the + /// class's. `out_of_body` records the syntactic origin for diagnostics only + /// — it must NOT influence resolution, dispatch, or coherence. + InClass { + class: LocalItemId, + out_of_body: bool, + }, + /// `implement<…> I for { … }`: an explicit for-type plus the + /// block's own generic parameters. + Free { + for_target: ast::TypeExpr, + generics: Vec, + }, +} + +/// A unified `implements` block (both kinds) stored in the `ItemTree`, keyed by +/// a stable `LocalItemId`. The interface target is kept as a raw +/// `TypeExpr` here; resolution to an `InterfaceLoc` + `Ty` happens lazily in the +/// `impl_data` query so HIR construction stays independent of name resolution. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplBlock { + pub subject: ImplSubject, + pub interface_target: ast::TypeExpr, + pub field_links: Vec, + pub associated_type_bindings: Vec, + pub methods: Vec>, + pub span: TextRange, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplementsBlock { + pub target: ast::TypeExpr, + pub field_links: Vec, + pub associated_type_bindings: Vec, + pub is_out_of_body: bool, + pub span: TextRange, +} + +/// A generic parameter on an out-of-body `implements` block, paired with its set +/// of `&`-separated interface bounds (`` → `bounds = []`; `` +/// → `bounds = [A, B]`). Pairing name and bounds makes a length mismatch between +/// the two unrepresentable. Bounds are unresolved `TypeExpr`s here; they resolve +/// to `baml_type::Interface` constraints downstream. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GenericParam { + pub name: Name, + pub bounds: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceFieldLink { + pub interface_field: Name, + pub class_field: Name, + pub span: TextRange, + pub interface_field_span: TextRange, + pub class_field_span: TextRange, +} + +impl InterfaceFieldLink { + pub(crate) fn from_ast(link: &ast::InterfaceFieldLinkDef) -> Self { + Self { + interface_field: link.interface_field.clone(), + class_field: link.class_field.clone(), + span: link.span, + interface_field_span: link.interface_field_span, + class_field_span: link.class_field_span, + } + } +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/lets.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/lets.rs new file mode 100644 index 00000000000..4bfab9cbbb8 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/lets.rs @@ -0,0 +1,14 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; +use text_size::TextRange; + +/// A top-level let binding stored in the `ItemTree`. +/// Carries the optional initializer `ExprBody` for body queries. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Let { + pub name: Name, + pub initializer: Option<(ast::ExprBody, ast::AstSourceMap)>, + pub origin: ast::LetOrigin, + pub span: TextRange, + pub name_span: TextRange, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/retry_policies.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/retry_policies.rs new file mode 100644 index 00000000000..c40cf546b91 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/retry_policies.rs @@ -0,0 +1,14 @@ +use baml_base::Name; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RetryPolicy { + pub name: Name, + /// Raw string value of `max_retries` (parsed at emit time). + pub max_retries: Option, + /// Raw string value of `initial_delay_ms`. + pub initial_delay_ms: Option, + /// Raw string value of multiplier. + pub multiplier: Option, + /// Raw string value of `max_delay_ms`. + pub max_delay_ms: Option, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/source_map.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/source_map.rs new file mode 100644 index 00000000000..2060efec8d9 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/source_map.rs @@ -0,0 +1,25 @@ +use rustc_hash::FxHashMap; +use text_size::TextRange; + +use crate::ids::{ClassMarker, EnumMarker, FunctionMarker, InterfaceMarker, LocalItemId}; + +/// Parallel source map for `ItemTree` — stores name spans that are +/// deliberately excluded from the semantic `ItemTree` to avoid polluting +/// Salsa's early-cutoff comparisons with position data. +/// +/// Follows the same body/signature source-map pattern used by +/// `function_body` / `function_body_source_map`. +#[derive(Debug, Clone, Default)] +pub struct ItemTreeSourceMap { + /// `name_span` for each class's fields, parallel to `Class::fields`. + pub class_field_spans: FxHashMap, Vec>, + /// `name_span` for each enum's variants, parallel to `Enum::variants`. + pub enum_variant_spans: FxHashMap, Vec>, + /// `name_span` for each function. + pub function_name_spans: FxHashMap, TextRange>, + /// `name_span` for each interface's fields, parallel to `Interface::fields`. + pub interface_field_spans: FxHashMap, Vec>, + /// `name_span` for each interface's required methods, parallel to + /// `Interface::required_methods`. + pub interface_method_spans: FxHashMap, Vec>, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/template_strings.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/template_strings.rs new file mode 100644 index 00000000000..ff97910176a --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/template_strings.rs @@ -0,0 +1,15 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::item_tree::FunctionParam; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TemplateString { + pub name: Name, + /// Template parameters with optional type annotations and spans. + pub params: Vec, + /// Template body text (Jinja template). + pub body: Option, + /// Full source span of the template string declaration. + pub span: TextRange, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/test_items.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/test_items.rs new file mode 100644 index 00000000000..d956b662c10 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/test_items.rs @@ -0,0 +1,11 @@ +use baml_base::Name; +pub use baml_compiler2_ast::ast::TestArgValue; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Test { + pub name: Name, + /// The function(s) this test exercises. + pub function_refs: Vec, + /// Test arguments as key-value pairs. + pub args: Vec<(Name, TestArgValue)>, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/type_aliases.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/type_aliases.rs new file mode 100644 index 00000000000..61ed69f75d0 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/type_aliases.rs @@ -0,0 +1,12 @@ +use baml_base::Name; +use baml_compiler2_ast::ast; +use text_size::TextRange; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TypeAlias { + pub name: Name, + /// The type expression on the RHS of the alias, if present. + pub type_expr: Option, + /// Full source span of the type alias declaration. + pub span: TextRange, +} From e1282e1010ba70c1287b57cffac2416bcd5562f1 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 15:12:04 -0700 Subject: [PATCH 02/20] Span-less HIR `TypeRef` The AST's `TypeExpr` holds spans which clobbers the salsa cache for whitespace changes. Based on rust-analyzer, `TypeRef` is the HIR's version, getting an item-scoped index that should be much more stable. The source map for the type refs is then a separate query. In this commit, `TypeRef` is currently unused. --- .../src/item_tree/common.rs | 4 +- .../crates/baml_compiler2_hir/src/lib.rs | 1 + .../crates/baml_compiler2_hir/src/type_ref.rs | 447 ++++++++++++++++++ 3 files changed, 450 insertions(+), 2 deletions(-) create mode 100644 baml_language/crates/baml_compiler2_hir/src/type_ref.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs index b9003ce80bb..b5ed47b593e 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/common.rs @@ -3,14 +3,14 @@ use baml_compiler2_ast::ast; /// A span-free attribute for position-independent storage in the `ItemTree`. /// Derived from `ast::RawAttribute` with all `TextRange`s stripped. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Attribute { pub args: Vec, pub name: Name, } /// A span-free attribute argument. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct AttributeArg { pub key: Option, pub value: String, diff --git a/baml_language/crates/baml_compiler2_hir/src/lib.rs b/baml_language/crates/baml_compiler2_hir/src/lib.rs index daed0d2839a..5d2f788e50a 100644 --- a/baml_language/crates/baml_compiler2_hir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_hir/src/lib.rs @@ -25,6 +25,7 @@ pub mod package; pub mod scope; pub mod semantic_index; pub mod signature; +pub mod type_ref; use std::sync::Arc; diff --git a/baml_language/crates/baml_compiler2_hir/src/type_ref.rs b/baml_language/crates/baml_compiler2_hir/src/type_ref.rs new file mode 100644 index 00000000000..87827b72245 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/type_ref.rs @@ -0,0 +1,447 @@ +//! Span-free type references — the HIR's own type-expression representation. +//! +//! `ast::TypeExpr` carries a `TextRange` on every node. That is correct for the +//! AST (spans are the syntax layer's job) but fatal for anything the HIR stores: +//! Salsa only overwrites a memoized value when it compares unequal, so a value +//! whose `PartialEq` ignores spans keeps the *old* spans forever after a +//! whitespace-only edit. A value cannot simultaneously be span-carrying, +//! compared semantically, and span-fresh. +//! +//! So the HIR stores `TypeRef` instead: the same tree, flattened into a +//! per-owner arena, with no spans anywhere. Spans live in a parallel +//! [`TypeRefSourceMap`] keyed by [`TypeRefId`], and are only reachable through +//! it. Attributes are the span-free [`Attribute`], not `ast::RawAttribute` +//! (whose span *is* part of its `PartialEq`, and which therefore leaks position +//! into type identity). +//! +//! The arena is scoped to **one item**. A file-wide arena would renumber every +//! later item's ids whenever an item was added, which would defeat the per-item +//! early cutoff this representation exists to enable. + +use std::ops::Index; + +use baml_base::{Literal, MediaKind, Name}; +use la_arena::{Arena, Idx}; +use text_size::TextRange; + +use crate::item_tree::Attribute; + +/// Identity of a type-reference node within its owning item's [`TypeRefStore`]. +pub type TypeRefId = Idx; + +/// A type reference before name resolution — structurally identical to +/// `ast::TypeExprKind`, but with children addressed by [`TypeRefId`] and with +/// every span removed. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct TypeRef { + pub kind: TypeRefKind, + /// Type-level attributes. Uniform across every kind, so — unlike + /// `ast::TypeExprKind` — it is stored once here rather than repeated in + /// all 24 variants. + pub attrs: Box<[Attribute]>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum TypeRefKind { + /// Named type path: `User`, `baml.http.Request`, `Stream`. + Path { + segments: Vec, + /// Generic type arguments (`` in `Stream`). Empty for non-generic paths. + generic_args: Box<[TypeRefId]>, + /// Named associated type bindings, e.g. `Iterator`. + associated_type_bindings: Box<[AssociatedTypeBindingRef]>, + }, + /// Associated type projection: `Base.Item` or `(Base as Interface).Item`. + AssociatedTypeProjection { + base: TypeRefId, + interface: Option, + member: Name, + }, + Int, + Bigint, + Float, + String, + Bool, + Null, + Never, + /// The `void` type — valid only as a function return type. + Void, + /// `Uint8Array` (binary data). + Uint8Array, + Media { + kind: MediaKind, + }, + /// `T?` + Optional { + inner: TypeRefId, + }, + /// `T[]` + List { + inner: TypeRefId, + }, + /// `map` + Map { + key: TypeRefId, + value: TypeRefId, + }, + /// `A | B | C` + Union { + variants: Box<[TypeRefId]>, + }, + /// Literal types in unions: `"user"`, `200`, `3.14`, `true`. + Literal { + value: Literal, + }, + /// Function type: `(params) -> return throws E`. + Function { + params: Box<[FunctionTypeParamRef]>, + ret: TypeRefId, + throws: Option, + }, + /// The `unknown` keyword type. + BuiltinUnknown, + /// The `type` meta-type keyword. + Type, + /// `$rust_type` — opaque Rust-managed state field type. + Rust, + /// Error-recovery sentinel. + Error, + /// Missing/omitted type. + Unknown, + /// The wildcard `_` — an inference hole. + Infer, +} + +/// A parameter of a function type. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct FunctionTypeParamRef { + pub name: Option, + pub optional: bool, + pub ty: TypeRefId, +} + +/// A named associated-type binding inside a type application. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AssociatedTypeBindingRef { + pub name: Name, + pub ty: TypeRefId, +} + +/// The span-free type-reference arena owned by a single item. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TypeRefStore { + types: Arena, +} + +impl TypeRefStore { + pub fn get(&self, id: TypeRefId) -> &TypeRef { + &self.types[id] + } + + pub fn iter(&self) -> impl Iterator { + self.types.iter() + } + + pub fn len(&self) -> usize { + self.types.len() + } + + pub fn is_empty(&self) -> bool { + self.types.is_empty() + } +} + +impl Index for TypeRefStore { + type Output = TypeRef; + fn index(&self, id: TypeRefId) -> &TypeRef { + &self.types[id] + } +} + +/// Spans for one item's [`TypeRefStore`], allocated in lockstep with it, so the +/// two arenas share indices. Deliberately kept out of `TypeRefStore` — see the +/// module docs. +/// +/// Unlike the store, this *does* compare spans: it is meant to invalidate on a +/// whitespace edit, precisely so the store doesn't have to. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TypeRefSourceMap { + spans: Arena, +} + +impl TypeRefSourceMap { + /// The source span of `id`. + /// + /// Synthesized nodes (which have no source text) are allocated with an empty + /// range; callers that need a user-visible anchor should fall back to the + /// owning item's span. + pub fn span(&self, id: TypeRefId) -> TextRange { + // The two arenas are filled in lockstep by `TypeRefBuilder::alloc`, so a + // `TypeRefId` indexes both. O(1) — do not reintroduce a scan here. + self.spans[Idx::from_raw(id.into_raw())] + } +} + +/// Lowers `ast::TypeExpr` trees into one item's [`TypeRefStore`], recording each +/// node's span in the parallel [`TypeRefSourceMap`]. +/// +/// One builder per item: ids are only meaningful relative to the store they were +/// allocated in. +#[derive(Debug, Default)] +pub struct TypeRefBuilder { + store: TypeRefStore, + source_map: TypeRefSourceMap, +} + +impl TypeRefBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Allocate a node and its span together. The lockstep is the invariant that + /// makes `TypeRefSourceMap::span` a plain index rather than a lookup. + fn alloc(&mut self, ty: TypeRef, span: TextRange) -> TypeRefId { + let id = self.store.types.alloc(ty); + let span_id = self.source_map.spans.alloc(span); + debug_assert_eq!( + id.into_raw(), + span_id.into_raw(), + "TypeRef and span arenas must stay in lockstep" + ); + id + } + + /// Lower a type expression and all of its children. Children are allocated + /// before their parent (post-order), so ids are a pure function of tree + /// shape — stable under whitespace edits. + pub fn lower(&mut self, te: &baml_compiler2_ast::ast::TypeExpr) -> TypeRefId { + use baml_compiler2_ast::ast::TypeExprKind as K; + + let attrs: Box<[Attribute]> = te.kind.attrs().iter().map(Attribute::from).collect(); + + let kind = match &te.kind { + K::Path { + segments, + generic_args, + associated_type_bindings, + .. + } => TypeRefKind::Path { + segments: segments.clone(), + generic_args: generic_args.iter().map(|arg| self.lower(arg)).collect(), + associated_type_bindings: associated_type_bindings + .iter() + .map(|binding| AssociatedTypeBindingRef { + name: binding.name.clone(), + ty: self.lower(&binding.ty), + }) + .collect(), + }, + K::AssociatedTypeProjection { + base, + interface, + member, + .. + } => TypeRefKind::AssociatedTypeProjection { + base: self.lower(base), + interface: interface.as_ref().map(|iface| self.lower(iface)), + member: member.clone(), + }, + K::Int { .. } => TypeRefKind::Int, + K::Bigint { .. } => TypeRefKind::Bigint, + K::Float { .. } => TypeRefKind::Float, + K::String { .. } => TypeRefKind::String, + K::Bool { .. } => TypeRefKind::Bool, + K::Null { .. } => TypeRefKind::Null, + K::Never { .. } => TypeRefKind::Never, + K::Void { .. } => TypeRefKind::Void, + K::Uint8Array { .. } => TypeRefKind::Uint8Array, + K::Media { kind, .. } => TypeRefKind::Media { kind: *kind }, + K::Optional { inner, .. } => TypeRefKind::Optional { + inner: self.lower(inner), + }, + K::List { inner, .. } => TypeRefKind::List { + inner: self.lower(inner), + }, + K::Map { key, value, .. } => TypeRefKind::Map { + key: self.lower(key), + value: self.lower(value), + }, + K::Union { variants, .. } => TypeRefKind::Union { + variants: variants.iter().map(|v| self.lower(v)).collect(), + }, + K::Literal { value, .. } => TypeRefKind::Literal { + value: value.clone(), + }, + K::Function { + params, + ret, + throws, + .. + } => TypeRefKind::Function { + params: params + .iter() + .map(|p| FunctionTypeParamRef { + name: p.name.clone(), + optional: p.optional, + ty: self.lower(&p.ty), + }) + .collect(), + ret: self.lower(ret), + throws: throws.as_ref().map(|t| self.lower(t)), + }, + K::BuiltinUnknown { .. } => TypeRefKind::BuiltinUnknown, + K::Type { .. } => TypeRefKind::Type, + K::Rust { .. } => TypeRefKind::Rust, + K::Error { .. } => TypeRefKind::Error, + K::Unknown { .. } => TypeRefKind::Unknown, + K::Infer { .. } => TypeRefKind::Infer, + }; + + self.alloc(TypeRef { kind, attrs }, te.span) + } + + /// Allocate a node that has no source text (a compiler-synthesized type). + /// Its span is empty; diagnostics should anchor to the owning item instead. + pub fn alloc_synthetic(&mut self, kind: TypeRefKind) -> TypeRefId { + self.alloc( + TypeRef { + kind, + attrs: Box::new([]), + }, + TextRange::default(), + ) + } + + pub fn finish(self) -> (TypeRefStore, TypeRefSourceMap) { + (self.store, self.source_map) + } +} + +#[cfg(test)] +mod tests { + use baml_compiler2_ast::ast::{RawAttribute, RawAttributeArg, TypeExpr, TypeExprKind}; + use text_size::TextSize; + + use super::*; + + fn span(start: u32, end: u32) -> TextRange { + TextRange::new(TextSize::from(start), TextSize::from(end)) + } + + fn attr(name: &str, value: &str, at: TextRange) -> RawAttribute { + RawAttribute { + name: Name::new(name), + args: vec![RawAttributeArg { + key: None, + value: value.to_string(), + span: at, + }], + span: at, + } + } + + /// `map`, with every node's span offset by `shift`. + fn map_of_string_to_int_list(shift: u32, attrs: Vec) -> TypeExpr { + let s = |a: u32, b: u32| span(a + shift, b + shift); + TypeExprKind::Map { + key: Box::new(TypeExprKind::String { attrs: vec![] }.at(s(4, 10))), + value: Box::new( + TypeExprKind::List { + inner: Box::new(TypeExprKind::Int { attrs: vec![] }.at(s(12, 15))), + attrs: vec![], + } + .at(s(12, 17)), + ), + attrs, + } + .at(s(0, 18)) + } + + fn lower(te: &TypeExpr) -> (TypeRefStore, TypeRefSourceMap, TypeRefId) { + let mut builder = TypeRefBuilder::new(); + let root = builder.lower(te); + let (store, source_map) = builder.finish(); + (store, source_map, root) + } + + #[test] + fn children_are_allocated_before_parents() { + let (store, _, root) = lower(&map_of_string_to_int_list(0, vec![])); + + // Post-order: the root is allocated last. + assert_eq!(root, store.iter().last().expect("non-empty store").0); + + // string, int, int[], map — post-order. + assert_eq!(store.len(), 4); + let kinds: Vec<_> = store.iter().map(|(_, t)| t.kind.clone()).collect(); + assert!(matches!(kinds[0], TypeRefKind::String)); + assert!(matches!(kinds[1], TypeRefKind::Int)); + assert!(matches!(kinds[2], TypeRefKind::List { .. })); + assert!(matches!(kinds[3], TypeRefKind::Map { .. })); + } + + /// The property the whole firewall rests on: shifting every span (a + /// whitespace-only edit) must leave the semantic value byte-for-byte equal, + /// so Salsa cuts off — while the source map still reports the *new* spans. + #[test] + fn shifting_spans_does_not_change_the_store() { + let (unshifted, map_a, root_a) = lower(&map_of_string_to_int_list(0, vec![])); + let (shifted, map_b, root_b) = lower(&map_of_string_to_int_list(100, vec![])); + + assert_eq!( + unshifted, shifted, + "a whitespace-only edit must not change the semantic type refs" + ); + + assert_eq!(map_a.span(root_a), span(0, 18)); + assert_eq!( + map_b.span(root_b), + span(100, 118), + "the source map must still track the new positions" + ); + } + + /// `ast::RawAttribute` puts its span in its own `PartialEq`, so attribute + /// spans leak into `TypeExpr` equality and silently destroy cutoff near any + /// `@description`. `TypeRef` carries the span-free `Attribute` instead. + #[test] + fn shifting_attribute_spans_does_not_change_the_store() { + let a = map_of_string_to_int_list(0, vec![attr("description", "hi", span(19, 40))]); + let b = map_of_string_to_int_list(0, vec![attr("description", "hi", span(219, 240))]); + + assert_ne!( + a, b, + "precondition: the AST does compare attribute spans (this is the leak)" + ); + + let (store_a, _, _) = lower(&a); + let (store_b, _, _) = lower(&b); + assert_eq!( + store_a, store_b, + "moving an attribute must not change the semantic type refs" + ); + } + + /// A real structural change must still compare unequal, or we would cut off + /// edits that actually matter. + #[test] + fn structural_changes_do_change_the_store() { + let (int_keyed, _, _) = lower(&map_of_string_to_int_list(0, vec![])); + let (bool_keyed, _, _) = lower( + &TypeExprKind::Map { + key: Box::new(TypeExprKind::Bool { attrs: vec![] }.at(span(4, 10))), + value: Box::new( + TypeExprKind::List { + inner: Box::new(TypeExprKind::Int { attrs: vec![] }.at(span(12, 15))), + attrs: vec![], + } + .at(span(12, 17)), + ), + attrs: vec![], + } + .at(span(0, 18)), + ); + + assert_ne!(int_keyed, bool_keyed); + } +} From 21c34d80acf1818a9f03eaca3422d22ae502084f Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 15:38:25 -0700 Subject: [PATCH 03/20] HIR type alias query firewall --- .../baml_compiler2_hir/src/item_data.rs | 31 +++++++ .../src/item_data/type_aliases.rs | 69 ++++++++++++++ .../crates/baml_compiler2_hir/src/lib.rs | 1 + .../baml_tests/src/incremental/scenarios.rs | 93 +++++++++++++++++++ 4 files changed, 194 insertions(+) create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data.rs b/baml_language/crates/baml_compiler2_hir/src/item_data.rs new file mode 100644 index 00000000000..8ac985dc37b --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data.rs @@ -0,0 +1,31 @@ +//! Firewall queries — the public way to read one item. +//! +//! `file_semantic_index` is `no_eq`: it re-runs on every edit to the file and +//! always reports "changed". Anything that reads the `ItemTree` through it +//! therefore re-runs on every keystroke, however unrelated the edit. These +//! queries are the firewall: they re-run cheaply, but their *results* are +//! compared with `PartialEq`, so a downstream query only re-runs when the item +//! it actually depends on changed. +//! +//! That only works if the result is span-free. Salsa keeps the old memoized +//! value whenever the new one compares equal (`salsa::update` — "this may cause +//! us not to update even if the value has changed"), so a result carrying spans +//! would hand out *stale* spans forever after a whitespace-only edit. Hence +//! every item is split in two: +//! +//! - `*_data` — semantic, span-free, uses [`TypeRef`](crate::type_ref::TypeRef). +//! This is what type checking reads. +//! - `*_source_map` — spans only. This is what diagnostics and the IDE read. +//! +//! The `ItemTree` itself may keep its spans: it lives behind the `no_eq` index, +//! which is overwritten wholesale on every revision, so its spans are always +//! fresh. It is memoized *results* that must be span-free. +//! +//! Each query takes a `*Loc` and nothing else — the file is inside the `Loc`, so +//! callers never thread a `SourceFile` through to reach an item. +//! +//! One submodule per item kind, mirroring `item_tree`. + +mod type_aliases; + +pub use type_aliases::*; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs new file mode 100644 index 00000000000..e425a3ff863 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs @@ -0,0 +1,69 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{ + loc::TypeAliasLoc, + type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, +}; + +/// Span-free semantic data for a `type X = ...` declaration. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TypeAliasData { + pub name: Name, + /// The type-reference arena owned by this alias. Scoped to the item, so + /// edits to sibling items cannot renumber these ids. + pub type_refs: TypeRefStore, + /// Root of the aliased type. `None` when the RHS was omitted or unparseable. + pub value: Option, +} + +/// Spans for a `TypeAlias`, parallel to [`TypeAliasData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TypeAliasSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Spans for every node in [`TypeAliasData::type_refs`]. + pub type_refs: TypeRefSourceMap, +} + +/// Semantic data for one type alias. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn type_alias_data<'db>(db: &'db dyn crate::Db, alias: TypeAliasLoc<'db>) -> TypeAliasData { + lower(db, alias).0 +} + +/// Spans for one type alias. Kept separate from [`type_alias_data`] so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn type_alias_source_map<'db>( + db: &'db dyn crate::Db, + alias: TypeAliasLoc<'db>, +) -> TypeAliasSourceMap { + lower(db, alias).1 +} + +/// Both halves share one lowering pass — the split into two queries is purely +/// about what each one lets downstream depend on. +fn lower<'db>( + db: &'db dyn crate::Db, + alias: TypeAliasLoc<'db>, +) -> (TypeAliasData, TypeAliasSourceMap) { + let item_tree = crate::file_item_tree(db, alias.file(db)); + let data = &item_tree[alias.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + let value = data.type_expr.as_ref().map(|te| type_refs.lower(te)); + let (store, spans) = type_refs.finish(); + + ( + TypeAliasData { + name: data.name.clone(), + type_refs: store, + value, + }, + TypeAliasSourceMap { + span: data.span, + type_refs: spans, + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/lib.rs b/baml_language/crates/baml_compiler2_hir/src/lib.rs index 5d2f788e50a..cbaaafcfc36 100644 --- a/baml_language/crates/baml_compiler2_hir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_hir/src/lib.rs @@ -18,6 +18,7 @@ pub mod contributions; pub mod diagnostic; pub mod file_package; pub mod ids; +pub mod item_data; pub mod item_tree; pub mod loc; pub mod namespace; diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index bc7a0234e5a..9a39224685c 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -343,3 +343,96 @@ function Greet(name: string) -> int { // Signature changes must invalidate queries test_db.assert_executed(|db| query_semantic_index(db, file), &[("lex_file", 1)]); } + +// ── Firewall queries ───────────────────────────────────────────────────────── +// +// `file_semantic_index` is `no_eq`, so it always reports "changed" and anything +// reading the `ItemTree` through it re-runs on every keystroke. The per-item +// firewall queries are what stop that from propagating: they re-run, but their +// results only *compare* unequal when the item genuinely changed, so Salsa cuts +// off there. +// +// That only holds because the semantic half is span-free. Salsa keeps the old +// memoized value whenever the new one compares equal, so if `*_data` carried +// spans it would either lose cutoff (spans in `PartialEq`) or hand out stale +// ones (spans ignored by `PartialEq`). These tests pin both halves of that. + +fn type_alias_loc<'db>( + db: &'db baml_project::ProjectDatabase, + file: SourceFile, + name: &str, +) -> baml_compiler2_hir::loc::TypeAliasLoc<'db> { + let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let (id, _) = item_tree + .type_aliases + .iter() + .find(|(_, alias)| alias.name.as_str() == name) + .unwrap_or_else(|| unreachable!("type alias `{name}` should exist")); + baml_compiler2_hir::loc::TypeAliasLoc::new(db, file, *id) +} + +/// A whitespace-only edit must leave the semantic data byte-for-byte equal (so +/// Salsa cuts off) while the source map still reports the *new* positions. +#[test] +fn whitespace_edit_preserves_item_data_but_moves_spans() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file("test.baml", "type Ids = int[]\n"); + + let (data_before, span_before) = { + let db = test_db.db(); + let loc = type_alias_loc(db, file, "Ids"); + ( + baml_compiler2_hir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_hir::item_data::type_alias_source_map(db, loc).clone(), + ) + }; + + // Push the declaration down a line. Semantics identical, every span shifted. + file.set_text(test_db.db_mut()) + .to("// a comment\ntype Ids = int[]\n".to_string()); + + let (data_after, span_after) = { + let db = test_db.db(); + let loc = type_alias_loc(db, file, "Ids"); + ( + baml_compiler2_hir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_hir::item_data::type_alias_source_map(db, loc).clone(), + ) + }; + + assert_eq!( + data_before, data_after, + "a whitespace-only edit must not change the semantic data, or nothing downstream can cut off" + ); + assert_ne!( + span_before, span_after, + "the source map must track the new positions — otherwise spans would go stale" + ); +} + +/// The converse: a real change to the aliased type must invalidate the semantic +/// data, or we would be cutting off edits that actually matter. +#[test] +fn semantic_edit_changes_item_data() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file("test.baml", "type Ids = int[]\n"); + + let data_before = { + let db = test_db.db(); + let loc = type_alias_loc(db, file, "Ids"); + baml_compiler2_hir::item_data::type_alias_data(db, loc).clone() + }; + + file.set_text(test_db.db_mut()) + .to("type Ids = string[]\n".to_string()); + + let data_after = { + let db = test_db.db(); + let loc = type_alias_loc(db, file, "Ids"); + baml_compiler2_hir::item_data::type_alias_data(db, loc).clone() + }; + + assert_ne!(data_before, data_after); +} From cb59f58edbda1381891090e35bcb29c77ba7c7d6 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 15:52:26 -0700 Subject: [PATCH 04/20] HIR class + function query firewalls --- .../baml_compiler2_hir/src/item_data.rs | 4 + .../src/item_data/classes.rs | 217 ++++++++++++++++++ .../src/item_data/functions.rs | 137 +++++++++++ .../baml_tests/src/incremental/scenarios.rs | 146 ++++++++++++ 4 files changed, 504 insertions(+) create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data.rs b/baml_language/crates/baml_compiler2_hir/src/item_data.rs index 8ac985dc37b..7106d418a3f 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_data.rs @@ -26,6 +26,10 @@ //! //! One submodule per item kind, mirroring `item_tree`. +mod classes; +mod functions; mod type_aliases; +pub use classes::*; +pub use functions::*; pub use type_aliases::*; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs new file mode 100644 index 00000000000..1f060a7d773 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs @@ -0,0 +1,217 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{ + item_tree::Attribute, + loc::{ClassLoc, FunctionLoc}, + type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, +}; + +/// Span-free semantic data for a `class` declaration. +/// +/// `methods` holds `FunctionLoc`s, not `LocalItemId`s: a `Loc` already carries +/// its file, so a consumer can go straight to `function_data` without knowing +/// which file the class came from. +#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)] +pub struct ClassData<'db> { + pub name: Name, + pub generic_params: Vec, + /// Every type reference in this class's *signature* — bounds, field types, + /// and `implements` targets. Scoped to the item, so edits to sibling items + /// cannot renumber these ids. + pub type_refs: TypeRefStore, + /// Parallel to `generic_params`. `Some` means `T extends `. + pub generic_param_bounds: Vec>, + pub fields: Vec, + pub methods: Vec>, + pub implements: Vec, + pub attributes: Vec, + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClassFieldData { + pub name: Name, + pub type_ref: Option, + pub attributes: Vec, + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplementsData { + pub target: TypeRefId, + pub field_links: Vec, + pub associated_type_bindings: Vec, + /// Syntactic origin, for diagnostics only — it must NOT influence + /// resolution, dispatch, or coherence. + pub is_out_of_body: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceFieldLinkData { + pub interface_field: Name, + pub class_field: Name, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeBindingData { + pub name: Name, + pub type_ref: Option, +} + +/// Spans for a `Class`, parallel to [`ClassData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClassSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Spans for every node in [`ClassData::type_refs`]. + pub type_refs: TypeRefSourceMap, + /// Name span per field, parallel to [`ClassData::fields`]. + pub field_name_spans: Vec, + /// Parallel to [`ClassData::implements`]. + pub implements: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplementsSourceMap { + pub span: TextRange, + pub field_links: Vec, + pub associated_type_bindings: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceFieldLinkSourceMap { + pub span: TextRange, + pub interface_field_span: TextRange, + pub class_field_span: TextRange, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeBindingSourceMap { + pub span: TextRange, + pub name_span: TextRange, +} + +/// Semantic data for one class. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn class_data<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> ClassData<'db> { + lower(db, class).0 +} + +/// Spans for one class. Kept separate from [`class_data`] so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn class_source_map<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> ClassSourceMap { + lower(db, class).1 +} + +/// Both halves share one lowering pass — the split into two queries is purely +/// about what each one lets downstream depend on. +/// +/// Type refs are allocated in a fixed order (bounds, then fields, then +/// `implements` targets) so that ids are a pure function of the class's shape. +fn lower<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> (ClassData<'db>, ClassSourceMap) { + let file = class.file(db); + let item_tree = crate::file_item_tree(db, file); + let item_source_map = crate::file_item_tree_source_map(db, file); + let data = &item_tree[class.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + + let generic_param_bounds = data + .generic_param_bounds + .iter() + .map(|bound| bound.as_ref().map(|te| type_refs.lower(te))) + .collect(); + + let fields = data + .fields + .iter() + .map(|field| ClassFieldData { + name: field.name.clone(), + type_ref: field.type_expr.as_ref().map(|te| type_refs.lower(te)), + attributes: field.attributes.clone(), + docstring: field.docstring.clone(), + }) + .collect(); + + let implements: Vec = data + .implements + .iter() + .map(|block| ImplementsData { + target: type_refs.lower(&block.target), + field_links: block + .field_links + .iter() + .map(|link| InterfaceFieldLinkData { + interface_field: link.interface_field.clone(), + class_field: link.class_field.clone(), + }) + .collect(), + associated_type_bindings: block + .associated_type_bindings + .iter() + .map(|binding| AssociatedTypeBindingData { + name: binding.name.clone(), + type_ref: binding.type_expr.as_ref().map(|te| type_refs.lower(te)), + }) + .collect(), + is_out_of_body: block.is_out_of_body, + }) + .collect(); + + let (store, spans) = type_refs.finish(); + + let implements_spans = data + .implements + .iter() + .map(|block| ImplementsSourceMap { + span: block.span, + field_links: block + .field_links + .iter() + .map(|link| InterfaceFieldLinkSourceMap { + span: link.span, + interface_field_span: link.interface_field_span, + class_field_span: link.class_field_span, + }) + .collect(), + associated_type_bindings: block + .associated_type_bindings + .iter() + .map(|binding| AssociatedTypeBindingSourceMap { + span: binding.span, + name_span: binding.name_span, + }) + .collect(), + }) + .collect(); + + ( + ClassData { + name: data.name.clone(), + generic_params: data.generic_params.clone(), + type_refs: store, + generic_param_bounds, + fields, + methods: data + .methods + .iter() + .map(|method| FunctionLoc::new(db, file, *method)) + .collect(), + implements, + attributes: data.attributes.clone(), + docstring: data.docstring.clone(), + }, + ClassSourceMap { + span: data.span, + type_refs: spans, + field_name_spans: item_source_map + .class_field_spans + .get(&class.id(db)) + .cloned() + .unwrap_or_default(), + implements: implements_spans, + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs new file mode 100644 index 00000000000..9a5d502319c --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs @@ -0,0 +1,137 @@ +use baml_base::Name; +use baml_compiler2_ast as ast; +use text_size::TextRange; + +use crate::{ + loc::FunctionLoc, + type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, +}; + +/// Span-free semantic data for a function's *signature*. +/// +/// The body is deliberately absent: it has its own query (`body::function_body`) +/// so that editing a body cannot invalidate anything that only depends on the +/// signature. +/// +/// This supersedes `signature::FunctionSignature`, which is memoized but still +/// carries spanned `ast::TypeExpr`s — meaning Salsa declines to overwrite it on a +/// whitespace-only edit and it serves *stale* spans from then on. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionData { + pub name: Name, + pub generic_params: Vec, + /// Every type reference in this function's signature — bounds, parameter + /// types, return type, throws. Scoped to the item, so edits to sibling items + /// cannot renumber these ids. + pub type_refs: TypeRefStore, + /// Parallel to `generic_params`. `Some` means `T extends `. + pub generic_param_bounds: Vec>, + pub params: Vec, + pub return_type: Option, + pub throws: Option, + pub origin: ast::FunctionOrigin, + pub docstring: Option, + /// Set when the fn def had a `//baml:tagged_string` marker. + pub is_tagged_template_tag: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionParamData { + pub name: Name, + pub type_ref: Option, + /// Whether a default expression was supplied. The expression itself lives in + /// `signature::function_parameter_defaults`. + pub has_default: bool, +} + +/// Spans for a `Function`, parallel to [`FunctionData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Span of just the function's name. + pub name_span: TextRange, + /// Spans for every node in [`FunctionData::type_refs`]. + pub type_refs: TypeRefSourceMap, + /// One span per parameter, parallel to [`FunctionData::params`]. + pub param_spans: Vec, +} + +/// Semantic data for one function signature. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn function_data<'db>(db: &'db dyn crate::Db, function: FunctionLoc<'db>) -> FunctionData { + lower(db, function).0 +} + +/// Spans for one function signature. Kept separate from [`function_data`] so that +/// a whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn function_source_map<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> FunctionSourceMap { + lower(db, function).1 +} + +/// Both halves share one lowering pass — the split into two queries is purely +/// about what each one lets downstream depend on. +/// +/// Type refs are allocated in a fixed order (bounds, params, return, throws) so +/// that ids are a pure function of the signature's shape. +fn lower<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> (FunctionData, FunctionSourceMap) { + let file = function.file(db); + let item_tree = crate::file_item_tree(db, file); + let item_source_map = crate::file_item_tree_source_map(db, file); + let data = &item_tree[function.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + + let generic_param_bounds = data + .generic_param_bounds + .iter() + .map(|bound| bound.as_ref().map(|te| type_refs.lower(te))) + .collect(); + + let params: Vec = data + .params + .iter() + .map(|param| FunctionParamData { + name: param.name.clone(), + type_ref: param.type_expr.as_ref().map(|te| type_refs.lower(te)), + has_default: param.default.is_some(), + }) + .collect(); + + let return_type = data.return_type.as_ref().map(|te| type_refs.lower(te)); + let throws = data.throws.as_ref().map(|te| type_refs.lower(te)); + + let (store, spans) = type_refs.finish(); + + ( + FunctionData { + name: data.name.clone(), + generic_params: data.generic_params.clone(), + type_refs: store, + generic_param_bounds, + params, + return_type, + throws, + origin: data.origin, + docstring: data.docstring.clone(), + is_tagged_template_tag: data.is_tagged_template_tag, + }, + FunctionSourceMap { + span: data.span, + name_span: item_source_map + .function_name_spans + .get(&function.id(db)) + .copied() + .unwrap_or_default(), + type_refs: spans, + param_spans: data.params.iter().map(|param| param.span).collect(), + }, + ) +} diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index 9a39224685c..8958befd873 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -436,3 +436,149 @@ fn semantic_edit_changes_item_data() { assert_ne!(data_before, data_after); } + +fn class_loc<'db>( + db: &'db baml_project::ProjectDatabase, + file: SourceFile, + name: &str, +) -> baml_compiler2_hir::loc::ClassLoc<'db> { + let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let (id, _) = item_tree + .classes + .iter() + .find(|(_, class)| class.name.as_str() == name) + .unwrap_or_else(|| unreachable!("class `{name}` should exist")); + baml_compiler2_hir::loc::ClassLoc::new(db, file, *id) +} + +fn function_loc<'db>( + db: &'db baml_project::ProjectDatabase, + file: SourceFile, + name: &str, +) -> baml_compiler2_hir::loc::FunctionLoc<'db> { + let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let (id, _) = item_tree + .functions + .iter() + .find(|(_, function)| function.name.as_str() == name) + .unwrap_or_else(|| unreachable!("function `{name}` should exist")); + baml_compiler2_hir::loc::FunctionLoc::new(db, file, *id) +} + +/// Everything span-bearing in a `ClassData`, as an owned value. +/// +/// `ClassData<'db>` holds `FunctionLoc<'db>`s, so even a clone keeps the `db` +/// borrow alive and cannot be held across an edit. `methods` is pure identity +/// and carries no spans, so projecting it away loses nothing these tests check. +type ClassFingerprint = ( + baml_base::Name, + Vec>, + baml_compiler2_hir::type_ref::TypeRefStore, + Vec, + Vec, + Vec, +); + +fn class_fingerprint( + db: &baml_project::ProjectDatabase, + file: SourceFile, + name: &str, +) -> ClassFingerprint { + let data = baml_compiler2_hir::item_data::class_data(db, class_loc(db, file, name)); + ( + data.name.clone(), + data.generic_param_bounds.clone(), + data.type_refs.clone(), + data.fields.clone(), + data.implements.clone(), + data.attributes.clone(), + ) +} + +/// The whole point of a per-item firewall: editing one item must leave every +/// *other* item's data untouched, so nothing downstream of them re-runs. +#[test] +fn editing_one_class_preserves_the_others_data() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "class Person {\n name string\n}\n\nclass Address {\n city string\n}\n", + ); + + let untouched_before = class_fingerprint(test_db.db(), file, "Address"); + + // Add a field to `Person`. `Address` is not touched — but it moves, so every + // span in it shifts. + file.set_text(test_db.db_mut()).to( + "class Person {\n name string\n age int\n}\n\nclass Address {\n city string\n}\n" + .to_string(), + ); + + let untouched_after = class_fingerprint(test_db.db(), file, "Address"); + let touched = class_fingerprint(test_db.db(), file, "Person"); + + assert_eq!( + untouched_before, untouched_after, + "editing `Person` must not invalidate `Address` — that is the firewall" + ); + assert_eq!(touched.3.len(), 2, "`Person` really did change"); +} + +/// `ast::RawAttribute` puts its span in its own `PartialEq`, and every +/// `TypeExprKind` variant holds `attrs`, so today a whitespace edit near any +/// `@description` makes the type compare unequal and silently destroys cutoff. +/// `ClassData` carries the span-free `Attribute` instead. +#[test] +fn moving_an_attribute_preserves_class_data() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "class Person {\n name string @description(\"who\")\n}\n", + ); + + let before = class_fingerprint(test_db.db(), file, "Person"); + + // Push the class down. The attribute's span moves with it. + file.set_text(test_db.db_mut()) + .to("// a comment\nclass Person {\n name string @description(\"who\")\n}\n".to_string()); + + let after = class_fingerprint(test_db.db(), file, "Person"); + + assert_eq!( + before, after, + "moving an attribute must not change the semantic data" + ); +} + +/// A function's signature data must not depend on its body — otherwise every +/// keystroke inside a body invalidates every caller's view of the signature. +#[test] +fn editing_a_function_body_preserves_its_signature_data() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Add(x: int, y: int) -> int {\n x + y\n}\n", + ); + + let before = { + let db = test_db.db(); + baml_compiler2_hir::item_data::function_data(db, function_loc(db, file, "Add")).clone() + }; + + file.set_text(test_db.db_mut()) + .to("function Add(x: int, y: int) -> int {\n y + x + 0\n}\n".to_string()); + + let after = { + let db = test_db.db(); + baml_compiler2_hir::item_data::function_data(db, function_loc(db, file, "Add")).clone() + }; + + assert_eq!( + before, after, + "a body edit must not invalidate the signature" + ); + assert_eq!(before.params.len(), 2); +} From 31d31d0f6b35479e281bbaa4cde36357a886d4b4 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 16:13:11 -0700 Subject: [PATCH 05/20] HIR remaining query firewalls --- .../baml_compiler2_hir/src/item_data.rs | 20 ++ .../src/item_data/classes.rs | 41 +--- .../src/item_data/clients.rs | 28 +++ .../src/item_data/common.rs | 52 +++++ .../src/item_data/enumeration.rs | 158 +++++++++++++ .../baml_compiler2_hir/src/item_data/enums.rs | 76 +++++++ .../src/item_data/functions.rs | 96 +++++++- .../baml_compiler2_hir/src/item_data/impls.rs | 179 +++++++++++++++ .../src/item_data/interfaces.rs | 214 ++++++++++++++++++ .../baml_compiler2_hir/src/item_data/lets.rs | 54 +++++ .../src/item_data/retry_policies.rs | 34 +++ .../src/item_data/template_strings.rs | 83 +++++++ .../src/item_data/test_items.rs | 27 +++ .../baml_tests/src/incremental/scenarios.rs | 2 +- 14 files changed, 1019 insertions(+), 45 deletions(-) create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/common.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data.rs b/baml_language/crates/baml_compiler2_hir/src/item_data.rs index 7106d418a3f..f5c9e32c1f1 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_data.rs @@ -27,9 +27,29 @@ //! One submodule per item kind, mirroring `item_tree`. mod classes; +mod clients; +mod common; +mod enumeration; +mod enums; mod functions; +mod impls; +mod interfaces; +mod lets; +mod retry_policies; +mod template_strings; +mod test_items; mod type_aliases; pub use classes::*; +pub use clients::*; +pub use common::*; +pub use enumeration::*; +pub use enums::*; pub use functions::*; +pub use impls::*; +pub use interfaces::*; +pub use lets::*; +pub use retry_policies::*; +pub use template_strings::*; +pub use test_items::*; pub use type_aliases::*; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs index 1f060a7d773..c089823e9be 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs @@ -2,6 +2,10 @@ use baml_base::Name; use text_size::TextRange; use crate::{ + item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FieldData, + InterfaceFieldLinkData, InterfaceFieldLinkSourceMap, + }, item_tree::Attribute, loc::{ClassLoc, FunctionLoc}, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, @@ -22,21 +26,13 @@ pub struct ClassData<'db> { pub type_refs: TypeRefStore, /// Parallel to `generic_params`. `Some` means `T extends `. pub generic_param_bounds: Vec>, - pub fields: Vec, + pub fields: Vec, pub methods: Vec>, pub implements: Vec, pub attributes: Vec, pub docstring: Option, } -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ClassFieldData { - pub name: Name, - pub type_ref: Option, - pub attributes: Vec, - pub docstring: Option, -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct ImplementsData { pub target: TypeRefId, @@ -47,18 +43,6 @@ pub struct ImplementsData { pub is_out_of_body: bool, } -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct InterfaceFieldLinkData { - pub interface_field: Name, - pub class_field: Name, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AssociatedTypeBindingData { - pub name: Name, - pub type_ref: Option, -} - /// Spans for a `Class`, parallel to [`ClassData`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ClassSourceMap { @@ -79,19 +63,6 @@ pub struct ImplementsSourceMap { pub associated_type_bindings: Vec, } -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct InterfaceFieldLinkSourceMap { - pub span: TextRange, - pub interface_field_span: TextRange, - pub class_field_span: TextRange, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AssociatedTypeBindingSourceMap { - pub span: TextRange, - pub name_span: TextRange, -} - /// Semantic data for one class. Span-free — see the module docs. #[salsa::tracked(returns(ref))] pub fn class_data<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> ClassData<'db> { @@ -127,7 +98,7 @@ fn lower<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> (ClassData<'db>, let fields = data .fields .iter() - .map(|field| ClassFieldData { + .map(|field| FieldData { name: field.name.clone(), type_ref: field.type_expr.as_ref().map(|te| type_refs.lower(te)), attributes: field.attributes.clone(), diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs new file mode 100644 index 00000000000..5b6a773c4d3 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs @@ -0,0 +1,28 @@ +use baml_base::Name; + +use crate::loc::ClientLoc; + +/// Semantic data for a `client` declaration. +/// +/// Clients carry neither type expressions nor spans, so there is no +/// `TypeRefStore` and no source-map twin — the data query alone is the firewall. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientData { + pub name: Name, + /// Provider name (e.g. "openai", "anthropic", "fallback", "round-robin"). + pub provider: Option, + /// Retry policy name, if configured. + pub retry_policy_name: Option, +} + +#[salsa::tracked(returns(ref))] +pub fn client_data<'db>(db: &'db dyn crate::Db, client: ClientLoc<'db>) -> ClientData { + let item_tree = crate::file_item_tree(db, client.file(db)); + let data = &item_tree[client.id(db)]; + + ClientData { + name: data.name.clone(), + provider: data.provider.clone(), + retry_policy_name: data.retry_policy_name.clone(), + } +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/common.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/common.rs new file mode 100644 index 00000000000..e6cd6370bc8 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/common.rs @@ -0,0 +1,52 @@ +//! Item-data pieces shared by more than one item kind. +//! +//! Mirrors `item_tree::common`. Interfaces reuse the class field shape and the +//! function parameter shape, exactly as they do in the `ItemTree`, so those live +//! here rather than in `classes`/`functions`. + +use baml_base::Name; +use text_size::TextRange; + +use crate::{item_tree::Attribute, type_ref::TypeRefId}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionParamData { + pub name: Name, + pub type_ref: Option, + /// Whether a default expression was supplied. The expression itself lives in + /// `signature::function_parameter_defaults`. + pub has_default: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FieldData { + pub name: Name, + pub type_ref: Option, + pub attributes: Vec, + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeBindingData { + pub name: Name, + pub type_ref: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeBindingSourceMap { + pub span: TextRange, + pub name_span: TextRange, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceFieldLinkData { + pub interface_field: Name, + pub class_field: Name, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceFieldLinkSourceMap { + pub span: TextRange, + pub interface_field_span: TextRange, + pub class_field_span: TextRange, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs new file mode 100644 index 00000000000..3a411cc6c01 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs @@ -0,0 +1,158 @@ +//! Enumeration queries — "which items does this file declare?" +//! +//! The other half of the firewall. Consumers get a list of `*Loc`s and then look +//! each one up with a `*_data` query; they never touch the `ItemTree`'s maps, and +//! they never mint a `Loc` themselves. +//! +//! Each list is sorted by the item's source position, so iteration order is +//! deterministic (the `ItemTree`'s maps are `FxHashMap`s and are not). Sorting by +//! position is safe *here* — the result is a set of `Loc`s, which are +//! position-independent, so moving an item around in the file does not change +//! this query's value unless it actually reorders declarations. + +use baml_base::SourceFile; + +use crate::loc::{ + ClassLoc, ClientLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, + TemplateStringLoc, TestLoc, TypeAliasLoc, +}; + +/// Declare an enumeration query: `file_(file) -> Vec`, ordered by +/// source position. +macro_rules! file_items { + ($(#[$meta:meta])* $name:ident, $map:ident, $loc:ident) => { + $(#[$meta])* + #[salsa::tracked(returns(ref))] + pub fn $name(db: &dyn crate::Db, file: SourceFile) -> Vec<$loc<'_>> { + let item_tree = crate::file_item_tree(db, file); + let mut items: Vec<_> = item_tree.$map.iter().collect(); + items.sort_by_key(|(_, item)| item.span.start()); + items + .into_iter() + .map(|(id, _)| $loc::new(db, file, *id)) + .collect() + } + }; +} + +file_items!( + /// Every class declared in `file`, in source order. + file_classes, + classes, + ClassLoc +); +file_items!( + /// Every function declared in `file`, in source order. Includes methods. + file_functions, + functions, + FunctionLoc +); +file_items!( + /// Every enum declared in `file`, in source order. + file_enums, + enums, + EnumLoc +); +file_items!( + /// Every interface declared in `file`, in source order. + file_interfaces, + interfaces, + InterfaceLoc +); +file_items!( + /// Every type alias declared in `file`, in source order. + file_type_aliases, + type_aliases, + TypeAliasLoc +); +file_items!( + /// Every template string declared in `file`, in source order. + file_template_strings, + template_strings, + TemplateStringLoc +); +file_items!( + /// Every top-level `let` declared in `file`, in source order. + file_lets, + lets, + LetLoc +); + +/// Every `implements` block in `file`, in source order — both in-body and +/// out-of-body. +/// +/// Source order is load-bearing here: coherence needs a deterministic order to +/// pick a winner among overlapping impls. +#[salsa::tracked(returns(ref))] +pub fn file_impls(db: &dyn crate::Db, file: SourceFile) -> Vec> { + let item_tree = crate::file_item_tree(db, file); + let mut items: Vec<_> = item_tree.impls.iter().collect(); + items.sort_by_key(|(_, block)| block.span.start()); + items + .into_iter() + .map(|(id, _)| ImplLoc::new(db, file, *id)) + .collect() +} + +/// The out-of-body (`implement I for T`) blocks in `file`, in source order. +#[salsa::tracked(returns(ref))] +pub fn file_free_impls(db: &dyn crate::Db, file: SourceFile) -> Vec> { + let item_tree = crate::file_item_tree(db, file); + item_tree + .free_impls + .iter() + .map(|id| ImplLoc::new(db, file, *id)) + .collect() +} + +/// The `implements` blocks whose subject is `class`, in source order. +#[salsa::tracked(returns(ref))] +pub fn class_impls<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> Vec> { + let file = class.file(db); + let item_tree = crate::file_item_tree(db, file); + item_tree + .class_to_impls + .get(&class.id(db)) + .map(|impls| impls.iter().map(|id| ImplLoc::new(db, file, *id)).collect()) + .unwrap_or_default() +} + +// `Client`, `RetryPolicy` and `Test` have no span in the `ItemTree`, so they +// cannot be ordered by source position like the rest. They are unordered sets; +// sort by name so iteration is at least deterministic. + +/// Every client declared in `file`, ordered by name. +#[salsa::tracked(returns(ref))] +pub fn file_clients(db: &dyn crate::Db, file: SourceFile) -> Vec> { + let item_tree = crate::file_item_tree(db, file); + let mut items: Vec<_> = item_tree.clients.iter().collect(); + items.sort_by(|(_, a), (_, b)| a.name.cmp(&b.name)); + items + .into_iter() + .map(|(id, _)| ClientLoc::new(db, file, *id)) + .collect() +} + +/// Every retry policy declared in `file`, ordered by name. +#[salsa::tracked(returns(ref))] +pub fn file_retry_policies(db: &dyn crate::Db, file: SourceFile) -> Vec> { + let item_tree = crate::file_item_tree(db, file); + let mut items: Vec<_> = item_tree.retry_policies.iter().collect(); + items.sort_by(|(_, a), (_, b)| a.name.cmp(&b.name)); + items + .into_iter() + .map(|(id, _)| RetryPolicyLoc::new(db, file, *id)) + .collect() +} + +/// Every test declared in `file`, ordered by name. +#[salsa::tracked(returns(ref))] +pub fn file_tests(db: &dyn crate::Db, file: SourceFile) -> Vec> { + let item_tree = crate::file_item_tree(db, file); + let mut items: Vec<_> = item_tree.tests.iter().collect(); + items.sort_by(|(_, a), (_, b)| a.name.cmp(&b.name)); + items + .into_iter() + .map(|(id, _)| TestLoc::new(db, file, *id)) + .collect() +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs new file mode 100644 index 00000000000..ff7a7da51dc --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs @@ -0,0 +1,76 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{item_tree::Attribute, loc::EnumLoc}; + +/// Span-free semantic data for an `enum` declaration. +/// +/// Enums carry no type expressions, so there is no `TypeRefStore` here. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EnumData { + pub name: Name, + pub variants: Vec, + pub attributes: Vec, + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EnumVariantData { + pub name: Name, + pub attributes: Vec, + pub docstring: Option, +} + +/// Spans for an `Enum`, parallel to [`EnumData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EnumSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Name span per variant, parallel to [`EnumData::variants`]. + pub variant_name_spans: Vec, +} + +/// Semantic data for one enum. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn enum_data<'db>(db: &'db dyn crate::Db, item: EnumLoc<'db>) -> EnumData { + lower(db, item).0 +} + +/// Spans for one enum. Kept separate from [`enum_data`] so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn enum_source_map<'db>(db: &'db dyn crate::Db, item: EnumLoc<'db>) -> EnumSourceMap { + lower(db, item).1 +} + +fn lower<'db>(db: &'db dyn crate::Db, item: EnumLoc<'db>) -> (EnumData, EnumSourceMap) { + let file = item.file(db); + let item_tree = crate::file_item_tree(db, file); + let item_source_map = crate::file_item_tree_source_map(db, file); + let data = &item_tree[item.id(db)]; + + ( + EnumData { + name: data.name.clone(), + variants: data + .variants + .iter() + .map(|variant| EnumVariantData { + name: variant.name.clone(), + attributes: variant.attributes.clone(), + docstring: variant.docstring.clone(), + }) + .collect(), + attributes: data.attributes.clone(), + docstring: data.docstring.clone(), + }, + EnumSourceMap { + span: data.span, + variant_name_spans: item_source_map + .enum_variant_spans + .get(&item.id(db)) + .cloned() + .unwrap_or_default(), + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs index 9a5d502319c..cd66a623a12 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs @@ -3,6 +3,9 @@ use baml_compiler2_ast as ast; use text_size::TextRange; use crate::{ + item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FunctionParamData, + }, loc::FunctionLoc, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; @@ -35,15 +38,6 @@ pub struct FunctionData { pub is_tagged_template_tag: bool, } -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FunctionParamData { - pub name: Name, - pub type_ref: Option, - /// Whether a default expression was supplied. The expression itself lives in - /// `signature::function_parameter_defaults`. - pub has_default: bool, -} - /// Spans for a `Function`, parallel to [`FunctionData`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct FunctionSourceMap { @@ -73,6 +67,90 @@ pub fn function_source_map<'db>( lower(db, function).1 } +/// The interface target a method was declared under, when it sits inside an +/// `implements I { … }` block — plus the associated-type bindings from that +/// block's header. `None` for class-level methods and for interface default +/// methods. +/// +/// The target stays an unresolved `TypeRef`; TIR resolves it to an +/// `InterfaceLoc` lazily, so HIR construction stays independent of name +/// resolution. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MethodInterfaceTarget { + pub type_refs: TypeRefStore, + pub target: TypeRefId, + pub associated_type_bindings: Vec, +} + +/// Spans for a [`MethodInterfaceTarget`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MethodInterfaceTargetSourceMap { + pub type_refs: TypeRefSourceMap, + /// Parallel to `MethodInterfaceTarget::associated_type_bindings`. + pub associated_type_bindings: Vec, +} + +/// See [`MethodInterfaceTarget`]. Span-free. +#[salsa::tracked(returns(ref))] +pub fn method_interface_target<'db>( + db: &'db dyn crate::Db, + method: FunctionLoc<'db>, +) -> Option { + lower_interface_target(db, method).map(|(data, _)| data) +} + +/// Spans for [`method_interface_target`]. Kept separate so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn method_interface_target_source_map<'db>( + db: &'db dyn crate::Db, + method: FunctionLoc<'db>, +) -> Option { + lower_interface_target(db, method).map(|(_, source_map)| source_map) +} + +fn lower_interface_target<'db>( + db: &'db dyn crate::Db, + method: FunctionLoc<'db>, +) -> Option<(MethodInterfaceTarget, MethodInterfaceTargetSourceMap)> { + let item_tree = crate::file_item_tree(db, method.file(db)); + let target_expr = item_tree.method_to_iface_target.get(&method.id(db))?; + let bindings = item_tree + .method_to_iface_associated_type_bindings + .get(&method.id(db)) + .map(Vec::as_slice) + .unwrap_or_default(); + + let mut type_refs = TypeRefBuilder::new(); + let target = type_refs.lower(target_expr); + let associated_type_bindings = bindings + .iter() + .map(|binding| AssociatedTypeBindingData { + name: binding.name.clone(), + type_ref: binding.type_expr.as_ref().map(|te| type_refs.lower(te)), + }) + .collect(); + let (store, spans) = type_refs.finish(); + + Some(( + MethodInterfaceTarget { + type_refs: store, + target, + associated_type_bindings, + }, + MethodInterfaceTargetSourceMap { + type_refs: spans, + associated_type_bindings: bindings + .iter() + .map(|binding| AssociatedTypeBindingSourceMap { + span: binding.span, + name_span: binding.name_span, + }) + .collect(), + }, + )) +} + /// Both halves share one lowering pass — the split into two queries is purely /// about what each one lets downstream depend on. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs new file mode 100644 index 00000000000..037c2db2455 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs @@ -0,0 +1,179 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{ + item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, InterfaceFieldLinkData, + InterfaceFieldLinkSourceMap, + }, + loc::{ClassLoc, FunctionLoc, ImplLoc}, + type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, +}; + +/// What an `implements` block applies to. +/// +/// Unifying the owner with the for-target makes "in-body with an explicit +/// for-target" and "out-of-body without one" both unrepresentable. +#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)] +pub enum ImplSubjectData<'db> { + /// `implements I { … }` in a class body. The for-type is the class itself. + /// `out_of_body` records the syntactic origin for diagnostics only — it must + /// NOT influence resolution, dispatch, or coherence. + InClass { + class: ClassLoc<'db>, + out_of_body: bool, + }, + /// `implement<…> I for { … }` — an explicit for-type plus the + /// block's own generic parameters. + Free { + for_target: TypeRefId, + generics: Vec, + }, +} + +/// A generic parameter on an out-of-body `implements` block, paired with its +/// `&`-separated interface bounds. Pairing name and bounds makes a length +/// mismatch between the two unrepresentable. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GenericParamData { + pub name: Name, + pub bounds: Vec, +} + +/// Span-free semantic data for one `implements` block (either kind). +#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)] +pub struct ImplBlockData<'db> { + pub subject: ImplSubjectData<'db>, + /// Every type reference in this block's header — the for-target, its generic + /// bounds, the interface target, and any associated-type bindings. + pub type_refs: TypeRefStore, + pub interface_target: TypeRefId, + pub field_links: Vec, + pub associated_type_bindings: Vec, + pub methods: Vec>, +} + +/// Spans for an `ImplBlock`, parallel to [`ImplBlockData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplBlockSourceMap { + /// Full source span of the block. + pub span: TextRange, + /// Spans for every node in [`ImplBlockData::type_refs`]. + pub type_refs: TypeRefSourceMap, + /// Parallel to [`ImplBlockData::field_links`]. + pub field_links: Vec, + /// Parallel to [`ImplBlockData::associated_type_bindings`]. + pub associated_type_bindings: Vec, +} + +/// Semantic data for one `implements` block. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn impl_block_data<'db>(db: &'db dyn crate::Db, block: ImplLoc<'db>) -> ImplBlockData<'db> { + lower(db, block).0 +} + +/// Spans for one `implements` block. Kept separate from [`impl_block_data`] so +/// that a whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn impl_block_source_map<'db>( + db: &'db dyn crate::Db, + block: ImplLoc<'db>, +) -> ImplBlockSourceMap { + lower(db, block).1 +} + +/// Type refs are allocated in a fixed order (for-target and its bounds, then the +/// interface target, then associated-type bindings) so that ids are a pure +/// function of the block's shape. +fn lower<'db>( + db: &'db dyn crate::Db, + block: ImplLoc<'db>, +) -> (ImplBlockData<'db>, ImplBlockSourceMap) { + use crate::item_tree::ImplSubject; + + let file = block.file(db); + let item_tree = crate::file_item_tree(db, file); + let data = &item_tree.impls[&block.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + + let subject = match &data.subject { + ImplSubject::InClass { class, out_of_body } => ImplSubjectData::InClass { + class: ClassLoc::new(db, file, *class), + out_of_body: *out_of_body, + }, + ImplSubject::Free { + for_target, + generics, + } => ImplSubjectData::Free { + for_target: type_refs.lower(for_target), + generics: generics + .iter() + .map(|param| GenericParamData { + name: param.name.clone(), + bounds: param + .bounds + .iter() + .map(|bound| type_refs.lower(bound)) + .collect(), + }) + .collect(), + }, + }; + + let interface_target = type_refs.lower(&data.interface_target); + + let associated_type_bindings = data + .associated_type_bindings + .iter() + .map(|binding| AssociatedTypeBindingData { + name: binding.name.clone(), + type_ref: binding.type_expr.as_ref().map(|te| type_refs.lower(te)), + }) + .collect(); + + let (store, spans) = type_refs.finish(); + + ( + ImplBlockData { + subject, + type_refs: store, + interface_target, + field_links: data + .field_links + .iter() + .map(|link| InterfaceFieldLinkData { + interface_field: link.interface_field.clone(), + class_field: link.class_field.clone(), + }) + .collect(), + associated_type_bindings, + methods: data + .methods + .iter() + .map(|method| FunctionLoc::new(db, file, *method)) + .collect(), + }, + ImplBlockSourceMap { + span: data.span, + type_refs: spans, + field_links: data + .field_links + .iter() + .map(|link| InterfaceFieldLinkSourceMap { + span: link.span, + interface_field_span: link.interface_field_span, + class_field_span: link.class_field_span, + }) + .collect(), + associated_type_bindings: data + .associated_type_bindings + .iter() + .map(|binding| AssociatedTypeBindingSourceMap { + span: binding.span, + name_span: binding.name_span, + }) + .collect(), + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs new file mode 100644 index 00000000000..b5c8290b56a --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs @@ -0,0 +1,214 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{ + item_data::common::{FieldData, FunctionParamData}, + item_tree::Attribute, + loc::{FunctionLoc, InterfaceLoc}, + type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, +}; + +/// Span-free semantic data for an `interface` declaration. +#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)] +pub struct InterfaceData<'db> { + pub name: Name, + pub generic_params: Vec, + /// Every type reference in this interface's signature — bounds, `requires` + /// targets, field types, associated-type bounds and defaults, and required + /// method signatures. Scoped to the item. + pub type_refs: TypeRefStore, + /// Parallel to `generic_params`. `Some` means `T extends `. + pub generic_param_bounds: Vec>, + /// Targets of `requires I1, I2, …`. + pub requires: Vec, + /// Field signatures. Interface fields cannot have default values. + pub fields: Vec, + pub associated_types: Vec, + /// Default methods (with bodies). Implementing classes inherit them. + pub default_methods: Vec>, + /// Required methods (no body). Implementing classes must provide one. + pub required_methods: Vec, + pub attributes: Vec, + pub docstring: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeData { + pub name: Name, + pub bound: Option, + pub default: Option, +} + +/// A required (no-body) method signature on an interface. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceMethodSigData { + pub name: Name, + pub generic_params: Vec, + pub generic_param_bounds: Vec>, + pub params: Vec, + pub return_type: Option, + pub throws: Option, + pub attributes: Vec, + pub docstring: Option, +} + +/// Spans for an `Interface`, parallel to [`InterfaceData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Spans for every node in [`InterfaceData::type_refs`]. + pub type_refs: TypeRefSourceMap, + /// Parallel to [`InterfaceData::associated_types`]. + pub associated_type_spans: Vec, + /// Parallel to [`InterfaceData::required_methods`]. + pub required_method_spans: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssociatedTypeSourceMap { + pub span: TextRange, + pub name_span: TextRange, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InterfaceMethodSigSourceMap { + pub span: TextRange, + /// One span per parameter, parallel to `InterfaceMethodSigData::params`. + pub param_spans: Vec, +} + +/// Semantic data for one interface. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn interface_data<'db>( + db: &'db dyn crate::Db, + interface: InterfaceLoc<'db>, +) -> InterfaceData<'db> { + lower(db, interface).0 +} + +/// Spans for one interface. Kept separate from [`interface_data`] so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn interface_source_map<'db>( + db: &'db dyn crate::Db, + interface: InterfaceLoc<'db>, +) -> InterfaceSourceMap { + lower(db, interface).1 +} + +/// Type refs are allocated in a fixed order (bounds, requires, fields, +/// associated types, then required methods) so that ids are a pure function of +/// the interface's shape. +fn lower<'db>( + db: &'db dyn crate::Db, + interface: InterfaceLoc<'db>, +) -> (InterfaceData<'db>, InterfaceSourceMap) { + let file = interface.file(db); + let item_tree = crate::file_item_tree(db, file); + let data = &item_tree[interface.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + + let generic_param_bounds = data + .generic_param_bounds + .iter() + .map(|bound| bound.as_ref().map(|te| type_refs.lower(te))) + .collect(); + + let requires = data + .requires + .iter() + .map(|target| type_refs.lower(target)) + .collect(); + + let fields = data + .fields + .iter() + .map(|field| FieldData { + name: field.name.clone(), + type_ref: field.type_expr.as_ref().map(|te| type_refs.lower(te)), + attributes: field.attributes.clone(), + docstring: field.docstring.clone(), + }) + .collect(); + + let associated_types = data + .associated_types + .iter() + .map(|assoc| AssociatedTypeData { + name: assoc.name.clone(), + bound: assoc.bound.as_ref().map(|te| type_refs.lower(te)), + default: assoc.default.as_ref().map(|te| type_refs.lower(te)), + }) + .collect(); + + let required_methods = data + .required_methods + .iter() + .map(|method| InterfaceMethodSigData { + name: method.name.clone(), + generic_params: method.generic_params.clone(), + generic_param_bounds: method + .generic_param_bounds + .iter() + .map(|bound| bound.as_ref().map(|te| type_refs.lower(te))) + .collect(), + params: method + .params + .iter() + .map(|param| FunctionParamData { + name: param.name.clone(), + type_ref: param.type_expr.as_ref().map(|te| type_refs.lower(te)), + has_default: param.default.is_some(), + }) + .collect(), + return_type: method.return_type.as_ref().map(|te| type_refs.lower(te)), + throws: method.throws.as_ref().map(|te| type_refs.lower(te)), + attributes: method.attributes.clone(), + docstring: method.docstring.clone(), + }) + .collect(); + + let (store, spans) = type_refs.finish(); + + ( + InterfaceData { + name: data.name.clone(), + generic_params: data.generic_params.clone(), + type_refs: store, + generic_param_bounds, + requires, + fields, + associated_types, + default_methods: data + .default_methods + .iter() + .map(|method| FunctionLoc::new(db, file, *method)) + .collect(), + required_methods, + attributes: data.attributes.clone(), + docstring: data.docstring.clone(), + }, + InterfaceSourceMap { + span: data.span, + type_refs: spans, + associated_type_spans: data + .associated_types + .iter() + .map(|assoc| AssociatedTypeSourceMap { + span: assoc.span, + name_span: assoc.name_span, + }) + .collect(), + required_method_spans: data + .required_methods + .iter() + .map(|method| InterfaceMethodSigSourceMap { + span: method.span, + param_spans: method.params.iter().map(|param| param.span).collect(), + }) + .collect(), + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs new file mode 100644 index 00000000000..6787985f1e2 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs @@ -0,0 +1,54 @@ +use baml_base::Name; +use baml_compiler2_ast as ast; +use text_size::TextRange; + +use crate::loc::LetLoc; + +/// Span-free semantic data for a top-level `let` binding. +/// +/// The initializer is deliberately absent: it has its own query +/// (`body::let_body`), so editing an initializer cannot invalidate anything that +/// only depends on the binding itself. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LetData { + pub name: Name, + pub origin: ast::LetOrigin, +} + +/// Spans for a `Let`, parallel to [`LetData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LetSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Span of just the binding's name. + pub name_span: TextRange, +} + +/// Semantic data for one `let` binding. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn let_data<'db>(db: &'db dyn crate::Db, binding: LetLoc<'db>) -> LetData { + lower(db, binding).0 +} + +/// Spans for one `let` binding. Kept separate from [`let_data`] so that a +/// whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn let_source_map<'db>(db: &'db dyn crate::Db, binding: LetLoc<'db>) -> LetSourceMap { + lower(db, binding).1 +} + +fn lower<'db>(db: &'db dyn crate::Db, binding: LetLoc<'db>) -> (LetData, LetSourceMap) { + let item_tree = crate::file_item_tree(db, binding.file(db)); + let data = &item_tree[binding.id(db)]; + + ( + LetData { + name: data.name.clone(), + origin: data.origin, + }, + LetSourceMap { + span: data.span, + name_span: data.name_span, + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs new file mode 100644 index 00000000000..5855dbd4576 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs @@ -0,0 +1,34 @@ +use baml_base::Name; + +use crate::loc::RetryPolicyLoc; + +/// Semantic data for a `retry_policy` declaration. +/// +/// Carries neither type expressions nor spans, so there is no source-map twin — +/// the data query alone is the firewall. Values stay as raw strings; they are +/// parsed at emit time. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RetryPolicyData { + pub name: Name, + pub max_retries: Option, + pub initial_delay_ms: Option, + pub multiplier: Option, + pub max_delay_ms: Option, +} + +#[salsa::tracked(returns(ref))] +pub fn retry_policy_data<'db>( + db: &'db dyn crate::Db, + policy: RetryPolicyLoc<'db>, +) -> RetryPolicyData { + let item_tree = crate::file_item_tree(db, policy.file(db)); + let data = &item_tree[policy.id(db)]; + + RetryPolicyData { + name: data.name.clone(), + max_retries: data.max_retries.clone(), + initial_delay_ms: data.initial_delay_ms.clone(), + multiplier: data.multiplier.clone(), + max_delay_ms: data.max_delay_ms.clone(), + } +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs new file mode 100644 index 00000000000..54a5a5f31da --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs @@ -0,0 +1,83 @@ +use baml_base::Name; +use text_size::TextRange; + +use crate::{ + item_data::common::FunctionParamData, + loc::TemplateStringLoc, + type_ref::{TypeRefBuilder, TypeRefSourceMap, TypeRefStore}, +}; + +/// Span-free semantic data for a `template_string` declaration. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TemplateStringData { + pub name: Name, + /// Type references in this template's parameter list. Scoped to the item. + pub type_refs: TypeRefStore, + pub params: Vec, + /// Template body text (Jinja). + pub body: Option, +} + +/// Spans for a `TemplateString`, parallel to [`TemplateStringData`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TemplateStringSourceMap { + /// Full source span of the declaration. + pub span: TextRange, + /// Spans for every node in [`TemplateStringData::type_refs`]. + pub type_refs: TypeRefSourceMap, + /// One span per parameter, parallel to [`TemplateStringData::params`]. + pub param_spans: Vec, +} + +/// Semantic data for one template string. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn template_string_data<'db>( + db: &'db dyn crate::Db, + template: TemplateStringLoc<'db>, +) -> TemplateStringData { + lower(db, template).0 +} + +/// Spans for one template string. Kept separate from [`template_string_data`] so +/// that a whitespace-only edit invalidates this but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn template_string_source_map<'db>( + db: &'db dyn crate::Db, + template: TemplateStringLoc<'db>, +) -> TemplateStringSourceMap { + lower(db, template).1 +} + +fn lower<'db>( + db: &'db dyn crate::Db, + template: TemplateStringLoc<'db>, +) -> (TemplateStringData, TemplateStringSourceMap) { + let item_tree = crate::file_item_tree(db, template.file(db)); + let data = &item_tree[template.id(db)]; + + let mut type_refs = TypeRefBuilder::new(); + let params = data + .params + .iter() + .map(|param| FunctionParamData { + name: param.name.clone(), + type_ref: param.type_expr.as_ref().map(|te| type_refs.lower(te)), + has_default: param.default.is_some(), + }) + .collect(); + let (store, spans) = type_refs.finish(); + + ( + TemplateStringData { + name: data.name.clone(), + type_refs: store, + params, + body: data.body.clone(), + }, + TemplateStringSourceMap { + span: data.span, + type_refs: spans, + param_spans: data.params.iter().map(|param| param.span).collect(), + }, + ) +} diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs new file mode 100644 index 00000000000..da3215183ca --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs @@ -0,0 +1,27 @@ +use baml_base::Name; + +use crate::{item_tree::TestArgValue, loc::TestLoc}; + +/// Semantic data for a `test` declaration. +/// +/// Carries neither type expressions nor spans, so there is no source-map twin — +/// the data query alone is the firewall. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TestData { + pub name: Name, + /// The function(s) this test exercises. + pub function_refs: Vec, + pub args: Vec<(Name, TestArgValue)>, +} + +#[salsa::tracked(returns(ref))] +pub fn test_data<'db>(db: &'db dyn crate::Db, test: TestLoc<'db>) -> TestData { + let item_tree = crate::file_item_tree(db, test.file(db)); + let data = &item_tree[test.id(db)]; + + TestData { + name: data.name.clone(), + function_refs: data.function_refs.clone(), + args: data.args.clone(), + } +} diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index 8958befd873..2879fdf98fb 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -474,7 +474,7 @@ type ClassFingerprint = ( baml_base::Name, Vec>, baml_compiler2_hir::type_ref::TypeRefStore, - Vec, + Vec, Vec, Vec, ); From d0fef032f01702080d51c4567c0ecb67cb580686 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 16:40:38 -0700 Subject: [PATCH 06/20] HIR item tree builder This separates the builder logic from the immutable `ItemTree` so the index counters should not clobber salsa Also record mappings between items and scopes instead of using span comparisons --- .../crates/baml_compiler2_hir/src/builder.rs | 83 +++- .../baml_compiler2_hir/src/item_data.rs | 2 + .../src/item_data/scopes.rs | 117 +++++ .../baml_compiler2_hir/src/item_tree.rs | 430 +----------------- .../src/item_tree/builder.rs | 412 +++++++++++++++++ .../crates/baml_compiler2_hir/src/scope.rs | 36 ++ .../baml_compiler2_hir/src/semantic_index.rs | 21 +- .../baml_tests/src/incremental/scenarios.rs | 106 +++++ 8 files changed, 759 insertions(+), 448 deletions(-) create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs create mode 100644 baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs diff --git a/baml_language/crates/baml_compiler2_hir/src/builder.rs b/baml_language/crates/baml_compiler2_hir/src/builder.rs index d1074a71d56..21e0a52adf8 100644 --- a/baml_language/crates/baml_compiler2_hir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_hir/src/builder.rs @@ -19,12 +19,12 @@ use crate::{ diagnostic::{Hir2Diagnostic, MemberSite}, file_package::file_package, ids::{FunctionMarker, LocalItemId}, - item_tree::{GenericParam, ImplBlock, ImplSubject, InterfaceFieldLink, ItemTree}, + item_tree::{GenericParam, ImplBlock, ImplSubject, InterfaceFieldLink}, loc::{ ClassLoc, ClientLoc, EnumLoc, FunctionLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, TemplateStringLoc, TestLoc, TypeAliasLoc, }, - scope::{FileScopeId, Scope, ScopeId, ScopeKind}, + scope::{FileScopeId, ItemScopeOwner, Scope, ScopeId, ScopeKind}, semantic_index::{ BindingId, DefinitionSite, FileSemanticIndex, LocalBinding, PathResolution, ScopeBindings, SemanticIndexExtra, visible_binding_at_in_scopes, @@ -60,6 +60,8 @@ pub struct SemanticIndexBuilder<'db> { scopes: Vec, scope_bindings: Vec, + /// Owning item -> its scope. Inverse of `Scope::owner`. + item_scopes: FxHashMap, /// Stack of currently-open scope IDs. scope_stack: Vec, /// Depth of class scopes we're inside (> 0 means methods shouldn't @@ -78,8 +80,7 @@ pub struct SemanticIndexBuilder<'db> { path_root_references: Vec, lambda_stack: Vec, - item_tree: ItemTree, - item_tree_source_map: crate::item_tree::ItemTreeSourceMap, + item_tree: crate::item_tree::builder::ItemTreeBuilder, type_contributions: Vec<(Name, Contribution<'db>)>, value_contributions: Vec<(Name, Contribution<'db>)>, diagnostics: Vec, @@ -94,14 +95,14 @@ impl<'db> SemanticIndexBuilder<'db> { file, scopes: Vec::new(), scope_bindings: Vec::new(), + item_scopes: FxHashMap::default(), scope_stack: Vec::new(), class_depth: 0, expr_scopes: Vec::new(), path_resolutions: Vec::new(), path_root_references: Vec::new(), lambda_stack: Vec::new(), - item_tree: ItemTree::new(), - item_tree_source_map: crate::item_tree::ItemTreeSourceMap::default(), + item_tree: crate::item_tree::builder::ItemTreeBuilder::new(), type_contributions: Vec::new(), value_contributions: Vec::new(), diagnostics: Vec::new(), @@ -185,13 +186,17 @@ impl<'db> SemanticIndexBuilder<'db> { })) }; + // Drops the collision counter — it has no meaning once the tree is built. + let (item_tree, item_tree_source_map) = self.item_tree.finish(); + FileSemanticIndex { scopes: self.scopes, expr_scopes: self.expr_scopes, scope_bindings: self.scope_bindings, scope_ids, - item_tree: Arc::new(self.item_tree), - item_tree_source_map: Arc::new(self.item_tree_source_map), + item_scopes: self.item_scopes, + item_tree: Arc::new(item_tree), + item_tree_source_map: Arc::new(item_tree_source_map), symbol_contributions: Arc::new(FileSymbolContributions { types: self.type_contributions, values: self.value_contributions, @@ -212,6 +217,7 @@ impl<'db> SemanticIndexBuilder<'db> { parent, kind, name, + owner: None, range, descendants: id.next()..id.next(), // empty initially; filled on pop is_template_body: false, @@ -220,6 +226,15 @@ impl<'db> SemanticIndexBuilder<'db> { self.scope_stack.push(id); } + /// Link a scope to the item it was opened for. + /// + /// Recorded here, at the one place that knows both, rather than recovered + /// later by comparing `item.span == scope.range`. + fn record_scope_owner(&mut self, scope: FileScopeId, owner: ItemScopeOwner) { + self.scopes[scope.index() as usize].owner = Some(owner); + self.item_scopes.insert(owner, scope); + } + fn pop_scope(&mut self) { let popped = self.scope_stack.pop().expect("scope stack underflow"); #[allow(clippy::cast_possible_truncation)] @@ -1173,7 +1188,6 @@ impl<'db> SemanticIndexBuilder<'db> { fn lower_function(&mut self, f: &ast::FunctionDef) -> LocalItemId { let local_id = self.item_tree.alloc_function(f); - ItemTree::collect_function_span(&mut self.item_tree_source_map, local_id, f); let loc = FunctionLoc::new(self.db, self.file, local_id); // Only contribute as a top-level symbol if not inside a class. @@ -1190,6 +1204,7 @@ impl<'db> SemanticIndexBuilder<'db> { self.push_scope(ScopeKind::Function, Some(f.name.clone()), f.span); let scope_id = self.current_scope_id(); + self.record_scope_owner(scope_id, ItemScopeOwner::Function(local_id)); for (idx, param) in f.params.iter().enumerate() { self.scope_bindings[scope_id.index() as usize] @@ -1218,7 +1233,6 @@ impl<'db> SemanticIndexBuilder<'db> { fn lower_class(&mut self, c: &ast::ClassDef) { let local_id = self.item_tree.alloc_class(c); - ItemTree::collect_class_spans(&mut self.item_tree_source_map, local_id, c); let loc = ClassLoc::new(self.db, self.file, local_id); self.type_contributions.push(( c.name.clone(), @@ -1229,6 +1243,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Class, Some(c.name.clone()), c.span); + let class_scope = self.current_scope_id(); + self.record_scope_owner(class_scope, ItemScopeOwner::Class(local_id)); // Unified per-scope duplicate detection: all members (fields, methods) // share one name-map so cross-kind collisions are also caught. @@ -1310,12 +1326,11 @@ impl<'db> SemanticIndexBuilder<'db> { // BEP-044: remember which interface this method came from so // `default.()` calls inside the body can resolve back // to the interface's default function. - self.item_tree - .method_to_iface_target - .insert(fid, impl_block.target.clone()); - self.item_tree - .method_to_iface_associated_type_bindings - .insert(fid, impl_block.associated_type_bindings.clone()); + self.item_tree.record_method_interface_target( + fid, + impl_block.target.clone(), + impl_block.associated_type_bindings.clone(), + ); block_method_ids.push(fid); } // Dual-write: also record this in-body impl under a stable `ImplId`. @@ -1351,6 +1366,7 @@ impl<'db> SemanticIndexBuilder<'db> { // For blanket impls (implements I for C), push a class-like scope // so TIR can resolve `self` and type variables in method bodies. let has_generic_params = !imp.generic_params.is_empty(); + let mut impl_scope: Option = None; if has_generic_params { // Derive a synthetic scope name from the for_target for `self` resolution. // Use the for_target's root name (e.g. "Container" from "Container"). @@ -1361,16 +1377,16 @@ impl<'db> SemanticIndexBuilder<'db> { _ => None, }; self.push_scope(ScopeKind::Class, scope_name, imp.span); + impl_scope = Some(self.current_scope_id()); } let mut method_ids = Vec::new(); for method in &imp.methods { let fid = self.lower_function(method); - self.item_tree - .method_to_iface_target - .insert(fid, imp.interface_target.clone()); - self.item_tree - .method_to_iface_associated_type_bindings - .insert(fid, imp.associated_type_bindings.clone()); + self.item_tree.record_method_interface_target( + fid, + imp.interface_target.clone(), + imp.associated_type_bindings.clone(), + ); method_ids.push(fid); } if has_generic_params { @@ -1403,7 +1419,10 @@ impl<'db> SemanticIndexBuilder<'db> { methods: method_ids, span: imp.span, }; - self.item_tree.alloc_impl(&iface_head, &for_head, block); + let impl_id = self.item_tree.alloc_impl(&iface_head, &for_head, block); + if let Some(scope) = impl_scope { + self.record_scope_owner(scope, ItemScopeOwner::Impl(impl_id)); + } } /// Lower an `interface I { ... }` declaration (BEP-044). @@ -1420,6 +1439,7 @@ impl<'db> SemanticIndexBuilder<'db> { // methods so we can record their `FunctionMarker` ids on the // interface. self.push_scope(ScopeKind::Class, Some(i.name.clone()), i.span); + let interface_scope = self.current_scope_id(); // BEP-044: default methods are lowered inside the interface's // `Class`-kind scope so the semantic index reports the interface @@ -1435,7 +1455,7 @@ impl<'db> SemanticIndexBuilder<'db> { self.class_depth -= 1; let local_id = self.item_tree.alloc_interface(i, default_method_ids); - ItemTree::collect_interface_spans(&mut self.item_tree_source_map, local_id, i); + self.record_scope_owner(interface_scope, ItemScopeOwner::Interface(local_id)); let loc = InterfaceLoc::new(self.db, self.file, local_id); self.type_contributions.push(( i.name.clone(), @@ -1483,7 +1503,6 @@ impl<'db> SemanticIndexBuilder<'db> { fn lower_enum(&mut self, e: &ast::EnumDef) { let local_id = self.item_tree.alloc_enum(e); - ItemTree::collect_enum_spans(&mut self.item_tree_source_map, local_id, e); let loc = EnumLoc::new(self.db, self.file, local_id); self.type_contributions.push(( e.name.clone(), @@ -1494,6 +1513,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Enum, Some(e.name.clone()), e.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::Enum(local_id)); let mut seen: FxHashMap> = FxHashMap::default(); for variant in &e.variants { @@ -1522,6 +1543,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::TypeAlias, Some(ta.name.clone()), ta.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::TypeAlias(local_id)); self.pop_scope(); } @@ -1537,6 +1560,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Item, Some(c.name.clone()), c.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::Client(local_id)); self.pop_scope(); } @@ -1552,6 +1577,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Item, Some(t.name.clone()), t.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::Test(local_id)); self.pop_scope(); } @@ -1567,6 +1594,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Function, Some(ts.name.clone()), ts.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::TemplateString(local_id)); self.pop_scope(); } @@ -1582,6 +1611,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Item, Some(rp.name.clone()), rp.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::RetryPolicy(local_id)); self.pop_scope(); } @@ -1597,6 +1628,8 @@ impl<'db> SemanticIndexBuilder<'db> { )); self.push_scope(ScopeKind::Let, Some(l.name.clone()), l.span); + let scope = self.current_scope_id(); + self.record_scope_owner(scope, ItemScopeOwner::Let(local_id)); if let Some((ref body, ref source_map)) = l.initializer { self.walk_expr_body(body, source_map); } diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data.rs b/baml_language/crates/baml_compiler2_hir/src/item_data.rs index f5c9e32c1f1..5f9105cf5e6 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_data.rs @@ -36,6 +36,7 @@ mod impls; mod interfaces; mod lets; mod retry_policies; +mod scopes; mod template_strings; mod test_items; mod type_aliases; @@ -50,6 +51,7 @@ pub use impls::*; pub use interfaces::*; pub use lets::*; pub use retry_policies::*; +pub use scopes::*; pub use template_strings::*; pub use test_items::*; pub use type_aliases::*; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs b/baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs new file mode 100644 index 00000000000..fb1b33f59c7 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs @@ -0,0 +1,117 @@ +//! Item ↔ scope queries. +//! +//! The builder opens an item's scope in the same step that allocates the item, so +//! the link is recorded directly ([`Scope::owner`](crate::scope::Scope::owner)). +//! Before that, consumers recovered it by scanning for a scope whose `range` +//! equalled the item's `span` — which made item spans load-bearing *semantic +//! identity* rather than diagnostic metadata, and blocked moving them into the +//! source map. These queries replace that scan. + +use crate::{ + loc::{ + ClassLoc, ClientLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, + TemplateStringLoc, TestLoc, TypeAliasLoc, + }, + scope::{ItemScopeOwner, ScopeId}, +}; + +/// The item a scope was opened for, as a `Loc`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, salsa::Update)] +pub enum ScopeOwner<'db> { + Function(FunctionLoc<'db>), + Class(ClassLoc<'db>), + Enum(EnumLoc<'db>), + Interface(InterfaceLoc<'db>), + TypeAlias(TypeAliasLoc<'db>), + TemplateString(TemplateStringLoc<'db>), + Client(ClientLoc<'db>), + Test(TestLoc<'db>), + RetryPolicy(RetryPolicyLoc<'db>), + Let(LetLoc<'db>), + Impl(ImplLoc<'db>), +} + +/// The item `scope` was opened for, or `None` for structural and expression +/// scopes (Project/Package/Namespace/File, Block/Lambda/MatchArm/…). +#[salsa::tracked] +pub fn scope_owner<'db>(db: &'db dyn crate::Db, scope: ScopeId<'db>) -> Option> { + let file = scope.file(db); + let index = crate::file_semantic_index(db, file); + let owner = index.scope_owner(scope.file_scope_id(db))?; + + Some(match owner { + ItemScopeOwner::Function(id) => ScopeOwner::Function(FunctionLoc::new(db, file, id)), + ItemScopeOwner::Class(id) => ScopeOwner::Class(ClassLoc::new(db, file, id)), + ItemScopeOwner::Enum(id) => ScopeOwner::Enum(EnumLoc::new(db, file, id)), + ItemScopeOwner::Interface(id) => ScopeOwner::Interface(InterfaceLoc::new(db, file, id)), + ItemScopeOwner::TypeAlias(id) => ScopeOwner::TypeAlias(TypeAliasLoc::new(db, file, id)), + ItemScopeOwner::TemplateString(id) => { + ScopeOwner::TemplateString(TemplateStringLoc::new(db, file, id)) + } + ItemScopeOwner::Client(id) => ScopeOwner::Client(ClientLoc::new(db, file, id)), + ItemScopeOwner::Test(id) => ScopeOwner::Test(TestLoc::new(db, file, id)), + ItemScopeOwner::RetryPolicy(id) => { + ScopeOwner::RetryPolicy(RetryPolicyLoc::new(db, file, id)) + } + ItemScopeOwner::Let(id) => ScopeOwner::Let(LetLoc::new(db, file, id)), + ItemScopeOwner::Impl(id) => ScopeOwner::Impl(ImplLoc::new(db, file, id)), + }) +} + +/// Declare the inverse: `_scope(loc) -> Option`. +macro_rules! item_scope { + ($(#[$meta:meta])* $name:ident, $loc:ident, $owner:ident) => { + $(#[$meta])* + #[salsa::tracked] + pub fn $name<'db>(db: &'db dyn crate::Db, item: $loc<'db>) -> Option> { + let file = item.file(db); + let index = crate::file_semantic_index(db, file); + let scope = index.item_scope(ItemScopeOwner::$owner(item.id(db)))?; + Some(index.scope_ids[scope.index() as usize]) + } + }; +} + +item_scope!( + /// The scope opened for `function`'s body. + function_scope, + FunctionLoc, + Function +); +item_scope!( + /// The scope opened for `class`'s members. + class_scope, + ClassLoc, + Class +); +item_scope!( + /// The scope opened for `interface`'s members. + interface_scope, + InterfaceLoc, + Interface +); +item_scope!( + /// The scope opened for `alias`'s right-hand side. + type_alias_scope, + TypeAliasLoc, + TypeAlias +); +item_scope!( + /// The scope owning `binding`'s initializer. + let_scope, + LetLoc, + Let +); +item_scope!( + /// The scope opened for `template`'s parameters and body. + template_string_scope, + TemplateStringLoc, + TemplateString +); +item_scope!( + /// The scope opened for an out-of-body `implement I for T` block's generic + /// parameters. `None` for blocks without generics — they open no scope. + impl_scope, + ImplLoc, + Impl +); diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs index a52510a5c15..2c3f33d2cf1 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs @@ -4,6 +4,7 @@ //! following the same scheme as `baml_compiler_hir::item_tree`. //! Items are indexed by name (not source position) for position-independence. +pub mod builder; mod classes; mod clients; mod common; @@ -19,7 +20,6 @@ mod type_aliases; use std::ops::Index; -use baml_base::Name; use baml_compiler2_ast as ast; pub use classes::*; pub use clients::*; @@ -33,23 +33,22 @@ use rustc_hash::FxHashMap; pub use source_map::*; pub use template_strings::*; pub use test_items::*; -use text_size::TextRange; pub use type_aliases::*; use crate::ids::{ - ClassMarker, ClientMarker, EnumMarker, FunctionMarker, ImplMarker, InterfaceMarker, ItemKind, - LetMarker, LocalItemId, RetryPolicyMarker, TemplateStringMarker, TestMarker, TypeAliasMarker, - hash_impl_key, hash_name, + ClassMarker, ClientMarker, EnumMarker, FunctionMarker, ImplMarker, InterfaceMarker, LetMarker, + LocalItemId, RetryPolicyMarker, TemplateStringMarker, TestMarker, TypeAliasMarker, }; // ── ItemTree ───────────────────────────────────────────────────────────────── /// Position-independent item storage for a single file. /// -/// Items are stored in hash maps keyed by name-based IDs. -/// The `next_index` map tracks the next available collision index -/// per `(ItemKind, hash)`. -#[derive(Debug, Clone, PartialEq, Eq)] +/// Items are stored in hash maps keyed by name-based IDs. This is a finished, +/// immutable value: everything needed to *build* it — the collision counter, the +/// source map — lives in [`ItemTreeBuilder`](builder::ItemTreeBuilder) and is +/// dropped once the tree is complete. +#[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct ItemTree { pub functions: FxHashMap, Function>, pub classes: FxHashMap, Class>, @@ -84,419 +83,6 @@ pub struct ItemTree { pub method_to_iface_target: FxHashMap, ast::TypeExpr>, pub method_to_iface_associated_type_bindings: FxHashMap, Vec>, - - /// Collision tracker: `(ItemKind, hash)` → next available index. - next_index: FxHashMap<(ItemKind, u16), u16>, -} - -impl Default for ItemTree { - fn default() -> Self { - Self::new() - } -} - -impl ItemTree { - pub fn new() -> Self { - Self { - functions: FxHashMap::default(), - classes: FxHashMap::default(), - enums: FxHashMap::default(), - interfaces: FxHashMap::default(), - type_aliases: FxHashMap::default(), - clients: FxHashMap::default(), - tests: FxHashMap::default(), - template_strings: FxHashMap::default(), - retry_policies: FxHashMap::default(), - lets: FxHashMap::default(), - impls: FxHashMap::default(), - class_to_impls: FxHashMap::default(), - free_impls: Vec::new(), - method_to_iface_target: FxHashMap::default(), - method_to_iface_associated_type_bindings: FxHashMap::default(), - next_index: FxHashMap::default(), - } - } - - /// Allocate a collision-resistant ID for an item. - fn alloc_id(&mut self, kind: ItemKind, name: &Name) -> LocalItemId { - let h = hash_name(name); - let index = self.next_index.entry((kind, h)).or_insert(0); - let id = LocalItemId::new(h, *index); - *index += 1; - id - } - - /// Allocate a function in the `ItemTree` with full AST data. - pub fn alloc_function(&mut self, f: &ast::FunctionDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Function, &f.name); - let params = f - .params - .iter() - .map(|p| FunctionParam { - name: p.name.clone(), - type_expr: p.type_expr.clone(), - default: p.default.map(|expr| DefaultExprRef { function: id, expr }), - span: p.span, - }) - .collect(); - self.functions.insert( - id, - Function { - name: f.name.clone(), - generic_params: f.generic_params.clone(), - generic_param_bounds: f.generic_param_bounds.clone(), - params, - defaults: f.defaults.clone(), - return_type: f.return_type.clone(), - throws: f.throws.clone(), - body: f.body.clone(), - declarative_meta: f.declarative_meta.clone(), - origin: f.origin, - docstring: f.docstring.clone(), - is_tagged_template_tag: f.is_tagged_template_tag, - span: f.span, - }, - ); - id - } - - pub fn alloc_class(&mut self, c: &ast::ClassDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Class, &c.name); - let fields = c - .fields - .iter() - .map(|f| ClassField { - name: f.name.clone(), - type_expr: f.type_expr.clone(), - attributes: f.attributes.iter().map(Attribute::from).collect(), - docstring: f.docstring.clone(), - }) - .collect(); - let implements = c - .implements - .iter() - .map(|b| ImplementsBlock { - target: b.target.clone(), - field_links: b - .field_links - .iter() - .map(InterfaceFieldLink::from_ast) - .collect(), - associated_type_bindings: b.associated_type_bindings.clone(), - is_out_of_body: b.is_out_of_body, - span: b.span, - }) - .collect(); - self.classes.insert( - id, - Class { - name: c.name.clone(), - generic_params: c.generic_params.clone(), - generic_param_bounds: c.generic_param_bounds.clone(), - fields, - methods: Vec::new(), - implements, - attributes: c.attributes.iter().map(Attribute::from).collect(), - docstring: c.docstring.clone(), - span: c.span, - }, - ); - id - } - - /// Attach method IDs to an already-allocated class. - pub fn set_class_methods( - &mut self, - class_id: LocalItemId, - methods: Vec>, - ) { - if let Some(class) = self.classes.get_mut(&class_id) { - class.methods = methods; - } - } - - /// Allocate a stable id for an `implements` block and store it in the - /// unified `impls` map. `iface_head`/`for_head` seed the position-independent - /// id (impls have no declared name); the collision index disambiguates impls - /// that share both heads. Records the `class_to_impls` edge for `InClass`. - pub fn alloc_impl( - &mut self, - iface_head: &Name, - for_head: &Name, - block: ImplBlock, - ) -> LocalItemId { - let h = hash_impl_key(iface_head, for_head); - let index = self.next_index.entry((ItemKind::Impl, h)).or_insert(0); - let id = LocalItemId::new(h, *index); - *index += 1; - match &block.subject { - ImplSubject::InClass { class, .. } => { - self.class_to_impls.entry(*class).or_default().push(id); - } - ImplSubject::Free { .. } => self.free_impls.push(id), - } - self.impls.insert(id, block); - id - } - - pub fn alloc_enum(&mut self, e: &ast::EnumDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Enum, &e.name); - let variants = e - .variants - .iter() - .map(|v| EnumVariant { - name: v.name.clone(), - attributes: v.attributes.iter().map(Attribute::from).collect(), - docstring: v.docstring.clone(), - }) - .collect(); - self.enums.insert( - id, - Enum { - name: e.name.clone(), - variants, - attributes: e.attributes.iter().map(Attribute::from).collect(), - docstring: e.docstring.clone(), - span: e.span, - }, - ); - id - } - - /// Populate source map spans for a class that was allocated via `alloc_class`. - pub fn collect_class_spans( - source_map: &mut ItemTreeSourceMap, - id: LocalItemId, - class_def: &ast::ClassDef, - ) { - let spans: Vec = class_def.fields.iter().map(|f| f.name_span).collect(); - source_map.class_field_spans.insert(id, spans); - } - - /// Populate source map name span for a function. - pub fn collect_function_span( - source_map: &mut ItemTreeSourceMap, - id: LocalItemId, - func_def: &ast::FunctionDef, - ) { - source_map - .function_name_spans - .insert(id, func_def.name_span); - } - - /// Populate source map spans for an enum that was allocated via `alloc_enum`. - pub fn collect_enum_spans( - source_map: &mut ItemTreeSourceMap, - id: LocalItemId, - enum_def: &ast::EnumDef, - ) { - let spans: Vec = enum_def.variants.iter().map(|v| v.name_span).collect(); - source_map.enum_variant_spans.insert(id, spans); - } - - /// Populate source map spans for an interface allocated via `alloc_interface`. - pub fn collect_interface_spans( - source_map: &mut ItemTreeSourceMap, - id: LocalItemId, - iface_def: &ast::InterfaceDef, - ) { - let field_spans: Vec = iface_def.fields.iter().map(|f| f.name_span).collect(); - source_map.interface_field_spans.insert(id, field_spans); - let method_spans: Vec = iface_def - .required_methods - .iter() - .map(|m| m.name_span) - .collect(); - source_map.interface_method_spans.insert(id, method_spans); - } - - /// Allocate an interface (BEP-044) in the `ItemTree`. - /// - /// `default_method_ids` are the `FunctionMarker` ids for any default - /// methods in the interface — those should be allocated separately via - /// `alloc_function` before this is called. - pub fn alloc_interface( - &mut self, - i: &ast::InterfaceDef, - default_method_ids: Vec>, - ) -> LocalItemId { - let id = self.alloc_id(ItemKind::Interface, &i.name); - let fields = i - .fields - .iter() - .map(|f| ClassField { - name: f.name.clone(), - type_expr: f.type_expr.clone(), - attributes: f.attributes.iter().map(Attribute::from).collect(), - docstring: f.docstring.clone(), - }) - .collect(); - let required_methods = i - .required_methods - .iter() - .map(|m| InterfaceMethodSig { - name: m.name.clone(), - generic_params: m.generic_params.clone(), - generic_param_bounds: m.generic_param_bounds.clone(), - params: m - .params - .iter() - .map(|p| FunctionParam { - name: p.name.clone(), - type_expr: p.type_expr.clone(), - default: None, - span: p.span, - }) - .collect(), - return_type: m.return_type.clone(), - throws: m.throws.clone(), - attributes: m.attributes.iter().map(Attribute::from).collect(), - docstring: m.docstring.clone(), - span: m.span, - }) - .collect(); - self.interfaces.insert( - id, - Interface { - name: i.name.clone(), - generic_params: i.generic_params.clone(), - generic_param_bounds: i.generic_param_bounds.clone(), - requires: i.requires.clone(), - fields, - associated_types: i.associated_types.clone(), - default_methods: default_method_ids, - required_methods, - attributes: i.attributes.iter().map(Attribute::from).collect(), - docstring: i.docstring.clone(), - span: i.span, - }, - ); - id - } - - pub fn alloc_type_alias(&mut self, ta: &ast::TypeAliasDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::TypeAlias, &ta.name); - self.type_aliases.insert( - id, - TypeAlias { - name: ta.name.clone(), - type_expr: ta.type_expr.clone(), - span: ta.span, - }, - ); - id - } - - pub fn alloc_client(&mut self, c: &ast::ClientDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Client, &c.name); - let provider = c - .config_items - .iter() - .find(|item| item.key.as_str() == "provider") - .map(|item| Name::new(item.value.trim().trim_matches('"'))); - let sub_client_names = c - .config_items - .iter() - .find(|item| item.key.as_str() == "options") - .map(|_| Vec::new()) // complex to parse; clients field is more relevant - .unwrap_or_default(); - let retry_policy_name = c - .config_items - .iter() - .find(|item| item.key.as_str() == "retry_policy") - .map(|item| Name::new(item.value.trim().trim_matches('"'))); - self.clients.insert( - id, - Client { - name: c.name.clone(), - provider, - sub_client_names, - retry_policy_name, - round_robin_start: None, - }, - ); - id - } - - pub fn alloc_test(&mut self, t: &ast::TestDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Test, &t.name); - self.tests.insert( - id, - Test { - name: t.name.clone(), - function_refs: t.function_refs.clone(), - args: t.args.clone(), - }, - ); - id - } - - pub fn alloc_template_string( - &mut self, - ts: &ast::TemplateStringDef, - ) -> LocalItemId { - let id = self.alloc_id(ItemKind::TemplateString, &ts.name); - let params = ts - .params - .iter() - .map(|p| FunctionParam { - name: p.name.clone(), - type_expr: p.type_expr.clone(), - default: None, - span: p.span, - }) - .collect(); - let body = ts.body.as_ref().map(|b| b.text.clone()); - self.template_strings.insert( - id, - TemplateString { - name: ts.name.clone(), - params, - body, - span: ts.span, - }, - ); - id - } - - pub fn alloc_retry_policy( - &mut self, - rp: &ast::RetryPolicyDef, - ) -> LocalItemId { - let id = self.alloc_id(ItemKind::RetryPolicy, &rp.name); - let get_field = |key: &str| -> Option { - rp.config_items - .iter() - .find(|item| item.key.as_str() == key) - .map(|item| item.value.trim().to_string()) - }; - self.retry_policies.insert( - id, - RetryPolicy { - name: rp.name.clone(), - max_retries: get_field("max_retries"), - initial_delay_ms: get_field("initial_delay_ms"), - multiplier: get_field("multiplier"), - max_delay_ms: get_field("max_delay_ms"), - }, - ); - id - } - - pub fn alloc_let(&mut self, l: &ast::LetDef) -> LocalItemId { - let id = self.alloc_id(ItemKind::Let, &l.name); - self.lets.insert( - id, - Let { - name: l.name.clone(), - initializer: l.initializer.clone(), - origin: l.origin, - span: l.span, - name_span: l.name_span, - }, - ); - id - } } // ── Index impls ─────────────────────────────────────────────────────────────── diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs new file mode 100644 index 00000000000..cadd6d88cb5 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs @@ -0,0 +1,412 @@ +//! Construction-time state for the `ItemTree`. +//! +//! `next_index` is the collision counter used while handing out `LocalItemId`s. +//! It is pure construction state — it has no meaning once the tree is built, and +//! leaving it on `ItemTree` would put mutable bookkeeping inside an immutable +//! value that derives `PartialEq`. It lives here instead. +//! +//! The builder also owns the `ItemTreeSourceMap`, so allocating an item and +//! recording its spans is a single call and the two cannot drift apart. + +use baml_base::Name; +use baml_compiler2_ast as ast; +use rustc_hash::FxHashMap; + +use crate::{ + ids::{ + ClassMarker, ClientMarker, EnumMarker, FunctionMarker, ImplMarker, InterfaceMarker, + ItemKind, LetMarker, LocalItemId, RetryPolicyMarker, TemplateStringMarker, TestMarker, + TypeAliasMarker, hash_impl_key, hash_name, + }, + item_tree::{ + Attribute, Class, ClassField, Client, DefaultExprRef, Enum, EnumVariant, Function, + FunctionParam, ImplBlock, ImplSubject, ImplementsBlock, Interface, InterfaceFieldLink, + InterfaceMethodSig, ItemTree, ItemTreeSourceMap, Let, RetryPolicy, TemplateString, Test, + TypeAlias, + }, +}; + +/// Builds an [`ItemTree`] and its [`ItemTreeSourceMap`] together. +#[derive(Debug, Default)] +pub struct ItemTreeBuilder { + tree: ItemTree, + source_map: ItemTreeSourceMap, + /// Collision tracker: `(ItemKind, hash)` -> next available index. + next_index: FxHashMap<(ItemKind, u16), u16>, +} + +impl ItemTreeBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Consume the builder, dropping `next_index` — it is meaningless once the + /// tree is complete. + pub fn finish(self) -> (ItemTree, ItemTreeSourceMap) { + (self.tree, self.source_map) + } + + /// Allocate a collision-resistant ID for an item. + fn alloc_id(&mut self, kind: ItemKind, name: &Name) -> LocalItemId { + let h = hash_name(name); + let index = self.next_index.entry((kind, h)).or_insert(0); + let id = LocalItemId::new(h, *index); + *index += 1; + id + } + + /// Allocate a function, recording its name span in the source map. + pub fn alloc_function(&mut self, f: &ast::FunctionDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Function, &f.name); + self.source_map.function_name_spans.insert(id, f.name_span); + let params = f + .params + .iter() + .map(|p| FunctionParam { + name: p.name.clone(), + type_expr: p.type_expr.clone(), + default: p.default.map(|expr| DefaultExprRef { function: id, expr }), + span: p.span, + }) + .collect(); + self.tree.functions.insert( + id, + Function { + name: f.name.clone(), + generic_params: f.generic_params.clone(), + generic_param_bounds: f.generic_param_bounds.clone(), + params, + defaults: f.defaults.clone(), + return_type: f.return_type.clone(), + throws: f.throws.clone(), + body: f.body.clone(), + declarative_meta: f.declarative_meta.clone(), + origin: f.origin, + docstring: f.docstring.clone(), + is_tagged_template_tag: f.is_tagged_template_tag, + span: f.span, + }, + ); + id + } + + /// Allocate a class, recording its field name spans in the source map. + pub fn alloc_class(&mut self, c: &ast::ClassDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Class, &c.name); + self.source_map + .class_field_spans + .insert(id, c.fields.iter().map(|f| f.name_span).collect()); + let fields = c + .fields + .iter() + .map(|f| ClassField { + name: f.name.clone(), + type_expr: f.type_expr.clone(), + attributes: f.attributes.iter().map(Attribute::from).collect(), + docstring: f.docstring.clone(), + }) + .collect(); + let implements = c + .implements + .iter() + .map(|b| ImplementsBlock { + target: b.target.clone(), + field_links: b + .field_links + .iter() + .map(InterfaceFieldLink::from_ast) + .collect(), + associated_type_bindings: b.associated_type_bindings.clone(), + is_out_of_body: b.is_out_of_body, + span: b.span, + }) + .collect(); + self.tree.classes.insert( + id, + Class { + name: c.name.clone(), + generic_params: c.generic_params.clone(), + generic_param_bounds: c.generic_param_bounds.clone(), + fields, + methods: Vec::new(), + implements, + attributes: c.attributes.iter().map(Attribute::from).collect(), + docstring: c.docstring.clone(), + span: c.span, + }, + ); + id + } + + /// Record the `implements I { … }` block a method was declared under, so + /// `default.()` calls in its body can resolve back to the interface. + /// Both halves are written together — they are always populated as a pair. + pub fn record_method_interface_target( + &mut self, + method: LocalItemId, + target: ast::TypeExpr, + associated_type_bindings: Vec, + ) { + self.tree.method_to_iface_target.insert(method, target); + self.tree + .method_to_iface_associated_type_bindings + .insert(method, associated_type_bindings); + } + + /// Attach method IDs to an already-allocated class. + pub fn set_class_methods( + &mut self, + class_id: LocalItemId, + methods: Vec>, + ) { + if let Some(class) = self.tree.classes.get_mut(&class_id) { + class.methods = methods; + } + } + + /// Allocate a stable id for an `implements` block and store it in the + /// unified `impls` map. `iface_head`/`for_head` seed the position-independent + /// id (impls have no declared name); the collision index disambiguates impls + /// that share both heads. Records the `class_to_impls` edge for `InClass`. + pub fn alloc_impl( + &mut self, + iface_head: &Name, + for_head: &Name, + block: ImplBlock, + ) -> LocalItemId { + let h = hash_impl_key(iface_head, for_head); + let index = self.next_index.entry((ItemKind::Impl, h)).or_insert(0); + let id = LocalItemId::new(h, *index); + *index += 1; + match &block.subject { + ImplSubject::InClass { class, .. } => { + self.tree.class_to_impls.entry(*class).or_default().push(id); + } + ImplSubject::Free { .. } => self.tree.free_impls.push(id), + } + self.tree.impls.insert(id, block); + id + } + + /// Allocate an enum, recording its variant name spans in the source map. + pub fn alloc_enum(&mut self, e: &ast::EnumDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Enum, &e.name); + self.source_map + .enum_variant_spans + .insert(id, e.variants.iter().map(|v| v.name_span).collect()); + let variants = e + .variants + .iter() + .map(|v| EnumVariant { + name: v.name.clone(), + attributes: v.attributes.iter().map(Attribute::from).collect(), + docstring: v.docstring.clone(), + }) + .collect(); + self.tree.enums.insert( + id, + Enum { + name: e.name.clone(), + variants, + attributes: e.attributes.iter().map(Attribute::from).collect(), + docstring: e.docstring.clone(), + span: e.span, + }, + ); + id + } + + /// Allocate an interface (BEP-044) in the `ItemTree`. + /// + /// `default_method_ids` are the `FunctionMarker` ids for any default + /// methods in the interface — those should be allocated separately via + /// `alloc_function` before this is called. + pub fn alloc_interface( + &mut self, + i: &ast::InterfaceDef, + default_method_ids: Vec>, + ) -> LocalItemId { + let id = self.alloc_id(ItemKind::Interface, &i.name); + self.source_map + .interface_field_spans + .insert(id, i.fields.iter().map(|f| f.name_span).collect()); + self.source_map + .interface_method_spans + .insert(id, i.required_methods.iter().map(|m| m.name_span).collect()); + let fields = i + .fields + .iter() + .map(|f| ClassField { + name: f.name.clone(), + type_expr: f.type_expr.clone(), + attributes: f.attributes.iter().map(Attribute::from).collect(), + docstring: f.docstring.clone(), + }) + .collect(); + let required_methods = i + .required_methods + .iter() + .map(|m| InterfaceMethodSig { + name: m.name.clone(), + generic_params: m.generic_params.clone(), + generic_param_bounds: m.generic_param_bounds.clone(), + params: m + .params + .iter() + .map(|p| FunctionParam { + name: p.name.clone(), + type_expr: p.type_expr.clone(), + default: None, + span: p.span, + }) + .collect(), + return_type: m.return_type.clone(), + throws: m.throws.clone(), + attributes: m.attributes.iter().map(Attribute::from).collect(), + docstring: m.docstring.clone(), + span: m.span, + }) + .collect(); + self.tree.interfaces.insert( + id, + Interface { + name: i.name.clone(), + generic_params: i.generic_params.clone(), + generic_param_bounds: i.generic_param_bounds.clone(), + requires: i.requires.clone(), + fields, + associated_types: i.associated_types.clone(), + default_methods: default_method_ids, + required_methods, + attributes: i.attributes.iter().map(Attribute::from).collect(), + docstring: i.docstring.clone(), + span: i.span, + }, + ); + id + } + + pub fn alloc_type_alias(&mut self, ta: &ast::TypeAliasDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::TypeAlias, &ta.name); + self.tree.type_aliases.insert( + id, + TypeAlias { + name: ta.name.clone(), + type_expr: ta.type_expr.clone(), + span: ta.span, + }, + ); + id + } + + pub fn alloc_client(&mut self, c: &ast::ClientDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Client, &c.name); + let provider = c + .config_items + .iter() + .find(|item| item.key.as_str() == "provider") + .map(|item| Name::new(item.value.trim().trim_matches('"'))); + let sub_client_names = c + .config_items + .iter() + .find(|item| item.key.as_str() == "options") + .map(|_| Vec::new()) // complex to parse; clients field is more relevant + .unwrap_or_default(); + let retry_policy_name = c + .config_items + .iter() + .find(|item| item.key.as_str() == "retry_policy") + .map(|item| Name::new(item.value.trim().trim_matches('"'))); + self.tree.clients.insert( + id, + Client { + name: c.name.clone(), + provider, + sub_client_names, + retry_policy_name, + round_robin_start: None, + }, + ); + id + } + + pub fn alloc_test(&mut self, t: &ast::TestDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Test, &t.name); + self.tree.tests.insert( + id, + Test { + name: t.name.clone(), + function_refs: t.function_refs.clone(), + args: t.args.clone(), + }, + ); + id + } + + pub fn alloc_template_string( + &mut self, + ts: &ast::TemplateStringDef, + ) -> LocalItemId { + let id = self.alloc_id(ItemKind::TemplateString, &ts.name); + let params = ts + .params + .iter() + .map(|p| FunctionParam { + name: p.name.clone(), + type_expr: p.type_expr.clone(), + default: None, + span: p.span, + }) + .collect(); + let body = ts.body.as_ref().map(|b| b.text.clone()); + self.tree.template_strings.insert( + id, + TemplateString { + name: ts.name.clone(), + params, + body, + span: ts.span, + }, + ); + id + } + + pub fn alloc_retry_policy( + &mut self, + rp: &ast::RetryPolicyDef, + ) -> LocalItemId { + let id = self.alloc_id(ItemKind::RetryPolicy, &rp.name); + let get_field = |key: &str| -> Option { + rp.config_items + .iter() + .find(|item| item.key.as_str() == key) + .map(|item| item.value.trim().to_string()) + }; + self.tree.retry_policies.insert( + id, + RetryPolicy { + name: rp.name.clone(), + max_retries: get_field("max_retries"), + initial_delay_ms: get_field("initial_delay_ms"), + multiplier: get_field("multiplier"), + max_delay_ms: get_field("max_delay_ms"), + }, + ); + id + } + + pub fn alloc_let(&mut self, l: &ast::LetDef) -> LocalItemId { + let id = self.alloc_id(ItemKind::Let, &l.name); + self.tree.lets.insert( + id, + Let { + name: l.name.clone(), + initializer: l.initializer.clone(), + origin: l.origin, + span: l.span, + name_span: l.name_span, + }, + ); + id + } +} diff --git a/baml_language/crates/baml_compiler2_hir/src/scope.rs b/baml_language/crates/baml_compiler2_hir/src/scope.rs index 874c475a041..bb0ae84c42b 100644 --- a/baml_language/crates/baml_compiler2_hir/src/scope.rs +++ b/baml_language/crates/baml_compiler2_hir/src/scope.rs @@ -9,6 +9,11 @@ use std::ops::Range; use baml_base::{Name, SourceFile}; use text_size::TextRange; +use crate::ids::{ + ClassMarker, ClientMarker, EnumMarker, FunctionMarker, ImplMarker, InterfaceMarker, LetMarker, + LocalItemId, RetryPolicyMarker, TemplateStringMarker, TestMarker, TypeAliasMarker, +}; + /// Dense sequential index into the per-file scope arena. /// `FileScopeId(0)` is always the Project scope (outermost). #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] @@ -82,6 +87,30 @@ pub enum ScopeKind { Let, } +/// The item a scope was opened for. +/// +/// The builder creates an item's scope in the same step that allocates the item, +/// so this link is recorded directly rather than reconstructed. Before it +/// existed, ~20 sites across TIR/MIR/LSP recovered it by comparing +/// `item.span == scope.range` — which made item spans load-bearing *semantic +/// identity* and blocked moving them into the source map. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ItemScopeOwner { + Function(LocalItemId), + Class(LocalItemId), + Enum(LocalItemId), + Interface(LocalItemId), + TypeAlias(LocalItemId), + TemplateString(LocalItemId), + Client(LocalItemId), + Test(LocalItemId), + RetryPolicy(LocalItemId), + Let(LocalItemId), + /// An out-of-body `implement I for T { … }` block, which opens a scope for + /// its own generic parameters. + Impl(LocalItemId), +} + /// A single scope node in the per-file scope tree. #[derive(Debug, Clone, PartialEq, Eq)] pub struct Scope { @@ -91,6 +120,13 @@ pub struct Scope { pub kind: ScopeKind, /// Optional name (packages, namespaces, items have names; blocks don't). pub name: Option, + /// The item this scope was opened for, if any. `None` for structural scopes + /// (Project/Package/Namespace/File) and for expression scopes (`Block`, + /// `Lambda`, `MatchArm`, …). + /// + /// Lives on `Scope` rather than in a parallel vec so it cannot drift out of + /// step with the scope arena. + pub owner: Option, /// Source range of this scope. Used by `scope_at_offset()` to find the /// innermost scope containing a cursor position. Structural scopes /// (Project, Package, Namespace) use the file's full range. diff --git a/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs b/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs index 4f2c7883b93..a69beffb431 100644 --- a/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs +++ b/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs @@ -34,11 +34,13 @@ pub enum PathResolution { Unknown, } +use rustc_hash::FxHashMap; + use crate::{ contributions::FileSymbolContributions, diagnostic::Hir2Diagnostic, item_tree::{ItemTree, ItemTreeSourceMap}, - scope::{FileScopeId, Scope, ScopeId, ScopeKind}, + scope::{FileScopeId, ItemScopeOwner, Scope, ScopeId, ScopeKind}, }; // ── DefinitionSite ─────────────────────────────────────────────────────────── @@ -222,6 +224,13 @@ pub struct FileSemanticIndex<'db> { /// Avoids repeated Salsa interning at query time. pub scope_ids: Vec>, + /// Item → the scope it opened. The inverse of `Scope::owner`. + /// + /// Recorded by the builder, which creates the scope in the same step that + /// allocates the item. Replaces the `item.span == scope.range` join that + /// consumers used to do, and is what lets item spans move into the source map. + pub item_scopes: FxHashMap, + /// Per-file item tree — maps `LocalItemId` to item data. pub item_tree: Arc, @@ -399,6 +408,16 @@ impl FileSemanticIndex<'_> { } } + /// The scope opened for `owner`, if it opened one. + pub fn item_scope(&self, owner: ItemScopeOwner) -> Option { + self.item_scopes.get(&owner).copied() + } + + /// The item `scope` was opened for, if it belongs to one. + pub fn scope_owner(&self, scope: FileScopeId) -> Option { + self.scopes.get(scope.index() as usize)?.owner + } + pub fn diagnostics(&self) -> &[Hir2Diagnostic] { self.extra .as_ref() diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index 2879fdf98fb..6c46839342f 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -582,3 +582,109 @@ fn editing_a_function_body_preserves_its_signature_data() { ); assert_eq!(before.params.len(), 2); } + +// ── Item ↔ scope index ─────────────────────────────────────────────────────── +// +// ~20 sites across TIR/MIR/LSP used to recover "the scope for this item" by +// scanning for a scope whose `range` equalled the item's `span`. That made item +// spans load-bearing *semantic identity*, which is what blocked moving them into +// the source map. The builder now records the link directly. + +/// The index must agree with the span-equality scan it replaces, for every +/// function in the file — otherwise migrating the call sites changes behavior. +#[test] +fn function_scope_index_agrees_with_the_span_join_it_replaces() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Add(x: int, y: int) -> int {\n x + y\n}\n\n\ + function Sub(x: int, y: int) -> int {\n x - y\n}\n\n\ + class Holder {\n n int\n\n function get(self) -> int {\n self.n\n }\n}\n", + ); + + let db = test_db.db(); + let index = baml_compiler2_hir::file_semantic_index(db, file); + let item_tree = baml_compiler2_hir::file_item_tree(db, file); + + assert!(!item_tree.functions.is_empty()); + + for (id, func) in item_tree.functions.iter() { + let loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *id); + + // The scan being retired. + let legacy = index + .scope_ids + .iter() + .copied() + .find(|scope_id| { + let scope = &index.scopes[scope_id.file_scope_id(db).index() as usize]; + matches!(scope.kind, baml_compiler2_hir::scope::ScopeKind::Function) + && scope.range == func.span + && scope.name.as_ref() == Some(&func.name) + }) + .map(|scope| scope.file_scope_id(db)); + + let indexed = baml_compiler2_hir::item_data::function_scope(db, loc) + .map(|scope| scope.file_scope_id(db)); + + assert_eq!( + legacy, indexed, + "index and span-join disagree for function `{}`", + func.name + ); + assert!(indexed.is_some(), "`{}` should have a scope", func.name); + } +} + +/// The point of the index: it is not derived from spans, so a whitespace edit +/// leaves it alone. +#[test] +fn function_scope_survives_a_whitespace_edit() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Add(x: int, y: int) -> int {\n x + y\n}\n", + ); + + let before = { + let db = test_db.db(); + let loc = function_loc(db, file, "Add"); + baml_compiler2_hir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) + }; + + file.set_text(test_db.db_mut()) + .to("// pushed down\nfunction Add(x: int, y: int) -> int {\n x + y\n}\n".to_string()); + + let after = { + let db = test_db.db(); + let loc = function_loc(db, file, "Add"); + baml_compiler2_hir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) + }; + + assert!(before.is_some()); + assert_eq!( + before, after, + "the item↔scope link must not depend on where the item sits in the file" + ); +} + +/// `scope_owner` is the inverse and must round-trip. +#[test] +fn scope_owner_round_trips() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db + .db_mut() + .add_file("test.baml", "function Add(x: int) -> int {\n x\n}\n"); + + let db = test_db.db(); + let loc = function_loc(db, file, "Add"); + let scope = baml_compiler2_hir::item_data::function_scope(db, loc).expect("scope"); + + assert_eq!( + baml_compiler2_hir::item_data::scope_owner(db, scope), + Some(baml_compiler2_hir::item_data::ScopeOwner::Function(loc)), + ); +} From 698e32a0bf8eac4db1b2d70f01aefa9bb8baf1a3 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 17:05:25 -0700 Subject: [PATCH 07/20] Make PPIR item tree canonical Nothing upstream should be using the HIR's item tree as it excludes synthetic items and thus has different indexing. All the firewall queries have been moved there for the downstream consumers. --- .../crates/baml_compiler2_hir/src/lib.rs | 1 - .../src/item_data.rs | 2 +- .../src/item_data/classes.rs | 14 ++-- .../src/item_data/clients.rs | 3 +- .../src/item_data/common.rs | 3 +- .../src/item_data/enumeration.rs | 3 +- .../src/item_data/enums.rs | 3 +- .../src/item_data/functions.rs | 12 ++-- .../src/item_data/impls.rs | 16 ++--- .../src/item_data/interfaces.rs | 8 +-- .../src/item_data/lets.rs | 3 +- .../src/item_data/retry_policies.rs | 3 +- .../src/item_data/scopes.rs | 4 +- .../src/item_data/template_strings.rs | 8 +-- .../src/item_data/test_items.rs | 3 +- .../src/item_data/type_aliases.rs | 5 +- .../crates/baml_compiler2_ppir/src/lib.rs | 9 ++- .../baml_tests/src/incremental/scenarios.rs | 66 ++++++++++++------- 18 files changed, 91 insertions(+), 75 deletions(-) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data.rs (95%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/classes.rs (97%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/clients.rs (95%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/common.rs (95%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/enumeration.rs (99%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/enums.rs (97%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/functions.rs (98%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/impls.rs (96%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/interfaces.rs (98%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/lets.rs (97%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/retry_policies.rs (95%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/scopes.rs (97%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/template_strings.rs (97%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/test_items.rs (92%) rename baml_language/crates/{baml_compiler2_hir => baml_compiler2_ppir}/src/item_data/type_aliases.rs (98%) diff --git a/baml_language/crates/baml_compiler2_hir/src/lib.rs b/baml_language/crates/baml_compiler2_hir/src/lib.rs index cbaaafcfc36..5d2f788e50a 100644 --- a/baml_language/crates/baml_compiler2_hir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_hir/src/lib.rs @@ -18,7 +18,6 @@ pub mod contributions; pub mod diagnostic; pub mod file_package; pub mod ids; -pub mod item_data; pub mod item_tree; pub mod loc; pub mod namespace; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data.rs similarity index 95% rename from baml_language/crates/baml_compiler2_hir/src/item_data.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data.rs index 5f9105cf5e6..245178be824 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data.rs @@ -13,7 +13,7 @@ //! would hand out *stale* spans forever after a whitespace-only edit. Hence //! every item is split in two: //! -//! - `*_data` — semantic, span-free, uses [`TypeRef`](crate::type_ref::TypeRef). +//! - `*_data` — semantic, span-free, uses [`TypeRef`](baml_compiler2_hir::type_ref::TypeRef). //! This is what type checking reads. //! - `*_source_map` — spans only. This is what diagnostics and the IDE read. //! diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs similarity index 97% rename from baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs index c089823e9be..47eddd292d2 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/classes.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs @@ -1,15 +1,15 @@ use baml_base::Name; -use text_size::TextRange; - -use crate::{ - item_data::common::{ - AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FieldData, - InterfaceFieldLinkData, InterfaceFieldLinkSourceMap, - }, +use baml_compiler2_hir::{ item_tree::Attribute, loc::{ClassLoc, FunctionLoc}, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; + +use crate::item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FieldData, InterfaceFieldLinkData, + InterfaceFieldLinkSourceMap, +}; /// Span-free semantic data for a `class` declaration. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/clients.rs similarity index 95% rename from baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/clients.rs index 5b6a773c4d3..602f457cc2c 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/clients.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/clients.rs @@ -1,6 +1,5 @@ use baml_base::Name; - -use crate::loc::ClientLoc; +use baml_compiler2_hir::loc::ClientLoc; /// Semantic data for a `client` declaration. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/common.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/common.rs similarity index 95% rename from baml_language/crates/baml_compiler2_hir/src/item_data/common.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/common.rs index e6cd6370bc8..2a63be1018f 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/common.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/common.rs @@ -5,10 +5,9 @@ //! here rather than in `classes`/`functions`. use baml_base::Name; +use baml_compiler2_hir::{item_tree::Attribute, type_ref::TypeRefId}; use text_size::TextRange; -use crate::{item_tree::Attribute, type_ref::TypeRefId}; - #[derive(Debug, Clone, PartialEq, Eq)] pub struct FunctionParamData { pub name: Name, diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs similarity index 99% rename from baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs index 3a411cc6c01..8e58b92d032 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/enumeration.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs @@ -11,8 +11,7 @@ //! this query's value unless it actually reorders declarations. use baml_base::SourceFile; - -use crate::loc::{ +use baml_compiler2_hir::loc::{ ClassLoc, ClientLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, TemplateStringLoc, TestLoc, TypeAliasLoc, }; diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enums.rs similarity index 97% rename from baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/enums.rs index ff7a7da51dc..b8846eafaf3 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/enums.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enums.rs @@ -1,8 +1,7 @@ use baml_base::Name; +use baml_compiler2_hir::{item_tree::Attribute, loc::EnumLoc}; use text_size::TextRange; -use crate::{item_tree::Attribute, loc::EnumLoc}; - /// Span-free semantic data for an `enum` declaration. /// /// Enums carry no type expressions, so there is no `TypeRefStore` here. diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs similarity index 98% rename from baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index cd66a623a12..b8b22f9c587 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -1,14 +1,14 @@ use baml_base::Name; use baml_compiler2_ast as ast; -use text_size::TextRange; - -use crate::{ - item_data::common::{ - AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FunctionParamData, - }, +use baml_compiler2_hir::{ loc::FunctionLoc, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; + +use crate::item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, FunctionParamData, +}; /// Span-free semantic data for a function's *signature*. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/impls.rs similarity index 96% rename from baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/impls.rs index 037c2db2455..52ad8b8d2c1 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/impls.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/impls.rs @@ -1,14 +1,14 @@ use baml_base::Name; -use text_size::TextRange; - -use crate::{ - item_data::common::{ - AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, InterfaceFieldLinkData, - InterfaceFieldLinkSourceMap, - }, +use baml_compiler2_hir::{ loc::{ClassLoc, FunctionLoc, ImplLoc}, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; + +use crate::item_data::common::{ + AssociatedTypeBindingData, AssociatedTypeBindingSourceMap, InterfaceFieldLinkData, + InterfaceFieldLinkSourceMap, +}; /// What an `implements` block applies to. /// @@ -89,7 +89,7 @@ fn lower<'db>( db: &'db dyn crate::Db, block: ImplLoc<'db>, ) -> (ImplBlockData<'db>, ImplBlockSourceMap) { - use crate::item_tree::ImplSubject; + use baml_compiler2_hir::item_tree::ImplSubject; let file = block.file(db); let item_tree = crate::file_item_tree(db, file); diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/interfaces.rs similarity index 98% rename from baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/interfaces.rs index b5c8290b56a..ae5890cbb03 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/interfaces.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/interfaces.rs @@ -1,12 +1,12 @@ use baml_base::Name; -use text_size::TextRange; - -use crate::{ - item_data::common::{FieldData, FunctionParamData}, +use baml_compiler2_hir::{ item_tree::Attribute, loc::{FunctionLoc, InterfaceLoc}, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; + +use crate::item_data::common::{FieldData, FunctionParamData}; /// Span-free semantic data for an `interface` declaration. #[derive(Debug, Clone, PartialEq, Eq, salsa::Update)] diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/lets.rs similarity index 97% rename from baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/lets.rs index 6787985f1e2..d63a85ec3ee 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/lets.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/lets.rs @@ -1,9 +1,8 @@ use baml_base::Name; use baml_compiler2_ast as ast; +use baml_compiler2_hir::loc::LetLoc; use text_size::TextRange; -use crate::loc::LetLoc; - /// Span-free semantic data for a top-level `let` binding. /// /// The initializer is deliberately absent: it has its own query diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/retry_policies.rs similarity index 95% rename from baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/retry_policies.rs index 5855dbd4576..08a9d83ce11 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/retry_policies.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/retry_policies.rs @@ -1,6 +1,5 @@ use baml_base::Name; - -use crate::loc::RetryPolicyLoc; +use baml_compiler2_hir::loc::RetryPolicyLoc; /// Semantic data for a `retry_policy` declaration. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/scopes.rs similarity index 97% rename from baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/scopes.rs index fb1b33f59c7..87b60d1136a 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/scopes.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/scopes.rs @@ -1,13 +1,13 @@ //! Item ↔ scope queries. //! //! The builder opens an item's scope in the same step that allocates the item, so -//! the link is recorded directly ([`Scope::owner`](crate::scope::Scope::owner)). +//! the link is recorded directly ([`Scope::owner`](baml_compiler2_hir::scope::Scope::owner)). //! Before that, consumers recovered it by scanning for a scope whose `range` //! equalled the item's `span` — which made item spans load-bearing *semantic //! identity* rather than diagnostic metadata, and blocked moving them into the //! source map. These queries replace that scan. -use crate::{ +use baml_compiler2_hir::{ loc::{ ClassLoc, ClientLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, TemplateStringLoc, TestLoc, TypeAliasLoc, diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/template_strings.rs similarity index 97% rename from baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/template_strings.rs index 54a5a5f31da..c5c27673977 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/template_strings.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/template_strings.rs @@ -1,11 +1,11 @@ use baml_base::Name; -use text_size::TextRange; - -use crate::{ - item_data::common::FunctionParamData, +use baml_compiler2_hir::{ loc::TemplateStringLoc, type_ref::{TypeRefBuilder, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; + +use crate::item_data::common::FunctionParamData; /// Span-free semantic data for a `template_string` declaration. #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/test_items.rs similarity index 92% rename from baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/test_items.rs index da3215183ca..3c47343efc0 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/test_items.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/test_items.rs @@ -1,6 +1,5 @@ use baml_base::Name; - -use crate::{item_tree::TestArgValue, loc::TestLoc}; +use baml_compiler2_hir::{item_tree::TestArgValue, loc::TestLoc}; /// Semantic data for a `test` declaration. /// diff --git a/baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs similarity index 98% rename from baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs rename to baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs index e425a3ff863..992c1c2df48 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_data/type_aliases.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs @@ -1,10 +1,9 @@ use baml_base::Name; -use text_size::TextRange; - -use crate::{ +use baml_compiler2_hir::{ loc::TypeAliasLoc, type_ref::{TypeRefBuilder, TypeRefId, TypeRefSourceMap, TypeRefStore}, }; +use text_size::TextRange; /// Span-free semantic data for a `type X = ...` declaration. #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/baml_language/crates/baml_compiler2_ppir/src/lib.rs b/baml_language/crates/baml_compiler2_ppir/src/lib.rs index d4ddfee48d4..d5da7e5d3bd 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/lib.rs @@ -9,6 +9,7 @@ //! **No union simplification in PPIR.** Deferred to TIR. pub mod expand; +pub mod item_data; pub mod ty; use std::sync::Arc; @@ -17,7 +18,7 @@ use baml_base::{Name, SourceFile, attr::TyAttrValue}; use baml_compiler2_ast as ast; use baml_compiler2_hir::{ contributions::FileSymbolContributions, - item_tree::ItemTree, + item_tree::{ItemTree, ItemTreeSourceMap}, namespace::{NameConflict, NamespaceId, NamespaceItems}, package::{PackageId, PackageItems, PackageItemsExtra}, scope::ScopeId, @@ -612,6 +613,12 @@ pub fn file_item_tree(db: &dyn Db, file: SourceFile) -> Arc { Arc::clone(&index.item_tree) } +/// Canonical item-tree source map (original + *$stream types). +pub fn file_item_tree_source_map(db: &dyn Db, file: SourceFile) -> Arc { + let index = file_semantic_index(db, file); + Arc::clone(&index.item_tree_source_map) +} + /// Canonical function body — uses PPIR's item tree (includes synthetic companions). /// /// TIR should call this instead of `baml_compiler2_hir::body::function_body` diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index 6c46839342f..2d19dc0cddb 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -362,7 +362,7 @@ fn type_alias_loc<'db>( file: SourceFile, name: &str, ) -> baml_compiler2_hir::loc::TypeAliasLoc<'db> { - let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); let (id, _) = item_tree .type_aliases .iter() @@ -383,8 +383,8 @@ fn whitespace_edit_preserves_item_data_but_moves_spans() { let db = test_db.db(); let loc = type_alias_loc(db, file, "Ids"); ( - baml_compiler2_hir::item_data::type_alias_data(db, loc).clone(), - baml_compiler2_hir::item_data::type_alias_source_map(db, loc).clone(), + baml_compiler2_ppir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_ppir::item_data::type_alias_source_map(db, loc).clone(), ) }; @@ -396,8 +396,8 @@ fn whitespace_edit_preserves_item_data_but_moves_spans() { let db = test_db.db(); let loc = type_alias_loc(db, file, "Ids"); ( - baml_compiler2_hir::item_data::type_alias_data(db, loc).clone(), - baml_compiler2_hir::item_data::type_alias_source_map(db, loc).clone(), + baml_compiler2_ppir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_ppir::item_data::type_alias_source_map(db, loc).clone(), ) }; @@ -422,7 +422,7 @@ fn semantic_edit_changes_item_data() { let data_before = { let db = test_db.db(); let loc = type_alias_loc(db, file, "Ids"); - baml_compiler2_hir::item_data::type_alias_data(db, loc).clone() + baml_compiler2_ppir::item_data::type_alias_data(db, loc).clone() }; file.set_text(test_db.db_mut()) @@ -431,7 +431,7 @@ fn semantic_edit_changes_item_data() { let data_after = { let db = test_db.db(); let loc = type_alias_loc(db, file, "Ids"); - baml_compiler2_hir::item_data::type_alias_data(db, loc).clone() + baml_compiler2_ppir::item_data::type_alias_data(db, loc).clone() }; assert_ne!(data_before, data_after); @@ -442,7 +442,7 @@ fn class_loc<'db>( file: SourceFile, name: &str, ) -> baml_compiler2_hir::loc::ClassLoc<'db> { - let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); let (id, _) = item_tree .classes .iter() @@ -456,7 +456,7 @@ fn function_loc<'db>( file: SourceFile, name: &str, ) -> baml_compiler2_hir::loc::FunctionLoc<'db> { - let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); let (id, _) = item_tree .functions .iter() @@ -474,8 +474,8 @@ type ClassFingerprint = ( baml_base::Name, Vec>, baml_compiler2_hir::type_ref::TypeRefStore, - Vec, - Vec, + Vec, + Vec, Vec, ); @@ -484,7 +484,7 @@ fn class_fingerprint( file: SourceFile, name: &str, ) -> ClassFingerprint { - let data = baml_compiler2_hir::item_data::class_data(db, class_loc(db, file, name)); + let data = baml_compiler2_ppir::item_data::class_data(db, class_loc(db, file, name)); ( data.name.clone(), data.generic_param_bounds.clone(), @@ -565,7 +565,7 @@ fn editing_a_function_body_preserves_its_signature_data() { let before = { let db = test_db.db(); - baml_compiler2_hir::item_data::function_data(db, function_loc(db, file, "Add")).clone() + baml_compiler2_ppir::item_data::function_data(db, function_loc(db, file, "Add")).clone() }; file.set_text(test_db.db_mut()) @@ -573,7 +573,7 @@ fn editing_a_function_body_preserves_its_signature_data() { let after = { let db = test_db.db(); - baml_compiler2_hir::item_data::function_data(db, function_loc(db, file, "Add")).clone() + baml_compiler2_ppir::item_data::function_data(db, function_loc(db, file, "Add")).clone() }; assert_eq!( @@ -596,18 +596,36 @@ fn editing_a_function_body_preserves_its_signature_data() { fn function_scope_index_agrees_with_the_span_join_it_replaces() { let mut test_db = IncrementalTestDb::new(); + // Includes a declarative LLM function: those synthesize companion functions, + // and `scope_at_offset`'s own docs note companions "share the same span as + // their parent" — so the span-join was ambiguous exactly there. If the index + // and the scan disagree for any of these, migrating the call sites is a + // behavior change and needs to be handled deliberately. let file = test_db.db_mut().add_file( "test.baml", "function Add(x: int, y: int) -> int {\n x + y\n}\n\n\ function Sub(x: int, y: int) -> int {\n x - y\n}\n\n\ - class Holder {\n n int\n\n function get(self) -> int {\n self.n\n }\n}\n", + class Holder {\n n int\n\n function get(self) -> int {\n self.n\n }\n}\n\n\ + function Greet(name: string) -> string {\n \ + client GPT4\n prompt #\"Hello {{name}}\"#\n}\n", ); let db = test_db.db(); - let index = baml_compiler2_hir::file_semantic_index(db, file); - let item_tree = baml_compiler2_hir::file_item_tree(db, file); + let index = baml_compiler2_ppir::file_semantic_index(db, file); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - assert!(!item_tree.functions.is_empty()); + // Guard against a vacuous test: the declarative `Greet` must actually have + // synthesized companions, or the ambiguous case is not being exercised. + let companions = item_tree + .functions + .values() + .filter(|f| !matches!(f.origin, baml_compiler2_ast::FunctionOrigin::UserDefined)) + .count(); + assert!( + companions > 0, + "fixture should synthesize companions; got {} functions, none synthetic", + item_tree.functions.len() + ); for (id, func) in item_tree.functions.iter() { let loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *id); @@ -625,7 +643,7 @@ fn function_scope_index_agrees_with_the_span_join_it_replaces() { }) .map(|scope| scope.file_scope_id(db)); - let indexed = baml_compiler2_hir::item_data::function_scope(db, loc) + let indexed = baml_compiler2_ppir::item_data::function_scope(db, loc) .map(|scope| scope.file_scope_id(db)); assert_eq!( @@ -651,7 +669,7 @@ fn function_scope_survives_a_whitespace_edit() { let before = { let db = test_db.db(); let loc = function_loc(db, file, "Add"); - baml_compiler2_hir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) + baml_compiler2_ppir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) }; file.set_text(test_db.db_mut()) @@ -660,7 +678,7 @@ fn function_scope_survives_a_whitespace_edit() { let after = { let db = test_db.db(); let loc = function_loc(db, file, "Add"); - baml_compiler2_hir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) + baml_compiler2_ppir::item_data::function_scope(db, loc).map(|scope| scope.file_scope_id(db)) }; assert!(before.is_some()); @@ -681,10 +699,10 @@ fn scope_owner_round_trips() { let db = test_db.db(); let loc = function_loc(db, file, "Add"); - let scope = baml_compiler2_hir::item_data::function_scope(db, loc).expect("scope"); + let scope = baml_compiler2_ppir::item_data::function_scope(db, loc).expect("scope"); assert_eq!( - baml_compiler2_hir::item_data::scope_owner(db, scope), - Some(baml_compiler2_hir::item_data::ScopeOwner::Function(loc)), + baml_compiler2_ppir::item_data::scope_owner(db, scope), + Some(baml_compiler2_ppir::item_data::ScopeOwner::Function(loc)), ); } From 576e9bf88d50d5dba3d7c397033b9b0115665cb1 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 21:00:44 -0700 Subject: [PATCH 08/20] Use firewall queries in TIR --- .../baml_compiler2_hir/src/item_tree.rs | 5 + .../src/item_tree/builder.rs | 33 +- .../src/item_tree/functions.rs | 21 + .../src/item_data/functions.rs | 42 + .../crates/baml_compiler2_tir/src/builder.rs | 12 +- .../crates/baml_compiler2_tir/src/callable.rs | 24 +- .../baml_compiler2_tir/src/inference.rs | 1048 ++++++++--------- .../baml_compiler2_tir/src/lower_type_expr.rs | 13 +- .../baml_tests/src/incremental/scenarios.rs | 146 +++ 9 files changed, 757 insertions(+), 587 deletions(-) diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs index 2c3f33d2cf1..dded064c96e 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs @@ -83,6 +83,11 @@ pub struct ItemTree { pub method_to_iface_target: FxHashMap, ast::TypeExpr>, pub method_to_iface_associated_type_bindings: FxHashMap, Vec>, + + /// Method → owning item. Inverse of `Class::methods` / + /// `Interface::default_methods` / a free impl's `ImplBlock::methods`; + /// absent for top-level functions. See [`MethodOwner`]. + pub method_owners: FxHashMap, MethodOwner>, } // ── Index impls ─────────────────────────────────────────────────────────────── diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs index cadd6d88cb5..f7268e25004 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs @@ -21,8 +21,8 @@ use crate::{ item_tree::{ Attribute, Class, ClassField, Client, DefaultExprRef, Enum, EnumVariant, Function, FunctionParam, ImplBlock, ImplSubject, ImplementsBlock, Interface, InterfaceFieldLink, - InterfaceMethodSig, ItemTree, ItemTreeSourceMap, Let, RetryPolicy, TemplateString, Test, - TypeAlias, + InterfaceMethodSig, ItemTree, ItemTreeSourceMap, Let, MethodOwner, RetryPolicy, + TemplateString, Test, TypeAlias, }, }; @@ -153,17 +153,32 @@ impl ItemTreeBuilder { .insert(method, associated_type_bindings); } - /// Attach method IDs to an already-allocated class. + /// Attach method IDs to an already-allocated class, recording each method's + /// owner. Membership and ownership are written by the same call so they + /// cannot drift apart. pub fn set_class_methods( &mut self, class_id: LocalItemId, methods: Vec>, ) { + for method in &methods { + self.record_method_owner(*method, MethodOwner::Class(class_id)); + } if let Some(class) = self.tree.classes.get_mut(&class_id) { class.methods = methods; } } + /// A method belongs to exactly one item; a second recording is a builder bug + /// (e.g. one function id handed to two owners). + fn record_method_owner(&mut self, method: LocalItemId, owner: MethodOwner) { + let previous = self.tree.method_owners.insert(method, owner); + debug_assert!( + previous.is_none(), + "method owner recorded twice: {previous:?} then {owner:?}" + ); + } + /// Allocate a stable id for an `implements` block and store it in the /// unified `impls` map. `iface_head`/`for_head` seed the position-independent /// id (impls have no declared name); the collision index disambiguates impls @@ -181,8 +196,15 @@ impl ItemTreeBuilder { match &block.subject { ImplSubject::InClass { class, .. } => { self.tree.class_to_impls.entry(*class).or_default().push(id); + // No owner recording: in-body impl methods are flattened into + // `Class::methods`, so `set_class_methods` owns them. + } + ImplSubject::Free { .. } => { + self.tree.free_impls.push(id); + for method in &block.methods { + self.record_method_owner(*method, MethodOwner::FreeImpl(id)); + } } - ImplSubject::Free { .. } => self.tree.free_impls.push(id), } self.tree.impls.insert(id, block); id @@ -233,6 +255,9 @@ impl ItemTreeBuilder { self.source_map .interface_method_spans .insert(id, i.required_methods.iter().map(|m| m.name_span).collect()); + for method in &default_method_ids { + self.record_method_owner(*method, MethodOwner::Interface(id)); + } let fields = i .fields .iter() diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs index 559ccc9ec49..ea5362fa7f9 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs @@ -40,6 +40,27 @@ pub struct Function { pub span: TextRange, } +/// The item a method belongs to. +/// +/// Recorded by the `ItemTreeBuilder` at the same call that establishes +/// membership (`set_class_methods` / `alloc_interface` / `alloc_impl`), so it +/// cannot drift from the forward lists. Before this existed, ~24 sites across +/// TIR/MIR/emit/HIR answered "who owns this method?" by scanning +/// `classes.values().find(|c| c.methods.contains(&id))` and friends — O(items) +/// per lookup, and with the class/interface halves drifting between copies. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MethodOwner { + /// A class-level method, *including* methods declared inside an in-body + /// `implements I { … }` block — those are flattened into `Class::methods`, + /// and their impl relationship lives in `method_to_iface_target` / + /// `class_to_impls`. + Class(LocalItemId), + /// An interface default method. + Interface(LocalItemId), + /// A method of an out-of-body `implements<…> I for T { … }` block. + FreeImpl(LocalItemId), +} + /// A function parameter entry in the `ItemTree`. #[derive(Debug, Clone, PartialEq, Eq)] pub struct FunctionParam { diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index b8b22f9c587..55028b0675f 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -67,6 +67,48 @@ pub fn function_source_map<'db>( lower(db, function).1 } +/// The item a method belongs to, as a `Loc`. +/// +/// Mirrors `item_tree::MethodOwner` (see its docs for the ownership rules — +/// notably, in-body `implements I { … }` methods are owned by their *class*). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, salsa::Update)] +pub enum MethodOwner<'db> { + Class(baml_compiler2_hir::loc::ClassLoc<'db>), + Interface(baml_compiler2_hir::loc::InterfaceLoc<'db>), + /// An out-of-body `implements<…> I for T { … }` block. + FreeImpl(baml_compiler2_hir::loc::ImplLoc<'db>), +} + +/// The item `method` belongs to, or `None` for a top-level function. +/// +/// Replaces the `classes.values().find(|c| c.methods.contains(&id))` scan +/// family — O(items) per lookup, with class-only and class-plus-interface +/// copies drifting apart across crates. +#[salsa::tracked] +pub fn method_owner<'db>( + db: &'db dyn crate::Db, + method: FunctionLoc<'db>, +) -> Option> { + use baml_compiler2_hir::item_tree; + + let file = method.file(db); + let owner = *crate::file_item_tree(db, file) + .method_owners + .get(&method.id(db))?; + + Some(match owner { + item_tree::MethodOwner::Class(id) => { + MethodOwner::Class(baml_compiler2_hir::loc::ClassLoc::new(db, file, id)) + } + item_tree::MethodOwner::Interface(id) => { + MethodOwner::Interface(baml_compiler2_hir::loc::InterfaceLoc::new(db, file, id)) + } + item_tree::MethodOwner::FreeImpl(id) => { + MethodOwner::FreeImpl(baml_compiler2_hir::loc::ImplLoc::new(db, file, id)) + } + }) +} + /// The interface target a method was declared under, when it sits inside an /// `implements I { … }` block — plus the associated-type bindings from that /// block's header. `None` for class-level methods and for interface default diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index 3025f0abc05..7f9370f680a 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -10057,12 +10057,12 @@ impl<'db> TypeInferenceBuilder<'db> { // `stream_llm_function`). Definition::Let(let_loc) => { let db = self.context.db(); - let file = let_loc.file(db); - let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let let_data = &item_tree[let_loc.id(db)]; - let index = baml_compiler2_ppir::file_semantic_index(db, file); - let fsi = index.scope_at_offset(let_data.span.start(), Some(&let_data.name)); - let scope_id = index.scope_ids[fsi.index() as usize]; + let Some(scope_id) = baml_compiler2_ppir::item_data::let_scope(db, let_loc) + else { + return Ty::Unknown { + attr: TyAttr::default(), + }; + }; let inference = crate::inference::infer_scope_types(db, scope_id); let body = baml_compiler2_hir::body::let_body(db, let_loc); if let baml_compiler2_hir::body::LetBody::Expr(expr_body) = body.as_ref() diff --git a/baml_language/crates/baml_compiler2_tir/src/callable.rs b/baml_language/crates/baml_compiler2_tir/src/callable.rs index 820baab0595..dc6e5ef64be 100644 --- a/baml_language/crates/baml_compiler2_tir/src/callable.rs +++ b/baml_language/crates/baml_compiler2_tir/src/callable.rs @@ -2,9 +2,7 @@ use std::collections::{BTreeSet, HashMap}; use baml_base::Name; use baml_compiler2_ast::ExprBody; -use baml_compiler2_hir::{ - body::FunctionBody, file_package, loc::FunctionLoc, package::PackageId, scope::ScopeKind, -}; +use baml_compiler2_hir::{body::FunctionBody, file_package, loc::FunctionLoc, package::PackageId}; use rustc_hash::FxHashMap; use crate::{ @@ -131,23 +129,6 @@ fn callable_key<'db>(db: &'db dyn crate::Db, function: FunctionLoc<'db>) -> Name throw_set_key(&namespace, &short_name) } -fn function_scope_id<'db>( - db: &'db dyn crate::Db, - function: FunctionLoc<'db>, -) -> Option> { - let file = function.file(db); - let index = baml_compiler2_ppir::file_semantic_index(db, file); - let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let func_data = &item_tree[function.id(db)]; - - index.scope_ids.iter().copied().find(|scope_id| { - let scope = &index.scopes[scope_id.file_scope_id(db).index() as usize]; - matches!(scope.kind, ScopeKind::Function) - && scope.range == func_data.span - && scope.name.as_ref() == Some(&func_data.name) - }) -} - fn named_callee_key<'db>( db: &'db dyn crate::Db, pkg_id: PackageId<'db>, @@ -411,7 +392,8 @@ pub fn callable_throws<'db>(db: &'db dyn crate::Db, function: FunctionLoc<'db>) match baml_compiler2_ppir::function_body(db, function).as_ref() { FunctionBody::Expr(body) => { - let Some(scope_id) = function_scope_id(db, function) else { + let Some(scope_id) = baml_compiler2_ppir::item_data::function_scope(db, function) + else { return Ty::Unknown { attr: TyAttr::default(), }; diff --git a/baml_language/crates/baml_compiler2_tir/src/inference.rs b/baml_language/crates/baml_compiler2_tir/src/inference.rs index 26f708a1c16..ff1ad516401 100644 --- a/baml_language/crates/baml_compiler2_tir/src/inference.rs +++ b/baml_language/crates/baml_compiler2_tir/src/inference.rs @@ -22,7 +22,7 @@ use baml_compiler2_ast::{ use baml_compiler2_hir::{ body::{FunctionBody, LetBody}, contributions::Definition, - loc::{ClassLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, TypeAliasLoc}, + loc::{ClassLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, TypeAliasLoc}, package::{PackageId, PackageItems}, scope::{FileScopeId, ScopeId, ScopeKind}, semantic_index::{BindingId, BindingKind}, @@ -728,6 +728,7 @@ fn apply_generic_env( #[derive(Clone, Copy)] struct GenericLookupContext<'a, 'db> { + db: &'db dyn crate::Db, index: &'a baml_compiler2_hir::semantic_index::FileSemanticIndex<'db>, item_tree: &'a baml_compiler2_hir::item_tree::ItemTree, } @@ -788,28 +789,30 @@ fn free_impl_generics( } } -/// The out-of-body `implement … for …` block whose method list contains `func_data`, as its -/// declared generics, if any. A method is identified by its source span (unique per declaration). -/// Reads the unified `impls` store via the `free_impls` index. +/// The out-of-body `implement … for …` block whose method list contains `func_id`, as its +/// declared generics, if any. Reads the unified `impls` store via the `free_impls` index. +/// +/// Membership is by id. It used to be by source span, which cannot distinguish a function +/// from its companions — all of `f`, `f$stream` and `f$parse_stream` carry one span. fn enclosing_impl_generics_for_func( item_tree: &baml_compiler2_hir::item_tree::ItemTree, - func_data: &baml_compiler2_hir::item_tree::Function, + func_id: baml_compiler2_hir::ids::LocalItemId, ) -> Option<(Vec, Vec>)> { item_tree.free_impls.iter().find_map(|impl_id| { let block = item_tree.impls.get(impl_id)?; block .methods - .iter() - .any(|&mid| item_tree[mid].span == func_data.span) + .contains(&func_id) .then(|| free_impl_generics(block)) }) } -fn generic_env_for_function_data( - ctx: GenericLookupContext<'_, '_>, +fn generic_env_for_function_data<'db>( + ctx: GenericLookupContext<'_, 'db>, function_scope: &baml_compiler2_hir::scope::Scope, - func_data: &baml_compiler2_hir::item_tree::Function, + func_loc: FunctionLoc<'db>, ) -> GenericEnv { + let func_data = &ctx.item_tree[func_loc.id(ctx.db)]; let mut env = GenericEnv::from_params(func_data.generic_params.clone()); env.add_bounds_for_declared_params(&func_data.generic_params, &func_data.generic_param_bounds); // A method in a generic `implements Iface for Target` block sees the @@ -818,7 +821,7 @@ fn generic_env_for_function_data( // scope. Mirror the `ScopeKind::Function` arm: impl generics take the place // of parent-type generics so a nested lambda body can resolve them too. if let Some((generic_params, generic_param_bounds)) = - enclosing_impl_generics_for_func(ctx.item_tree, func_data) + enclosing_impl_generics_for_func(ctx.item_tree, func_loc.id(ctx.db)) { env.prepend_declared(&generic_params, &generic_param_bounds); } else { @@ -836,11 +839,14 @@ fn enclosing_function_generic_env_from_let( let scope = &ctx.index.scopes[fsi.index() as usize]; match scope.kind { ScopeKind::Function => { - let func_data = - ctx.item_tree.functions.values().find(|fd| { - fd.span == scope.range && scope.name.as_ref() == Some(&fd.name) - })?; - return Some(generic_env_for_function_data(ctx, scope, func_data)); + let owner = baml_compiler2_ppir::item_data::scope_owner( + ctx.db, + ctx.index.scope_ids[fsi.index() as usize], + )?; + let baml_compiler2_ppir::item_data::ScopeOwner::Function(func_loc) = owner else { + return None; + }; + return Some(generic_env_for_function_data(ctx, scope, func_loc)); } ScopeKind::Let => current = scope.parent, _ => current = scope.parent, @@ -1655,468 +1661,456 @@ pub fn infer_scope_types<'db>( // Both checks are required to disambiguate companion functions that // share the parent's span. let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let mut found = false; - for (local_id, func_data) in &item_tree.functions { - if func_data.span == scope.range && scope.name.as_ref() == Some(&func_data.name) { - let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *local_id); - let body = baml_compiler2_ppir::function_body(db, func_loc); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); - - let enclosing_impl = item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - block.methods.contains(local_id).then_some(block) - }); + // The HIR builder records which item opened each scope, so the owner is + // read directly rather than recovered by matching `item.span == scope.range` + // — a join that could not tell a function from its companions, which share + // one span. Template strings also open `ScopeKind::Function` scopes; those + // have a non-Function owner and fall through to the `else`. + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Function(func_loc)) = + baml_compiler2_ppir::item_data::scope_owner(db, scope_id) + { + let local_id = &func_loc.id(db); + let func_data = &item_tree[func_loc.id(db)]; + let body = baml_compiler2_ppir::function_body(db, func_loc); + let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + + let enclosing_impl = item_tree.free_impls.iter().find_map(|impl_id| { + let block = item_tree.impls.get(impl_id)?; + block.methods.contains(local_id).then_some(block) + }); - let mut env = GenericEnv::from_params(sig.user_generic_params.clone()); - env.params - .extend(sig.synthetic_effect_params.iter().cloned()); - report_duplicate_generic_params( - &builder, - &sig.user_generic_params, - func_data.span, - ); - if let Some(imp) = enclosing_impl { - let (impl_generic_params, impl_generic_bounds) = free_impl_generics(imp); - for mp in &sig.user_generic_params { - if impl_generic_params.iter().any(|cp| cp == mp) { - builder.report_at_span( - crate::infer_context::TirTypeError::TypeParamShadowedImplParam { - param_name: mp.clone(), - }, - func_data.span, - ); - } - } - env.prepend_declared(&impl_generic_params, &impl_generic_bounds); - } else if let Some(parent) = parent_type_generic_env( - GenericLookupContext { - index, - item_tree: &item_tree, - }, - scope.parent, - ) { - for mp in &sig.user_generic_params { - if parent.params.iter().any(|cp| cp == mp) { - builder.report_at_span( - crate::infer_context::TirTypeError::TypeParamShadowed { - param_name: mp.clone(), - type_name: parent.type_name.clone(), - owner: parent.owner, - }, - func_data.span, - ); - } + let mut env = GenericEnv::from_params(sig.user_generic_params.clone()); + env.params + .extend(sig.synthetic_effect_params.iter().cloned()); + report_duplicate_generic_params(&builder, &sig.user_generic_params, func_data.span); + if let Some(imp) = enclosing_impl { + let (impl_generic_params, impl_generic_bounds) = free_impl_generics(imp); + for mp in &sig.user_generic_params { + if impl_generic_params.iter().any(|cp| cp == mp) { + builder.report_at_span( + crate::infer_context::TirTypeError::TypeParamShadowedImplParam { + param_name: mp.clone(), + }, + func_data.span, + ); } - env.prepend_declared(&parent.params, &parent.bounds); } - // BEP-044 (Self-as-type-variable): inside an interface's own - // method, `self` is a `Self` type variable bound by the - // interface instead of the interface existential. Register it - // in the shared generic env so the normal bound-aware member - // resolution path handles `Self`. - let interface_self_bound: Option = if enclosing_impl - .is_none() - { - scope.parent.and_then(|parent_idx| { - let parent = &index.scopes[parent_idx.index() as usize]; - if !matches!(parent.kind, ScopeKind::Class) { - return None; - } - let cn = parent.name.as_ref()?; - let def = pkg_items.lookup_type(&pkg_info.namespace_path, cn)?; - if !matches!( - def, - baml_compiler2_hir::contributions::Definition::Interface(_) - ) { - return None; - } - let qtn = crate::lower_type_expr::qualify_def(db, def, cn); - let iface = item_tree.interfaces.values().find(|i| &i.name == cn)?; - let args = iface - .generic_params - .iter() - .map(|p| Ty::TypeVar(p.clone(), TyAttr::default())) - .collect(); - // `Self`'s bound is the interface as a *constraint* (a - // `baml_type::Interface`, not a `Ty::Interface` existential): it pins - // none of its associated types — they are `Self`'s own, abstract in a - // default body, so `Self.Item` stays symbolic. Empty bindings, never - // `Ty::Error` placeholders which would wrongly pin `Item = Error`. - Some(baml_type::Interface::new(qtn, args, Vec::new())) - }) - } else { - None - }; - if let Some(bound) = interface_self_bound.clone() { - let self_param = Name::new("Self"); - if !env.params.iter().any(|p| p == &self_param) { - env.params.push(self_param.clone()); + env.prepend_declared(&impl_generic_params, &impl_generic_bounds); + } else if let Some(parent) = parent_type_generic_env( + GenericLookupContext { + db, + index, + item_tree: &item_tree, + }, + scope.parent, + ) { + for mp in &sig.user_generic_params { + if parent.params.iter().any(|cp| cp == mp) { + builder.report_at_span( + crate::infer_context::TirTypeError::TypeParamShadowed { + param_name: mp.clone(), + type_name: parent.type_name.clone(), + owner: parent.owner, + }, + func_data.span, + ); } - env.add_concrete_bound(self_param, bound); } - env.add_bounds_for_declared_params( - &func_data.generic_params, - &func_data.generic_param_bounds, - ); - apply_generic_env( - db, - &mut builder, - pkg_items, - &pkg_info.namespace_path, - &env, - func_data.span, - ); - if let Some(sm) = baml_compiler2_ppir::function_body_source_map(db, func_loc) { - builder.set_body_source_map(sm); + env.prepend_declared(&parent.params, &parent.bounds); + } + // BEP-044 (Self-as-type-variable): inside an interface's own + // method, `self` is a `Self` type variable bound by the + // interface instead of the interface existential. Register it + // in the shared generic env so the normal bound-aware member + // resolution path handles `Self`. + let interface_self_bound: Option = if enclosing_impl.is_none() + { + scope.parent.and_then(|parent_idx| { + let parent = &index.scopes[parent_idx.index() as usize]; + if !matches!(parent.kind, ScopeKind::Class) { + return None; + } + let cn = parent.name.as_ref()?; + let def = pkg_items.lookup_type(&pkg_info.namespace_path, cn)?; + if !matches!( + def, + baml_compiler2_hir::contributions::Definition::Interface(_) + ) { + return None; + } + let qtn = crate::lower_type_expr::qualify_def(db, def, cn); + let iface = item_tree.interfaces.values().find(|i| &i.name == cn)?; + let args = iface + .generic_params + .iter() + .map(|p| Ty::TypeVar(p.clone(), TyAttr::default())) + .collect(); + // `Self`'s bound is the interface as a *constraint* (a + // `baml_type::Interface`, not a `Ty::Interface` existential): it pins + // none of its associated types — they are `Self`'s own, abstract in a + // default body, so `Self.Item` stays symbolic. Empty bindings, never + // `Ty::Error` placeholders which would wrongly pin `Item = Error`. + Some(baml_type::Interface::new(qtn, args, Vec::new())) + }) + } else { + None + }; + if let Some(bound) = interface_self_bound.clone() { + let self_param = Name::new("Self"); + if !env.params.iter().any(|p| p == &self_param) { + env.params.push(self_param.clone()); } - builder.set_auto_derived(matches!( - func_data.origin, - ast::FunctionOrigin::AutoDerive - )); - // BEP-044: if this function lives inside an - // `implements I { ... }` block, attach `I`'s QTN so - // `default.(...)` resolves against I's - // contract. - if let Some(target) = item_tree.method_to_iface_target.get(local_id) - && let baml_compiler2_ast::TypeExprKind::Path { segments, .. } = - &target.kind - && let Some((head, name)) = segments - .split_last() - .map(|(last, head)| (head, last.clone())) + env.add_concrete_bound(self_param, bound); + } + env.add_bounds_for_declared_params( + &func_data.generic_params, + &func_data.generic_param_bounds, + ); + apply_generic_env( + db, + &mut builder, + pkg_items, + &pkg_info.namespace_path, + &env, + func_data.span, + ); + if let Some(sm) = baml_compiler2_ppir::function_body_source_map(db, func_loc) { + builder.set_body_source_map(sm); + } + builder + .set_auto_derived(matches!(func_data.origin, ast::FunctionOrigin::AutoDerive)); + // BEP-044: if this function lives inside an + // `implements I { ... }` block, attach `I`'s QTN so + // `default.(...)` resolves against I's + // contract. + if let Some(target) = item_tree.method_to_iface_target.get(local_id) + && let baml_compiler2_ast::TypeExprKind::Path { segments, .. } = &target.kind + && let Some((head, name)) = segments + .split_last() + .map(|(last, head)| (head, last.clone())) + { + let lookup_ns: &[Name] = if head.is_empty() { + &pkg_info.namespace_path + } else { + head + }; + if let Some(def) = pkg_items.lookup_type(lookup_ns, &name) + && let baml_compiler2_hir::contributions::Definition::Interface(_) = def { - let lookup_ns: &[Name] = if head.is_empty() { - &pkg_info.namespace_path - } else { - head - }; - if let Some(def) = pkg_items.lookup_type(lookup_ns, &name) - && let baml_compiler2_hir::contributions::Definition::Interface(_) = def - { - let qtn = crate::lower_type_expr::qualify_def(db, def, &name); - builder.set_implements_block_interface(Some(qtn)); - } + let qtn = crate::lower_type_expr::qualify_def(db, def, &name); + builder.set_implements_block_interface(Some(qtn)); } + } - if let FunctionBody::Expr(expr_body) = body.as_ref() { - // The method's in-scope interface bounds (its own params, the - // enclosing class/interface's, and `Self`'s constraint) so a - // `T.member` / `Self.member` projection in the receiver pattern, - // interface arguments, or a binding value resolves nominally. - let env_bounds = - env_interface_bounds(db, pkg_items, &pkg_info.namespace_path, &env); - // Determine enclosing class name for `self` parameter - // resolution and BEP-044 `Self`-type substitution. - let enclosing_class_name: Option = - scope.parent.and_then(|parent_idx| { - let parent = &index.scopes[parent_idx.index() as usize]; - if matches!(parent.kind, ScopeKind::Class) { - parent.name.clone() - } else { - None - } - }); - // `Self`'s type for this body — resolved through the lowering context - // below, never a bare-name substitution: the rigid `Self` type variable - // (interface's own default method), the impl's receiver pattern, or the - // enclosing class's full receiver type (`Foo`, carrying its generics). - let self_ty: Option = if interface_self_bound.is_some() { - Some(crate::self_type::self_type_for_interface_default()) - } else if let Some(imp) = enclosing_impl - && let baml_compiler2_hir::item_tree::ImplSubject::Free { - for_target, - .. - } = &imp.subject - { - let mut diags = Vec::new(); - Some(crate::lower_type_expr::lower_type_expr( - for_target, - &crate::lower_type_expr::ScopeCtx { - db, - package_items: pkg_items, - ns_context: &pkg_info.namespace_path, - generic_params: &env.params, - bounds: &env_bounds, - self_ty: None, - }, - &mut diags, - )) + if let FunctionBody::Expr(expr_body) = body.as_ref() { + // The method's in-scope interface bounds (its own params, the + // enclosing class/interface's, and `Self`'s constraint) so a + // `T.member` / `Self.member` projection in the receiver pattern, + // interface arguments, or a binding value resolves nominally. + let env_bounds = + env_interface_bounds(db, pkg_items, &pkg_info.namespace_path, &env); + // Determine enclosing class name for `self` parameter + // resolution and BEP-044 `Self`-type substitution. + let enclosing_class_name: Option = scope.parent.and_then(|parent_idx| { + let parent = &index.scopes[parent_idx.index() as usize]; + if matches!(parent.kind, ScopeKind::Class) { + parent.name.clone() } else { - enclosing_class_name.as_ref().and_then(|cn| { - let baml_compiler2_hir::contributions::Definition::Class(class_loc) = - pkg_items.lookup_type(&pkg_info.namespace_path, cn)? - else { - return None; - }; - let class_file = class_loc.file(db); - let class_pkg = - baml_compiler2_hir::file_package::file_package(db, class_file); - let class_tree = baml_compiler2_hir::file_item_tree(db, class_file); - let class_data = class_tree.classes.get(&class_loc.id(db))?; - Some(crate::self_type::self_type_for_class( - class_data, - &class_pkg.namespace_path, - class_pkg.package.clone(), - )) - }) - }; - // The interface bound `Self.Assoc` projects through, as a constraint - // (interface's own default method); `None` when `Self` is concrete. - let self_bound: Option = interface_self_bound; - - let sig_sm = baml_compiler2_ppir::elaborated_function_signature_source_map( - db, func_loc, - ); - let mut type_bindings: FxHashMap = - type_bindings_for_params(&env.params); - if let Some(target) = item_tree.method_to_iface_target.get(local_id) - && let Some(iface_loc) = crate::interfaces::resolve_path_to_interface( + None + } + }); + // `Self`'s type for this body — resolved through the lowering context + // below, never a bare-name substitution: the rigid `Self` type variable + // (interface's own default method), the impl's receiver pattern, or the + // enclosing class's full receiver type (`Foo`, carrying its generics). + let self_ty: Option = if interface_self_bound.is_some() { + Some(crate::self_type::self_type_for_interface_default()) + } else if let Some(imp) = enclosing_impl + && let baml_compiler2_hir::item_tree::ImplSubject::Free { + for_target, .. + } = &imp.subject + { + let mut diags = Vec::new(); + Some(crate::lower_type_expr::lower_type_expr( + for_target, + &crate::lower_type_expr::ScopeCtx { db, - target, - pkg_items, - &pkg_info.namespace_path, - ) - { - let iface_file = iface_loc.file(db); - let iface_tree = baml_compiler2_hir::file_item_tree(db, iface_file); - if let Some(iface_data) = iface_tree.interfaces.get(&iface_loc.id(db)) { - let mut iface_type_bindings = type_bindings.clone(); - if let baml_compiler2_ast::TypeExprKind::Path { - generic_args, .. - } = &target.kind + package_items: pkg_items, + ns_context: &pkg_info.namespace_path, + generic_params: &env.params, + bounds: &env_bounds, + self_ty: None, + }, + &mut diags, + )) + } else { + enclosing_class_name.as_ref().and_then(|cn| { + let baml_compiler2_hir::contributions::Definition::Class(class_loc) = + pkg_items.lookup_type(&pkg_info.namespace_path, cn)? + else { + return None; + }; + let class_file = class_loc.file(db); + let class_pkg = + baml_compiler2_hir::file_package::file_package(db, class_file); + let class_tree = baml_compiler2_hir::file_item_tree(db, class_file); + let class_data = class_tree.classes.get(&class_loc.id(db))?; + Some(crate::self_type::self_type_for_class( + class_data, + &class_pkg.namespace_path, + class_pkg.package.clone(), + )) + }) + }; + // The interface bound `Self.Assoc` projects through, as a constraint + // (interface's own default method); `None` when `Self` is concrete. + let self_bound: Option = interface_self_bound; + + let sig_sm = + baml_compiler2_ppir::elaborated_function_signature_source_map(db, func_loc); + let mut type_bindings: FxHashMap = + type_bindings_for_params(&env.params); + if let Some(target) = item_tree.method_to_iface_target.get(local_id) + && let Some(iface_loc) = crate::interfaces::resolve_path_to_interface( + db, + target, + pkg_items, + &pkg_info.namespace_path, + ) + { + let iface_file = iface_loc.file(db); + let iface_tree = baml_compiler2_hir::file_item_tree(db, iface_file); + if let Some(iface_data) = iface_tree.interfaces.get(&iface_loc.id(db)) { + let mut iface_type_bindings = type_bindings.clone(); + if let baml_compiler2_ast::TypeExprKind::Path { generic_args, .. } = + &target.kind + { + for (param, arg) in + iface_data.generic_params.iter().zip(generic_args) { - for (param, arg) in - iface_data.generic_params.iter().zip(generic_args) - { - let mut arg_diags = Vec::new(); - let ty = { - let generic_params: Vec<_> = - type_bindings.keys().cloned().collect(); - crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr( - arg, - &crate::lower_type_expr::ScopeCtx { - db, - package_items: pkg_items, - ns_context: &pkg_info.namespace_path, - generic_params: &generic_params, - bounds: &env_bounds, - self_ty: None, - }, - &mut arg_diags, - ), - &type_bindings, - ) - }; - for diag in arg_diags { - builder.report_at_span(diag, target.span); - } - iface_type_bindings.insert(param.clone(), ty.clone()); - type_bindings.entry(param.clone()).or_insert(ty); - } - } - let explicit_bindings = item_tree - .method_to_iface_associated_type_bindings - .get(local_id) - .cloned() - .unwrap_or_default(); - for assoc in &iface_data.associated_types { - if let Some(binding) = - explicit_bindings.iter().find(|b| b.name == assoc.name) - && let Some(te) = &binding.type_expr - { - // Lower the binding's type through a context that resolves - // `Self`, then substitute the in-scope generics / - // associated types accumulated so far. - let binding_generic_params: Vec = + let mut arg_diags = Vec::new(); + let ty = { + let generic_params: Vec<_> = type_bindings.keys().cloned().collect(); - let binding_ctx = crate::lower_type_expr::ScopeCtx { - db, - package_items: pkg_items, - ns_context: &pkg_info.namespace_path, - generic_params: &binding_generic_params, - bounds: &env_bounds, - self_ty: self_ty.clone(), - }; - let mut binding_diags = Vec::new(); - let ty = crate::generics::substitute_ty( + crate::generics::substitute_ty( &crate::lower_type_expr::lower_type_expr( - te, - &binding_ctx, - &mut binding_diags, + arg, + &crate::lower_type_expr::ScopeCtx { + db, + package_items: pkg_items, + ns_context: &pkg_info.namespace_path, + generic_params: &generic_params, + bounds: &env_bounds, + self_ty: None, + }, + &mut arg_diags, ), &type_bindings, - ); - for diag in binding_diags { - builder.report_at_span(diag, te.span); - } - // Into `iface_type_bindings` only — the binding realizes - // later defaults, but the bare name is NOT in scope - // (banned: the method must write `Self.Item`). - iface_type_bindings.insert(assoc.name.clone(), ty); - continue; - } - if let Some((default_ty, _diags)) = - crate::interfaces::interface_associated_type_default( - db, - iface_loc, - assoc.name.clone(), ) - { - // The default is lowered once (symbolic `Self`) by the - // shared query; substitute the generics / associated - // types accumulated so far. `Self` stays symbolic here - // (an interface method body's receiver is rigid), so a - // Self-referencing default resolves to `(Self as I).X`. - // Diagnostics surface at the interface declaration. - // Into `iface_type_bindings` only (see the explicit- - // binding arm above): bare names are banned. - let ty = crate::generics::substitute_ty( - &default_ty, - &iface_type_bindings, - ); - iface_type_bindings.insert(assoc.name.clone(), ty); + }; + for diag in arg_diags { + builder.report_at_span(diag, target.span); } + iface_type_bindings.insert(param.clone(), ty.clone()); + type_bindings.entry(param.clone()).or_insert(ty); } } - } - // Inside an interface's own default body, associated types are - // deliberately NOT registered as bare names: a bare `Item` is - // banned everywhere — the body writes `Self.Item`, which lowers - // through the `Self` bound installed below. - builder.set_type_bindings(type_bindings.clone()); - // Lower body type annotations through the shared context: `Self` and - // `Self.Assoc` resolve via `self_ty` and the `Self` bound; interface / - // method generics and associated types then substitute in. A bare - // associated name is an in-scope type variable (it is a `type_bindings` - // key) substituted to its symbolic projection, matching `Self.Assoc`. - let body_generic_params: Vec = - type_bindings.keys().cloned().collect(); - let mut body_bounds = - crate::lower_type_expr::function_in_scope_generic_param_bounds( - db, func_loc, - ) - .clone(); - if let Some(bound) = &self_bound { - body_bounds.insert(Name::new("Self"), vec![bound.clone()]); - } - let ctx = crate::lower_type_expr::ScopeCtx { - db, - package_items: pkg_items, - ns_context: &pkg_info.namespace_path, - generic_params: &body_generic_params, - bounds: &body_bounds, - self_ty: self_ty.clone(), - }; - let lower_with_self = |te: &baml_compiler2_ast::TypeExpr, - diags: &mut Vec< - crate::infer_context::TirTypeError, - >| { - crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr(te, &ctx, diags), - &type_bindings, - ) - }; - - // Get declared return type - let return_ty = sig - .return_type - .as_ref() - .map(|te| { - let span = sig_sm.return_type_span.unwrap_or(func_data.span); - let mut diags = Vec::new(); - let ty = lower_with_self(te, &mut diags); - for diag in diags { - builder.report_at_span(diag, span); - } - builder.validate_type_generic_bounds_at_span(span, &ty); - ty - }) - .unwrap_or(Ty::Unknown { - attr: TyAttr::default(), - }); - - // Set declared return type for return statement checking - builder.set_return_type(return_ty.clone()); - - // Add parameter bindings as locals - for (i, param) in sig.params.iter().enumerate() { - let param_type_span = sig_sm - .param_type_spans - .get(i) - .copied() - .flatten() - .or_else(|| sig_sm.param_spans.get(i).copied()) + let explicit_bindings = item_tree + .method_to_iface_associated_type_bindings + .get(local_id) + .cloned() .unwrap_or_default(); - let mut param_ty_validated = false; - let param_ty = if param.name.as_str() == "self" - && matches!( - param.ty.kind, - baml_compiler2_ast::TypeExprKind::Unknown { .. } - ) { - // `self`'s type is the method's `Self` receiver, resolved once - // above (rigid `Self` var for an interface's own default method, - // the impl's receiver pattern, or the enclosing class's `Foo`). - self_ty.clone().unwrap_or(Ty::Unknown { - attr: TyAttr::default(), - }) - } else { - param_ty_validated = true; - let mut param_diags = Vec::new(); - let ty = lower_with_self(¶m.ty, &mut param_diags); - for diag in param_diags { - builder.report_at_span(diag, param_type_span); + for assoc in &iface_data.associated_types { + if let Some(binding) = + explicit_bindings.iter().find(|b| b.name == assoc.name) + && let Some(te) = &binding.type_expr + { + // Lower the binding's type through a context that resolves + // `Self`, then substitute the in-scope generics / + // associated types accumulated so far. + let binding_generic_params: Vec = + type_bindings.keys().cloned().collect(); + let binding_ctx = crate::lower_type_expr::ScopeCtx { + db, + package_items: pkg_items, + ns_context: &pkg_info.namespace_path, + generic_params: &binding_generic_params, + bounds: &env_bounds, + self_ty: self_ty.clone(), + }; + let mut binding_diags = Vec::new(); + let ty = crate::generics::substitute_ty( + &crate::lower_type_expr::lower_type_expr( + te, + &binding_ctx, + &mut binding_diags, + ), + &type_bindings, + ); + for diag in binding_diags { + builder.report_at_span(diag, te.span); + } + // Into `iface_type_bindings` only — the binding realizes + // later defaults, but the bare name is NOT in scope + // (banned: the method must write `Self.Item`). + iface_type_bindings.insert(assoc.name.clone(), ty); + continue; + } + if let Some((default_ty, _diags)) = + crate::interfaces::interface_associated_type_default( + db, + iface_loc, + assoc.name.clone(), + ) + { + // The default is lowered once (symbolic `Self`) by the + // shared query; substitute the generics / associated + // types accumulated so far. `Self` stays symbolic here + // (an interface method body's receiver is rigid), so a + // Self-referencing default resolves to `(Self as I).X`. + // Diagnostics surface at the interface declaration. + // Into `iface_type_bindings` only (see the explicit- + // binding arm above): bare names are banned. + let ty = crate::generics::substitute_ty( + &default_ty, + &iface_type_bindings, + ); + iface_type_bindings.insert(assoc.name.clone(), ty); } - builder.validate_type_generic_bounds_at_span(param_type_span, &ty); - ty - }; - if !param_ty_validated { - builder.validate_type_generic_bounds_at_span( - param_type_span, - ¶m_ty, - ); } - builder.add_local(param.name.clone(), param_ty.clone()); - builder.param_types.push((param.name.clone(), param_ty)); } + } + // Inside an interface's own default body, associated types are + // deliberately NOT registered as bare names: a bare `Item` is + // banned everywhere — the body writes `Self.Item`, which lowers + // through the `Self` bound installed below. + builder.set_type_bindings(type_bindings.clone()); + // Lower body type annotations through the shared context: `Self` and + // `Self.Assoc` resolve via `self_ty` and the `Self` bound; interface / + // method generics and associated types then substitute in. A bare + // associated name is an in-scope type variable (it is a `type_bindings` + // key) substituted to its symbolic projection, matching `Self.Assoc`. + let body_generic_params: Vec = type_bindings.keys().cloned().collect(); + let mut body_bounds = + crate::lower_type_expr::function_in_scope_generic_param_bounds( + db, func_loc, + ) + .clone(); + if let Some(bound) = &self_bound { + body_bounds.insert(Name::new("Self"), vec![bound.clone()]); + } + let ctx = crate::lower_type_expr::ScopeCtx { + db, + package_items: pkg_items, + ns_context: &pkg_info.namespace_path, + generic_params: &body_generic_params, + bounds: &body_bounds, + self_ty: self_ty.clone(), + }; + let lower_with_self = |te: &baml_compiler2_ast::TypeExpr, + diags: &mut Vec< + crate::infer_context::TirTypeError, + >| { + crate::generics::substitute_ty( + &crate::lower_type_expr::lower_type_expr(te, &ctx, diags), + &type_bindings, + ) + }; - let param_types = builder.param_types.clone(); - let parameter_defaults = - baml_compiler2_ppir::function_parameter_defaults(db, func_loc); - builder.check_function_parameter_defaults( - &func_data.params, - ¶meter_defaults, - ¶m_types, - ); + // Get declared return type + let return_ty = sig + .return_type + .as_ref() + .map(|te| { + let span = sig_sm.return_type_span.unwrap_or(func_data.span); + let mut diags = Vec::new(); + let ty = lower_with_self(te, &mut diags); + for diag in diags { + builder.report_at_span(diag, span); + } + builder.validate_type_generic_bounds_at_span(span, &ty); + ty + }) + .unwrap_or(Ty::Unknown { + attr: TyAttr::default(), + }); - // Check root expression against declared return type - if let Some(root_expr) = expr_body.root_expr { - builder.check_expr(root_expr, expr_body, &return_ty); + // Set declared return type for return statement checking + builder.set_return_type(return_ty.clone()); + + // Add parameter bindings as locals + for (i, param) in sig.params.iter().enumerate() { + let param_type_span = sig_sm + .param_type_spans + .get(i) + .copied() + .flatten() + .or_else(|| sig_sm.param_spans.get(i).copied()) + .unwrap_or_default(); + let mut param_ty_validated = false; + let param_ty = if param.name.as_str() == "self" + && matches!( + param.ty.kind, + baml_compiler2_ast::TypeExprKind::Unknown { .. } + ) { + // `self`'s type is the method's `Self` receiver, resolved once + // above (rigid `Self` var for an interface's own default method, + // the impl's receiver pattern, or the enclosing class's `Foo`). + self_ty.clone().unwrap_or(Ty::Unknown { + attr: TyAttr::default(), + }) + } else { + param_ty_validated = true; + let mut param_diags = Vec::new(); + let ty = lower_with_self(¶m.ty, &mut param_diags); + for diag in param_diags { + builder.report_at_span(diag, param_type_span); + } + builder.validate_type_generic_bounds_at_span(param_type_span, &ty); + ty + }; + if !param_ty_validated { + builder + .validate_type_generic_bounds_at_span(param_type_span, ¶m_ty); } + builder.add_local(param.name.clone(), param_ty.clone()); + builder.param_types.push((param.name.clone(), param_ty)); + } - // Validate declared `throws` against effective escaping throws. - // Auto-derived methods (synthesized `to_json` / - // `from_json`) use a conservative pre-baked throws - // clause; the body's actual escaping throws can be - // wider when fields have malformed/unknown types. - // The user can't fix the synthesized contract, so - // skip the entire check for auto-derive bodies. - let is_auto_derive = - matches!(func_data.origin, ast::FunctionOrigin::AutoDerive); - if !is_auto_derive { - builder.check_throws_contract( - expr_body, - sig.throws.as_ref(), - sig_sm.throws_type_span, - func_data.span, - true, - ); - } + let param_types = builder.param_types.clone(); + let parameter_defaults = + baml_compiler2_ppir::function_parameter_defaults(db, func_loc); + builder.check_function_parameter_defaults( + &func_data.params, + ¶meter_defaults, + ¶m_types, + ); + + // Check root expression against declared return type + if let Some(root_expr) = expr_body.root_expr { + builder.check_expr(root_expr, expr_body, &return_ty); + } + + // Validate declared `throws` against effective escaping throws. + // Auto-derived methods (synthesized `to_json` / + // `from_json`) use a conservative pre-baked throws + // clause; the body's actual escaping throws can be + // wider when fields have malformed/unknown types. + // The user can't fix the synthesized contract, so + // skip the entire check for auto-derive bodies. + let is_auto_derive = + matches!(func_data.origin, ast::FunctionOrigin::AutoDerive); + if !is_auto_derive { + builder.check_throws_contract( + expr_body, + sig.throws.as_ref(), + sig_sm.throws_type_span, + func_data.span, + true, + ); } - found = true; - break; } - } - if !found { + } else { // Template strings create ScopeKind::Function scopes but are // stored in item_tree.template_strings, not item_tree.functions. // They have no expression body to type-check, so skip silently. @@ -2225,16 +2219,17 @@ pub fn infer_scope_types<'db>( // Walk ancestors to find a Function or Let scope that has a body. 'ancestor_walk: for ancestor_fsi in index.ancestor_scopes(file_scope) { let ancestor_scope = &index.scopes[ancestor_fsi.index() as usize]; + let ancestor_owner = baml_compiler2_ppir::item_data::scope_owner( + db, + index.scope_ids[ancestor_fsi.index() as usize], + ); match &ancestor_scope.kind { ScopeKind::Function => { - // Find the function by span + name in the item tree - for func_data in item_tree.functions.values() { - if func_data.span != ancestor_scope.range { - continue; - } - if ancestor_scope.name.as_ref() != Some(&func_data.name) { - continue; - } + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Function( + ancestor_func, + )) = ancestor_owner + { + let func_data = &item_tree[ancestor_func.id(db)]; // Get the function body from item_tree (includes source map) if let Some(baml_compiler2_ast::FunctionBodyDef::Expr( ref func_body, @@ -2265,11 +2260,12 @@ pub fn infer_scope_types<'db>( let env = extend_env_with_lambda_generics( generic_env_for_function_data( GenericLookupContext { + db, index, item_tree: &item_tree, }, ancestor_scope, - func_data, + ancestor_func, ), func_def, ); @@ -2306,15 +2302,9 @@ pub fn infer_scope_types<'db>( } } ScopeKind::Let => { - // Find the let binding by span + name in the item tree - for (local_id, let_data) in &item_tree.lets { - if let_data.span != ancestor_scope.range { - continue; - } - if ancestor_scope.name.as_ref() != Some(&let_data.name) { - continue; - } - let let_loc = LetLoc::new(db, file, *local_id); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Let(let_loc)) = + ancestor_owner + { let body = baml_compiler2_hir::body::let_body(db, let_loc); let source_map_opt = baml_compiler2_hir::body::let_body_source_map(db, let_loc); @@ -2341,6 +2331,7 @@ pub fn infer_scope_types<'db>( let env = extend_env_with_lambda_generics( enclosing_function_generic_env_from_let( GenericLookupContext { + db, index, item_tree: &item_tree, }, @@ -2388,10 +2379,9 @@ pub fn infer_scope_types<'db>( } ScopeKind::Class => { let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - for (local_id, class_data) in &item_tree.classes { - if class_data.span != scope.range || scope.name.as_ref() != Some(&class_data.name) { - continue; - } + let scope_item = baml_compiler2_ppir::item_data::scope_owner(db, scope_id); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Class(class_loc)) = scope_item { + let class_data = &item_tree[class_loc.id(db)]; report_duplicate_generic_params( &builder, &class_data.generic_params, @@ -2410,20 +2400,17 @@ pub fn infer_scope_types<'db>( &env, class_data.span, ); - let class_loc = ClassLoc::new(db, file, *local_id); let resolved = resolve_class_fields(db, class_loc); for (field, (_, ty, _)) in class_data.fields.iter().zip(resolved.fields.iter()) { if let Some(type_expr) = &field.type_expr { builder.validate_type_generic_bounds_at_span(type_expr.span, ty); } } - break; } - for (local_id, iface_data) in &item_tree.interfaces { - if iface_data.span != scope.range || scope.name.as_ref() != Some(&iface_data.name) { - continue; - } - let iface_loc = InterfaceLoc::new(db, file, *local_id); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Interface(iface_loc)) = + scope_item + { + let iface_data = &item_tree[iface_loc.id(db)]; report_duplicate_generic_params( &builder, &iface_data.generic_params, @@ -2787,16 +2774,15 @@ pub fn infer_scope_types<'db>( // the impl site when one class field is forced to satisfy two of them (E0116), and // ambiguous unqualified access is a use-site check. There is deliberately no // declaration-level inherited-field-conflict check. - break; } } ScopeKind::Let => { // Top-level let binding — find the matching let in the item tree // and type-infer its initializer expression. - let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - for (local_id, let_data) in &item_tree.lets { - if let_data.span == scope.range && scope.name.as_ref() == Some(&let_data.name) { - let let_loc = LetLoc::new(db, file, *local_id); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Let(let_loc)) = + baml_compiler2_ppir::item_data::scope_owner(db, scope_id) + { + { let body = baml_compiler2_hir::body::let_body(db, let_loc); if let LetBody::Expr(expr_body) = body.as_ref() { @@ -2805,22 +2791,19 @@ pub fn infer_scope_types<'db>( builder.infer_expr(root_expr, expr_body); } } - break; } } } ScopeKind::TypeAlias => { let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - for (local_id, alias_data) in &item_tree.type_aliases { - if alias_data.span != scope.range || scope.name.as_ref() != Some(&alias_data.name) { - continue; - } - let alias_loc = TypeAliasLoc::new(db, file, *local_id); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::TypeAlias(alias_loc)) = + baml_compiler2_ppir::item_data::scope_owner(db, scope_id) + { + let alias_data = &item_tree[alias_loc.id(db)]; let resolved = resolve_type_alias(db, alias_loc); if let Some(type_expr) = &alias_data.type_expr { builder.validate_type_generic_bounds_at_span(type_expr.span, &resolved.ty); } - break; } } _ => { @@ -3240,72 +3223,47 @@ pub fn render_scope_diagnostics<'db>( // then use find_lambda_by_span to get the lambda's own source map. let lambda_span = scope.range; let mut found_sm = None; - 'ancestor: for ancestor_fsi in index.ancestor_scopes(file_scope) { - let ancestor = &index.scopes[ancestor_fsi.index() as usize]; - match &ancestor.kind { - ScopeKind::Function => { - for func_data in item_tree.functions.values() { - if func_data.span == ancestor.range - && ancestor.name.as_ref() == Some(&func_data.name) - { - if let Some(baml_compiler2_ast::FunctionBodyDef::Expr( - ref body, - ref sm, - )) = func_data.body - { - if let Some((_, _, lambda_sm, _)) = - find_lambda_by_span(body, sm, lambda_span) - { - found_sm = Some(lambda_sm.clone()); - } - } - break 'ancestor; + for ancestor_fsi in index.ancestor_scopes(file_scope) { + let ancestor_scope = index.scope_ids[ancestor_fsi.index() as usize]; + let Some(owner) = baml_compiler2_ppir::item_data::scope_owner(db, ancestor_scope) + else { + continue; + }; + // The first Function/Let ancestor owns the body the lambda lives in. + let body_and_map = match owner { + baml_compiler2_ppir::item_data::ScopeOwner::Function(func_loc) => { + match &item_tree[func_loc.id(db)].body { + Some(baml_compiler2_ast::FunctionBodyDef::Expr(body, sm)) => { + Some((body.clone(), sm.clone())) } + _ => None, } } - ScopeKind::Let => { - for let_data in item_tree.lets.values() { - if let_data.span == ancestor.range - && ancestor.name.as_ref() == Some(&let_data.name) - { - if let Some((ref body, ref sm)) = let_data.initializer { - if let Some((_, _, lambda_sm, _)) = - find_lambda_by_span(body, sm, lambda_span) - { - found_sm = Some(lambda_sm.clone()); - } - } - break 'ancestor; - } - } + baml_compiler2_ppir::item_data::ScopeOwner::Let(let_loc) => { + item_tree[let_loc.id(db)].initializer.clone() } - _ => {} + _ => continue, + }; + if let Some((body, sm)) = body_and_map + && let Some((_, _, lambda_sm, _)) = find_lambda_by_span(&body, &sm, lambda_span) + { + found_sm = Some(lambda_sm.clone()); } + break; } found_sm } _ => { - // For Function/Let scopes, find the source map directly. - // Use PPIR's canonical version so PPIR-synthesized functions are found. - item_tree - .functions - .iter() - .find(|(_, f)| f.span == scope.range && scope.name.as_ref() == Some(&f.name)) - .and_then(|(local_id, _)| { - let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *local_id); + // Function/Let scopes: the owner is recorded, so no scan is needed. + match baml_compiler2_ppir::item_data::scope_owner(db, scope_id) { + Some(baml_compiler2_ppir::item_data::ScopeOwner::Function(func_loc)) => { baml_compiler2_ppir::function_body_source_map(db, func_loc) - }) - .or_else(|| { - // Also search let bindings. - item_tree - .lets - .iter() - .find(|(_, l)| l.span == scope.range) - .and_then(|(local_id, _)| { - let let_loc = baml_compiler2_hir::loc::LetLoc::new(db, file, *local_id); - baml_compiler2_hir::body::let_body_source_map(db, let_loc) - }) - }) + } + Some(baml_compiler2_ppir::item_data::ScopeOwner::Let(let_loc)) => { + baml_compiler2_hir::body::let_body_source_map(db, let_loc) + } + _ => None, + } } }; diff --git a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs index 0d424171233..1bdce62579e 100644 --- a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs +++ b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs @@ -1715,17 +1715,8 @@ mod tests { .copied() .find(|&id| tree[id].name.as_str() == "first") .expect("`first` default method"); - let func_data = &tree[fn_id]; - - let index = baml_compiler2_ppir::file_semantic_index(&db, file); - let scope_id = index - .scope_ids - .iter() - .copied() - .find(|scope_id| { - let scope = &index.scopes[scope_id.file_scope_id(&db).index() as usize]; - scope.range == func_data.span && scope.name.as_ref() == Some(&func_data.name) - }) + let fn_loc = baml_compiler2_hir::loc::FunctionLoc::new(&db, file, fn_id); + let scope_id = baml_compiler2_ppir::item_data::function_scope(&db, fn_loc) .expect("`first`'s body scope"); let inference = crate::inference::infer_scope_types(&db, scope_id); diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index 2d19dc0cddb..cbe66ad0bef 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -706,3 +706,149 @@ fn scope_owner_round_trips() { Some(baml_compiler2_ppir::item_data::ScopeOwner::Function(loc)), ); } + +// ── Method → owner index ───────────────────────────────────────────────────── + +/// `method_owner` must agree with the three scans it replaces (classes by +/// `methods`, interfaces by `default_methods`, free impls by `methods`), for +/// every function in a fixture covering all the ownership cases: a plain class +/// method, an in-body `implements` method (owned by the *class*), an interface +/// default method, an out-of-body impl method, and a top-level function. +#[test] +fn method_owner_index_agrees_with_the_scans_it_replaces() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + r#" +interface Greeter { + function greet(self) -> string + function shout(self) -> string throws never { + "HI" + } +} + +class Person { + name string + + function rename(self, name: string) -> string throws never { + name + } + + implements Greeter { + function greet(self) -> string throws never { + self.name + } + } +} + +class Robot { + id string +} + +// A simple `implements I for C` is merged onto the class during AST lowering +// (its method becomes class-owned); only a *generic* out-of-body impl stays a +// free impl block. +implements Greeter for Robot { + function greet(self) -> string throws never { + self.id + } +} + +interface Valued { + function get(self) -> T +} + +class Box { + value T +} + +implements Valued for Box { + function get(self) -> T throws never { + self.value + } +} + +function free_standing(x: int) -> int throws never { + x +} +"#, + ); + + let db = test_db.db(); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); + assert!(!item_tree.functions.is_empty()); + + let mut cases = (0usize, 0usize, 0usize, 0usize); + for id in item_tree.functions.keys() { + let loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *id); + + // The scans being retired. + let by_class = item_tree + .classes + .iter() + .find(|(_, class)| class.methods.contains(id)) + .map(|(class_id, _)| *class_id); + let by_interface = item_tree + .interfaces + .iter() + .find(|(_, iface)| iface.default_methods.contains(id)) + .map(|(iface_id, _)| *iface_id); + let by_free_impl = item_tree + .free_impls + .iter() + .find(|impl_id| item_tree.impls[impl_id].methods.contains(id)) + .copied(); + + let indexed = baml_compiler2_ppir::item_data::method_owner(db, loc); + + use baml_compiler2_ppir::item_data::MethodOwner; + match (by_class, by_interface, by_free_impl) { + (Some(class_id), None, None) => { + cases.0 += 1; + assert!( + matches!(indexed, Some(MethodOwner::Class(c)) if c.id(db) == class_id), + "class scan and index disagree for {:?}", + item_tree[*id].name + ); + } + (None, Some(iface_id), None) => { + cases.1 += 1; + assert!( + matches!(indexed, Some(MethodOwner::Interface(i)) if i.id(db) == iface_id), + "interface scan and index disagree for {:?}", + item_tree[*id].name + ); + } + (None, None, Some(impl_id)) => { + cases.2 += 1; + assert!( + matches!(indexed, Some(MethodOwner::FreeImpl(b)) if b.id(db) == impl_id), + "free-impl scan and index disagree for {:?}", + item_tree[*id].name + ); + } + (None, None, None) => { + cases.3 += 1; + assert_eq!( + indexed, None, + "top-level function {:?} should have no owner", + item_tree[*id].name + ); + } + other => unreachable!( + "a method can only have one owner; scans returned {other:?} for {:?}", + item_tree[*id].name + ), + } + } + + // Guard against a vacuous fixture: every ownership case must be present. + assert!( + cases.0 >= 2, + "expected class methods (plain + in-body impl)" + ); + assert!(cases.1 >= 1, "expected an interface default method"); + assert!(cases.2 >= 1, "expected a free-impl method"); + assert!(cases.3 >= 1, "expected a top-level function"); +} From 7b177a4e5e3d0f889dcd532061eb2ad3fae4a2ac Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Mon, 13 Jul 2026 22:17:34 -0700 Subject: [PATCH 09/20] Unified class generic params lookup And some other conversion to the firewalled queries --- .../crates/baml_compiler2_emit/src/lib.rs | 30 ++--- .../baml_compiler2_hir/src/item_tree.rs | 22 ++++ .../baml_compiler2_hir/src/signature.rs | 16 +-- .../crates/baml_compiler2_mir/src/lower.rs | 122 +++++++++--------- .../crates/baml_compiler2_ppir/src/lib.rs | 25 +--- .../crates/baml_compiler2_tir/src/builder.rs | 64 ++++----- .../crates/baml_compiler2_tir/src/callable.rs | 31 ++--- .../baml_compiler2_tir/src/inference.rs | 24 ++-- .../baml_compiler2_tir/src/lower_type_expr.rs | 59 ++++----- 9 files changed, 177 insertions(+), 216 deletions(-) diff --git a/baml_language/crates/baml_compiler2_emit/src/lib.rs b/baml_language/crates/baml_compiler2_emit/src/lib.rs index fd5618afbba..c7609da67af 100644 --- a/baml_language/crates/baml_compiler2_emit/src/lib.rs +++ b/baml_language/crates/baml_compiler2_emit/src/lib.rs @@ -3474,6 +3474,8 @@ fn compute_function_metadata_from_item_tree( parameter_defaults: &baml_compiler2_hir::signature::FunctionParameterDefaults, cache: &ResolvedAliases, ) -> FunctionSignatureMetadata { + use baml_compiler2_hir::item_tree::MethodOwner; + let param_names: Vec = func_data .params .iter() @@ -3498,26 +3500,24 @@ fn compute_function_metadata_from_item_tree( // both its `for_target` (for `Self` substitution) and its declared // generics/bounds are recoverable — replacing the removed // `item_tree.implements_for`, which exposed those as flat fields. - let enclosing_free_impl = item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - block - .methods - .contains(&func_id) - .then_some((*impl_id, block)) - }); + let owner = item_tree.method_owners.get(&func_id); + let enclosing_free_impl = match owner { + Some(MethodOwner::FreeImpl(impl_id)) => Some((*impl_id, &item_tree.impls[impl_id])), + _ => None, + }; let enclosing_impl_for_target = enclosing_free_impl.and_then(|(_, block)| match &block.subject { baml_compiler2_hir::item_tree::ImplSubject::Free { for_target, .. } => Some(for_target), baml_compiler2_hir::item_tree::ImplSubject::InClass { .. } => None, }); - let enclosing_class = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&func_id)); - let enclosing_interface = item_tree - .interfaces - .values() - .find(|iface_data| iface_data.default_methods.contains(&func_id)); + let enclosing_class = match owner { + Some(MethodOwner::Class(class_id)) => Some(&item_tree[*class_id]), + _ => None, + }; + let enclosing_interface = match owner { + Some(MethodOwner::Interface(iface_id)) => Some(&item_tree[*iface_id]), + _ => None, + }; let self_replacement = enclosing_impl_for_target .cloned() .or_else(|| { diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs index dded064c96e..07bb4cfb966 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs @@ -20,6 +20,7 @@ mod type_aliases; use std::ops::Index; +use baml_base::Name; use baml_compiler2_ast as ast; pub use classes::*; pub use clients::*; @@ -90,6 +91,27 @@ pub struct ItemTree { pub method_owners: FxHashMap, MethodOwner>, } +impl ItemTree { + /// Generic parameters of the type declaration enclosing `method` — the + /// class's for a class method (BEP-044: a generic interface's default + /// method likewise sees the interface's), empty for top-level functions. + /// + /// Also empty for a *free-impl* method: an out-of-body block's generics + /// live on the `ImplBlock` and are threaded by the impl-specific paths, + /// not treated as enclosing-type parameters. + /// + /// The single successor of the `classes.values().find(|c| + /// c.methods.contains(…))` scans that used to be copied (divergently) + /// across HIR, PPIR and TIR. + pub fn enclosing_type_generic_params(&self, method: LocalItemId) -> &[Name] { + match self.method_owners.get(&method) { + Some(MethodOwner::Class(id)) => &self[*id].generic_params, + Some(MethodOwner::Interface(id)) => &self[*id].generic_params, + Some(MethodOwner::FreeImpl(_)) | None => &[], + } + } +} + // ── Index impls ─────────────────────────────────────────────────────────────── impl Index> for ItemTree { diff --git a/baml_language/crates/baml_compiler2_hir/src/signature.rs b/baml_language/crates/baml_compiler2_hir/src/signature.rs index 24ebbcd1309..9619c2b28a5 100644 --- a/baml_language/crates/baml_compiler2_hir/src/signature.rs +++ b/baml_language/crates/baml_compiler2_hir/src/signature.rs @@ -438,18 +438,6 @@ pub fn elaborate_function_signature_parts( } } -fn enclosing_class_generic_params( - item_tree: &crate::item_tree::ItemTree, - function_id: crate::ids::LocalItemId, -) -> Vec { - item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&function_id)) - .map(|class_data| class_data.generic_params.clone()) - .unwrap_or_default() -} - fn elaborated_function_signature_with_source_map<'db>( db: &'db dyn crate::Db, function: FunctionLoc<'db>, @@ -475,7 +463,9 @@ fn elaborated_function_signature_with_source_map<'db>( let return_type = func_data.return_type.clone(); let throws = func_data.throws.clone(); - let reserved_effect_param_names = enclosing_class_generic_params(&item_tree, function.id(db)); + let reserved_effect_param_names = item_tree + .enclosing_type_generic_params(function.id(db)) + .to_vec(); let signature = Arc::new(elaborate_function_signature_parts( func_data.name.clone(), func_data.generic_params.clone(), diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 0698c3c594f..929549dccb0 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -681,10 +681,13 @@ fn enclosing_free_impl( item_tree: &baml_compiler2_hir::item_tree::ItemTree, func_id: baml_compiler2_hir::ids::LocalItemId, ) -> Option<&baml_compiler2_hir::item_tree::ImplBlock> { - item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - block.methods.contains(&func_id).then_some(block) - }) + match item_tree.method_owners.get(&func_id)? { + baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(impl_id) => { + item_tree.impls.get(impl_id) + } + baml_compiler2_hir::item_tree::MethodOwner::Class(_) + | baml_compiler2_hir::item_tree::MethodOwner::Interface(_) => None, + } } /// An out-of-body impl block's declared generics as `(param names, each's single optional @@ -1335,45 +1338,38 @@ pub fn def_to_item_ref<'db>(db: &'db dyn crate::Db, def: Definition<'db>) -> Ite Definition::Let(loc) => item_tree[loc.id(db)].name.clone(), }; - // For function definitions, check if this is a class method by searching - // the item tree's class methods lists. + // Function definitions: a method needs a Method-shaped ItemRef so it gets a + // distinct global slot keyed on its owner's name (instead of colliding with + // same-named free functions in the package). if let Definition::Function(func_loc) = def { - let func_local_id = func_loc.id(db); - for (class_id, class_data) in &item_tree.classes { - if class_data.methods.contains(&func_local_id) { + use baml_compiler2_hir::item_tree::MethodOwner; + match item_tree.method_owners.get(&func_loc.id(db)) { + Some(MethodOwner::Class(class_id)) => { let class_loc = baml_compiler2_hir::loc::ClassLoc::new(db, file, *class_id); return method_item_ref(db, class_loc, func_loc); } - } - // BEP-044: interface default methods are also stored in - // `item_tree.functions` and need a Method-shaped ItemRef so they - // get a distinct global slot keyed on the interface name (instead - // of colliding with same-named free functions in the package). - for iface_data in item_tree.interfaces.values() { - if iface_data.default_methods.contains(&func_local_id) { + Some(MethodOwner::Interface(iface_id)) => { return ItemRef::Method { package: pkg_info.package.clone(), namespace: pkg_info.namespace_path, - class: iface_data.name.clone(), + class: item_tree[*iface_id].name.clone(), name, }; } - } - for impl_id in &item_tree.free_impls { - let Some(block) = item_tree.impls.get(impl_id) else { - continue; - }; - if block.methods.contains(&func_local_id) - && let baml_compiler2_hir::item_tree::ImplSubject::Free { for_target, .. } = + Some(MethodOwner::FreeImpl(impl_id)) => { + let block = &item_tree.impls[impl_id]; + if let baml_compiler2_hir::item_tree::ImplSubject::Free { for_target, .. } = &block.subject - { - return ItemRef::Method { - package: pkg_info.package.clone(), - namespace: pkg_info.namespace_path, - class: Name::new(format!("{}$for${}", block.interface_target, for_target)), - name, - }; + { + return ItemRef::Method { + package: pkg_info.package.clone(), + namespace: pkg_info.namespace_path, + class: Name::new(format!("{}$for${}", block.interface_target, for_target)), + name, + }; + } } + None => {} } } @@ -2633,10 +2629,12 @@ impl<'db> LoweringContext<'db> { // `self` dispatch through the interface — `interface_dispatch_target_for_tir_ty` // already follows type-var bounds — so default methods keep dispatching // through the concrete implementor. - for iface in item_tree.interfaces.values() { - if iface.default_methods.contains(&func_loc.id(db)) - && let Some(def) = - pkg_items_for_bounds.lookup_type(&pkg_info.namespace_path, &iface.name) + if let Some(baml_compiler2_hir::item_tree::MethodOwner::Interface(iface_id)) = + item_tree.method_owners.get(&func_loc.id(db)) + { + let iface = &item_tree[*iface_id]; + if let Some(def) = + pkg_items_for_bounds.lookup_type(&pkg_info.namespace_path, &iface.name) { let qtn = baml_compiler2_tir::lower_type_expr::qualify_def(db, def, &iface.name); let args = iface @@ -2666,7 +2664,6 @@ impl<'db> LoweringContext<'db> { baml_compiler2_tir::ty::TyAttr::default(), ), ); - break; } } @@ -3854,17 +3851,18 @@ impl<'db> LoweringContext<'db> { let fl = self.func_loc?; let item_tree = file_item_tree(self.db, fl.file(self.db)); let func_id = fl.id(self.db); - item_tree - .interfaces - .values() - .find(|iface_data| iface_data.default_methods.contains(&func_id)) - .map(|iface_data| { - iface_data - .associated_types - .iter() - .map(|assoc| assoc.name.clone()) - .collect() - }) + let baml_compiler2_hir::item_tree::MethodOwner::Interface(iface_id) = + item_tree.method_owners.get(&func_id)? + else { + return None; + }; + Some(&item_tree[*iface_id]).map(|iface_data| { + iface_data + .associated_types + .iter() + .map(|assoc| assoc.name.clone()) + .collect() + }) } /// Rewrite every `Self.X` path (where `X` is one of `assoc_names`) to the bare `X` @@ -8871,11 +8869,10 @@ impl LoweringContext<'_> { // class-method convention — interface params first, then fn params — so // `TypeVar(T)` lowers to `TypeArgRef(N)` against the frame type args the // interface-dispatch switch seeds (see `emit_method_candidate_switch`). - if let Some(iface_data) = item_tree - .interfaces - .values() - .find(|iface_data| iface_data.default_methods.contains(&func_id)) + if let Some(baml_compiler2_hir::item_tree::MethodOwner::Interface(iface_id)) = + item_tree.method_owners.get(&func_id) { + let iface_data = &item_tree[*iface_id]; let mut params = iface_data.generic_params.clone(); params.extend( iface_data @@ -8886,12 +8883,12 @@ impl LoweringContext<'_> { params.extend(item_tree[func_id].generic_params.iter().cloned()); return params; } - let mut params: Vec = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&func_id)) - .map(|class_data| class_data.generic_params.clone()) - .unwrap_or_default(); + let mut params: Vec = match item_tree.method_owners.get(&func_id) { + Some(baml_compiler2_hir::item_tree::MethodOwner::Class(class_id)) => { + item_tree[*class_id].generic_params.clone() + } + _ => Vec::new(), + }; params.extend(item_tree[func_id].generic_params.iter().cloned()); // Inside a (possibly nested) generic lambda body, the lambda's own // type params follow the enclosing function's, matching the runtime @@ -8919,14 +8916,11 @@ impl LoweringContext<'_> { .clone(); let item_tree = file_item_tree(self.db, fl.file(self.db)); let func_id = fl.id(self.db); - if let Some(impl_id) = item_tree.free_impls.iter().copied().find(|impl_id| { - item_tree - .impls - .get(impl_id) - .is_some_and(|block| block.methods.contains(&func_id)) - }) { + if let Some(baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(impl_id)) = + item_tree.method_owners.get(&func_id) + { let impl_loc = - baml_compiler2_hir::loc::ImplLoc::new(self.db, fl.file(self.db), impl_id); + baml_compiler2_hir::loc::ImplLoc::new(self.db, fl.file(self.db), *impl_id); for (name, conjunction) in baml_compiler2_tir::lower_type_expr::impl_generic_param_bounds(self.db, impl_loc) { diff --git a/baml_language/crates/baml_compiler2_ppir/src/lib.rs b/baml_language/crates/baml_compiler2_ppir/src/lib.rs index d5da7e5d3bd..c2d59889ea0 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/lib.rs @@ -694,27 +694,6 @@ pub fn function_signature<'db>( }) } -fn enclosing_class_generic_params( - item_tree: &ItemTree, - function_id: baml_compiler2_hir::ids::LocalItemId, -) -> Vec { - if let Some(class_data) = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&function_id)) - { - return class_data.generic_params.clone(); - } - // BEP-044: a generic interface's default method sees the interface's type - // params (`interface Container { function f(self) -> T { ... } }`). - item_tree - .interfaces - .values() - .find(|iface_data| iface_data.default_methods.contains(&function_id)) - .map(|iface_data| iface_data.generic_params.clone()) - .unwrap_or_default() -} - /// Canonical elaborated callable signature — uses PPIR's item tree. pub fn elaborated_function_signature<'db>( db: &'db dyn Db, @@ -742,7 +721,9 @@ pub fn elaborated_function_signature<'db>( let return_type = func_data.return_type.clone(); let throws = func_data.throws.clone(); - let reserved_effect_param_names = enclosing_class_generic_params(&item_tree, function.id(db)); + let reserved_effect_param_names = item_tree + .enclosing_type_generic_params(function.id(db)) + .to_vec(); Arc::new( baml_compiler2_hir::signature::elaborate_function_signature_parts( diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index 7f9370f680a..4359c4a9216 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -2009,12 +2009,13 @@ impl<'db> TypeInferenceBuilder<'db> { // must be visible or `C` would resolve to `unknown`. let file = func_loc.file(db); let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let enclosing_class_params: Vec = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&func_loc.id(db))) - .map(|class_data| class_data.generic_params.clone()) - .unwrap_or_default(); + let enclosing_class_params: Vec = + match baml_compiler2_ppir::item_data::method_owner(db, func_loc) { + Some(baml_compiler2_ppir::item_data::MethodOwner::Class(class_loc)) => { + item_tree[class_loc.id(db)].generic_params.clone() + } + _ => Vec::new(), + }; // Only user-declared generic params are *supplied at the call site*; // synthetic effect params are always inferred. For static-method-on- // generic-class calls, the class params are also supplied (`Class<...>.m`) @@ -2511,37 +2512,28 @@ impl<'db> TypeInferenceBuilder<'db> { let db = self.context.db(); let item_tree = baml_compiler2_ppir::file_item_tree(db, func_loc.file(db)); let func_id = func_loc.id(db); - // The enclosing out-of-body impl's declared generic-param names, from the unified `impls` - // store via the `free_impls` index (replacing the removed `implements_for`). - if let Some(impl_params) = item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - if !block.methods.contains(&func_id) { - return None; - } - let baml_compiler2_hir::item_tree::ImplSubject::Free { generics, .. } = &block.subject - else { - return None; - }; - Some(generics.iter().map(|g| g.name.clone()).collect::>()) - }) { - return (impl_params, item_tree[func_id].generic_params.clone()); - } - let fn_params = item_tree[func_id].generic_params.clone(); - if let Some(iface_data) = item_tree - .interfaces - .values() - .find(|iface_data| iface_data.default_methods.contains(&func_id)) - { - return (iface_data.generic_params.clone(), fn_params); - } - - let owner_params = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&func_id)) - .map(|class_data| class_data.generic_params.clone()) - .unwrap_or_default(); + // The owner's declared generic-param names: an out-of-body impl's block + // generics, an interface's params for a default method, or the class's. + let owner_params = match baml_compiler2_ppir::item_data::method_owner(db, func_loc) { + Some(baml_compiler2_ppir::item_data::MethodOwner::FreeImpl(impl_loc)) => { + match &item_tree.impls[&impl_loc.id(db)].subject { + baml_compiler2_hir::item_tree::ImplSubject::Free { generics, .. } => { + generics.iter().map(|g| g.name.clone()).collect() + } + baml_compiler2_hir::item_tree::ImplSubject::InClass { .. } => unreachable!( + "MethodOwner::FreeImpl always names an ImplSubject::Free block" + ), + } + } + Some(baml_compiler2_ppir::item_data::MethodOwner::Interface(iface_loc)) => { + item_tree[iface_loc.id(db)].generic_params.clone() + } + Some(baml_compiler2_ppir::item_data::MethodOwner::Class(class_loc)) => { + item_tree[class_loc.id(db)].generic_params.clone() + } + None => Vec::new(), + }; (owner_params, fn_params) } diff --git a/baml_language/crates/baml_compiler2_tir/src/callable.rs b/baml_language/crates/baml_compiler2_tir/src/callable.rs index dc6e5ef64be..b7700f3fab2 100644 --- a/baml_language/crates/baml_compiler2_tir/src/callable.rs +++ b/baml_language/crates/baml_compiler2_tir/src/callable.rs @@ -37,7 +37,9 @@ fn lowered_declared_callable_throws<'db>( let pkg_id = PackageId::new(db, pkg_info.package.clone()); let pkg_items = baml_compiler2_ppir::package_items(db, pkg_id); - let mut generic_params = enclosing_class_generic_params(&item_tree, function.id(db)); + let mut generic_params = item_tree + .enclosing_type_generic_params(function.id(db)) + .to_vec(); generic_params.extend(sig.user_generic_params.iter().cloned()); generic_params.extend(sig.synthetic_effect_params.iter().cloned()); @@ -71,7 +73,9 @@ fn signature_cycle_initial_callable_throws<'db>( let pkg_id = PackageId::new(db, pkg_info.package.clone()); let pkg_items = baml_compiler2_ppir::package_items(db, pkg_id); - let mut generic_params = enclosing_class_generic_params(&item_tree, function.id(db)); + let mut generic_params = item_tree + .enclosing_type_generic_params(function.id(db)) + .to_vec(); generic_params.extend(sig.user_generic_params.iter().cloned()); generic_params.extend(sig.synthetic_effect_params.iter().cloned()); @@ -95,29 +99,18 @@ fn signature_cycle_initial_callable_throws<'db>( join_throw_facts(&facts) } -fn enclosing_class_generic_params( - item_tree: &baml_compiler2_hir::item_tree::ItemTree, - function_id: baml_compiler2_hir::ids::LocalItemId, -) -> Vec { - item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&function_id)) - .map(|class_data| class_data.generic_params.clone()) - .unwrap_or_default() -} - fn callable_short_name<'db>(db: &'db dyn crate::Db, function: FunctionLoc<'db>) -> Name { let file = function.file(db); let item_tree = baml_compiler2_ppir::file_item_tree(db, file); let func_data = &item_tree[function.id(db)]; - if let Some(class_data) = item_tree - .classes - .values() - .find(|class_data| class_data.methods.contains(&function.id(db))) + // Only a class owner qualifies the name — interface default methods and + // free-impl methods keep their bare name, preserving the throw-set key + // format the scan produced. + if let Some(baml_compiler2_hir::item_tree::MethodOwner::Class(class_id)) = + item_tree.method_owners.get(&function.id(db)) { - Name::new(format!("{}.{}", class_data.name, func_data.name)) + Name::new(format!("{}.{}", item_tree[*class_id].name, func_data.name)) } else { func_data.name.clone() } diff --git a/baml_language/crates/baml_compiler2_tir/src/inference.rs b/baml_language/crates/baml_compiler2_tir/src/inference.rs index ff1ad516401..9179e74b08c 100644 --- a/baml_language/crates/baml_compiler2_tir/src/inference.rs +++ b/baml_language/crates/baml_compiler2_tir/src/inference.rs @@ -798,13 +798,13 @@ fn enclosing_impl_generics_for_func( item_tree: &baml_compiler2_hir::item_tree::ItemTree, func_id: baml_compiler2_hir::ids::LocalItemId, ) -> Option<(Vec, Vec>)> { - item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - block - .methods - .contains(&func_id) - .then(|| free_impl_generics(block)) - }) + match item_tree.method_owners.get(&func_id)? { + baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(impl_id) => { + Some(free_impl_generics(&item_tree.impls[impl_id])) + } + baml_compiler2_hir::item_tree::MethodOwner::Class(_) + | baml_compiler2_hir::item_tree::MethodOwner::Interface(_) => None, + } } fn generic_env_for_function_data<'db>( @@ -1674,10 +1674,12 @@ pub fn infer_scope_types<'db>( let body = baml_compiler2_ppir::function_body(db, func_loc); let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); - let enclosing_impl = item_tree.free_impls.iter().find_map(|impl_id| { - let block = item_tree.impls.get(impl_id)?; - block.methods.contains(local_id).then_some(block) - }); + let enclosing_impl = match item_tree.method_owners.get(local_id) { + Some(baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(impl_id)) => { + item_tree.impls.get(impl_id) + } + _ => None, + }; let mut env = GenericEnv::from_params(sig.user_generic_params.clone()); env.params diff --git a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs index 1bdce62579e..2184f53061b 100644 --- a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs +++ b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs @@ -1106,30 +1106,25 @@ pub fn function_in_scope_generic_param_bounds<'db>( function_loc: baml_compiler2_hir::loc::FunctionLoc<'db>, ) -> TypeVarBoundsMap { let file = function_loc.file(db); - let item_tree = baml_compiler2_hir::file_item_tree(db, file); + // PPIR's tree is the canonical one (it includes synthesized companions); + // reading HIR's here while the owner index reads PPIR's would mix the two. + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); let function_id = function_loc.id(db); let mut bounds = TypeVarBoundsMap::default(); - for (class_id, class) in &item_tree.classes { - if class.methods.contains(&function_id) { + match baml_compiler2_ppir::item_data::method_owner(db, function_loc) { + Some(baml_compiler2_ppir::item_data::MethodOwner::Class(class_loc)) => { bounds.extend( - class_generic_param_bounds( - db, - baml_compiler2_hir::loc::ClassLoc::new(db, file, *class_id), - ) - .iter() - .map(|(name, conjunction)| (name.clone(), conjunction.clone())), + class_generic_param_bounds(db, class_loc) + .iter() + .map(|(name, conjunction)| (name.clone(), conjunction.clone())), ); } - } - for (interface_id, interface) in &item_tree.interfaces { - if interface.default_methods.contains(&function_id) { + Some(baml_compiler2_ppir::item_data::MethodOwner::Interface(iface_loc)) => { + let interface = &item_tree[iface_loc.id(db)]; bounds.extend( - interface_generic_param_bounds( - db, - baml_compiler2_hir::loc::InterfaceLoc::new(db, file, *interface_id), - ) - .iter() - .map(|(name, conjunction)| (name.clone(), conjunction.clone())), + interface_generic_param_bounds(db, iface_loc) + .iter() + .map(|(name, conjunction)| (name.clone(), conjunction.clone())), ); // Inside an interface's own default method `Self` is a rigid type // variable bounded by the interface itself (as a constraint, no @@ -1141,9 +1136,7 @@ pub fn function_in_scope_generic_param_bounds<'db>( // slot at template lowering. let iface_qtn = qualify_def( db, - baml_compiler2_hir::contributions::Definition::Interface( - baml_compiler2_hir::loc::InterfaceLoc::new(db, file, *interface_id), - ), + baml_compiler2_hir::contributions::Definition::Interface(iface_loc), &interface.name, ); let iface_args = interface @@ -1156,24 +1149,18 @@ pub fn function_in_scope_generic_param_bounds<'db>( vec![baml_type::Interface::new(iface_qtn, iface_args, Vec::new())], ); } - } - // A method of an out-of-body `implements` block sees the block's generics (`implements Renderable for Wrapped`) — so a `T.member` / `Self.member` projection - // in its signature or body resolves through `T`'s declared bound. (In-body impl methods sit - // in the class, already covered above.) - for impl_id in &item_tree.free_impls { - if let Some(block) = item_tree.impls.get(impl_id) - && block.methods.contains(&function_id) - { + // A method of an out-of-body `implements` block sees the block's generics + // (`implements Renderable for Wrapped`) — so a + // `T.member` / `Self.member` projection in its signature or body resolves + // through `T`'s declared bound. (In-body impl methods are class-owned.) + Some(baml_compiler2_ppir::item_data::MethodOwner::FreeImpl(impl_loc)) => { bounds.extend( - impl_generic_param_bounds( - db, - baml_compiler2_hir::loc::ImplLoc::new(db, file, *impl_id), - ) - .iter() - .map(|(name, conjunction)| (name.clone(), conjunction.clone())), + impl_generic_param_bounds(db, impl_loc) + .iter() + .map(|(name, conjunction)| (name.clone(), conjunction.clone())), ); } + None => {} } if let Some(function) = item_tree.functions.get(&function_id) { let pkg_info = baml_compiler2_hir::file_package::file_package(db, file); From 854e4f02c27231c4de7a0188acbed4b09a59834c Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Tue, 14 Jul 2026 13:26:44 -0700 Subject: [PATCH 10/20] TIR lower type refs We should now lower the salsa-stable type refs instead of type exprs (`lower_type_expr` is now a shim) --- .../baml_compiler2_tir/src/lower_type_expr.rs | 791 +++++++++--------- 1 file changed, 407 insertions(+), 384 deletions(-) diff --git a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs index 2184f53061b..0de2c3568f5 100644 --- a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs +++ b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs @@ -498,379 +498,67 @@ pub fn lower_type_expr( ctx: &dyn TypeExprContext<'_>, diagnostics: &mut Vec, ) -> Ty { - let db = ctx.db(); - match &type_expr.kind { - TypeExprKind::Path { + // Transitional shim: lower the AST node into a scratch span-free arena and + // run the native lowering. Diagnostics carry no spans (callers position + // them from their own `TypeExpr`), so nothing is lost by the round-trip. + let mut type_refs = baml_compiler2_hir::type_ref::TypeRefBuilder::new(); + let id = type_refs.lower(type_expr); + let (store, _spans) = type_refs.finish(); + lower_type_ref(&store, id, ctx, diagnostics) +} + +/// Resolve a span-free [`TypeRef`](baml_compiler2_hir::type_ref::TypeRef) to a +/// `Ty` — the native form of [`lower_type_expr`], for callers holding firewall +/// data (`function_data` / `class_data` / …) rather than AST nodes. +pub fn lower_type_ref( + store: &baml_compiler2_hir::type_ref::TypeRefStore, + id: baml_compiler2_hir::type_ref::TypeRefId, + ctx: &dyn TypeExprContext<'_>, + diagnostics: &mut Vec, +) -> Ty { + use baml_compiler2_hir::type_ref::TypeRefKind; + + match &store[id].kind { + TypeRefKind::Path { segments, generic_args, associated_type_bindings, - .. - } => { - // `Self` / `Self.Member…`: resolve the receiver through the context, but only - // when `Self` is in scope (`lower_self` is `Some`). A bare `Self` is the - // receiver type; `Self.Item` is an associated-type projection rooted at it. - // Falls through otherwise, so a stray `Self` still takes the unresolved path. - if segments[0].as_str() == "Self" - && generic_args.is_empty() - && associated_type_bindings.is_empty() - && let Some(self_ty) = ctx.lower_self() - { - if segments.len() == 1 { - return self_ty; - } - let mut ty = self_ty; - for member in &segments[1..] { - let lowered = crate::builder::associated_projection::lower_projection( - ctx, - ty, - None, - member.clone(), - ); - diagnostics.extend(lowered.diagnostics); - ty = lowered.ty; - } - return ty; - } - match ctx.resolve_type(segments) { - Ok(def) => { - let short = segments.last().expect("non-empty path"); - match def { - Definition::Class(class_loc) => { - // Collect and lower generic args, storing them in Ty::Class. - let lowered_args: Vec = generic_args - .iter() - .map(|ga| lower_type_expr(ga, ctx, diagnostics)) - .collect(); - - let class_tree = - baml_compiler2_ppir::file_item_tree(db, class_loc.file(db)); - let expected_type_args = class_tree - .classes - .get(&class_loc.id(db)) - .map(|class| class.generic_params.len()) - .unwrap_or(0); - // A bare generic name (`generic_args.is_empty()`) is left - // unchecked here: it is a deliberate wildcard in several - // positions (`reflect.type_of()`, construction - // `Box { .. }` where args infer from fields, an interface's - // own `Self` type). Object construction that cannot infer - // its args is reported by `infer_object_expr` - // (`CannotInferTypeParameter`) instead. - if !generic_args.is_empty() && generic_args.len() != expected_type_args - { - diagnostics.push(TirTypeError::WrongNumberOfTypeArgs { - type_name: short.clone(), - expected: expected_type_args, - got: generic_args.len(), - }); - } - - // BEP-034: `baml.future.Future` resolves to the - // dedicated `Ty::Future` variant rather than the - // generic `Ty::Class`. The class declaration exists - // as a regular `.baml` file (for method dispatch via - // the standard PackageBaml path), but `spawn` / - // `await` keys off `Ty::Future` directly. Mirrors - // how `int[]` resolves to `Ty::List` even though - // `class Array` is a regular class declaration. - let qtn = qualify_def(db, def, short); - if qtn.name().as_str() == "Future" - && qtn.package().as_str() == "baml" - && qtn.namespace().len() == 1 - && qtn.namespace()[0].as_str() == "future" - && lowered_args.len() == 2 - { - return Ty::Future( - Box::new(lowered_args[0].clone()), - Box::new(lowered_args[1].clone()), - TyAttr::default(), - ); - } - - Ty::Class(qtn, lowered_args, TyAttr::default()) - } - Definition::Interface(iface_loc) => { - // Same generic-arg handling as `Class` — interface - // parameters are valid in the same positions. - let lowered_args: Vec = generic_args - .iter() - .map(|ga| lower_type_expr(ga, ctx, diagnostics)) - .collect(); - let iface_tree = - baml_compiler2_ppir::file_item_tree(db, iface_loc.file(db)); - let expected_type_args = iface_tree - .interfaces - .get(&iface_loc.id(db)) - .map(|iface| iface.generic_params.len()) - .unwrap_or(0); - if !generic_args.is_empty() && generic_args.len() != expected_type_args - { - diagnostics.push(TirTypeError::WrongNumberOfTypeArgs { - type_name: short.clone(), - expected: expected_type_args, - got: generic_args.len(), - }); - } - let known_associated_types: FxHashSet = iface_tree - .interfaces - .get(&iface_loc.id(db)) - .map(|iface| { - iface - .associated_types - .iter() - .map(|assoc| assoc.name.clone()) - .collect() - }) - .unwrap_or_default(); - let iface_qtn = qualify_def(db, def, short); - let mut seen_associated_bindings = FxHashSet::default(); - let lowered_associated_bindings: Vec<(baml_base::Name, Ty)> = - associated_type_bindings - .iter() - .map(|binding| { - if !known_associated_types.contains(&binding.name) { - diagnostics.push(TirTypeError::UnresolvedType { - name: binding.name.clone(), - suggestions: known_associated_types - .iter() - .cloned() - .collect(), - }); - } - // The same associated type specified twice in an - // existential's arguments (`I`) — the - // existential counterpart of the impl-side duplicate. - if !seen_associated_bindings.insert(binding.name.clone()) { - diagnostics.push( - TirTypeError::DuplicateAssociatedTypeBinding { - interface: iface_qtn.clone(), - name: binding.name.clone(), - }, - ); - } - ( - binding.name.clone(), - lower_type_expr(&binding.ty, ctx, diagnostics), - ) - }) - .collect(); - // §1.7(a): eagerly fill each omitted, defaulted associated type at - // this existential. `Self` is the existential itself (its explicit - // pins plus the defaults filled so far), so a Self-referencing - // default (`type Items = Self.Item[]`) reduces against them. The - // default is lowered once — with a symbolic `Self` — by - // `interface_associated_type_default`, and substituted here. - let (iface_generic_params, iface_assoc_names): (Vec<_>, Vec<_>) = - iface_tree - .interfaces - .get(&iface_loc.id(db)) - .map(|iface| { - ( - iface.generic_params.clone(), - iface - .associated_types - .iter() - .map(|assoc| assoc.name.clone()) - .collect(), - ) - }) - .unwrap_or_default(); - let mut associated_bindings = lowered_associated_bindings; - for assoc_name in iface_assoc_names { - if associated_bindings.iter().any(|(n, _)| *n == assoc_name) { - continue; - } - if let Some((default, _)) = - crate::interfaces::interface_associated_type_default( - db, - iface_loc, - assoc_name.clone(), - ) - { - let self_ty = Ty::Interface( - iface_qtn.clone(), - lowered_args.clone(), - associated_bindings.clone(), - TyAttr::default(), - ); - let filled = crate::interfaces::realize_associated_default( - &default, - &iface_generic_params, - &lowered_args, - &self_ty, - ); - associated_bindings.push((assoc_name, filled)); - } - } - Ty::Interface( - iface_qtn, - lowered_args, - associated_bindings, - TyAttr::default(), - ) - } - Definition::Enum(_) => { - // Enums are not generic — validate args and emit a diagnostic if any were supplied. - for ga in generic_args { - let _ = lower_type_expr(ga, ctx, diagnostics); - } - if !generic_args.is_empty() { - diagnostics.push(TirTypeError::TypeIsNotGeneric { - type_name: short.clone(), - kind: "enum", - }); - } - Ty::Enum(qualify_def(db, def, short), TyAttr::default()) - } - Definition::TypeAlias(_) => { - // Type aliases are not generic — validate args and emit a diagnostic if any were supplied. - for ga in generic_args { - let _ = lower_type_expr(ga, ctx, diagnostics); - } - if !generic_args.is_empty() { - diagnostics.push(TirTypeError::TypeIsNotGeneric { - type_name: short.clone(), - kind: "type alias", - }); - } - Ty::TypeAlias(qualify_def(db, def, short), TyAttr::default()) - } - // Unreachable in practice: the package `types` namespace only holds - // class/enum/interface/alias contributions (hir builder). Defensive — - // if a non-type definition ever lands here, report it rather than - // silently producing a compatible-with-everything sentinel. - _ => { - diagnostics.push(TirTypeError::UnresolvedType { - name: short.clone(), - suggestions: Box::new([]), - }); - Ty::Error { - attr: TyAttr::default(), - } - } - } - } - Err(suggestions) => { - // A single-segment name that is an in-scope type variable - // (e.g. T, K, V) lowers to `Ty::TypeVar`, not an error. - if segments.len() == 1 && ctx.type_var_bounds(&segments[0]).is_some() { - return Ty::TypeVar(segments[0].clone(), TyAttr::default()); - } - // Enum-variant fallback: a path like `Status.Active` (or - // `pkg.ns.Status.Active`) won't resolve as a type — `Active` - // isn't a type, it's a variant. Try interpreting the last - // segment as the variant and the rest as the enum's path. - if segments.len() >= 2 { - let (variant, enum_path) = segments.split_last().unwrap(); - let enum_short = enum_path.last().unwrap(); - if let Ok(def @ Definition::Enum(enum_loc)) = ctx.resolve_type(enum_path) { - // Verify the variant actually exists on the enum; - // otherwise `Status.Typo` would silently produce a - // bogus `Ty::EnumVariant` and downstream code would - // never see `UnresolvedType`. - let item_tree = - baml_compiler2_ppir::file_item_tree(db, enum_loc.file(db)); - let enum_data = &item_tree[enum_loc.id(db)]; - if enum_data.variants.iter().any(|v| v.name == *variant) { - return Ty::EnumVariant( - qualify_def(db, def, enum_short), - variant.clone(), - TyAttr::default(), - ); - } - } - } - // Associated type projection fallback: after ordinary type - // paths and enum variants have had first refusal, treat - // `Base.Member` as shorthand for an associated type projection. - // This preserves enum disambiguation (`Status.Active`) and - // still accepts aliases, type variables, concrete classes, - // interfaces, and nested projections as projection bases. - if segments.len() >= 2 - && generic_args.is_empty() - && associated_type_bindings.is_empty() - { - let base_expr = TypeExprKind::Path { - segments: segments[..segments.len() - 1].to_vec(), - generic_args: Vec::new(), - associated_type_bindings: Vec::new(), - attrs: Vec::new(), - } - .at(type_expr.span); - let mut base_diags = Vec::new(); - let base_ty = lower_type_expr(&base_expr, ctx, &mut base_diags); - if base_diags.is_empty() && can_be_associated_type_projection_base(&base_ty) - { - let member = segments.last().expect("non-empty path").clone(); - let lowered = crate::builder::associated_projection::lower_projection( - ctx, base_ty, None, member, - ); - diagnostics.extend(lowered.diagnostics); - return lowered.ty; - } - } - let name_str = segments - .iter() - .map(smol_str::SmolStr::as_str) - .collect::>() - .join("."); - // If this otherwise-unresolved single-segment name is an associated - // type of the enclosing interface's `Self`, the user almost certainly - // meant the projection `Self.` (bare associated-type names are - // illegal). Surface that as the leading "did you mean" on the ordinary - // unresolved-type error — a real in-scope type of the same name resolves - // in the `Ok` branch above and never reaches here. - let suggestions = if segments.len() == 1 - && generic_args.is_empty() - && associated_type_bindings.is_empty() - && ctx.self_associated_type_names().contains(&segments[0]) - { - std::iter::once(baml_base::Name::new(format!("Self.{name_str}"))) - .chain(suggestions.iter().cloned()) - .collect() - } else { - suggestions - }; - diagnostics.push(TirTypeError::UnresolvedType { - name: baml_base::Name::new(&name_str), - suggestions, - }); - // An unresolved name is unrecoverable — `Ty::Error` (diagnosed, poisons - // downstream), never `Ty::Unknown` (the missing/inferable sentinel). - Ty::Error { - attr: TyAttr::default(), - } - } - } - } - TypeExprKind::Int { .. } => Ty::Int { + } => lower_path( + store, + segments, + generic_args, + associated_type_bindings, + ctx, + diagnostics, + ), + TypeRefKind::Int => Ty::Int { attr: TyAttr::default(), }, - TypeExprKind::Bigint { .. } => Ty::Bigint { + TypeRefKind::Bigint => Ty::Bigint { attr: TyAttr::default(), }, - TypeExprKind::Float { .. } => Ty::Float { + TypeRefKind::Float => Ty::Float { attr: TyAttr::default(), }, - TypeExprKind::String { .. } => Ty::String { + TypeRefKind::String => Ty::String { attr: TyAttr::default(), }, - TypeExprKind::Bool { .. } => Ty::Bool { + TypeRefKind::Bool => Ty::Bool { attr: TyAttr::default(), }, - TypeExprKind::Null { .. } => Ty::Null { + TypeRefKind::Null => Ty::Null { attr: TyAttr::default(), }, - TypeExprKind::Never { .. } => Ty::Never { + TypeRefKind::Never => Ty::Never { attr: TyAttr::default(), }, - TypeExprKind::Void { .. } => Ty::Void { + TypeRefKind::Void => Ty::Void { attr: TyAttr::default(), }, - TypeExprKind::Uint8Array { .. } => Ty::Uint8Array { + TypeRefKind::Uint8Array => Ty::Uint8Array { attr: TyAttr::default(), }, - TypeExprKind::Media { kind, .. } => match kind { + TypeRefKind::Media { kind } => match kind { baml_base::MediaKind::Image => Ty::Media(MediaKind::Image, TyAttr::default()), baml_base::MediaKind::Audio => Ty::Media(MediaKind::Audio, TyAttr::default()), baml_base::MediaKind::Video => Ty::Media(MediaKind::Video, TyAttr::default()), @@ -881,19 +569,19 @@ pub fn lower_type_expr( }, }, // `T?` is sugar for `T | null` — lower it directly to a nullable union. - TypeExprKind::Optional { inner, .. } => { - Ty::optional(lower_type_expr(inner, ctx, diagnostics)) + TypeRefKind::Optional { inner } => { + Ty::optional(lower_type_ref(store, *inner, ctx, diagnostics)) } - TypeExprKind::List { inner, .. } => Ty::List( - Box::new(lower_type_expr(inner, ctx, diagnostics)), + TypeRefKind::List { inner } => Ty::List( + Box::new(lower_type_ref(store, *inner, ctx, diagnostics)), TyAttr::default(), ), - TypeExprKind::Map { key, value, .. } => { - let key_ty = lower_type_expr(key, ctx, diagnostics); - let value_ty = lower_type_expr(value, ctx, diagnostics); + TypeRefKind::Map { key, value } => { + let key_ty = lower_type_ref(store, *key, ctx, diagnostics); + let value_ty = lower_type_ref(store, *value, ctx, diagnostics); // Map keys are strings at runtime, so the key type must denote `string` // (or a subset of it) — E0067. - if !is_supported_map_key_type(db, &key_ty, &mut FxHashSet::default()) { + if !is_supported_map_key_type(ctx.db(), &key_ty, &mut FxHashSet::default()) { diagnostics.push(TirTypeError::InvalidMapKeyType { key: key_ty.clone(), }); @@ -904,20 +592,17 @@ pub fn lower_type_expr( attr: TyAttr::default(), } } - TypeExprKind::Union { - variants: members, .. - } => Ty::Union( - members + TypeRefKind::Union { variants } => Ty::Union( + variants .iter() - .map(|m| lower_type_expr(m, ctx, diagnostics)) + .map(|&member| lower_type_ref(store, member, ctx, diagnostics)) .collect(), TyAttr::default(), ), - TypeExprKind::Function { + TypeRefKind::Function { params, ret, throws, - .. } => { // A function type carries no generics of its own; its type variables // come from the enclosing context. @@ -926,7 +611,7 @@ pub fn lower_type_expr( .iter() .map(|p| FunctionParamTy { name: p.name.clone(), - ty: lower_type_expr(&p.ty, ctx, diagnostics), + ty: lower_type_ref(store, p.ty, ctx, diagnostics), mode: if p.optional { FunctionParamMode::Optional } else { @@ -934,11 +619,10 @@ pub fn lower_type_expr( }, }) .collect(), - ret: Box::new(lower_type_expr(ret, ctx, diagnostics)), + ret: Box::new(lower_type_ref(store, *ret, ctx, diagnostics)), throws: Box::new( throws - .as_deref() - .map(|throws| lower_type_expr(throws, ctx, diagnostics)) + .map(|throws| lower_type_ref(store, throws, ctx, diagnostics)) .unwrap_or(Ty::Never { attr: TyAttr::default(), }), @@ -946,16 +630,14 @@ pub fn lower_type_expr( attr: TyAttr::default(), } } - TypeExprKind::AssociatedTypeProjection { + TypeRefKind::AssociatedTypeProjection { base, interface, member, - .. } => { - let base_ty = lower_type_expr(base, ctx, diagnostics); - let explicit_interface = interface - .as_ref() - .map(|interface| lower_type_expr(interface, ctx, diagnostics)); + let base_ty = lower_type_ref(store, *base, ctx, diagnostics); + let explicit_interface = + interface.map(|interface| lower_type_ref(store, interface, ctx, diagnostics)); let lowered = crate::builder::associated_projection::lower_projection( ctx, base_ty, @@ -965,34 +647,34 @@ pub fn lower_type_expr( diagnostics.extend(lowered.diagnostics); lowered.ty } - TypeExprKind::Literal { value: lit, .. } => { + TypeRefKind::Literal { value: lit } => { Ty::Literal(lit.clone(), Freshness::Regular, TyAttr::default()) } - TypeExprKind::BuiltinUnknown { .. } => Ty::BuiltinUnknown { + TypeRefKind::BuiltinUnknown => Ty::BuiltinUnknown { attr: TyAttr::default(), }, // Parse recovery: the parser already reported the syntax error — unrecoverable. - TypeExprKind::Error { .. } => Ty::Error { + TypeRefKind::Error => Ty::Error { attr: TyAttr::default(), }, // A missing type (an unannotated lambda param, an elided receiver type): genuinely // "not yet known", to be filled by inference — NOT an error sentinel. - TypeExprKind::Unknown { .. } => Ty::Unknown { + TypeRefKind::Unknown => Ty::Unknown { attr: TyAttr::default(), }, // Dedicated Ty::Type variant — see ty.rs doc comment for design rationale. - TypeExprKind::Type { .. } => Ty::Type { + TypeRefKind::Type => Ty::Type { attr: TyAttr::default(), }, // `$rust_type` — opaque Rust-managed state field type. - TypeExprKind::Rust { .. } => Ty::RustType { + TypeRefKind::Rust => Ty::RustType { attr: TyAttr::default(), }, // The wildcard `_` (a type-inference placeholder) cannot be inferred — inference for // `_` is unimplemented. Reject it with a diagnostic and lower to `Ty::Error` — never // `Ty::Infer`, which the canonical normalizer treats as `unreachable!`. The user must // write the type explicitly. - TypeExprKind::Infer { .. } => { + TypeRefKind::Infer => { diagnostics.push(TirTypeError::CannotInferType); Ty::Error { attr: TyAttr::default(), @@ -1001,6 +683,347 @@ pub fn lower_type_expr( } } +/// The `Path` arm of [`lower_type_ref`], factored out so the associated-type +/// projection fallback can re-lower a path *prefix* by recursing with a shorter +/// segment list — where the AST form had to synthesize a fresh `TypeExpr` node. +fn lower_path( + store: &baml_compiler2_hir::type_ref::TypeRefStore, + segments: &[baml_base::Name], + generic_args: &[baml_compiler2_hir::type_ref::TypeRefId], + associated_type_bindings: &[baml_compiler2_hir::type_ref::AssociatedTypeBindingRef], + ctx: &dyn TypeExprContext<'_>, + diagnostics: &mut Vec, +) -> Ty { + let db = ctx.db(); + // `Self` / `Self.Member…`: resolve the receiver through the context, but only + // when `Self` is in scope (`lower_self` is `Some`). A bare `Self` is the + // receiver type; `Self.Item` is an associated-type projection rooted at it. + // Falls through otherwise, so a stray `Self` still takes the unresolved path. + if segments[0].as_str() == "Self" + && generic_args.is_empty() + && associated_type_bindings.is_empty() + && let Some(self_ty) = ctx.lower_self() + { + if segments.len() == 1 { + return self_ty; + } + let mut ty = self_ty; + for member in &segments[1..] { + let lowered = crate::builder::associated_projection::lower_projection( + ctx, + ty, + None, + member.clone(), + ); + diagnostics.extend(lowered.diagnostics); + ty = lowered.ty; + } + return ty; + } + match ctx.resolve_type(segments) { + Ok(def) => { + let short = segments.last().expect("non-empty path"); + match def { + Definition::Class(class_loc) => { + // Collect and lower generic args, storing them in Ty::Class. + let lowered_args: Vec = generic_args + .iter() + .map(|&ga| lower_type_ref(store, ga, ctx, diagnostics)) + .collect(); + + let class_tree = baml_compiler2_ppir::file_item_tree(db, class_loc.file(db)); + let expected_type_args = class_tree + .classes + .get(&class_loc.id(db)) + .map(|class| class.generic_params.len()) + .unwrap_or(0); + // A bare generic name (`generic_args.is_empty()`) is left + // unchecked here: it is a deliberate wildcard in several + // positions (`reflect.type_of()`, construction + // `Box { .. }` where args infer from fields, an interface's + // own `Self` type). Object construction that cannot infer + // its args is reported by `infer_object_expr` + // (`CannotInferTypeParameter`) instead. + if !generic_args.is_empty() && generic_args.len() != expected_type_args { + diagnostics.push(TirTypeError::WrongNumberOfTypeArgs { + type_name: short.clone(), + expected: expected_type_args, + got: generic_args.len(), + }); + } + + // BEP-034: `baml.future.Future` resolves to the + // dedicated `Ty::Future` variant rather than the + // generic `Ty::Class`. The class declaration exists + // as a regular `.baml` file (for method dispatch via + // the standard PackageBaml path), but `spawn` / + // `await` keys off `Ty::Future` directly. Mirrors + // how `int[]` resolves to `Ty::List` even though + // `class Array` is a regular class declaration. + let qtn = qualify_def(db, def, short); + if qtn.name().as_str() == "Future" + && qtn.package().as_str() == "baml" + && qtn.namespace().len() == 1 + && qtn.namespace()[0].as_str() == "future" + && lowered_args.len() == 2 + { + return Ty::Future( + Box::new(lowered_args[0].clone()), + Box::new(lowered_args[1].clone()), + TyAttr::default(), + ); + } + + Ty::Class(qtn, lowered_args, TyAttr::default()) + } + Definition::Interface(iface_loc) => { + // Same generic-arg handling as `Class` — interface + // parameters are valid in the same positions. + let lowered_args: Vec = generic_args + .iter() + .map(|&ga| lower_type_ref(store, ga, ctx, diagnostics)) + .collect(); + let iface_tree = baml_compiler2_ppir::file_item_tree(db, iface_loc.file(db)); + let expected_type_args = iface_tree + .interfaces + .get(&iface_loc.id(db)) + .map(|iface| iface.generic_params.len()) + .unwrap_or(0); + if !generic_args.is_empty() && generic_args.len() != expected_type_args { + diagnostics.push(TirTypeError::WrongNumberOfTypeArgs { + type_name: short.clone(), + expected: expected_type_args, + got: generic_args.len(), + }); + } + let known_associated_types: FxHashSet = iface_tree + .interfaces + .get(&iface_loc.id(db)) + .map(|iface| { + iface + .associated_types + .iter() + .map(|assoc| assoc.name.clone()) + .collect() + }) + .unwrap_or_default(); + let iface_qtn = qualify_def(db, def, short); + let mut seen_associated_bindings = FxHashSet::default(); + let lowered_associated_bindings: Vec<(baml_base::Name, Ty)> = + associated_type_bindings + .iter() + .map(|binding| { + if !known_associated_types.contains(&binding.name) { + diagnostics.push(TirTypeError::UnresolvedType { + name: binding.name.clone(), + suggestions: known_associated_types + .iter() + .cloned() + .collect(), + }); + } + // The same associated type specified twice in an + // existential's arguments (`I`) — the + // existential counterpart of the impl-side duplicate. + if !seen_associated_bindings.insert(binding.name.clone()) { + diagnostics.push( + TirTypeError::DuplicateAssociatedTypeBinding { + interface: iface_qtn.clone(), + name: binding.name.clone(), + }, + ); + } + ( + binding.name.clone(), + lower_type_ref(store, binding.ty, ctx, diagnostics), + ) + }) + .collect(); + // §1.7(a): eagerly fill each omitted, defaulted associated type at + // this existential. `Self` is the existential itself (its explicit + // pins plus the defaults filled so far), so a Self-referencing + // default (`type Items = Self.Item[]`) reduces against them. The + // default is lowered once — with a symbolic `Self` — by + // `interface_associated_type_default`, and substituted here. + let (iface_generic_params, iface_assoc_names): (Vec<_>, Vec<_>) = iface_tree + .interfaces + .get(&iface_loc.id(db)) + .map(|iface| { + ( + iface.generic_params.clone(), + iface + .associated_types + .iter() + .map(|assoc| assoc.name.clone()) + .collect(), + ) + }) + .unwrap_or_default(); + let mut associated_bindings = lowered_associated_bindings; + for assoc_name in iface_assoc_names { + if associated_bindings.iter().any(|(n, _)| *n == assoc_name) { + continue; + } + if let Some((default, _)) = + crate::interfaces::interface_associated_type_default( + db, + iface_loc, + assoc_name.clone(), + ) + { + let self_ty = Ty::Interface( + iface_qtn.clone(), + lowered_args.clone(), + associated_bindings.clone(), + TyAttr::default(), + ); + let filled = crate::interfaces::realize_associated_default( + &default, + &iface_generic_params, + &lowered_args, + &self_ty, + ); + associated_bindings.push((assoc_name, filled)); + } + } + Ty::Interface( + iface_qtn, + lowered_args, + associated_bindings, + TyAttr::default(), + ) + } + Definition::Enum(_) => { + // Enums are not generic — validate args and emit a diagnostic if any were supplied. + for &ga in generic_args { + let _ = lower_type_ref(store, ga, ctx, diagnostics); + } + if !generic_args.is_empty() { + diagnostics.push(TirTypeError::TypeIsNotGeneric { + type_name: short.clone(), + kind: "enum", + }); + } + Ty::Enum(qualify_def(db, def, short), TyAttr::default()) + } + Definition::TypeAlias(_) => { + // Type aliases are not generic — validate args and emit a diagnostic if any were supplied. + for &ga in generic_args { + let _ = lower_type_ref(store, ga, ctx, diagnostics); + } + if !generic_args.is_empty() { + diagnostics.push(TirTypeError::TypeIsNotGeneric { + type_name: short.clone(), + kind: "type alias", + }); + } + Ty::TypeAlias(qualify_def(db, def, short), TyAttr::default()) + } + // Unreachable in practice: the package `types` namespace only holds + // class/enum/interface/alias contributions (hir builder). Defensive — + // if a non-type definition ever lands here, report it rather than + // silently producing a compatible-with-everything sentinel. + _ => { + diagnostics.push(TirTypeError::UnresolvedType { + name: short.clone(), + suggestions: Box::new([]), + }); + Ty::Error { + attr: TyAttr::default(), + } + } + } + } + Err(suggestions) => { + // A single-segment name that is an in-scope type variable + // (e.g. T, K, V) lowers to `Ty::TypeVar`, not an error. + if segments.len() == 1 && ctx.type_var_bounds(&segments[0]).is_some() { + return Ty::TypeVar(segments[0].clone(), TyAttr::default()); + } + // Enum-variant fallback: a path like `Status.Active` (or + // `pkg.ns.Status.Active`) won't resolve as a type — `Active` + // isn't a type, it's a variant. Try interpreting the last + // segment as the variant and the rest as the enum's path. + if segments.len() >= 2 { + let (variant, enum_path) = segments.split_last().unwrap(); + let enum_short = enum_path.last().unwrap(); + if let Ok(def @ Definition::Enum(enum_loc)) = ctx.resolve_type(enum_path) { + // Verify the variant actually exists on the enum; + // otherwise `Status.Typo` would silently produce a + // bogus `Ty::EnumVariant` and downstream code would + // never see `UnresolvedType`. + let item_tree = baml_compiler2_ppir::file_item_tree(db, enum_loc.file(db)); + let enum_data = &item_tree[enum_loc.id(db)]; + if enum_data.variants.iter().any(|v| v.name == *variant) { + return Ty::EnumVariant( + qualify_def(db, def, enum_short), + variant.clone(), + TyAttr::default(), + ); + } + } + } + // Associated type projection fallback: after ordinary type + // paths and enum variants have had first refusal, treat + // `Base.Member` as shorthand for an associated type projection. + // This preserves enum disambiguation (`Status.Active`) and + // still accepts aliases, type variables, concrete classes, + // interfaces, and nested projections as projection bases. + if segments.len() >= 2 && generic_args.is_empty() && associated_type_bindings.is_empty() + { + let mut base_diags = Vec::new(); + let base_ty = lower_path( + store, + &segments[..segments.len() - 1], + &[], + &[], + ctx, + &mut base_diags, + ); + if base_diags.is_empty() && can_be_associated_type_projection_base(&base_ty) { + let member = segments.last().expect("non-empty path").clone(); + let lowered = crate::builder::associated_projection::lower_projection( + ctx, base_ty, None, member, + ); + diagnostics.extend(lowered.diagnostics); + return lowered.ty; + } + } + let name_str = segments + .iter() + .map(smol_str::SmolStr::as_str) + .collect::>() + .join("."); + // If this otherwise-unresolved single-segment name is an associated + // type of the enclosing interface's `Self`, the user almost certainly + // meant the projection `Self.` (bare associated-type names are + // illegal). Surface that as the leading "did you mean" on the ordinary + // unresolved-type error — a real in-scope type of the same name resolves + // in the `Ok` branch above and never reaches here. + let suggestions = if segments.len() == 1 + && generic_args.is_empty() + && associated_type_bindings.is_empty() + && ctx.self_associated_type_names().contains(&segments[0]) + { + std::iter::once(baml_base::Name::new(format!("Self.{name_str}"))) + .chain(suggestions.iter().cloned()) + .collect() + } else { + suggestions + }; + diagnostics.push(TirTypeError::UnresolvedType { + name: baml_base::Name::new(&name_str), + suggestions, + }); + // An unresolved name is unrecoverable — `Ty::Error` (diagnosed, poisons + // downstream), never `Ty::Unknown` (the missing/inferable sentinel). + Ty::Error { + attr: TyAttr::default(), + } + } + } +} + /// Derive the qualified name for a type from its Definition's file location. pub fn qualify_def( db: &dyn crate::Db, From 2b366fb9fb9d97bdc32643ff15f07f72f4f98024 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Tue, 14 Jul 2026 13:56:58 -0700 Subject: [PATCH 11/20] Span-free elaborated function data --- .../src/item_data/functions.rs | 94 +++++++++++++++++++ .../crates/baml_compiler2_tir/src/callable.rs | 19 +++- .../baml_tests/src/incremental/scenarios.rs | 46 +++++++++ 3 files changed, 154 insertions(+), 5 deletions(-) diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index 55028b0675f..7a8c13cb5a4 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -67,6 +67,100 @@ pub fn function_source_map<'db>( lower(db, function).1 } +/// Span-free semantic data for a function's *elaborated* signature — the +/// canonical callable view TIR consumes. +/// +/// Elaboration (performed by the shared `hir::signature` machinery, which this +/// wraps) keeps the user-written top-level throws contract optional but makes +/// every nested function-type throws surface explicit: +/// - immediate callback parameter roots with omitted throws are opened to a +/// fresh synthetic effect parameter +/// - immediate function-valued return roots derive their omitted throws from +/// any immediate callback parameters they expose +/// - every other omitted nested function-type throws becomes `never` +/// +/// This is the **tracked** successor of `ppir::elaborated_function_signature` +/// (an untracked fn returning spanned `TypeExpr`s — zero memoization, and +/// unsafe to memoize as-is because Salsa would retain its stale spans). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ElaboratedFunctionData { + pub name: Name, + pub user_generic_params: Vec, + /// Fresh `__effect_param_N` names minted for callback params with omitted + /// throws, in parameter order. + pub synthetic_effect_params: Vec, + /// The arena for the *elaborated* types below — distinct from + /// [`FunctionData::type_refs`], which holds the raw signature. + pub type_refs: TypeRefStore, + pub params: Vec, + pub return_type: Option, + pub throws: Option, +} + +/// Spans for an [`ElaboratedFunctionData`]. Synthetic nodes (effect params, +/// filled `never`s) carry empty ranges — anchor to the owning item instead. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ElaboratedFunctionSourceMap { + pub type_refs: TypeRefSourceMap, +} + +/// Elaborated signature for one function. Span-free — see the module docs. +#[salsa::tracked(returns(ref))] +pub fn elaborated_function_data<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> ElaboratedFunctionData { + lower_elaborated(db, function).0 +} + +/// Spans for one elaborated signature. Kept separate from +/// [`elaborated_function_data`] so that a whitespace-only edit invalidates this +/// but not the semantic data. +#[salsa::tracked(returns(ref))] +pub fn elaborated_function_source_map<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> ElaboratedFunctionSourceMap { + lower_elaborated(db, function).1 +} + +/// Runs the existing (TypeExpr-based, shared-with-HIR) elaboration, then lowers +/// its output into a span-free arena. Only the representation changes here; the +/// elaboration semantics live in one place. +fn lower_elaborated<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> (ElaboratedFunctionData, ElaboratedFunctionSourceMap) { + let sig = crate::elaborated_function_signature(db, function); + + let mut type_refs = TypeRefBuilder::new(); + let params = sig + .params + .iter() + .map(|param| FunctionParamData { + name: param.name.clone(), + type_ref: Some(type_refs.lower(¶m.ty)), + has_default: param.has_default, + }) + .collect(); + let return_type = sig.return_type.as_ref().map(|te| type_refs.lower(te)); + let throws = sig.throws.as_ref().map(|te| type_refs.lower(te)); + let (store, spans) = type_refs.finish(); + + ( + ElaboratedFunctionData { + name: sig.name.clone(), + user_generic_params: sig.user_generic_params.clone(), + synthetic_effect_params: sig.synthetic_effect_params.clone(), + type_refs: store, + params, + return_type, + throws, + }, + ElaboratedFunctionSourceMap { type_refs: spans }, + ) +} + /// The item a method belongs to, as a `Loc`. /// /// Mirrors `item_tree::MethodOwner` (see its docs for the ownership rules — diff --git a/baml_language/crates/baml_compiler2_tir/src/callable.rs b/baml_language/crates/baml_compiler2_tir/src/callable.rs index b7700f3fab2..8751cc5b0a1 100644 --- a/baml_language/crates/baml_compiler2_tir/src/callable.rs +++ b/baml_language/crates/baml_compiler2_tir/src/callable.rs @@ -32,7 +32,7 @@ fn lowered_declared_callable_throws<'db>( ) -> Option { let file = function.file(db); let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, function); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, function); let pkg_info = file_package::file_package(db, file); let pkg_id = PackageId::new(db, pkg_info.package.clone()); let pkg_items = baml_compiler2_ppir::package_items(db, pkg_id); @@ -43,9 +43,10 @@ fn lowered_declared_callable_throws<'db>( generic_params.extend(sig.user_generic_params.iter().cloned()); generic_params.extend(sig.synthetic_effect_params.iter().cloned()); - sig.throws.as_ref().map(|declared_throws| { + sig.throws.map(|declared_throws| { let mut diags = Vec::new(); - crate::lower_type_expr::lower_type_expr( + crate::lower_type_expr::lower_type_ref( + &sig.type_refs, declared_throws, &crate::lower_type_expr::ScopeCtx { db, @@ -68,7 +69,7 @@ fn signature_cycle_initial_callable_throws<'db>( ) -> Ty { let file = function.file(db); let item_tree = baml_compiler2_ppir::file_item_tree(db, file); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, function); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, function); let pkg_info = file_package::file_package(db, file); let pkg_id = PackageId::new(db, pkg_info.package.clone()); let pkg_items = baml_compiler2_ppir::package_items(db, pkg_id); @@ -89,8 +90,16 @@ fn signature_cycle_initial_callable_throws<'db>( }; let mut facts = BTreeSet::new(); for param in &sig.params { + let Some(param_ty) = param.type_ref else { + continue; + }; let mut diags = Vec::new(); - let lowered = crate::lower_type_expr::lower_type_expr(¶m.ty, ¶m_scope, &mut diags); + let lowered = crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param_ty, + ¶m_scope, + &mut diags, + ); if let Ty::Function { throws, .. } = lowered { facts.extend(crate::throw_inference::flatten_ty_to_facts(&throws)); } diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index cbe66ad0bef..e69e719c789 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -852,3 +852,49 @@ function free_standing(x: int) -> int throws never { assert!(cases.2 >= 1, "expected a free-impl method"); assert!(cases.3 >= 1, "expected a top-level function"); } + +/// The elaborated signature is the canonical callable view TIR consumes; its +/// tracked, span-free form must survive whitespace *and* body edits untouched, +/// and must still perform the elaboration (callback params with omitted throws +/// get synthetic effect parameters). +#[test] +fn elaborated_function_data_cuts_off_and_still_elaborates() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Apply(x: int, f: (int) -> int) -> int throws never {\n f(x)\n}\n", + ); + + let before = { + let db = test_db.db(); + let loc = function_loc(db, file, "Apply"); + baml_compiler2_ppir::item_data::elaborated_function_data(db, loc).clone() + }; + + // The callback param `f` omits its throws — elaboration must have opened a + // synthetic effect parameter for it. + assert_eq!( + before.synthetic_effect_params.len(), + 1, + "callback with omitted throws should mint one effect param" + ); + assert_eq!(before.params.len(), 2); + + // Whitespace + body edit together: signature semantics unchanged. + file.set_text(test_db.db_mut()).to( + "// moved\nfunction Apply(x: int, f: (int) -> int) -> int throws never {\n f(x) + 0\n}\n" + .to_string(), + ); + + let after = { + let db = test_db.db(); + let loc = function_loc(db, file, "Apply"); + baml_compiler2_ppir::item_data::elaborated_function_data(db, loc).clone() + }; + + assert_eq!( + before, after, + "whitespace/body edits must not change the elaborated signature" + ); +} From c9c2fb3325b01d5670ed4239dd94971b384dfa83 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Tue, 14 Jul 2026 14:59:58 -0700 Subject: [PATCH 12/20] Use elaborated function param data --- .../src/item_data/functions.rs | 19 +++- .../crates/baml_compiler2_tir/src/builder.rs | 95 ++++++++++++------- .../crates/baml_compiler2_tir/src/callable.rs | 5 +- 3 files changed, 76 insertions(+), 43 deletions(-) diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index 7a8c13cb5a4..8a299313665 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -92,11 +92,24 @@ pub struct ElaboratedFunctionData { /// The arena for the *elaborated* types below — distinct from /// [`FunctionData::type_refs`], which holds the raw signature. pub type_refs: TypeRefStore, - pub params: Vec, + pub params: Vec, pub return_type: Option, pub throws: Option, } +/// A parameter of an *elaborated* signature. +/// +/// Unlike [`FunctionParamData`](crate::item_data::FunctionParamData) — where +/// `type_ref: Option` mirrors the user writing or omitting an annotation — +/// `type_ref` here is total: elaboration substitutes `Unknown` for a missing +/// annotation before this struct is built, so "no type" is unrepresentable. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ElaboratedParamData { + pub name: Name, + pub type_ref: TypeRefId, + pub has_default: bool, +} + /// Spans for an [`ElaboratedFunctionData`]. Synthetic nodes (effect params, /// filled `never`s) carry empty ranges — anchor to the owning item instead. #[derive(Debug, Clone, PartialEq, Eq)] @@ -137,9 +150,9 @@ fn lower_elaborated<'db>( let params = sig .params .iter() - .map(|param| FunctionParamData { + .map(|param| ElaboratedParamData { name: param.name.clone(), - type_ref: Some(type_refs.lower(¶m.ty)), + type_ref: type_refs.lower(¶m.ty), has_default: param.has_default, }) .collect(); diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index 4359c4a9216..b9e7b405c3b 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -1914,7 +1914,7 @@ impl<'db> TypeInferenceBuilder<'db> { if let Some(crate::inference::MemberResolution::Free { func_loc }) = self.resolutions.get(&expr) { - !baml_compiler2_ppir::elaborated_function_signature(self.context.db(), *func_loc) + !baml_compiler2_ppir::item_data::elaborated_function_data(self.context.db(), *func_loc) .user_generic_params .is_empty() } else { @@ -2001,7 +2001,7 @@ impl<'db> TypeInferenceBuilder<'db> { _ => return None, }; let db = self.context.db(); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); // The enclosing class's generic params (`[]` for a free function). These // are always in scope when lowering the method's *bounds* — a bound may // reference a class generic (`>` on a method of @@ -9825,7 +9825,7 @@ impl<'db> TypeInferenceBuilder<'db> { expr_id, crate::inference::MemberResolution::Free { func_loc }, ); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let function_generic_params: Vec = sig .user_generic_params .iter() @@ -9849,8 +9849,11 @@ impl<'db> TypeInferenceBuilder<'db> { .iter() .map(|param| FunctionParamTy { name: Some(param.name.clone()), - ty: crate::lower_type_expr::lower_type_expr( - ¶m.ty, &sig_scope, &mut diags, + ty: crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param.type_ref, + &sig_scope, + &mut diags, ), mode: if param.has_default { FunctionParamMode::Optional @@ -9861,9 +9864,13 @@ impl<'db> TypeInferenceBuilder<'db> { .collect(), ret: Box::new( sig.return_type - .as_ref() - .map(|te| { - crate::lower_type_expr::lower_type_expr(te, &sig_scope, &mut diags) + .map(|id| { + crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + id, + &sig_scope, + &mut diags, + ) }) .unwrap_or(Ty::Unknown { attr: TyAttr::default(), @@ -9984,7 +9991,8 @@ impl<'db> TypeInferenceBuilder<'db> { Definition::Function(func_loc) => { // Get function signature to build the function type let db = self.context.db(); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = + baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let function_generic_params: Vec = sig .user_generic_params .iter() @@ -10014,8 +10022,11 @@ impl<'db> TypeInferenceBuilder<'db> { .iter() .map(|param| FunctionParamTy { name: Some(param.name.clone()), - ty: crate::lower_type_expr::lower_type_expr( - ¶m.ty, &sig_scope, &mut diags, + ty: crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param.type_ref, + &sig_scope, + &mut diags, ), mode: if param.has_default { FunctionParamMode::Optional @@ -10026,10 +10037,12 @@ impl<'db> TypeInferenceBuilder<'db> { .collect(), ret: Box::new( sig.return_type - .as_ref() - .map(|te| { - crate::lower_type_expr::lower_type_expr( - te, &sig_scope, &mut diags, + .map(|id| { + crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + id, + &sig_scope, + &mut diags, ) }) .unwrap_or(Ty::Unknown { @@ -12063,7 +12076,7 @@ impl<'db> TypeInferenceBuilder<'db> { } let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, method_id); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let mut diags = Vec::new(); // Build the self type WITH concrete type args, or TypeVars for @@ -12196,17 +12209,21 @@ impl<'db> TypeInferenceBuilder<'db> { .params .iter() .map(|param| { - let param_ty = if param.name.as_str() == "self" + let is_unannotated_self = param.name.as_str() == "self" && matches!( - param.ty.kind, - baml_compiler2_ast::TypeExprKind::Unknown { .. } - ) { + sig.type_refs[param.type_ref].kind, + baml_compiler2_hir::type_ref::TypeRefKind::Unknown + ); + let param_ty = if is_unannotated_self { // self with no annotation → use the enclosing class type class_ty.clone() } else { crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr( - ¶m.ty, &ctx, &mut diags, + &crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param.type_ref, + &ctx, + &mut diags, ), &bindings, ) @@ -12224,10 +12241,14 @@ impl<'db> TypeInferenceBuilder<'db> { .collect(), ret: Box::new( sig.return_type - .as_ref() - .map(|te| { + .map(|id| { crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr(te, &ctx, &mut diags), + &crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + id, + &ctx, + &mut diags, + ), &bindings, ) }) @@ -12384,7 +12405,7 @@ impl<'db> TypeInferenceBuilder<'db> { // every param and the return type below. let generic_params: Vec<_> = bindings.keys().cloned().collect(); let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, method_id); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let mut diags = Vec::new(); // Build the class type for self parameter resolution. // For generics, apply type_args (e.g. Array). @@ -12459,16 +12480,18 @@ impl<'db> TypeInferenceBuilder<'db> { .params .iter() .map(|param| { - let ty = if param.name.as_str() == "self" + let is_unannotated_self = param.name.as_str() == "self" && matches!( - param.ty.kind, - baml_compiler2_ast::TypeExprKind::Unknown { .. } - ) { + sig.type_refs[param.type_ref].kind, + baml_compiler2_hir::type_ref::TypeRefKind::Unknown + ); + let ty = if is_unannotated_self { builtin_class_ty.clone() } else { crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr( - ¶m.ty, + &crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param.type_ref, &crate::lower_type_expr::ScopeCtx { db, package_items: self.package_items, @@ -12495,11 +12518,11 @@ impl<'db> TypeInferenceBuilder<'db> { .collect(); let ret = sig .return_type - .as_ref() - .map(|te| { + .map(|id| { crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr( - te, + &crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + id, &crate::lower_type_expr::ScopeCtx { db, package_items: self.package_items, diff --git a/baml_language/crates/baml_compiler2_tir/src/callable.rs b/baml_language/crates/baml_compiler2_tir/src/callable.rs index 8751cc5b0a1..56e16175e14 100644 --- a/baml_language/crates/baml_compiler2_tir/src/callable.rs +++ b/baml_language/crates/baml_compiler2_tir/src/callable.rs @@ -90,13 +90,10 @@ fn signature_cycle_initial_callable_throws<'db>( }; let mut facts = BTreeSet::new(); for param in &sig.params { - let Some(param_ty) = param.type_ref else { - continue; - }; let mut diags = Vec::new(); let lowered = crate::lower_type_expr::lower_type_ref( &sig.type_refs, - param_ty, + param.type_ref, ¶m_scope, &mut diags, ); From 3c78bbdb9b3b65f6c13359b8c796e953c5ceb721 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Tue, 14 Jul 2026 15:43:01 -0700 Subject: [PATCH 13/20] Use type refs for more tir stuff --- .../crates/baml_compiler2_tir/src/builder.rs | 10 +++-- .../baml_compiler2_tir/src/inference.rs | 20 +++++----- .../src/package_interface.rs | 39 +++++++++++-------- 3 files changed, 38 insertions(+), 31 deletions(-) diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index b9e7b405c3b..c9b1ffabfb0 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -8187,18 +8187,20 @@ impl<'db> TypeInferenceBuilder<'db> { pub fn check_throws_contract( &mut self, body: &ExprBody, - declared_throws: Option<&TypeExpr>, + type_refs: &baml_compiler2_hir::type_ref::TypeRefStore, + declared_throws: Option, throws_span: Option, fallback_span: TextRange, warn_extraneous: bool, ) { - let Some(declared_expr) = declared_throws else { + let Some(declared_id) = declared_throws else { return; }; let mut diags = Vec::new(); - let declared_ty = crate::lower_type_expr::lower_type_expr( - declared_expr, + let declared_ty = crate::lower_type_expr::lower_type_ref( + type_refs, + declared_id, &crate::lower_type_expr::ScopeCtx { db: self.context.db(), package_items: self.package_items, diff --git a/baml_language/crates/baml_compiler2_tir/src/inference.rs b/baml_language/crates/baml_compiler2_tir/src/inference.rs index 9179e74b08c..d17d27cdba7 100644 --- a/baml_language/crates/baml_compiler2_tir/src/inference.rs +++ b/baml_language/crates/baml_compiler2_tir/src/inference.rs @@ -1672,7 +1672,7 @@ pub fn infer_scope_types<'db>( let local_id = &func_loc.id(db); let func_data = &item_tree[func_loc.id(db)]; let body = baml_compiler2_ppir::function_body(db, func_loc); - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let enclosing_impl = match item_tree.method_owners.get(local_id) { Some(baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(impl_id)) => { @@ -2009,12 +2009,12 @@ pub fn infer_scope_types<'db>( bounds: &body_bounds, self_ty: self_ty.clone(), }; - let lower_with_self = |te: &baml_compiler2_ast::TypeExpr, + let lower_with_self = |id: baml_compiler2_hir::type_ref::TypeRefId, diags: &mut Vec< crate::infer_context::TirTypeError, >| { crate::generics::substitute_ty( - &crate::lower_type_expr::lower_type_expr(te, &ctx, diags), + &crate::lower_type_expr::lower_type_ref(&sig.type_refs, id, &ctx, diags), &type_bindings, ) }; @@ -2022,11 +2022,10 @@ pub fn infer_scope_types<'db>( // Get declared return type let return_ty = sig .return_type - .as_ref() - .map(|te| { + .map(|id| { let span = sig_sm.return_type_span.unwrap_or(func_data.span); let mut diags = Vec::new(); - let ty = lower_with_self(te, &mut diags); + let ty = lower_with_self(id, &mut diags); for diag in diags { builder.report_at_span(diag, span); } @@ -2052,8 +2051,8 @@ pub fn infer_scope_types<'db>( let mut param_ty_validated = false; let param_ty = if param.name.as_str() == "self" && matches!( - param.ty.kind, - baml_compiler2_ast::TypeExprKind::Unknown { .. } + sig.type_refs[param.type_ref].kind, + baml_compiler2_hir::type_ref::TypeRefKind::Unknown ) { // `self`'s type is the method's `Self` receiver, resolved once // above (rigid `Self` var for an interface's own default method, @@ -2064,7 +2063,7 @@ pub fn infer_scope_types<'db>( } else { param_ty_validated = true; let mut param_diags = Vec::new(); - let ty = lower_with_self(¶m.ty, &mut param_diags); + let ty = lower_with_self(param.type_ref, &mut param_diags); for diag in param_diags { builder.report_at_span(diag, param_type_span); } @@ -2105,7 +2104,8 @@ pub fn infer_scope_types<'db>( if !is_auto_derive { builder.check_throws_contract( expr_body, - sig.throws.as_ref(), + &sig.type_refs, + sig.throws, sig_sm.throws_type_span, func_data.span, true, diff --git a/baml_language/crates/baml_compiler2_tir/src/package_interface.rs b/baml_language/crates/baml_compiler2_tir/src/package_interface.rs index 8c6e2987a26..620c52e5210 100644 --- a/baml_language/crates/baml_compiler2_tir/src/package_interface.rs +++ b/baml_language/crates/baml_compiler2_tir/src/package_interface.rs @@ -295,7 +295,7 @@ fn lower_class_method_signature<'db>( ns_path: &[Name], diags: &mut Vec, ) -> LoweredClassMethodSignature { - let sig = baml_compiler2_ppir::elaborated_function_signature(db, method_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, method_loc); let body = baml_compiler2_ppir::function_body(db, method_loc); let mut all_generic_params = class_data.generic_params.clone(); @@ -323,20 +323,21 @@ fn lower_class_method_signature<'db>( bounds: method_bounds, self_ty: Some(self_ty.clone()), }; - let lower_with_self = |te: &baml_compiler2_ast::TypeExpr, diags: &mut Vec| { - crate::lower_type_expr::lower_type_expr(te, &ctx, diags) + let lower_with_self = |id: baml_compiler2_hir::type_ref::TypeRefId, + diags: &mut Vec| { + crate::lower_type_expr::lower_type_ref(&sig.type_refs, id, &ctx, diags) }; let mut params = Vec::new(); for param in &sig.params { let param_ty = if param.name.as_str() == "self" && matches!( - param.ty.kind, - baml_compiler2_ast::TypeExprKind::Unknown { .. } + sig.type_refs[param.type_ref].kind, + baml_compiler2_hir::type_ref::TypeRefKind::Unknown ) { self_ty.clone() } else { - lower_with_self(¶m.ty, diags) + lower_with_self(param.type_ref, diags) }; params.push(exported_function_param( param.name.clone(), @@ -345,14 +346,14 @@ fn lower_class_method_signature<'db>( )); } - let return_type = sig.return_type.as_ref().map_or( + let return_type = sig.return_type.map_or( Ty::Unknown { attr: TyAttr::default(), }, - |te| lower_with_self(te, diags), + |id| lower_with_self(id, diags), ); - let declared_throws = sig.throws.as_ref().map(|te| lower_with_self(te, diags)); + let declared_throws = sig.throws.map(|id| lower_with_self(id, diags)); let callable_throws = crate::callable::callable_throws(db, method_loc).clone(); let builtin_kind = match body.as_ref() { @@ -504,7 +505,7 @@ fn lower_function_export<'db>( name: &Name, ) -> ExportedFunction { let func_ns = file_package::file_package(db, func_loc.file(db)).namespace_path; - let sig = baml_compiler2_ppir::elaborated_function_signature(db, func_loc); + let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); let body = baml_compiler2_ppir::function_body(db, func_loc); let mut diags = Vec::new(); let function_generic_params: Vec = sig @@ -528,7 +529,12 @@ fn lower_function_export<'db>( let mut params = Vec::new(); for param in &sig.params { - let param_ty = crate::lower_type_expr::lower_type_expr(¶m.ty, &sig_scope, &mut diags); + let param_ty = crate::lower_type_expr::lower_type_ref( + &sig.type_refs, + param.type_ref, + &sig_scope, + &mut diags, + ); params.push(exported_function_param( param.name.clone(), param_ty, @@ -536,17 +542,16 @@ fn lower_function_export<'db>( )); } - let return_type = sig.return_type.as_ref().map_or( + let return_type = sig.return_type.map_or( Ty::Unknown { attr: TyAttr::default(), }, - |te| crate::lower_type_expr::lower_type_expr(te, &sig_scope, &mut diags), + |id| crate::lower_type_expr::lower_type_ref(&sig.type_refs, id, &sig_scope, &mut diags), ); - let declared_throws = sig - .throws - .as_ref() - .map(|te| crate::lower_type_expr::lower_type_expr(te, &sig_scope, &mut diags)); + let declared_throws = sig.throws.map(|id| { + crate::lower_type_expr::lower_type_ref(&sig.type_refs, id, &sig_scope, &mut diags) + }); let callable_throws = crate::callable::callable_throws(db, func_loc).clone(); let builtin_kind = match body.as_ref() { From 03a91241976aa60e17fe7ec54320adaae4fe086b Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Tue, 14 Jul 2026 17:15:30 -0700 Subject: [PATCH 14/20] Move span for fn+class+iface+let to source map Span changes with whitespace so is not a cache-friendly strategy (and it can be ambiguous for synthetic items). For this reason we also change several points where we were looking up by span to instead use the much more stable id. --- .../crates/baml_compiler2_mir/src/lower.rs | 50 ++++---------- .../baml_compiler2_tir/src/inference.rs | 36 +++++----- .../baml_lsp2_actions/src/annotations.rs | 30 ++------ .../crates/baml_lsp2_actions/src/check.rs | 19 +---- baml_language/crates/baml_project/src/db.rs | 69 +++++++++++-------- 5 files changed, 81 insertions(+), 123 deletions(-) diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 929549dccb0..497de8b0aee 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -2503,34 +2503,15 @@ impl<'db> LoweringContext<'db> { ) -> Self { let file = func_loc.file(db); - // --- Resolve FunctionLoc → FileScopeId via span --- let item_tree = file_item_tree(db, file); let func_data = &item_tree[func_loc.id(db)]; - let func_span = func_data.span; - let index = file_semantic_index(db, file); - // For synthesized functions whose span is `0..0` (e.g. `$init_test_N`), - // `scope_at_offset` may return a descendant Lambda scope instead of the - // Function scope itself, because all synthesized expressions share span - // `0..0` and the descendant search finds a matching lambda first. - // Avoid this by searching explicitly for a `ScopeKind::Function` scope - // with the correct name and span before falling back to `scope_at_offset`. - let func_scope_id: FileScopeId = index - .scopes - .iter() - .enumerate() - .find_map(|(i, scope)| { - if scope.kind == baml_compiler2_hir::scope::ScopeKind::Function - && scope.range == func_span - && scope.name.as_ref() == Some(&func_data.name) - { - #[allow(clippy::cast_possible_truncation)] - Some(FileScopeId::new(i as u32)) - } else { - None - } - }) - .unwrap_or_else(|| index.scope_at_offset(func_span.start(), Some(&func_data.name))); + // The scope this function opened, from the recorded item↔scope index. + // Exact — no span match, so companion functions and synthesized `0..0` + // functions (which the old scan special-cased) resolve correctly. + let func_scope_id = baml_compiler2_ppir::item_data::function_scope(db, func_loc) + .expect("every item-tree function has a recorded scope") + .file_scope_id(db); // --- Collect per-scope TIR inference views (func + all descendants) --- // Borrows the Salsa-cached `infer_scope_types` results instead of @@ -2753,14 +2734,13 @@ impl<'db> LoweringContext<'db> { ) -> Self { let file = let_loc.file(db); - // --- Resolve LetLoc → FileScopeId via span --- let item_tree = file_item_tree(db, file); let let_data = &item_tree[let_loc.id(db)]; - let let_span = let_data.span; let let_name = let_data.name.clone(); - let index = file_semantic_index(db, file); - let let_scope_id: FileScopeId = index.scope_at_offset(let_span.start(), Some(&let_name)); + let let_scope_id = baml_compiler2_ppir::item_data::let_scope(db, let_loc) + .expect("every item-tree let has a recorded scope") + .file_scope_id(db); // --- Collect per-scope TIR inference views (let + all descendants) --- // Borrows the Salsa-cached `infer_scope_types` results instead of @@ -3994,13 +3974,13 @@ impl<'db> LoweringContext<'db> { let index = file_semantic_index(self.db, self.file); let item_tree = file_item_tree(self.db, self.file); let func_data = &item_tree[func_loc.id(self.db)]; + let func_span = baml_compiler2_ppir::item_data::function_source_map(self.db, func_loc).span; // Set the function-level span on the builder so MirFunction::span is populated. - self.builder.set_span(baml_base::Span::new( - self.file.file_id(self.db), - func_data.span, - )); - let func_scope_id: FileScopeId = - index.scope_at_offset(func_data.span.start(), Some(&func_data.name)); + self.builder + .set_span(baml_base::Span::new(self.file.file_id(self.db), func_span)); + let func_scope_id = baml_compiler2_ppir::item_data::function_scope(self.db, func_loc) + .expect("every item-tree function has a recorded scope") + .file_scope_id(self.db); let func_scope = &index.scopes[func_scope_id.index() as usize]; let enclosing_class_name: Option = func_scope.parent.and_then(|parent_idx| { let parent = &index.scopes[parent_idx.index() as usize]; diff --git a/baml_language/crates/baml_compiler2_tir/src/inference.rs b/baml_language/crates/baml_compiler2_tir/src/inference.rs index d17d27cdba7..e4b9c5b679d 100644 --- a/baml_language/crates/baml_compiler2_tir/src/inference.rs +++ b/baml_language/crates/baml_compiler2_tir/src/inference.rs @@ -1671,6 +1671,8 @@ pub fn infer_scope_types<'db>( { let local_id = &func_loc.id(db); let func_data = &item_tree[func_loc.id(db)]; + let func_span = + baml_compiler2_ppir::item_data::function_source_map(db, func_loc).span; let body = baml_compiler2_ppir::function_body(db, func_loc); let sig = baml_compiler2_ppir::item_data::elaborated_function_data(db, func_loc); @@ -1684,7 +1686,7 @@ pub fn infer_scope_types<'db>( let mut env = GenericEnv::from_params(sig.user_generic_params.clone()); env.params .extend(sig.synthetic_effect_params.iter().cloned()); - report_duplicate_generic_params(&builder, &sig.user_generic_params, func_data.span); + report_duplicate_generic_params(&builder, &sig.user_generic_params, func_span); if let Some(imp) = enclosing_impl { let (impl_generic_params, impl_generic_bounds) = free_impl_generics(imp); for mp in &sig.user_generic_params { @@ -1693,7 +1695,7 @@ pub fn infer_scope_types<'db>( crate::infer_context::TirTypeError::TypeParamShadowedImplParam { param_name: mp.clone(), }, - func_data.span, + func_span, ); } } @@ -1714,7 +1716,7 @@ pub fn infer_scope_types<'db>( type_name: parent.type_name.clone(), owner: parent.owner, }, - func_data.span, + func_span, ); } } @@ -1774,7 +1776,7 @@ pub fn infer_scope_types<'db>( pkg_items, &pkg_info.namespace_path, &env, - func_data.span, + func_span, ); if let Some(sm) = baml_compiler2_ppir::function_body_source_map(db, func_loc) { builder.set_body_source_map(sm); @@ -2023,7 +2025,7 @@ pub fn infer_scope_types<'db>( let return_ty = sig .return_type .map(|id| { - let span = sig_sm.return_type_span.unwrap_or(func_data.span); + let span = sig_sm.return_type_span.unwrap_or(func_span); let mut diags = Vec::new(); let ty = lower_with_self(id, &mut diags); for diag in diags { @@ -2107,7 +2109,7 @@ pub fn infer_scope_types<'db>( &sig.type_refs, sig.throws, sig_sm.throws_type_span, - func_data.span, + func_span, true, ); } @@ -2384,11 +2386,9 @@ pub fn infer_scope_types<'db>( let scope_item = baml_compiler2_ppir::item_data::scope_owner(db, scope_id); if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Class(class_loc)) = scope_item { let class_data = &item_tree[class_loc.id(db)]; - report_duplicate_generic_params( - &builder, - &class_data.generic_params, - class_data.span, - ); + let class_span = + baml_compiler2_ppir::item_data::class_source_map(db, class_loc).span; + report_duplicate_generic_params(&builder, &class_data.generic_params, class_span); let mut env = GenericEnv::from_params(class_data.generic_params.clone()); env.add_bounds_for_declared_params( &class_data.generic_params, @@ -2400,7 +2400,7 @@ pub fn infer_scope_types<'db>( pkg_items, &pkg_info.namespace_path, &env, - class_data.span, + class_span, ); let resolved = resolve_class_fields(db, class_loc); for (field, (_, ty, _)) in class_data.fields.iter().zip(resolved.fields.iter()) { @@ -2413,11 +2413,9 @@ pub fn infer_scope_types<'db>( scope_item { let iface_data = &item_tree[iface_loc.id(db)]; - report_duplicate_generic_params( - &builder, - &iface_data.generic_params, - iface_data.span, - ); + let iface_span = + baml_compiler2_ppir::item_data::interface_source_map(db, iface_loc).span; + report_duplicate_generic_params(&builder, &iface_data.generic_params, iface_span); // Associated types share the interface's type-level namespace with its // generic parameters (a bare `Assoc` reference lowers as a type variable), // so a name collision — with a parameter or another associated type — @@ -2454,7 +2452,7 @@ pub fn infer_scope_types<'db>( if let Some(chain) = crate::interfaces::interface_requires_cycle(db, iface_loc) { builder.report_at_span( crate::infer_context::TirTypeError::InterfaceRequiresCycle { chain }, - iface_data.span, + iface_span, ); } // E0136: a field type may not name the *bare* `Self` type (a recursive field @@ -2588,7 +2586,7 @@ pub fn infer_scope_types<'db>( pkg_items, &pkg_info.namespace_path, &iface_env, - iface_data.span, + iface_span, ); // Computed once per env — every signature type expr below shares it. let iface_env_bounds = diff --git a/baml_language/crates/baml_lsp2_actions/src/annotations.rs b/baml_language/crates/baml_lsp2_actions/src/annotations.rs index e9d23589433..aaabc8be478 100644 --- a/baml_language/crates/baml_lsp2_actions/src/annotations.rs +++ b/baml_language/crates/baml_lsp2_actions/src/annotations.rs @@ -57,11 +57,7 @@ use baml_compiler2_ast::{ Expr, ExprId, Stmt, ast::{AstSourceMap, DeclarativeMeta, ExprBody, FunctionBodyDef, FunctionOrigin}, }; -use baml_compiler2_hir::{ - body::FunctionBody, - loc::FunctionLoc, - scope::{FileScopeId, ScopeKind}, -}; +use baml_compiler2_hir::{body::FunctionBody, loc::FunctionLoc, scope::FileScopeId}; use baml_compiler2_tir::{inference::infer_scope_types, ty::Ty}; use text_size::TextSize; @@ -143,9 +139,12 @@ pub fn file_annotations(db: &dyn Db, file: SourceFile) -> Vec continue; }; - let owner_scope = function_scope_for(index, func_data.span, &func_data.name) + let owner_scope = baml_compiler2_ppir::item_data::function_scope(db, func_loc) + .map(|scope| scope.file_scope_id(db)) .unwrap_or_else(|| { - index.scope_at_offset(func_data.span.start(), Some(&func_data.name)) + let func_span = + baml_compiler2_ppir::item_data::function_source_map(db, func_loc).span; + index.scope_at_offset(func_span.start(), Some(&func_data.name)) }); process_body( db, @@ -354,23 +353,6 @@ fn is_synthetic_registration(body: &ExprBody, callee: ExprId) -> bool { matches!(name, "register_test" | "register_test_set") } -fn function_scope_for( - index: &SemanticIndex<'_>, - span: text_size::TextRange, - name: &baml_base::Name, -) -> Option { - index - .scopes - .iter() - .enumerate() - .find(|(_, scope)| { - matches!(scope.kind, ScopeKind::Function) - && scope.range == span - && scope.name.as_ref() == Some(name) - }) - .map(|(idx, _)| FileScopeId::new(u32::try_from(idx).expect("scope index fits in u32"))) -} - fn scope_at_offset_within_body( index: &SemanticIndex<'_>, offset: TextSize, diff --git a/baml_language/crates/baml_lsp2_actions/src/check.rs b/baml_language/crates/baml_lsp2_actions/src/check.rs index 4a8013d5a66..50733b5290e 100644 --- a/baml_language/crates/baml_lsp2_actions/src/check.rs +++ b/baml_language/crates/baml_lsp2_actions/src/check.rs @@ -396,7 +396,8 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec { param_types.push((param.name.clone(), param_ty)); } - if let Some(scope_id) = function_scope_id(index, func_data) { + let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(db, file, *local_id); + if let Some(scope_id) = baml_compiler2_ppir::item_data::function_scope(db, func_loc) { let context = baml_compiler2_tir::infer_context::InferContext::new(db, scope_id); let mut builder = baml_compiler2_tir::builder::TypeInferenceBuilder::new( context, @@ -1648,22 +1649,6 @@ fn collect_type_aliases_for_resolution_context<'db>( aliases } -fn function_scope_id<'db>( - index: &baml_compiler2_hir::semantic_index::FileSemanticIndex<'db>, - func_data: &baml_compiler2_hir::item_tree::Function, -) -> Option> { - index - .scopes - .iter() - .zip(index.scope_ids.iter()) - .find_map(|(scope, scope_id)| { - (matches!(scope.kind, ScopeKind::Function) - && scope.range == func_data.span - && scope.name.as_ref() == Some(&func_data.name)) - .then_some(*scope_id) - }) -} - fn is_function_default_signature_diagnostic( diag: &baml_compiler2_tir::infer_context::TirDiagnostic<'_>, ) -> bool { diff --git a/baml_language/crates/baml_project/src/db.rs b/baml_language/crates/baml_project/src/db.rs index 6d396a16d43..a8a7cbc9d23 100644 --- a/baml_language/crates/baml_project/src/db.rs +++ b/baml_language/crates/baml_project/src/db.rs @@ -173,6 +173,19 @@ pub struct ProjectDatabase { removed_file_tombstones: HashMap, } +/// Origin-preference order for disambiguating functions that share one +/// declaration span (a declarative LLM function and its `$stream` / +/// `$parse_stream` companions): the user-authored function sorts first. +fn func_origin_rank(origin: baml_compiler2_ast::ast::FunctionOrigin) -> u8 { + use baml_compiler2_ast::ast::FunctionOrigin; + match origin { + FunctionOrigin::UserDefined => 0, + FunctionOrigin::Companion => 1, + FunctionOrigin::Internal => 2, + FunctionOrigin::AutoDerive => 3, + } +} + #[salsa::db] impl salsa::Database for ProjectDatabase {} @@ -660,6 +673,8 @@ impl ProjectDatabase { let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, *local_id); + let func_span = + baml_compiler2_ppir::item_data::function_source_map(self, func_loc).span; let body = baml_compiler2_ppir::function_body(self, func_loc); // Check if this is an LLM function via declarative_meta (not body variant, @@ -682,9 +697,7 @@ impl ProjectDatabase { "unknown".to_string() }; let mut graph = build_llm_control_flow_graph(function_name, &client_name); - if let Some(source_span) = - self.source_span_for_range(source_file, func_data.span) - { + if let Some(source_span) = self.source_span_for_range(source_file, func_span) { if let Some(node) = graph.nodes.values_mut().next() { node.source_span = Some(source_span); } @@ -709,7 +722,7 @@ impl ProjectDatabase { // whole function declaration so clicking the root in the // playground selects the function (mirrors the LLM path above). if let Some(root_span) = - self.source_span_for_range(source_file, func_data.span) + self.source_span_for_range(source_file, func_span) { if let Some(root) = graph.nodes.values_mut().find(|node| { node.node_type @@ -831,7 +844,7 @@ impl ProjectDatabase { let mut unique_title = None; for source_file in self.file_map.values().copied() { let index = baml_compiler2_ppir::file_semantic_index(self, source_file); - for func_data in index.item_tree.functions.values() { + for (local_id, func_data) in &index.item_tree.functions { if !self.function_name_matches_source_name( source_file, &func_data.name, @@ -839,8 +852,12 @@ impl ProjectDatabase { ) { continue; } + let func_loc = + baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, *local_id); + let func_span = + baml_compiler2_ppir::item_data::function_source_map(self, func_loc).span; let text = source_file.text(self); - let start = usize::from(func_data.span.start()).min(text.len()); + let start = usize::from(func_span.start()).min(text.len()); if let Some(title) = header_title_above(&text[..start]) { match &unique_title { Some(existing) if existing != &title => return None, @@ -1263,7 +1280,6 @@ impl ProjectDatabase { source_file: SourceFile, offset: text_size::TextSize, ) -> Option<(String, bool)> { - use baml_compiler2_ast::ast::FunctionOrigin; use baml_compiler2_hir::scope::ScopeKind; let index = baml_compiler2_ppir::file_semantic_index(self, source_file); @@ -1278,19 +1294,19 @@ impl ProjectDatabase { let func_scope_range = index.scopes[func_scope_id.index() as usize].range; - // Match against item tree functions by span - let item_tree = &index.item_tree; - let (local_id, _) = item_tree + // A declarative LLM function and its `$stream`/`$parse_stream` companions + // share one declaration span, hence one scope range — so multiple + // functions match here. Prefer the user-authored one (origin order). + let (local_id, _) = index + .item_tree .functions .iter() - .filter(|(_, func_data)| func_data.span == func_scope_range) - .min_by_key(|(_, func_data)| match func_data.origin { - FunctionOrigin::UserDefined => 0, - FunctionOrigin::Companion => 1, - FunctionOrigin::Internal => 2, - FunctionOrigin::AutoDerive => 3, - })?; - + .filter(|(id, _)| { + let loc = baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, **id); + baml_compiler2_ppir::item_data::function_source_map(self, loc).span + == func_scope_range + }) + .min_by_key(|(_, func_data)| func_origin_rank(func_data.origin))?; let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, *local_id); let sig = baml_compiler2_ppir::function_signature(self, func_loc); let body = baml_compiler2_ppir::function_body(self, func_loc); @@ -1386,7 +1402,6 @@ impl ProjectDatabase { source_file: SourceFile, offset: text_size::TextSize, ) -> (Option, Vec) { - use baml_compiler2_ast::ast::FunctionOrigin; use baml_compiler2_hir::scope::ScopeKind; let index = baml_compiler2_ppir::file_semantic_index(self, source_file); @@ -1401,18 +1416,16 @@ impl ProjectDatabase { }; let func_scope_range = index.scopes[func_scope_id.index() as usize].range; - - let item_tree = &index.item_tree; - if let Some((local_id, _)) = item_tree + if let Some((local_id, _)) = index + .item_tree .functions .iter() - .filter(|(_, func_data)| func_data.span == func_scope_range) - .min_by_key(|(_, func_data)| match func_data.origin { - FunctionOrigin::UserDefined => 0, - FunctionOrigin::Companion => 1, - FunctionOrigin::Internal => 2, - FunctionOrigin::AutoDerive => 3, + .filter(|(id, _)| { + let loc = baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, **id); + baml_compiler2_ppir::item_data::function_source_map(self, loc).span + == func_scope_range }) + .min_by_key(|(_, func_data)| func_origin_rank(func_data.origin)) { let func_loc = baml_compiler2_hir::loc::FunctionLoc::new(self, source_file, *local_id); let Some(source_map) = baml_compiler2_ppir::function_body_source_map(self, func_loc) From 6a1ecd6f04dda9eb39cf3a1691650e11d4fcbd4d Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Wed, 15 Jul 2026 10:26:40 -0700 Subject: [PATCH 15/20] Delete old `Self`-substitution logic --- .../crates/baml_compiler2_emit/src/lib.rs | 44 ++-- .../crates/baml_compiler2_mir/src/lower.rs | 21 +- .../baml_compiler2_tir/src/lower_type_expr.rs | 208 ------------------ 3 files changed, 41 insertions(+), 232 deletions(-) diff --git a/baml_language/crates/baml_compiler2_emit/src/lib.rs b/baml_language/crates/baml_compiler2_emit/src/lib.rs index c7609da67af..eb862fc4ec5 100644 --- a/baml_language/crates/baml_compiler2_emit/src/lib.rs +++ b/baml_language/crates/baml_compiler2_emit/src/lib.rs @@ -3674,14 +3674,39 @@ fn compute_function_metadata_from_item_tree( None => rustc_hash::FxHashMap::default(), }; + // The concrete receiver (`ClassName` or a free-impl `for` target), lowered + // once. A non-interface method's `Self` / `Self.Assoc` then root at it through the + // `self_ty` channel — the same projection path the interface branch drives with its + // rigid `Self` type variable. `None` for a free function (no receiver in scope) and + // unused for interface methods (which bind `Self` via `interface_signature_bindings`). + let receiver_ty: Option = if enclosing_interface.is_none() { + self_replacement.as_ref().map(|replacement| { + let mut receiver_diags = Vec::new(); + baml_compiler2_tir::lower_type_expr::lower_type_expr( + replacement, + &baml_compiler2_tir::lower_type_expr::ScopeCtx { + db, + package_items: pkg_items, + ns_context: &pkg_info.namespace_path, + generic_params: &enclosing_generics, + bounds: &scope_bounds, + self_ty: None, + }, + &mut receiver_diags, + ) + }) + } else { + None + }; + // Lower a signature type expression against this method's scope, recording any // diagnostics into `diags`. For an interface method the associated-type / `Self` // bindings are applied by lowering with their names in scope (so `Item` lowers // to `TypeVar(Item)`) and then substituting; for every other method `Self` is - // first replaced by the concrete receiver type. In both cases `scope_bounds` - // drives projection resolution. Namespace-relative resolution (e.g. `MyLorem` - // in a signature under `ns_lorem/`) uses the file's namespace so a non-root-ns - // class does not erase to `unknown`. + // bound to the concrete `receiver_ty` via the `self_ty` channel. In both cases + // `scope_bounds` drives projection resolution. Namespace-relative resolution + // (e.g. `MyLorem` in a signature under `ns_lorem/`) uses the file's namespace so + // a non-root-ns class does not erase to `unknown`. let lower_scoped = |te: &TypeExpr, diags: &mut Vec| -> baml_compiler2_tir::ty::Ty { @@ -3704,22 +3729,15 @@ fn compute_function_metadata_from_item_tree( &interface_signature_bindings, ); } - let resolved_te = match &self_replacement { - Some(replacement) => baml_compiler2_tir::lower_type_expr::substitute_paths_in( - te, - &std::collections::HashMap::from([(Name::new("Self"), replacement.clone())]), - ), - None => te.clone(), - }; baml_compiler2_tir::lower_type_expr::lower_type_expr( - &resolved_te, + te, &baml_compiler2_tir::lower_type_expr::ScopeCtx { db, package_items: pkg_items, ns_context: &pkg_info.namespace_path, generic_params: &enclosing_generics, bounds: &scope_bounds, - self_ty: None, + self_ty: receiver_ty.clone(), }, diags, ) diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 497de8b0aee..23b99259415 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -3256,34 +3256,33 @@ impl<'db> LoweringContext<'db> { /// Lower a method-signature type expression (a parameter or return type) to /// a runtime type. In a method signature `Self` is the receiver type - /// variable and `Self.Assoc` is an associated-type projection onto it. - /// A bare `lower_type_expr_in_ns` has neither in scope and would erase both - /// to `Ty::Unknown`, tripping the runtime lowering boundary — so rewrite - /// `Self.Assoc` paths into projections and bind `Self` as a type variable. + /// variable and `Self.Assoc` is an associated-type projection onto it. A bare + /// `lower_type_expr_in_ns` has neither in scope and would erase both to + /// `Ty::Unknown`, tripping the runtime lowering boundary — so bind `Self` to + /// its rigid type variable through the `self_ty` channel, which roots both a + /// bare `Self` and each `Self.Assoc` projection at it. fn lower_signature_runtime_ty( &self, te: &baml_compiler2_ast::TypeExpr, pkg_items: &baml_compiler2_hir::package::PackageItems<'_>, ns_context: &[baml_base::Name], ) -> RuntimeTy { - let self_subst = std::collections::HashMap::from([( - baml_base::Name::new("Self"), - baml_compiler2_tir::lower_type_expr::type_expr_for_name(baml_base::Name::new("Self")), - )]); - let te = baml_compiler2_tir::lower_type_expr::substitute_paths_in(te, &self_subst); let mut generic_params = self.enclosing_generic_params(); generic_params.push(baml_base::Name::new("Self")); let generic_param_bounds = self.enclosing_generic_param_bounds(); let mut diags = Vec::new(); let tir_ty = baml_compiler2_tir::lower_type_expr::lower_type_expr( - &te, + te, &baml_compiler2_tir::lower_type_expr::ScopeCtx { db: self.db, package_items: pkg_items, ns_context, generic_params: &generic_params, bounds: &generic_param_bounds, - self_ty: None, + self_ty: Some(Tir2Ty::TypeVar( + baml_base::Name::new("Self"), + TyAttr::default(), + )), }, &mut diags, ); diff --git a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs index 0de2c3568f5..0f7063444b5 100644 --- a/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs +++ b/baml_language/crates/baml_compiler2_tir/src/lower_type_expr.rs @@ -230,151 +230,6 @@ impl<'db> TypeExprContext<'db> for ScopeCtx<'_, 'db> { } } -/// Like [`lower_type_expr`], but resolves unqualified names relative to -/// `ns_context` first (e.g. `["fs"]`), falling back to the package root. -/// -/// Use this when lowering type expressions from a function/class signature -/// that lives in a sub-namespace of its package. For example, `File` in -/// `baml/fs.baml` (namespace `["fs"]`) resolves via `lookup_type(&["fs", "File"])`. -fn projection_chain_from_root( - root: TypeExpr, - members: &[baml_base::Name], - attrs: &[baml_compiler2_ast::RawAttribute], -) -> TypeExpr { - members - .iter() - .enumerate() - .fold(root, |base, (idx, member)| { - let span = base.span; - TypeExprKind::AssociatedTypeProjection { - base: Box::new(base), - interface: None, - member: member.clone(), - attrs: if idx + 1 == members.len() { - attrs.to_vec() - } else { - Vec::new() - }, - } - .at(span) - }) -} - -/// BEP-044 generic bounds: walk `type_expr` and replace any single- -/// segment `Path` whose name appears in `subst` with the matching -/// replacement `TypeExpr`. Used at MIR lowering to substitute bounded -/// type-vars (`T extends Named`) with their bound (`Named`) so the -/// runtime sees a concrete type to dispatch on. -pub fn substitute_paths_in( - type_expr: &TypeExpr, - subst: &std::collections::HashMap, -) -> TypeExpr { - if subst.is_empty() { - return type_expr.clone(); - } - substitute_paths_walk(type_expr, subst) -} - -fn substitute_paths_walk( - ty: &TypeExpr, - subst: &std::collections::HashMap, -) -> TypeExpr { - match &ty.kind { - TypeExprKind::Path { - segments, - generic_args, - associated_type_bindings, - attrs, - } => { - if generic_args.is_empty() - && associated_type_bindings.is_empty() - && let Some(replacement) = subst.get(&segments[0]) - { - if segments.len() == 1 { - return replacement.clone(); - } - return projection_chain_from_root(replacement.clone(), &segments[1..], attrs); - } - TypeExprKind::Path { - segments: segments.clone(), - generic_args: generic_args - .iter() - .map(|a| substitute_paths_walk(a, subst)) - .collect(), - associated_type_bindings: associated_type_bindings - .iter() - .map(|binding| baml_compiler2_ast::AssociatedTypeBinding { - name: binding.name.clone(), - ty: Box::new(substitute_paths_walk(&binding.ty, subst)), - }) - .collect(), - attrs: attrs.clone(), - } - .at(ty.span) - } - TypeExprKind::AssociatedTypeProjection { - base, - interface, - member, - attrs, - } => TypeExprKind::AssociatedTypeProjection { - base: Box::new(substitute_paths_walk(base, subst)), - interface: interface - .as_ref() - .map(|interface| Box::new(substitute_paths_walk(interface, subst))), - member: member.clone(), - attrs: attrs.clone(), - } - .at(ty.span), - TypeExprKind::List { inner, attrs } => TypeExprKind::List { - inner: Box::new(substitute_paths_walk(inner, subst)), - attrs: attrs.clone(), - } - .at(ty.span), - TypeExprKind::Optional { inner, attrs } => TypeExprKind::Optional { - inner: Box::new(substitute_paths_walk(inner, subst)), - attrs: attrs.clone(), - } - .at(ty.span), - TypeExprKind::Map { key, value, attrs } => TypeExprKind::Map { - key: Box::new(substitute_paths_walk(key, subst)), - value: Box::new(substitute_paths_walk(value, subst)), - attrs: attrs.clone(), - } - .at(ty.span), - TypeExprKind::Union { variants, attrs } => TypeExprKind::Union { - variants: variants - .iter() - .map(|v| substitute_paths_walk(v, subst)) - .collect(), - attrs: attrs.clone(), - } - .at(ty.span), - TypeExprKind::Function { - params, - ret, - throws, - attrs, - } => TypeExprKind::Function { - params: params - .iter() - .map(|param| { - let mut param = param.clone(); - param.ty = substitute_paths_walk(¶m.ty, subst); - param - }) - .collect(), - ret: Box::new(substitute_paths_walk(ret, subst)), - throws: throws - .as_ref() - .map(|ty| Box::new(substitute_paths_walk(ty, subst))), - attrs: attrs.clone(), - } - .at(ty.span), - _ => ty.clone(), - } -} - /// Build a `TypeExprKind::Path` referring to a single named type — used to desugar the bare /// `self` receiver to a `Self` path so it flows through normal parameter lowering. pub fn type_expr_for_name(name: baml_base::Name) -> TypeExpr { @@ -3016,69 +2871,6 @@ mod tests { } } - #[test] - fn substitute_paths_recurses_into_function_type() { - let type_expr = TypeExprKind::Function { - params: vec![FunctionTypeParam { - name: Some(Name::new("value")), - optional: false, - ty: path("T"), - }], - ret: Box::new(path("T")), - throws: Some(Box::new(path("E"))), - attrs: vec![], - } - .at(text_size::TextRange::default()); - let mut subst = std::collections::HashMap::new(); - subst.insert( - Name::new("T"), - TypeExprKind::String { attrs: vec![] }.at(text_size::TextRange::default()), - ); - subst.insert( - Name::new("E"), - TypeExprKind::Bool { attrs: vec![] }.at(text_size::TextRange::default()), - ); - - let TypeExprKind::Function { - params, - ret, - throws, - .. - } = substitute_paths_in(&type_expr, &subst).kind - else { - panic!("expected function type"); - }; - - assert!(matches!(params[0].ty.kind, TypeExprKind::String { .. })); - assert!(matches!(ret.kind, TypeExprKind::String { .. })); - assert!(matches!( - throws.as_deref().map(|t| &t.kind), - Some(TypeExprKind::Bool { .. }) - )); - } - - #[test] - fn substitute_paths_rewrites_substituted_root_member_path_to_projection() { - let type_expr = path_segments(&["T", "Item"]); - let replacement = path("Box"); - let mut subst = std::collections::HashMap::new(); - subst.insert(Name::new("T"), replacement.clone()); - - let TypeExprKind::AssociatedTypeProjection { - base, - interface, - member, - .. - } = substitute_paths_in(&type_expr, &subst).kind - else { - panic!("expected associated type projection"); - }; - - assert_eq!(*base, replacement); - assert!(interface.is_none()); - assert_eq!(member, Name::new("Item")); - } - #[test] fn lower_function_type_preserves_parameter_optionality() { let mut db = TestDb::default(); From c4c221e3b33beeb6251570379634690c336c37fa Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Wed, 15 Jul 2026 11:50:39 -0700 Subject: [PATCH 16/20] Add `function_llm_meta` query Acts as a firewall for llm function metadata that does not require full information or span data --- .../crates/baml_cli/src/test_command.rs | 6 ++-- .../src/item_data/functions.rs | 30 +++++++++++++++++++ .../baml_lsp2_actions/src/annotations.rs | 8 ++--- baml_language/crates/baml_project/src/db.rs | 30 +++++++------------ .../crates/baml_project/src/symbols.rs | 28 ++++++----------- 5 files changed, 57 insertions(+), 45 deletions(-) diff --git a/baml_language/crates/baml_cli/src/test_command.rs b/baml_language/crates/baml_cli/src/test_command.rs index a6c36bb214d..919328fcabe 100644 --- a/baml_language/crates/baml_cli/src/test_command.rs +++ b/baml_language/crates/baml_cli/src/test_command.rs @@ -499,20 +499,20 @@ fn discover_legacy_tests( db: &ProjectDatabase, project: baml_workspace::Project, ) -> Vec { - use baml_db::baml_compiler2_hir; + use baml_db::baml_compiler2_ppir::item_data::{file_tests, test_data}; let mut tests = Vec::new(); let root = project.root(db); for source_file in db.get_source_files() { - let item_tree = baml_compiler2_hir::file_item_tree(db, source_file); // Root-relative for display, matching how emit records source paths — // keeps `--list` output identical between compiled and // bytecode-cache-served runs. let file_path = source_file.path(db); let file_path = file_path.strip_prefix(&root).unwrap_or(&file_path); - for test in item_tree.tests.values() { + for test_loc in file_tests(db, source_file) { + let test = test_data(db, *test_loc); for func_ref in &test.function_refs { tests.push(LegacyTest { function_name: func_ref.to_string(), diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index 8a299313665..e8b687a3609 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -67,6 +67,36 @@ pub fn function_source_map<'db>( lower(db, function).1 } +/// The span-free facts a function's `declarative_meta` exposes for an LLM +/// (`{ client …; prompt … }`) function: its declared client name (itself optional). +/// [`function_llm_meta`] wraps this in an `Option`, where `None` marks a non-LLM +/// function — so a client name can never be attached to one. +/// +/// The full [`ast::LlmBodyDef`] (prompt template, interpolation spans, companion +/// bodies) carries spans, so it stays behind a body-ish read. Metadata consumers +/// that only need these facts front the item tree through this projection and get +/// early cutoff — editing an unrelated function no longer invalidates them. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FunctionLlmMeta { + pub client_name: Option, +} + +/// The [`FunctionLlmMeta`] projection for one function, or `None` when it is not an +/// LLM function. See [`FunctionLlmMeta`] for why the full LLM body is excluded. +#[salsa::tracked(returns(ref))] +pub fn function_llm_meta<'db>( + db: &'db dyn crate::Db, + function: FunctionLoc<'db>, +) -> Option { + let item_tree = crate::file_item_tree(db, function.file(db)); + item_tree[function.id(db)] + .declarative_meta + .as_ref() + .map(|ast::DeclarativeMeta::Llm(llm)| FunctionLlmMeta { + client_name: llm.client.clone(), + }) +} + /// Span-free semantic data for a function's *elaborated* signature — the /// canonical callable view TIR consumes. /// diff --git a/baml_language/crates/baml_lsp2_actions/src/annotations.rs b/baml_language/crates/baml_lsp2_actions/src/annotations.rs index aaabc8be478..bf3c2be328f 100644 --- a/baml_language/crates/baml_lsp2_actions/src/annotations.rs +++ b/baml_language/crates/baml_lsp2_actions/src/annotations.rs @@ -55,7 +55,7 @@ use baml_base::SourceFile; use baml_compiler2_ast::{ Expr, ExprId, Stmt, - ast::{AstSourceMap, DeclarativeMeta, ExprBody, FunctionBodyDef, FunctionOrigin}, + ast::{AstSourceMap, ExprBody, FunctionBodyDef, FunctionOrigin}, }; use baml_compiler2_hir::{body::FunctionBody, loc::FunctionLoc, scope::FileScopeId}; use baml_compiler2_tir::{inference::infer_scope_types, ty::Ty}; @@ -114,6 +114,8 @@ pub fn file_annotations(db: &dyn Db, file: SourceFile) -> Vec let mut out: Vec = Vec::new(); for (func_local_id, func_data) in &item_tree.functions { + let func_loc = FunctionLoc::new(db, file, *func_local_id); + // Process user-written functions and methods, plus the synthesized // `$init_test*` registration functions (so test/testset bodies — which // lower to lambdas — get hints). Skip LLM declarative functions: we must @@ -123,13 +125,11 @@ pub fn file_annotations(db: &dyn Db, file: SourceFile) -> Vec let is_user = func_data.origin == FunctionOrigin::UserDefined; let is_test_init = func_data.name.as_str().starts_with("$init_test"); if (!is_user && !is_test_init) - || matches!(func_data.declarative_meta, Some(DeclarativeMeta::Llm(_))) + || baml_compiler2_ppir::item_data::function_llm_meta(db, func_loc).is_some() { continue; } - let func_loc = FunctionLoc::new(db, file, *func_local_id); - let body = baml_compiler2_hir::body::function_body(db, func_loc); let FunctionBody::Expr(expr_body) = body.as_ref() else { continue; diff --git a/baml_language/crates/baml_project/src/db.rs b/baml_language/crates/baml_project/src/db.rs index a8a7cbc9d23..5b400050b34 100644 --- a/baml_language/crates/baml_project/src/db.rs +++ b/baml_language/crates/baml_project/src/db.rs @@ -677,25 +677,17 @@ impl ProjectDatabase { baml_compiler2_ppir::item_data::function_source_map(self, func_loc).span; let body = baml_compiler2_ppir::function_body(self, func_loc); - // Check if this is an LLM function via declarative_meta (not body variant, - // since compiler2 desugars LLM functions to Expr bodies). - let is_llm = matches!( - func_data.declarative_meta, - Some(baml_compiler2_ast::ast::DeclarativeMeta::Llm(_)) - ); - - result = if is_llm { - let client_name = - if let Some(baml_compiler2_ast::ast::DeclarativeMeta::Llm(ref llm)) = - func_data.declarative_meta - { - llm.client - .as_ref() - .map(|c: &baml_db::Name| c.to_string()) - .unwrap_or_else(|| "unknown".to_string()) - } else { - "unknown".to_string() - }; + // LLM functions desugar to Expr bodies, so it is `declarative_meta` + // (surfaced span-free by `function_llm_meta`) — not the body variant — + // that marks them. + result = if let Some(llm_meta) = + baml_compiler2_ppir::item_data::function_llm_meta(self, func_loc) + { + let client_name = llm_meta + .client_name + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| "unknown".to_string()); let mut graph = build_llm_control_flow_graph(function_name, &client_name); if let Some(source_span) = self.source_span_for_range(source_file, func_span) { if let Some(node) = graph.nodes.values_mut().next() { diff --git a/baml_language/crates/baml_project/src/symbols.rs b/baml_language/crates/baml_project/src/symbols.rs index d8a86400c69..a110dbcf4c0 100644 --- a/baml_language/crates/baml_project/src/symbols.rs +++ b/baml_language/crates/baml_project/src/symbols.rs @@ -5,9 +5,9 @@ use baml_compiler2_hir::{ contributions::{Definition, DefinitionKind}, - file_item_tree, package::{PackageId, package_items}, }; +use baml_compiler2_ppir::item_data::{function_data, function_llm_meta, test_data}; use baml_compiler2_tir::package_interface::package_interface; use baml_db::Name; @@ -138,26 +138,17 @@ pub fn list_functions_with_metadata(db: &ProjectDatabase) -> FunctionListing { for (namespace_path, ns_items) in &pkg.namespaces { for (name, defn) in &ns_items.values { if let Definition::Function(func_loc) = defn { - let item_tree = file_item_tree(db, func_loc.file(db)); - let func = &item_tree[func_loc.id(db)]; - - let is_llm = matches!( - func.declarative_meta, - Some(baml_compiler2_ast::ast::DeclarativeMeta::Llm(_)) - ); - let client_name = - if let Some(baml_compiler2_ast::ast::DeclarativeMeta::Llm(ref llm)) = - func.declarative_meta - { - llm.client.as_ref().map(std::string::ToString::to_string) - } else { - None - }; + let llm_meta = function_llm_meta(db, *func_loc); + let is_llm = llm_meta.is_some(); + let client_name = llm_meta + .as_ref() + .and_then(|meta| meta.client_name.as_ref()) + .map(std::string::ToString::to_string); // Sub-functions have names with '$' (e.g. MyFunc$render_prompt) let is_sub_function = name.as_str().contains('$'); - let origin: FunctionOrigin = func.origin.into(); + let origin: FunctionOrigin = function_data(db, *func_loc).origin.into(); // Companions clone parent params verbatim and non-userDefined // functions are hidden by default — extracting schemas for // them only duplicates payload. The UI degrades to raw mode. @@ -210,8 +201,7 @@ pub fn list_tests_with_metadata(db: &ProjectDatabase) -> Vec { for (namespace_path, ns_items) in &pkg.namespaces { for (name, defn) in &ns_items.values { if let Definition::Test(test_loc) = defn { - let item_tree = file_item_tree(db, test_loc.file(db)); - let test = &item_tree[test_loc.id(db)]; + let test = test_data(db, *test_loc); let args_json = serialize_test_args(&test.args); result.extend(test.function_refs.iter().map(|function_name| { From 72ed7f777a19db92d8f0ccb3d735ce862badb65a Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Wed, 15 Jul 2026 12:08:26 -0700 Subject: [PATCH 17/20] Tests to validate firewall queries --- .../baml_tests/src/incremental/scenarios.rs | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index e69e719c789..c3b9bae496b 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -6,6 +6,7 @@ //! NOTE: These scenarios were ported from the legacy HIR (baml_compiler_hir) //! to compiler2 HIR (baml_compiler2_hir) as part of the compiler2 migration. +use baml_compiler2_tir::inference::infer_scope_types; use baml_db::{SourceFile, baml_compiler2_hir}; use salsa::Setter; @@ -898,3 +899,87 @@ fn elaborated_function_data_cuts_off_and_still_elaborates() { "whitespace/body edits must not change the elaborated signature" ); } + +/// `function_llm_meta` is a *projection*: it exposes only whether a function is an +/// LLM function and its client, deliberately excluding the prompt template (which +/// carries spans and changes constantly). Editing the prompt must therefore leave +/// the projection equal, so consumers that only care about `is_llm`/client cut off +/// even though the underlying `declarative_meta` changed. +#[test] +fn editing_a_function_prompt_preserves_its_llm_meta() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Greet(name: string) -> string {\n client GPT4\n prompt #\"Hi {{name}}\"#\n}\n", + ); + + let before = { + let db = test_db.db(); + baml_compiler2_ppir::item_data::function_llm_meta(db, function_loc(db, file, "Greet")) + .clone() + }; + + // Rewrite only the prompt — the client (the one fact the projection keeps) is + // untouched. + file.set_text(test_db.db_mut()).to( + "function Greet(name: string) -> string {\n client GPT4\n prompt #\"Hello there {{name}}!\"#\n}\n" + .to_string(), + ); + + let after = { + let db = test_db.db(); + baml_compiler2_ppir::item_data::function_llm_meta(db, function_loc(db, file, "Greet")) + .clone() + }; + + assert!(before.is_some(), "`Greet` is an LLM function"); + assert_eq!( + before, after, + "a prompt-only edit must not change is_llm/client — the projection exists to exclude the prompt" + ); +} + +/// PLANNED (plan Phase 5). The whole refactor exists so that a cosmetic edit does +/// not re-run type inference — the exact claim ARCHITECTURE.md:728 makes but that +/// is *false* today: `infer_scope_types` reads the `no_eq` `file_semantic_index` +/// directly (`inference.rs:1524`), so any edit re-executes it. Un-ignore once +/// inference consumes the per-item firewall queries instead of the coarse index. +#[test] +#[ignore = "infer_scope_types still reads the no_eq file_semantic_index directly; un-ignore after it migrates to the firewall queries (plan Phase 5)"] +fn comment_edit_does_not_reexecute_type_inference() { + let mut test_db = IncrementalTestDb::new(); + + let file = test_db.db_mut().add_file( + "test.baml", + "function Add(x: int, y: int) -> int {\n x + y\n}\n", + ); + + // Prime inference for `Add`'s body scope. + let scope_id = { + let db = test_db.db(); + baml_compiler2_ppir::item_data::function_scope(db, function_loc(db, file, "Add")) + .expect("Add has a scope") + }; + let _ = test_db.log_executed(|db| { + let _ = infer_scope_types(db, scope_id); + }); + + // Add a comment: semantically a no-op for the function body. + file.set_text(test_db.db_mut()) + .to("// a comment\nfunction Add(x: int, y: int) -> int {\n x + y\n}\n".to_string()); + + // Re-fetch the scope (its tracked-struct id may have been re-minted by the + // no_eq index) and assert inference is served from cache. + let scope_id = { + let db = test_db.db(); + baml_compiler2_ppir::item_data::function_scope(db, function_loc(db, file, "Add")) + .expect("Add has a scope") + }; + test_db.assert_not_executed( + |db| { + let _ = infer_scope_types(db, scope_id); + }, + &["infer_scope_types"], + ); +} From bc22c086498205a23b769560441ef30e6ffce796 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Thu, 16 Jul 2026 10:50:57 -0700 Subject: [PATCH 18/20] Cleanup & bugfixes --- baml_language/ARCHITECTURE.md | 43 +++++++++++++------ .../crates/baml_cli/src/test_command.rs | 16 ++++++- .../crates/baml_compiler2_emit/src/lib.rs | 6 +-- .../src/item_tree/builder.rs | 2 +- .../baml_compiler2_hir/src/signature.rs | 11 +++-- .../crates/baml_compiler2_mir/src/lower.rs | 8 +++- .../src/item_data/classes.rs | 7 ++- .../src/item_data/enumeration.rs | 9 ++-- .../src/item_data/functions.rs | 7 ++- .../src/item_data/type_aliases.rs | 7 ++- .../baml_lsp2_actions/src/definition.rs | 10 ++++- .../baml_tests/src/incremental/scenarios.rs | 12 +++--- 12 files changed, 97 insertions(+), 41 deletions(-) diff --git a/baml_language/ARCHITECTURE.md b/baml_language/ARCHITECTURE.md index dee4b36f5ec..49795644960 100644 --- a/baml_language/ARCHITECTURE.md +++ b/baml_language/ARCHITECTURE.md @@ -169,6 +169,7 @@ The parser produces a **CST** (Concrete Syntax Tree), which is a lossless, error - **Shadowing rules** — The HIR decides where shadowing is allowed (e.g., a match arm variable may shadow a function parameter). - **Lambda capture analysis** — Which variables a lambda captures is determined here. You don't need to know the type to know *what* is captured, only what names are in scope. - **Package and namespace aggregation** — Cross-file symbol merging happens here. +- **Item tree + type-reference arena** — A span-free `ItemTree` (items keyed by position-independent `LocalItemId`) with a parallel `ItemTreeSourceMap`, plus a flat span-free `TypeRef` arena that replaces inline `ast::TypeExpr`. `ItemTreeBuilder` constructs both together and records the item↔scope and method→owner indices. This is the substrate the PPIR firewall queries front (see [Salsa Early Cutoff](#salsa-early-cutoff-how-edits-stay-local)). **What does NOT live here:** - Node transformations. The HIR should NOT construct new AST nodes. If you find yourself doing that, the work belongs in the AST layer. @@ -199,11 +200,13 @@ The parser produces a **CST** (Concrete Syntax Tree), which is a lossless, error - Synthesis of `*$stream` variants for classes and type aliases - Stream expansion logic (`stream_expand`, `expand_partial`) - SAP (streaming attribute propagation) attributes -- Canonical queries that merge original items with synthetic stream items +- **The canonical item layer** — `ppir::file_item_tree` merges original items with the synthetic stream items, and the per-item **firewall queries** (`item_data`: enumeration + `*_data` + `*_source_map`) that front the item tree live here. -**How it works:** The PPIR generates synthetic AST items (the stream variants) and feeds them **back into the HIR**. This means the flow is actually: HIR → PPIR → back to HIR (with expanded symbols) → TIR. The TIR then consumes the enriched HIR index that includes both original and stream-expanded items. +**How it works:** The PPIR generates synthetic AST items (the stream variants) and re-runs the HIR builder over the *merged* list — originals first, then synthetics appended. The flow is HIR → PPIR → (re-run the HIR builder on merged items) → TIR. PPIR's tree is **canonical**, but it is an internal substrate: downstream layers read it **only through the firewall queries** in `item_data` (enumeration + `*_data` + `*_source_map`) — never `file_item_tree` directly, and never the HIR pre-expansion tree. A `ClassLoc`/`FunctionLoc` therefore unambiguously means a canonical item. (Originals keep identical `LocalItemId`s in the pre-expansion and canonical trees because they are allocated first, in the same order — which is what makes a HIR-derived `*Loc` safe to pass to a firewall query. Migrating the remaining direct `file_item_tree` readers onto the firewall queries is in progress.) -**Key Salsa query:** `ppir_expansion_items(db, file)` — Synthetic stream items per file +**Key Salsa queries:** +- `ppir_expansion_items(db, file)` — Synthetic stream items per file +- Firewall queries (`item_data`) — the consumer API over the canonical item tree: enumeration (`file_classes` / `file_functions` / …), lookup (`class_data` / `function_data` / …), and spans (`class_source_map` / …). `file_item_tree` itself is the internal substrate these are built on, not for direct downstream use. --- @@ -569,9 +572,18 @@ baml_test! { ## Rules for Adding Spans to Data Structures -**Do not add `TextRange` or span fields to your data structures.** There is a dedicated mechanism for associating spans with nodes. If you add `TextRange` directly to a data structure, you break Salsa incrementality for everything downstream — a change to whitespace (which changes spans but not semantics) will unnecessarily invalidate all dependent queries. +**The load-bearing invariant: anything a memoized, `PartialEq`-compared query *returns* must be span-free.** Salsa only overwrites a memoized value when the new one is *not* `Eq` to the old, so a span inside a query result is a trap either way: + +- If the span is inside `PartialEq`, a whitespace edit makes the value compare unequal and **destroys early cutoff** for everything downstream. +- If the span is *ignored* by `PartialEq` (as `ast::TypeExpr` does with its own `span`), the value compares equal, Salsa keeps the old memo, and the query **serves a stale span forever**. + +Concretely: -Use expression IDs and the span lookup infrastructure instead. If you're unsure how to associate span information with a new construct, ask before implementing. +- **Do not add `TextRange`/span fields to data that flows out of a tracked query.** Type references use the span-free `TypeRef` arena (ids into a per-item store), not `ast::TypeExpr`. Attributes use the span-free `item_tree::Attribute`, not `ast::RawAttribute` (whose `PartialEq` *includes* its span — a subtle transitive leak through any `TypeExpr`). +- **Spans live in a parallel `*_source_map` query**, keyed by the same id (`LocalItemId` / `TypeRefId`), never inline with the data. +- The one place spans may remain inline is the coarse `no_eq` index (`file_semantic_index` and its `ItemTree`): because that query is `no_eq` it is always recomputed, so its spans are always fresh — never stale. Everything *fronting* it (the firewall `*_data` queries) must be span-free. + +If you're unsure how to associate span information with a new construct, ask before implementing. --- @@ -713,22 +725,27 @@ The Salsa query model has one critical optimization beyond basic memoization: ** ### How it works in practice -Every item is physically split into **two** tracked queries: one for semantic data (span-free), one for source maps (spans only). For example, `function_signature` returns names and `TypeExpr`s with no `TextRange`, while `function_signature_source_map` returns only spans. The type checker reads the semantic query but never the source map query. +The item tree is produced by a single coarse query, `file_semantic_index`, marked `no_eq` — it always reports "changed", so its spans are always fresh but it provides no cutoff itself. In front of it sit fine-grained **firewall queries** (in `baml_compiler2_ppir::item_data`), one family per item kind: + +- *Enumeration* — `file_classes(file)` / `file_functions(file)` / … return a `Vec` of interned `*Loc` handles (a `ClassLoc` carries its own file plus a position-independent `LocalItemId`). +- *Lookup* — `class_data(ClassLoc)` / `function_data(FunctionLoc)` / … return **span-free** semantic data. Type references inside them are ids into a per-item `TypeRef` arena — a flat, span-free replacement for `ast::TypeExpr` — never `TextRange`s. +- *Spans* — `class_source_map(ClassLoc)` / `function_source_map(FunctionLoc)` / … are separate queries holding the spans, keyed by the same `*Loc`. The type checker reads the data query; only diagnostics read the source-map query. + +Because the `*_data` queries are tracked and `PartialEq`-compared, a whitespace or comment edit that shifts spans but not semantics leaves them **byte-identical** → `PartialEq` returns `true` → anything depending on a `*_data` query cuts off. **This is why the data must be span-free** (see [Rules for Adding Spans](#rules-for-adding-spans-to-data-structures)): a span inside a `PartialEq`-compared result would either destroy the cutoff or, if `PartialEq` ignored it, serve a *stale* span forever. -Items are keyed by **position-independent IDs** — a hash of the item's name, not its position in the file. Adding a blank line before `function Greet(...)` doesn't change the hash of `"Greet"`, so the Salsa query key stays the same and cached results survive. +Items are keyed by **position-independent IDs** — `LocalItemId` is a hash of the item's name plus a collision index, not its position in the file. Adding a blank line before `function Greet(...)` doesn't change the hash of `"Greet"`, so the `FunctionLoc` key stays the same and cached results survive. ### Concrete trace: adding a comment to a file User adds `// comment` to `file_a.baml`. File B is untouched. 1. `file_a.text` is marked changed. -2. `file_semantic_index(file_a)` re-runs (it's marked `no_eq`, so always reports "changed"). -3. `namespace_items(user_root)` re-runs — re-collects contributions from all files. But the result is identical: same names, same definition handles. Its `PartialEq` returns `true`. **Early cutoff fires.** -4. `package_items` — NOT re-run (its dependency didn't change). -5. `infer_scope_types` for any scope — NOT re-run. -6. `file_semantic_index(file_b)` — NOT re-run (its input `file_b.text` is unchanged). +2. `file_semantic_index(file_a)` re-runs (`no_eq` — always "changed"); its spans are refreshed. +3. Each per-item firewall query for `file_a` (`class_data`, `function_data`, …) re-runs but **early-cuts**: the semantic data is identical, so its `PartialEq` returns `true`. Its `*_source_map` twin re-runs and *does* change (new spans) — correctly, since positions moved. +4. `namespace_items(user_root)` re-runs and early-cuts when the name set is unchanged. (Exception: a file already holding a *duplicate-name* conflict currently carries a `name_span` in the conflict record, so a cosmetic edit there loses cutoff — a known gap being closed.) +5. `file_semantic_index(file_b)` — NOT re-run (its input `file_b.text` is unchanged), so nothing about file B recomputes. -Result: a comment addition re-runs the lexer and HIR for that one file, then stops. A whitespace edit shifts spans but leaves `TypeExpr` trees identical → `function_signature` early-cuts → type inference stays cached. +**Status / caveat.** The firewall queries above exist and cut off, but many downstream consumers are **not yet migrated onto them** — notably `infer_scope_types`, which still reads the coarse `no_eq` `file_semantic_index` directly and therefore re-runs on *any* edit to its file today. Realizing end-to-end cutoff (a comment edit not re-running type inference) requires migrating those consumers to the firewall queries; that work is in progress. Treat this section as the mechanism, not a fully-realized guarantee — the incremental tests in `baml_tests` pin what actually holds today, and one of them (`comment_edit_does_not_reexecute_type_inference`) is deliberately `#[ignore]`d precisely because inference is not yet migrated. --- diff --git a/baml_language/crates/baml_cli/src/test_command.rs b/baml_language/crates/baml_cli/src/test_command.rs index 919328fcabe..edc1f6cc443 100644 --- a/baml_language/crates/baml_cli/src/test_command.rs +++ b/baml_language/crates/baml_cli/src/test_command.rs @@ -470,7 +470,21 @@ fn render_test_list( ); // Indented list under the cargo-style status line. These are content (the // actual list), not status updates, so they go to stdout as plain prints. - for t in legacy_selected { + // + // Canonicalize the order at this single render point — both the fresh-compile + // and bytecode-cache paths flow through here — so `--list` output is + // byte-identical between them regardless of upstream map/enumeration order. + // Sort by (function, test) so each function's tests group together (a legacy + // `test` is function-attached), with the path as a final deterministic tiebreak. + let mut legacy_sorted: Vec<&crate::bytecode_cache::CachedLegacyTest> = + legacy_selected.iter().collect(); + legacy_sorted.sort_by(|a, b| { + a.function_name + .cmp(&b.function_name) + .then_with(|| a.test_name.cmp(&b.test_name)) + .then_with(|| a.file_path.cmp(&b.file_path)) + }); + for t in legacy_sorted { println!(" {}::{} ({})", t.function_name, t.test_name, t.file_path); } for name in testset_names { diff --git a/baml_language/crates/baml_compiler2_emit/src/lib.rs b/baml_language/crates/baml_compiler2_emit/src/lib.rs index eb862fc4ec5..78c8b4e1efe 100644 --- a/baml_language/crates/baml_compiler2_emit/src/lib.rs +++ b/baml_language/crates/baml_compiler2_emit/src/lib.rs @@ -3503,7 +3503,7 @@ fn compute_function_metadata_from_item_tree( let owner = item_tree.method_owners.get(&func_id); let enclosing_free_impl = match owner { Some(MethodOwner::FreeImpl(impl_id)) => Some((*impl_id, &item_tree.impls[impl_id])), - _ => None, + Some(MethodOwner::Class(_) | MethodOwner::Interface(_)) | None => None, }; let enclosing_impl_for_target = enclosing_free_impl.and_then(|(_, block)| match &block.subject { @@ -3512,11 +3512,11 @@ fn compute_function_metadata_from_item_tree( }); let enclosing_class = match owner { Some(MethodOwner::Class(class_id)) => Some(&item_tree[*class_id]), - _ => None, + Some(MethodOwner::FreeImpl(_) | MethodOwner::Interface(_)) | None => None, }; let enclosing_interface = match owner { Some(MethodOwner::Interface(iface_id)) => Some(&item_tree[*iface_id]), - _ => None, + Some(MethodOwner::Class(_) | MethodOwner::FreeImpl(_)) | None => None, }; let self_replacement = enclosing_impl_for_target .cloned() diff --git a/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs index f7268e25004..5c3c48cb4d8 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs @@ -335,7 +335,7 @@ impl ItemTreeBuilder { .config_items .iter() .find(|item| item.key.as_str() == "options") - .map(|_| Vec::new()) // complex to parse; clients field is more relevant + .map(|_| Vec::new()) // sub-clients are not parsed from `options`; left empty (unused downstream) .unwrap_or_default(); let retry_policy_name = c .config_items diff --git a/baml_language/crates/baml_compiler2_hir/src/signature.rs b/baml_language/crates/baml_compiler2_hir/src/signature.rs index 9619c2b28a5..1f72025003c 100644 --- a/baml_language/crates/baml_compiler2_hir/src/signature.rs +++ b/baml_language/crates/baml_compiler2_hir/src/signature.rs @@ -14,11 +14,14 @@ use text_size::TextRange; use crate::{item_tree::DefaultExprRef, loc::FunctionLoc}; -/// Compiler2 function signature — param names + unresolved `TypeExpr`. +/// Compiler2 function signature — param names + unresolved `ast::TypeExpr`. /// -/// No spans — those live in `SignatureSourceMap`. -/// `TypeExpr` is already span-free (spans live in `TypeExpr` at the -/// AST layer and are split out here). +/// The `SignatureSourceMap` twin holds the item-level spans. This struct is NOT +/// fully span-free, though: `ast::TypeExpr` carries its own `span` inline (its +/// `PartialEq` ignores that span but transitively compares `RawAttribute` spans), +/// so a whitespace edit near an attribute can still bust this query's cutoff. The +/// span-free successor is `ppir::function_data` over the `TypeRef` arena; this +/// struct is retained until its consumers migrate. #[derive(Debug, Clone, PartialEq, Eq)] pub struct FunctionSignature { pub name: baml_base::Name, diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 23b99259415..d84940f64aa 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -8866,7 +8866,13 @@ impl LoweringContext<'_> { Some(baml_compiler2_hir::item_tree::MethodOwner::Class(class_id)) => { item_tree[*class_id].generic_params.clone() } - _ => Vec::new(), + // `Interface` owners are handled by the early return above; a free + // impl's own generics are not enclosing-type parameters. + Some( + baml_compiler2_hir::item_tree::MethodOwner::Interface(_) + | baml_compiler2_hir::item_tree::MethodOwner::FreeImpl(_), + ) + | None => Vec::new(), }; params.extend(item_tree[func_id].generic_params.iter().cloned()); // Inside a (possibly nested) generic lambda body, the lambda's own diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs index 47eddd292d2..15e63936c0d 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs @@ -76,8 +76,11 @@ pub fn class_source_map<'db>(db: &'db dyn crate::Db, class: ClassLoc<'db>) -> Cl lower(db, class).1 } -/// Both halves share one lowering pass — the split into two queries is purely -/// about what each one lets downstream depend on. +/// The `*_data` and `*_source_map` queries each call this and keep one half, so +/// `lower` runs once per query — not a single shared pass. It is deterministic, +/// so the `TypeRefId`s the data half hands out validly index the source map's +/// arena. The split into two queries is purely about what each one lets +/// downstream depend on. /// /// Type refs are allocated in a fixed order (bounds, then fields, then /// `implements` targets) so that ids are a pure function of the class's shape. diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs index 8e58b92d032..618a62b10c9 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs @@ -17,7 +17,8 @@ use baml_compiler2_hir::loc::{ }; /// Declare an enumeration query: `file_(file) -> Vec`, ordered by -/// source position. +/// source position. Synthetic `*$stream` companions carry no source span, so they +/// sort first (at offset 0); user-declared items follow in source order. macro_rules! file_items { ($(#[$meta:meta])* $name:ident, $map:ident, $loc:ident) => { $(#[$meta])* @@ -35,13 +36,13 @@ macro_rules! file_items { } file_items!( - /// Every class declared in `file`, in source order. + /// Every class in `file`, in source order (synthetic `*$stream` companions first). file_classes, classes, ClassLoc ); file_items!( - /// Every function declared in `file`, in source order. Includes methods. + /// Every function in `file`, in source order (synthetic companions first). Includes methods. file_functions, functions, FunctionLoc @@ -59,7 +60,7 @@ file_items!( InterfaceLoc ); file_items!( - /// Every type alias declared in `file`, in source order. + /// Every type alias in `file`, in source order (synthetic `*$stream` companions first). file_type_aliases, type_aliases, TypeAliasLoc diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs index e8b687a3609..5fa8647c650 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -330,8 +330,11 @@ fn lower_interface_target<'db>( )) } -/// Both halves share one lowering pass — the split into two queries is purely -/// about what each one lets downstream depend on. +/// The `*_data` and `*_source_map` queries each call this and keep one half, so +/// `lower` runs once per query — not a single shared pass. It is deterministic, +/// so the `TypeRefId`s the data half hands out validly index the source map's +/// arena. The split into two queries is purely about what each one lets +/// downstream depend on. /// /// Type refs are allocated in a fixed order (bounds, params, return, throws) so /// that ids are a pure function of the signature's shape. diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs index 992c1c2df48..7a859016b1f 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs @@ -41,8 +41,11 @@ pub fn type_alias_source_map<'db>( lower(db, alias).1 } -/// Both halves share one lowering pass — the split into two queries is purely -/// about what each one lets downstream depend on. +/// The `*_data` and `*_source_map` queries each call this and keep one half, so +/// `lower` runs once per query — not a single shared pass. It is deterministic, +/// so the `TypeRefId`s the data half hands out validly index the source map's +/// arena. The split into two queries is purely about what each one lets +/// downstream depend on. fn lower<'db>( db: &'db dyn crate::Db, alias: TypeAliasLoc<'db>, diff --git a/baml_language/crates/baml_lsp2_actions/src/definition.rs b/baml_language/crates/baml_lsp2_actions/src/definition.rs index 466734c82f0..bbc2f567b88 100644 --- a/baml_language/crates/baml_lsp2_actions/src/definition.rs +++ b/baml_language/crates/baml_lsp2_actions/src/definition.rs @@ -373,8 +373,14 @@ fn resolve_field_access_at( field_name, } => { let target_file = class_loc.file(db); - let target_item_tree = baml_compiler2_hir::file_item_tree(db, target_file); - let target_source_map = baml_compiler2_hir::file_item_tree_source_map(db, target_file); + // Read the canonical (PPIR) tree, not the HIR pre-expansion tree: this + // `class_loc` is inferred from a member access, so it can be a synthetic + // `$stream` class (the type of a streamed partial), which is absent from + // the pre-expansion tree and would panic on index here. PPIR carries it, + // and a `$stream` field's name-span is the original class field's, so + // goto-definition still lands on the user-authored field. + let target_item_tree = baml_compiler2_ppir::file_item_tree(db, target_file); + let target_source_map = baml_compiler2_ppir::file_item_tree_source_map(db, target_file); let class = &target_item_tree[class_loc.id(db)]; let field_idx = class.fields.iter().position(|f| f.name == *field_name)?; let field_spans = target_source_map.class_field_spans.get(&class_loc.id(db))?; diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index c3b9bae496b..3b007f324dd 100644 --- a/baml_language/crates/baml_tests/src/incremental/scenarios.rs +++ b/baml_language/crates/baml_tests/src/incremental/scenarios.rs @@ -940,13 +940,13 @@ fn editing_a_function_prompt_preserves_its_llm_meta() { ); } -/// PLANNED (plan Phase 5). The whole refactor exists so that a cosmetic edit does -/// not re-run type inference — the exact claim ARCHITECTURE.md:728 makes but that -/// is *false* today: `infer_scope_types` reads the `no_eq` `file_semantic_index` -/// directly (`inference.rs:1524`), so any edit re-executes it. Un-ignore once -/// inference consumes the per-item firewall queries instead of the coarse index. +/// The whole refactor exists so that a cosmetic edit does not re-run type +/// inference. That does not hold yet: `infer_scope_types` reads the `no_eq` +/// `file_semantic_index` directly, so any edit to its file re-executes it. +/// Un-ignore once inference consumes the per-item firewall queries instead of the +/// coarse index. #[test] -#[ignore = "infer_scope_types still reads the no_eq file_semantic_index directly; un-ignore after it migrates to the firewall queries (plan Phase 5)"] +#[ignore = "infer_scope_types still reads the no_eq file_semantic_index directly; un-ignore once it consumes the firewall queries"] fn comment_edit_does_not_reexecute_type_inference() { let mut test_db = IncrementalTestDb::new(); From e03d422b669b8d487cb5fe26c8ae1279002333da Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Thu, 16 Jul 2026 14:19:41 -0700 Subject: [PATCH 19/20] Item enumeration deterministic tie breaker --- .../src/item_data/enumeration.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs index 618a62b10c9..ce5d640b5d7 100644 --- a/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs @@ -26,7 +26,18 @@ macro_rules! file_items { pub fn $name(db: &dyn crate::Db, file: SourceFile) -> Vec<$loc<'_>> { let item_tree = crate::file_item_tree(db, file); let mut items: Vec<_> = item_tree.$map.iter().collect(); - items.sort_by_key(|(_, item)| item.span.start()); + // Source position, then a stable tiebreaker for items that share an + // offset: synthetic `*$stream` companions all sit at offset 0, and + // `$map` is an `FxHashMap` (nondeterministic iteration), so ties must + // be broken by a position-independent key — otherwise this query's + // value (and its early-cutoff) would be unstable across rebuilds. + items.sort_by(|(left_id, left), (right_id, right)| { + left.span + .start() + .cmp(&right.span.start()) + .then_with(|| left.name.cmp(&right.name)) + .then_with(|| left_id.as_u32().cmp(&right_id.as_u32())) + }); items .into_iter() .map(|(id, _)| $loc::new(db, file, *id)) From a81dec3d0f90537fec5f77c111670ed0b06a54c9 Mon Sep 17 00:00:00 2001 From: 2kai2kai2 Date: Thu, 16 Jul 2026 17:36:18 -0700 Subject: [PATCH 20/20] Bump size gate --- .github/workflows/cargo-tests.reusable.yaml | 10 ++++++++ baml_language/.cargo/size-gate.toml | 24 +++++++++---------- .../.ci/size-gate/aarch64-apple-darwin.toml | 14 +++++------ .../.ci/size-gate/x86_64-pc-windows-msvc.toml | 14 +++++------ .../size-gate/x86_64-unknown-linux-gnu.toml | 14 +++++------ 5 files changed, 43 insertions(+), 33 deletions(-) diff --git a/.github/workflows/cargo-tests.reusable.yaml b/.github/workflows/cargo-tests.reusable.yaml index faa09158499..e5f61abb4a0 100644 --- a/.github/workflows/cargo-tests.reusable.yaml +++ b/.github/workflows/cargo-tests.reusable.yaml @@ -258,6 +258,16 @@ jobs: runs-on: blacksmith-8vcpu-windows-2025 env: CARGO_BUILD_JOBS: "16" + # Windows debug builds emit very large PDBs; full debuginfo across the + # whole `--all-features` test build overruns the runner's ~130 GB disk + # ("no space on device" during codegen, before any test runs). + # `line-tables-only` keeps file:line backtraces for test failures while + # dropping the bulk of the debug-info size (variable/type DIEs). Pinned on + # both `dev` (dependencies) and `test` (the `--test` binaries, which + # inherit `dev` but are pinned explicitly so the reduction is unambiguous). + # Drop to "0" if this is still too tight. + CARGO_PROFILE_DEV_DEBUG: "line-tables-only" + CARGO_PROFILE_TEST_DEBUG: "line-tables-only" # TODO: Reduce this once our changes to aws-sdk-rust are upstreamed # Right now it needs to download our repo every time which is very slow timeout-minutes: 50 diff --git a/baml_language/.cargo/size-gate.toml b/baml_language/.cargo/size-gate.toml index 69adbe0d4e7..4fa22b87dab 100644 --- a/baml_language/.cargo/size-gate.toml +++ b/baml_language/.cargo/size-gate.toml @@ -42,29 +42,29 @@ policy.max_delta_pct = 3.0 # Baselines re-recorded after merging canary (~+4%). Reclaiming the `syn` # cost would mean dropping prettyplease (raw, unformatted generated SDK # output). -# macOS arm64: baseline 18.30 MB file (18_300_864), ceiling = baseline x1.03. +# macOS arm64: baseline 18.90 MB file (18_900_160), ceiling = baseline x1.03. [artifacts.baml-cli.platform.aarch64-apple-darwin] -max_file_bytes = 18_849_890 -# Linux x86_64: baseline 23.75 MB file (23_753_040), ceiling = baseline x1.03. +max_file_bytes = 19_467_165 +# Linux x86_64: baseline 24.51 MB file (24_507_664), ceiling = baseline x1.03. [artifacts.baml-cli.platform.x86_64-unknown-linux-gnu] -max_file_bytes = 24_465_631 -# Windows x86_64: baseline 19.80 MB file (19_798_528), ceiling = baseline x1.03. +max_file_bytes = 25_242_894 +# Windows x86_64: baseline 20.43 MB file (20_426_752), ceiling = baseline x1.03. [artifacts.baml-cli.platform.x86_64-pc-windows-msvc] -max_file_bytes = 20_392_484 +max_file_bytes = 21_039_555 [artifacts.packed-program] kind = "pack" policy.max_delta_pct = 3.0 -# macOS arm64: baseline 12.64 MB file. +# macOS arm64: baseline 13.21 MB file. [artifacts.packed-program.platform.aarch64-apple-darwin] -max_file_bytes = 13_019_384 -# Linux x86_64: baseline 16.26 MB file. +max_file_bytes = 13_602_990 +# Linux x86_64: baseline 17.03 MB file. [artifacts.packed-program.platform.x86_64-unknown-linux-gnu] -max_file_bytes = 16_771_886 -# Windows x86_64: baseline 13.53 MB file. +max_file_bytes = 17_542_251 +# Windows x86_64: baseline 14.14 MB file. [artifacts.packed-program.platform.x86_64-pc-windows-msvc] -max_file_bytes = 13_944_063 +max_file_bytes = 14_567_300 [artifacts.packed-program.pack] cli_package = "baml_cli" diff --git a/baml_language/.ci/size-gate/aarch64-apple-darwin.toml b/baml_language/.ci/size-gate/aarch64-apple-darwin.toml index b7bb35afe26..460105fef06 100644 --- a/baml_language/.ci/size-gate/aarch64-apple-darwin.toml +++ b/baml_language/.ci/size-gate/aarch64-apple-darwin.toml @@ -1,12 +1,12 @@ version = 1 -recorded_at = "2026-07-16T10:00:08Z" -git_sha = "79aa0f44e323e356a172e16e6b0f4a42e569959d" +recorded_at = "2026-07-17T00:35:21Z" +git_sha = "b01ee9c671a00b44a7a9a3832e149e3c2ec93c13" [artifacts.baml-cli] -file_bytes = 18300864 -stripped_bytes = 18300912 -gzip_bytes = 8798867 +file_bytes = 18900160 +stripped_bytes = 18900208 +gzip_bytes = 8955790 [artifacts.packed-program] -file_bytes = 12640178 -gzip_bytes = 5993393 +file_bytes = 13206786 +gzip_bytes = 6154161 diff --git a/baml_language/.ci/size-gate/x86_64-pc-windows-msvc.toml b/baml_language/.ci/size-gate/x86_64-pc-windows-msvc.toml index 99d8cdd85e8..c49191d5455 100644 --- a/baml_language/.ci/size-gate/x86_64-pc-windows-msvc.toml +++ b/baml_language/.ci/size-gate/x86_64-pc-windows-msvc.toml @@ -1,12 +1,12 @@ version = 1 -recorded_at = "2026-07-16T10:00:08Z" -git_sha = "79aa0f44e323e356a172e16e6b0f4a42e569959d" +recorded_at = "2026-07-17T00:35:21Z" +git_sha = "b01ee9c671a00b44a7a9a3832e149e3c2ec93c13" [artifacts.baml-cli] -file_bytes = 19798528 -stripped_bytes = 19798528 -gzip_bytes = 8992285 +file_bytes = 20426752 +stripped_bytes = 20426752 +gzip_bytes = 9169338 [artifacts.packed-program] -file_bytes = 13537925 -gzip_bytes = 6082881 +file_bytes = 14143010 +gzip_bytes = 6248500 diff --git a/baml_language/.ci/size-gate/x86_64-unknown-linux-gnu.toml b/baml_language/.ci/size-gate/x86_64-unknown-linux-gnu.toml index f38352d31c5..06df65b9f78 100644 --- a/baml_language/.ci/size-gate/x86_64-unknown-linux-gnu.toml +++ b/baml_language/.ci/size-gate/x86_64-unknown-linux-gnu.toml @@ -1,12 +1,12 @@ version = 1 -recorded_at = "2026-07-16T10:00:08Z" -git_sha = "79aa0f44e323e356a172e16e6b0f4a42e569959d" +recorded_at = "2026-07-17T00:35:21Z" +git_sha = "b01ee9c671a00b44a7a9a3832e149e3c2ec93c13" [artifacts.baml-cli] -file_bytes = 23753040 -stripped_bytes = 23753032 -gzip_bytes = 10113500 +file_bytes = 24507664 +stripped_bytes = 24507656 +gzip_bytes = 10325653 [artifacts.packed-program] -file_bytes = 16283384 -gzip_bytes = 6846279 +file_bytes = 17031312 +gzip_bytes = 7039119