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 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 a6c36bb214d..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 { @@ -499,20 +513,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_emit/src/lib.rs b/baml_language/crates/baml_compiler2_emit/src/lib.rs index fd5618afbba..78c8b4e1efe 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])), + Some(MethodOwner::Class(_) | MethodOwner::Interface(_)) | None => 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]), + Some(MethodOwner::FreeImpl(_) | MethodOwner::Interface(_)) | None => None, + }; + let enclosing_interface = match owner { + Some(MethodOwner::Interface(iface_id)) => Some(&item_tree[*iface_id]), + Some(MethodOwner::Class(_) | MethodOwner::FreeImpl(_)) | None => None, + }; let self_replacement = enclosing_impl_for_target .cloned() .or_else(|| { @@ -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_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_tree.rs b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs index cd52ada06ec..07bb4cfb966 100644 --- a/baml_language/crates/baml_compiler2_hir/src/item_tree.rs +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree.rs @@ -4,387 +4,52 @@ //! 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; +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; -use text_size::TextRange; +pub use source_map::*; +pub use template_strings::*; +pub use test_items::*; +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, }; -// ── 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. /// -/// 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>, @@ -420,417 +85,30 @@ pub struct ItemTree { 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() - } + /// 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>, } 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`. + /// 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. /// - /// `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 + /// 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 => &[], + } } } 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..5c3c48cb4d8 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/builder.rs @@ -0,0 +1,437 @@ +//! 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, MethodOwner, 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, 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 + /// 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); + // 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)); + } + } + } + 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()); + for method in &default_method_ids { + self.record_method_owner(*method, MethodOwner::Interface(id)); + } + 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()) // sub-clients are not parsed from `options`; left empty (unused downstream) + .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/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..b5ed47b593e --- /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, Hash)] +pub struct Attribute { + pub args: Vec, + pub name: Name, +} + +/// A span-free attribute argument. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +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..ea5362fa7f9 --- /dev/null +++ b/baml_language/crates/baml_compiler2_hir/src/item_tree/functions.rs @@ -0,0 +1,77 @@ +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, +} + +/// 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 { + 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, +} diff --git a/baml_language/crates/baml_compiler2_hir/src/lib.rs b/baml_language/crates/baml_compiler2_hir/src/lib.rs index d4eb8f8daae..67148709370 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/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_compiler2_hir/src/signature.rs b/baml_language/crates/baml_compiler2_hir/src/signature.rs index 46eabbccdee..2e87512f20a 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, @@ -220,18 +223,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>, @@ -257,7 +248,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_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); + } +} diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index d472692a686..01520b8af9f 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 => {} } } @@ -2601,34 +2597,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 @@ -2727,10 +2704,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 @@ -2760,7 +2739,6 @@ impl<'db> LoweringContext<'db> { baml_compiler2_tir::ty::TyAttr::default(), ), ); - break; } } @@ -2852,14 +2830,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 @@ -3377,34 +3354,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, ); @@ -3965,17 +3941,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` @@ -4107,13 +4084,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]; @@ -9162,11 +9139,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 @@ -9177,12 +9153,18 @@ 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() + } + // `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 // type params follow the enclosing function's, matching the runtime @@ -9210,14 +9192,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/item_data.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data.rs new file mode 100644 index 00000000000..245178be824 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data.rs @@ -0,0 +1,57 @@ +//! 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`](baml_compiler2_hir::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 classes; +mod clients; +mod common; +mod enumeration; +mod enums; +mod functions; +mod impls; +mod interfaces; +mod lets; +mod retry_policies; +mod scopes; +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 scopes::*; +pub use template_strings::*; +pub use test_items::*; +pub use type_aliases::*; 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 new file mode 100644 index 00000000000..15e63936c0d --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/classes.rs @@ -0,0 +1,191 @@ +use baml_base::Name; +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. +/// +/// `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 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, +} + +/// 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, +} + +/// 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 +} + +/// 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. +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| 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 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_ppir/src/item_data/clients.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/clients.rs new file mode 100644 index 00000000000..602f457cc2c --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/clients.rs @@ -0,0 +1,27 @@ +use baml_base::Name; +use baml_compiler2_hir::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_ppir/src/item_data/common.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/common.rs new file mode 100644 index 00000000000..2a63be1018f --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/common.rs @@ -0,0 +1,51 @@ +//! 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 baml_compiler2_hir::{item_tree::Attribute, type_ref::TypeRefId}; +use text_size::TextRange; + +#[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_ppir/src/item_data/enumeration.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs new file mode 100644 index 00000000000..ce5d640b5d7 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enumeration.rs @@ -0,0 +1,169 @@ +//! 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 baml_compiler2_hir::loc::{ + ClassLoc, ClientLoc, EnumLoc, FunctionLoc, ImplLoc, InterfaceLoc, LetLoc, RetryPolicyLoc, + TemplateStringLoc, TestLoc, TypeAliasLoc, +}; + +/// Declare an enumeration query: `file_(file) -> Vec`, ordered by +/// 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])* + #[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(); + // 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)) + .collect() + } + }; +} + +file_items!( + /// Every class in `file`, in source order (synthetic `*$stream` companions first). + file_classes, + classes, + ClassLoc +); +file_items!( + /// Every function in `file`, in source order (synthetic companions first). 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 in `file`, in source order (synthetic `*$stream` companions first). + 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_ppir/src/item_data/enums.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/enums.rs new file mode 100644 index 00000000000..b8846eafaf3 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/enums.rs @@ -0,0 +1,75 @@ +use baml_base::Name; +use baml_compiler2_hir::{item_tree::Attribute, loc::EnumLoc}; +use text_size::TextRange; + +/// 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_ppir/src/item_data/functions.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs new file mode 100644 index 00000000000..5fa8647c650 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/functions.rs @@ -0,0 +1,397 @@ +use baml_base::Name; +use baml_compiler2_ast as ast; +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*. +/// +/// 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, +} + +/// 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 +} + +/// 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. +/// +/// 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, +} + +/// 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)] +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| ElaboratedParamData { + name: param.name.clone(), + type_ref: 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 — +/// 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 +/// 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(), + }, + )) +} + +/// 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. +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_compiler2_ppir/src/item_data/impls.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/impls.rs new file mode 100644 index 00000000000..52ad8b8d2c1 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/impls.rs @@ -0,0 +1,179 @@ +use baml_base::Name; +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. +/// +/// 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 baml_compiler2_hir::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_ppir/src/item_data/interfaces.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/interfaces.rs new file mode 100644 index 00000000000..ae5890cbb03 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/interfaces.rs @@ -0,0 +1,214 @@ +use baml_base::Name; +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)] +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_ppir/src/item_data/lets.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/lets.rs new file mode 100644 index 00000000000..d63a85ec3ee --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/lets.rs @@ -0,0 +1,53 @@ +use baml_base::Name; +use baml_compiler2_ast as ast; +use baml_compiler2_hir::loc::LetLoc; +use text_size::TextRange; + +/// 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_ppir/src/item_data/retry_policies.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/retry_policies.rs new file mode 100644 index 00000000000..08a9d83ce11 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/retry_policies.rs @@ -0,0 +1,33 @@ +use baml_base::Name; +use baml_compiler2_hir::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_ppir/src/item_data/scopes.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/scopes.rs new file mode 100644 index 00000000000..87b60d1136a --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/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`](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 baml_compiler2_hir::{ + 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_ppir/src/item_data/template_strings.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/template_strings.rs new file mode 100644 index 00000000000..c5c27673977 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/template_strings.rs @@ -0,0 +1,83 @@ +use baml_base::Name; +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)] +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_ppir/src/item_data/test_items.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/test_items.rs new file mode 100644 index 00000000000..3c47343efc0 --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/test_items.rs @@ -0,0 +1,26 @@ +use baml_base::Name; +use baml_compiler2_hir::{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_compiler2_ppir/src/item_data/type_aliases.rs b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs new file mode 100644 index 00000000000..7a859016b1f --- /dev/null +++ b/baml_language/crates/baml_compiler2_ppir/src/item_data/type_aliases.rs @@ -0,0 +1,71 @@ +use baml_base::Name; +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)] +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 +} + +/// 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>, +) -> (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_ppir/src/lib.rs b/baml_language/crates/baml_compiler2_ppir/src/lib.rs index c86e28280b3..89809e256f6 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, @@ -627,6 +628,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` @@ -708,27 +715,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, @@ -756,7 +742,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 00f304162f1..ef34d65eb47 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -188,12 +188,15 @@ pub(crate) fn callee_generics_for_func<'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 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(); + // Enclosing class params via the firewall `method_owner` index (O(1)) rather + // than an O(classes) `methods.contains` scan — identical result (a class + // method's owner class's params, else none). + let 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(), + }; // Lower the user generic params' interface bounds in the callee's own // package/namespace, with the class params in scope (see above). The @@ -2015,7 +2018,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 { @@ -2582,37 +2585,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) } @@ -8292,18 +8286,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, @@ -9930,7 +9926,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() @@ -9954,8 +9950,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 @@ -9966,9 +9965,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(), @@ -10089,7 +10092,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() @@ -10119,8 +10123,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 @@ -10131,10 +10138,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 { @@ -10154,12 +10163,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() @@ -12168,7 +12177,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 @@ -12301,17 +12310,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, ) @@ -12329,10 +12342,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, ) }) @@ -12489,7 +12506,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). @@ -12564,16 +12581,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, @@ -12600,11 +12619,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 7cbe79b217d..ffb98722f6c 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::{ @@ -34,18 +32,21 @@ 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); - 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()); - 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,12 +69,14 @@ 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); - 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()); @@ -88,7 +91,12 @@ fn signature_cycle_initial_callable_throws<'db>( let mut facts = BTreeSet::new(); for param in &sig.params { 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.type_ref, + ¶m_scope, + &mut diags, + ); if let Ty::Function { throws, .. } = lowered { facts.extend(crate::throw_inference::flatten_ty_to_facts(&throws)); } @@ -97,29 +105,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() } @@ -131,23 +128,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 +391,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 72176af10cd..1cd3cebc61c 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) - .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( - 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, @@ -1703,468 +1709,460 @@ 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 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); + + 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 - .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_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_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_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_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 = |id: baml_compiler2_hir::type_ref::TypeRefId, + diags: &mut Vec< + crate::infer_context::TirTypeError, + >| { + crate::generics::substitute_ty( + &crate::lower_type_expr::lower_type_ref(&sig.type_refs, id, &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 + .map(|id| { + 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 { + 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!( + 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, + // 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(param.type_ref, &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.type_refs, + sig.throws, + sig_sm.throws_type_span, + func_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. @@ -2327,16 +2325,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, @@ -2367,11 +2366,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, ); @@ -2408,15 +2408,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); @@ -2443,6 +2437,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, }, @@ -2490,15 +2485,12 @@ 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; - } - report_duplicate_generic_params( - &builder, - &class_data.generic_params, - class_data.span, - ); + 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)]; + 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, @@ -2510,27 +2502,22 @@ pub fn infer_scope_types<'db>( pkg_items, &pkg_info.namespace_path, &env, - class_data.span, + class_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); - report_duplicate_generic_params( - &builder, - &iface_data.generic_params, - iface_data.span, - ); + if let Some(baml_compiler2_ppir::item_data::ScopeOwner::Interface(iface_loc)) = + scope_item + { + let iface_data = &item_tree[iface_loc.id(db)]; + 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 — @@ -2567,7 +2554,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 @@ -2701,7 +2688,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 = @@ -2889,16 +2876,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() { @@ -2907,22 +2893,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; } } _ => { @@ -3372,72 +3355,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 92cd82638da..7dc242b82a0 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 { @@ -498,379 +353,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 +424,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 +447,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. @@ -928,21 +468,12 @@ pub fn lower_type_expr( // parameter by HIR signature elaboration, and lambda literals infer their // own throws from the body. Recover with `never` so downstream checking // doesn't cascade. - let throws_ty = match throws.as_deref() { - Some(throws) => lower_type_expr(throws, ctx, diagnostics), - None => { - diagnostics.push(TirTypeError::FunctionTypeMissingThrows); - Ty::Never { - attr: TyAttr::default(), - } - } - }; Ty::Function { params: params .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 { @@ -950,21 +481,33 @@ pub fn lower_type_expr( }, }) .collect(), - ret: Box::new(lower_type_expr(ret, ctx, diagnostics)), - throws: Box::new(throws_ty), + ret: Box::new(lower_type_ref(store, *ret, ctx, diagnostics)), + // #4034 (E0151): a function TYPE must declare its throws. The + // positions that legitimately omit it never reach here with `None` + // (declaration callback params are opened to a synthetic effect + // param by signature elaboration; lambda literals infer their own), + // so a `None` here is a user error — report it and recover with + // `never` so downstream checking doesn't cascade. + throws: Box::new(match throws { + Some(throws) => lower_type_ref(store, *throws, ctx, diagnostics), + None => { + diagnostics.push(TirTypeError::FunctionTypeMissingThrows); + Ty::Never { + attr: TyAttr::default(), + } + } + }), 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, @@ -974,34 +517,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(), @@ -1010,6 +553,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, @@ -1115,30 +999,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 @@ -1150,9 +1029,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 @@ -1165,24 +1042,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); @@ -1724,17 +1595,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); @@ -3024,69 +2886,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(); 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() { diff --git a/baml_language/crates/baml_lsp2_actions/src/annotations.rs b/baml_language/crates/baml_lsp2_actions/src/annotations.rs index e9d23589433..bf3c2be328f 100644 --- a/baml_language/crates/baml_lsp2_actions/src/annotations.rs +++ b/baml_language/crates/baml_lsp2_actions/src/annotations.rs @@ -55,13 +55,9 @@ use baml_base::SourceFile; use baml_compiler2_ast::{ Expr, ExprId, Stmt, - ast::{AstSourceMap, DeclarativeMeta, ExprBody, FunctionBodyDef, FunctionOrigin}, -}; -use baml_compiler2_hir::{ - body::FunctionBody, - loc::FunctionLoc, - scope::{FileScopeId, ScopeKind}, + 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}; use text_size::TextSize; @@ -118,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 @@ -127,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; @@ -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 deccfb0fe3d..f3469fc671a 100644 --- a/baml_language/crates/baml_lsp2_actions/src/check.rs +++ b/baml_language/crates/baml_lsp2_actions/src/check.rs @@ -395,7 +395,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, res_ctx, pkg_id, scope_id, aliases, @@ -1622,22 +1623,6 @@ fn jinja_diagnostic_id(message: &str) -> DiagnosticId { } } -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_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_project/src/db.rs b/baml_language/crates/baml_project/src/db.rs index 6d396a16d43..5b400050b34 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,31 +673,23 @@ 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, - // 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_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 +714,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 +836,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 +844,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 +1272,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 +1286,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 +1394,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 +1408,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) 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| { diff --git a/baml_language/crates/baml_tests/src/incremental/scenarios.rs b/baml_language/crates/baml_tests/src/incremental/scenarios.rs index bc7a0234e5a..3b007f324dd 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; @@ -343,3 +344,642 @@ 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_ppir::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_ppir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_ppir::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_ppir::item_data::type_alias_data(db, loc).clone(), + baml_compiler2_ppir::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_ppir::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_ppir::item_data::type_alias_data(db, loc).clone() + }; + + 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_ppir::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_ppir::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_ppir::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_ppir::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_ppir::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); +} + +// ── 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(); + + // 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\n\ + function Greet(name: string) -> string {\n \ + client GPT4\n prompt #\"Hello {{name}}\"#\n}\n", + ); + + let db = test_db.db(); + let index = baml_compiler2_ppir::file_semantic_index(db, file); + let item_tree = baml_compiler2_ppir::file_item_tree(db, file); + + // 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); + + // 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_ppir::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_ppir::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_ppir::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_ppir::item_data::function_scope(db, loc).expect("scope"); + + assert_eq!( + baml_compiler2_ppir::item_data::scope_owner(db, scope), + 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"); +} + +/// 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" + ); +} + +/// `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" + ); +} + +/// 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 once it consumes the firewall queries"] +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"], + ); +}