diff --git a/baml_language/crates/baml_compiler2_emit/src/lib.rs b/baml_language/crates/baml_compiler2_emit/src/lib.rs index d2415c84afc..9a80d001a7e 100644 --- a/baml_language/crates/baml_compiler2_emit/src/lib.rs +++ b/baml_language/crates/baml_compiler2_emit/src/lib.rs @@ -35,12 +35,14 @@ use bex_vm_types::{ ObjectIndex, ObjectPool, Program, }; -/// Build a per-package `ResolvedAliases` cache, keyed by package name. -fn build_alias_caches( - db: &dyn baml_compiler2_mir::Db, +/// Per-package `ResolvedAliases` refs, keyed by package name — borrowed +/// straight from the tracked `resolved_aliases_for_package` query (computed +/// once per package for the whole compile). +fn build_alias_caches<'db>( + db: &'db dyn baml_compiler2_mir::Db, all_files: &[baml_base::SourceFile], -) -> HashMap { - let mut caches: HashMap = HashMap::new(); +) -> HashMap { + let mut caches: HashMap = HashMap::new(); for file in all_files { let pkg_info = file_package(db, *file); caches.entry(pkg_info.package.clone()).or_insert_with(|| { @@ -240,7 +242,7 @@ fn build_interface_def( fn build_packages( db: &dyn baml_compiler2_mir::Db, all_files: &[baml_base::SourceFile], - alias_caches: &HashMap, + alias_caches: &HashMap, function_indices: &HashMap, interface_indices: &HashMap, program_packages: &mut indexmap::IndexMap, diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 8c5733132df..6a40efeeb53 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -60,14 +60,24 @@ use baml_compiler2_tir::ty::{ FunctionParamMode, FunctionParamTy as Tir2FunctionParamTy, QualifiedTypeName, Ty as Tir2Ty, }; -/// Build the [`ResolvedAliases`] type-alias environment for a package, +/// Salsa query: the [`ResolvedAliases`] type-alias environment for a package, /// including dependency packages. The pure erasure that consumes it lives in /// `baml_type` ([`ResolvedAliases::convert`]), wrapped compiler-side by /// `convert_tir_ty_for_runtime`; only this db-querying constructor stays /// compiler-side. -pub fn resolved_aliases_for_package( - db: &dyn crate::Db, - pkg_id: baml_compiler2_hir::package::PackageId, +/// +/// Memoized per package: the emit driver and `package_lowering_data` both +/// demand it, and each computation re-collects (and clones) every alias of +/// the package *and all its dependencies* — so before tracking, dependency +/// packages' aliases were re-collected once per dependent, per caller. +/// +/// Intentionally distinct from TIR's `package_alias_env`, which sees only +/// dependency *interface exports*; MIR/emit need every dependency alias for +/// runtime type erasure. +#[salsa::tracked(returns(ref))] +pub fn resolved_aliases_for_package<'db>( + db: &'db dyn crate::Db, + pkg_id: baml_compiler2_hir::package::PackageId<'db>, ) -> ResolvedAliases { use baml_compiler2_hir::package::{package_dependencies, package_items}; @@ -1443,7 +1453,9 @@ fn package_lowering_data<'db>( ) -> PackageLoweringData { use baml_compiler2_hir::package::{package_dependencies, package_items}; - let resolved_aliases = resolved_aliases_for_package(db, pkg_id); + // Cloned out of the tracked query's cached value once per package: this + // struct is itself a per-package Salsa value and owns its env. + let resolved_aliases = resolved_aliases_for_package(db, pkg_id).clone(); let mut class_fields = ClassFieldIndices::default(); let mut class_field_types = ClassFieldTypes::default(); diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index c894d4e2374..97d7034d5d8 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -13,7 +13,7 @@ //! NOT recurse into it — lambda bodies are separate scopes with their own //! `infer_scope_types` Salsa query. -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeSet; use baml_base::{Name, SourceFile}; use baml_compiler2_ast::{ @@ -529,7 +529,7 @@ impl baml_type::normalize::TypeContext for NormalizeCtx<'_, '_> { b.context.db(), concrete, interface, - &b.aliases, + b.aliases, |actual, bound| b.is_subtype(actual, bound), ) } @@ -650,9 +650,10 @@ pub struct TypeInferenceBuilder<'db> { /// Declared return type for the function (used to check return statements). declared_return_ty: Option, /// Resolved type-alias environment (alias map + precomputed recursive - /// set), built once per scope so per-comparison normalization never - /// re-derives the recursive-alias analysis. - aliases: crate::normalize::ResolvedAliases, + /// set), borrowed from the per-package Salsa query + /// (`inference::package_alias_env`) — computed once per package, not per + /// scope, and never re-derived per comparison. + aliases: &'db crate::normalize::ResolvedAliases, /// Namespace path for the file being analyzed (e.g. `["env"]` for `baml/env.baml`). ns_context: Vec, /// BEP-044: when this body is the override of an interface method @@ -952,7 +953,7 @@ impl<'db> TypeInferenceBuilder<'db> { res_ctx: &'db PackageResolutionContext<'db>, package_id: PackageId<'db>, scope: ScopeId<'db>, - aliases: HashMap, + aliases: &'db crate::normalize::ResolvedAliases, ) -> Self { let db = context.db(); let package_items = &res_ctx.own_items; @@ -973,7 +974,7 @@ impl<'db> TypeInferenceBuilder<'db> { package_id, scope, declared_return_ty: None, - aliases: crate::normalize::resolved_aliases_from_map(aliases), + aliases, ns_context, implements_block_interface: None, generic_param_bounds: rustc_hash::FxHashMap::default(), @@ -1790,7 +1791,7 @@ impl<'db> TypeInferenceBuilder<'db> { } fn container_arg_subtype_without_nominal(&self, actual: &Ty, expected: &Ty) -> bool { - crate::normalize::is_subtype_of(actual, expected, &self.aliases) + crate::normalize::is_subtype_of(actual, expected, self.aliases) } fn function_coercion_for( @@ -4863,7 +4864,7 @@ impl<'db> TypeInferenceBuilder<'db> { crate::associated_projection::AssociatedProjectionResolver::with_resolution_context( self.context.db(), self.res_ctx, - &self.aliases, + self.aliases, &self.generic_param_bounds, ); let completed_bindings: Vec<(Name, Ty)> = self @@ -5896,7 +5897,7 @@ impl<'db> TypeInferenceBuilder<'db> { let db = self.context.db(); let registry = crate::interfaces::package_implements_registry(db, self.package_id); - registry.first_failing_bound(&inferred, expected, &self.aliases, |a, b| { + registry.first_failing_bound(&inferred, expected, self.aliases, |a, b| { self.is_subtype(a, b) }) } else { @@ -11163,7 +11164,7 @@ impl<'db> TypeInferenceBuilder<'db> { crate::associated_projection::AssociatedProjectionResolver::with_resolution_context( self.context.db(), self.res_ctx, - &self.aliases, + self.aliases, &self.generic_param_bounds, ); if let Some(projection_bound) = @@ -11244,7 +11245,7 @@ impl<'db> TypeInferenceBuilder<'db> { crate::associated_projection::AssociatedProjectionResolver::with_resolution_context( self.context.db(), self.res_ctx, - &self.aliases, + self.aliases, &self.generic_param_bounds, ) .types_equivalent(a, b) @@ -11471,7 +11472,7 @@ impl<'db> TypeInferenceBuilder<'db> { &rule.for_ty_pattern, &actual_ty, &rule.generic_params, - &self.aliases, + self.aliases, ) else { continue; }; @@ -11485,7 +11486,7 @@ impl<'db> TypeInferenceBuilder<'db> { if !registry.type_implements_interface_via_rule( &actual_ty, &implemented_iface, - &self.aliases, + self.aliases, |actual, bound| self.is_subtype(actual, bound), ) { continue; @@ -14359,9 +14360,10 @@ impl<'db> TypeInferenceBuilder<'db> { .map(|a| crate::generics::substitute_ty(a, &bindings)) .collect(); if subst.len() == target_iface_args.len() - && subst.iter().zip(target_iface_args).all(|(a, b)| { - crate::normalize::is_same_normalized_type(a, b, &self.aliases) - }) + && subst + .iter() + .zip(target_iface_args) + .all(|(a, b)| crate::normalize::is_same_normalized_type(a, b, self.aliases)) { count += 1; } @@ -14473,7 +14475,7 @@ impl<'db> TypeInferenceBuilder<'db> { &rule.for_ty_pattern, base_ty, &rule.generic_params, - &self.aliases, + self.aliases, ) else { continue; }; @@ -14495,7 +14497,7 @@ impl<'db> TypeInferenceBuilder<'db> { if !registry.type_implements_interface_via_rule( base_ty, &requested, - &self.aliases, + self.aliases, |a, c| self.is_subtype(a, c), ) { continue; @@ -15611,7 +15613,7 @@ impl<'db> TypeInferenceBuilder<'db> { crate::associated_projection::AssociatedProjectionResolver::with_resolution_context( self.context.db(), self.res_ctx, - &self.aliases, + self.aliases, &self.generic_param_bounds, ) .resolve_deep(ty) @@ -15663,7 +15665,7 @@ impl<'db> TypeInferenceBuilder<'db> { crate::associated_projection::AssociatedProjectionResolver::with_resolution_context( self.context.db(), self.res_ctx, - &self.aliases, + self.aliases, &self.generic_param_bounds, ) .projection_views_equivalent(a, b) @@ -15812,7 +15814,7 @@ impl<'db> TypeInferenceBuilder<'db> { return registry.type_implements_interface_via_rule( sub, &requested_iface_ty, - &self.aliases, + self.aliases, |actual, bound| self.is_subtype(actual, bound), ); } @@ -15898,7 +15900,7 @@ impl<'db> TypeInferenceBuilder<'db> { // and the `extends` chain is reflected in the per-class data already. // For a pure interface-to-interface check, fall through to structural // equality (matches today's behaviour for unrelated interfaces). - crate::normalize::is_subtype_of(sub, sup, &self.aliases) + crate::normalize::is_subtype_of(sub, sup, self.aliases) } fn structural_subtype_with_nominal_interfaces(&self, sub: &Ty, sup: &Ty) -> Option { @@ -15910,10 +15912,10 @@ impl<'db> TypeInferenceBuilder<'db> { match (sub, sup) { (Ty::List(..) | Ty::EvolvingList(..), Ty::List(..) | Ty::EvolvingList(..)) => { - Some(crate::normalize::is_subtype_of(sub, sup, &self.aliases)) + Some(crate::normalize::is_subtype_of(sub, sup, self.aliases)) } (Ty::Map { .. } | Ty::EvolvingMap(..), Ty::Map { .. } | Ty::EvolvingMap(..)) => { - Some(crate::normalize::is_subtype_of(sub, sup, &self.aliases)) + Some(crate::normalize::is_subtype_of(sub, sup, self.aliases)) } (Ty::Future(sub_value, sub_error, _), Ty::Future(sup_value, sup_error, _)) => { Some(self.is_subtype(sub_value, sup_value) && self.is_subtype(sub_error, sup_error)) diff --git a/baml_language/crates/baml_compiler2_tir/src/callable.rs b/baml_language/crates/baml_compiler2_tir/src/callable.rs index 8225d6b80ae..2c754d1f15a 100644 --- a/baml_language/crates/baml_compiler2_tir/src/callable.rs +++ b/baml_language/crates/baml_compiler2_tir/src/callable.rs @@ -404,17 +404,14 @@ pub fn callable_throws<'db>(db: &'db dyn crate::Db, function: FunctionLoc<'db>) }; }; let inference = infer_scope_types(db, scope_id); - let res_ctx = package_resolution_context(db, pkg_id); - let aliases = crate::normalize::resolved_aliases_from_map( - crate::inference::package_alias_map(db, res_ctx), - ); + let aliases = crate::inference::package_alias_env(db, pkg_id); crate::throws_analysis::collect_escaping_throws( &CallableThrowsAnalysis { db, pkg_id, ns_context: &pkg_info.namespace_path, inference, - aliases: &aliases, + aliases, }, body, ) diff --git a/baml_language/crates/baml_compiler2_tir/src/inference.rs b/baml_language/crates/baml_compiler2_tir/src/inference.rs index 1d779bf93e8..9ad9a256059 100644 --- a/baml_language/crates/baml_compiler2_tir/src/inference.rs +++ b/baml_language/crates/baml_compiler2_tir/src/inference.rs @@ -1226,7 +1226,7 @@ pub fn infer_scope_types<'db>( let res_ctx = crate::package_interface::package_resolution_context(db, pkg_id); let pkg_items = &res_ctx.own_items; - let aliases = package_alias_map(db, res_ctx); + let aliases = package_alias_env(db, pkg_id); let context = InferContext::new(db, scope_id); let mut builder = TypeInferenceBuilder::new(context, res_ctx, pkg_id, scope_id, aliases); @@ -2319,6 +2319,26 @@ pub fn collect_type_aliases<'db>( aliases } +/// Salsa query: the full type-alias environment visible from a package — the +/// alias map ([`package_alias_map`]) plus its precomputed recursive-alias set. +/// +/// This is THE alias environment for per-scope inference and throws analysis: +/// computed once per package instead of once per scope (the old scheme +/// rebuilt the map — re-collecting every dependency's exports and cloning +/// every alias `Ty` — for every `infer_scope_types` call). +/// +/// Note this is intentionally distinct from the MIR-side +/// `resolved_aliases_for_package`, which collects *all* dependency aliases +/// from their items; this one sees only dependency *interface exports*. +#[salsa::tracked(returns(ref))] +pub fn package_alias_env<'db>( + db: &'db dyn crate::Db, + pkg_id: PackageId<'db>, +) -> crate::normalize::ResolvedAliases { + let res_ctx = crate::package_interface::package_resolution_context(db, pkg_id); + crate::normalize::resolved_aliases_from_map(package_alias_map(db, res_ctx)) +} + /// Build the type-alias map visible from a package: its own aliases plus those /// re-exported by its dependencies. Shared by per-scope inference and throws /// analysis so both expand the same aliases (e.g. `testing.TestSetBody`). diff --git a/baml_language/crates/baml_lsp2_actions/src/check.rs b/baml_language/crates/baml_lsp2_actions/src/check.rs index 0eb11ac164e..f6111739046 100644 --- a/baml_language/crates/baml_lsp2_actions/src/check.rs +++ b/baml_language/crates/baml_lsp2_actions/src/check.rs @@ -153,7 +153,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec { let pkg_id = baml_compiler2_hir::package::PackageId::new(db, pkg_info.package.clone()); let res_ctx = baml_compiler2_tir::package_interface::package_resolution_context(db, pkg_id); let pkg_items = &res_ctx.own_items; - let aliases = collect_type_aliases_for_resolution_context(db, res_ctx); + let aliases = baml_compiler2_tir::inference::package_alias_env(db, pkg_id); let ast_items = { let tree = baml_compiler_parser::syntax_tree(db, file); let (items, _, _) = baml_compiler2_ast::lower_file(&tree); @@ -165,7 +165,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec { &ast_items, pkg_items, &pkg_info.namespace_path, - &aliases, + aliases, )); // ── 5. Jinja prompt/template diagnostics ──────────────────────────────── @@ -370,11 +370,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec { if let Some(scope_id) = function_scope_id(index, func_data) { 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.aliases.clone(), + context, res_ctx, pkg_id, scope_id, aliases, ); builder.set_generic_params(generic_params); for (name, ty) in ¶m_types { @@ -432,7 +428,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec { &ast_items, pkg_items, &pkg_info.namespace_path, - &aliases, + aliases, )); // Deduplicate: multiple steps can produce the same diagnostic (e.g. scope @@ -5901,27 +5897,6 @@ fn jinja_diagnostic_id(message: &str) -> DiagnosticId { } } -fn collect_type_aliases_for_resolution_context<'db>( - db: &'db dyn Db, - res_ctx: &'db baml_compiler2_tir::package_interface::PackageResolutionContext<'db>, -) -> baml_compiler2_tir::normalize::ResolvedAliases { - let mut aliases = baml_compiler2_tir::inference::collect_type_aliases(db, &res_ctx.own_items); - for (_dep_name, dep_iface) in &res_ctx.dep_interfaces { - for types_in_ns in dep_iface.types.values() { - for exported in types_in_ns.values() { - if let baml_compiler2_tir::package_interface::ExportedType::TypeAlias { - qtn, - resolved, - } = exported - { - aliases.insert(qtn.clone(), resolved.clone()); - } - } - } - } - baml_compiler2_tir::normalize::resolved_aliases_from_map(aliases) -} - fn function_scope_id<'db>( index: &baml_compiler2_hir::semantic_index::FileSemanticIndex<'db>, func_data: &baml_compiler2_hir::item_tree::Function,