Compiler perf: track package alias environments in Salsa (1,167 rebuilds → ~10)

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<Name, ResolvedAliases> {
-    let mut caches: HashMap<Name, ResolvedAliases> = HashMap::new();
+) -> HashMap<Name, &'db ResolvedAliases> {
+    let mut caches: HashMap<Name, &'db ResolvedAliases> = 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<Name, ResolvedAliases>,
+    alias_caches: &HashMap<Name, &ResolvedAliases>,
     function_indices: &HashMap<String, usize>,
     interface_indices: &HashMap<baml_type::TypeName, usize>,
     program_packages: &mut indexmap::IndexMap<Name, bex_vm_types::types::ProgramPackage>,

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();

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<Ty>,
     /// 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<Name>,
     /// 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<crate::ty::QualifiedTypeName, Ty>,
+        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<bool> {
@@ -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))

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,
             )

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`).

baml_language/crates/baml_lsp2_actions/src/check.rs

@@ -153,7 +153,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec<Diagnostic> {
     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<Diagnostic> {
         &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<Diagnostic> {
         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 &param_types {
@@ -432,7 +428,7 @@ pub fn check_file(db: &dyn Db, file: SourceFile) -> Vec<Diagnostic> {
         &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,