Consolidate name resolution through PackageResolutionContext

baml_language/crates/baml_compiler2_hir/src/package.rs

@@ -134,54 +134,20 @@ unsafe impl salsa::Update for PackageItems<'_> {
 }
 
 impl<'db> PackageItems<'db> {
-    /// Look up a type by path segments (e.g., `["llm", "render_prompt"]`).
+    /// Look up a type by explicit namespace and item name.
     ///
-    /// Tries progressively longer namespace prefixes: for path `["a", "B"]`,
-    /// first tries namespace `["a"]` + name `"B"`, then namespace `[]` + name `"a"`.
-    pub fn lookup_type(&self, path: &[Name]) -> Option<Definition<'db>> {
-        if path.is_empty() {
-            return None;
-        }
-        for split in (0..path.len()).rev() {
-            let ns_path = &path[..split];
-            let item_name = &path[split];
-            if let Some(ns) = self.namespaces.get(ns_path) {
-                if let Some(def) = ns.types.get(item_name) {
-                    return Some(*def);
-                }
-            }
-        }
-        None
+    /// Single hash lookup — no split-loop ambiguity.
+    /// `namespace` is the namespace path (e.g. `["llm"]` or `[]` for root).
+    /// `item` is the unqualified item name (e.g. `"Response"`).
+    pub fn lookup_type(&self, namespace: &[Name], item: &Name) -> Option<Definition<'db>> {
+        self.namespaces.get(namespace)?.types.get(item).copied()
     }
 
-    /// Look up a type by short name, searching across ALL namespaces.
+    /// Look up a value by explicit namespace and item name.
     ///
-    /// Use when the caller only has a short name (e.g. `"File"`) without
-    /// knowing which namespace it lives in. Returns the first match found.
-    pub fn lookup_type_any_ns(&self, name: &Name) -> Option<Definition<'db>> {
-        for ns in self.namespaces.values() {
-            if let Some(def) = ns.types.get(name) {
-                return Some(*def);
-            }
-        }
-        None
-    }
-
-    /// Look up a value by path segments.
-    pub fn lookup_value(&self, path: &[Name]) -> Option<Definition<'db>> {
-        if path.is_empty() {
-            return None;
-        }
-        for split in (0..path.len()).rev() {
-            let ns_path = &path[..split];
-            let item_name = &path[split];
-            if let Some(ns) = self.namespaces.get(ns_path) {
-                if let Some(def) = ns.values.get(item_name) {
-                    return Some(*def);
-                }
-            }
-        }
-        None
+    /// Single hash lookup — no split-loop ambiguity.
+    pub fn lookup_value(&self, namespace: &[Name], item: &Name) -> Option<Definition<'db>> {
+        self.namespaces.get(namespace)?.values.get(item).copied()
     }
 }
 

baml_language/crates/baml_compiler2_hir/src/signature.rs

@@ -37,6 +37,9 @@ pub struct FunctionSignature {
 pub struct SignatureSourceMap {
     /// One span per parameter, parallel to `FunctionSignature::params`.
     pub param_spans: Vec<TextRange>,
+    /// One span per parameter's type expression (just the type, not the name).
+    /// `None` when the parameter has no explicit type annotation.
+    pub param_type_spans: Vec<Option<TextRange>>,
     /// Span of the return type annotation, if present.
     pub return_type_span: Option<TextRange>,
     /// Span of the throws type annotation, if present.
@@ -79,6 +82,11 @@ fn function_signature_with_source_map<'db>(
     // Build source map — spans only (separate for early-cutoff)
     let source_map = SignatureSourceMap {
         param_spans: func_data.params.iter().map(|p| p.span).collect(),
+        param_type_spans: func_data
+            .params
+            .iter()
+            .map(|p| p.type_expr.as_ref().map(|te| te.span))
+            .collect(),
         return_type_span: func_data.return_type.as_ref().map(|te| te.span),
         throws_type_span: func_data.throws.as_ref().map(|te| te.span),
     };

baml_language/crates/baml_compiler2_mir/src/lower.rs

@@ -786,12 +786,19 @@ impl<'db> LoweringContext<'db> {
     /// Used for `TypedBinding` patterns where TIR may not have populated the
     /// bindings map (e.g. catch arm and match arm patterns).
     fn resolve_type_annotation(&self, ty_expr: &baml_compiler2_ast::TypeExpr) -> Ty {
-        use baml_compiler2_tir::lower_type_expr::lower_type_expr;
+        use baml_compiler2_tir::lower_type_expr::lower_type_expr_in_ns;
         let pkg_info = file_package(self.db, self.file);
         let pkg_id = PackageId::new(self.db, pkg_info.package);
         let pkg_items = package_items(self.db, pkg_id);
         let mut diags = Vec::new();
-        let tir_ty = lower_type_expr(self.db, ty_expr, pkg_items, &[], &mut diags);
+        let tir_ty = lower_type_expr_in_ns(
+            self.db,
+            ty_expr,
+            pkg_items,
+            &pkg_info.namespace_path,
+            &[],
+            &mut diags,
+        );
         convert_tir2_ty(&tir_ty, &self.type_aliases, &self.recursive_aliases)
     }
 }
@@ -800,7 +807,7 @@ impl<'db> LoweringContext<'db> {
 
 impl LoweringContext<'_> {
     fn lower_function_body(&mut self) -> MirFunction {
-        use baml_compiler2_tir::lower_type_expr::lower_type_expr;
+        use baml_compiler2_tir::lower_type_expr::lower_type_expr_in_ns;
 
         let func_loc = self
             .func_loc
@@ -817,7 +824,14 @@ impl LoweringContext<'_> {
             .as_ref()
             .map(|te| {
                 let mut diags = Vec::new();
-                let tir_ty = lower_type_expr(self.db, te, pkg_items, &[], &mut diags);
+                let tir_ty = lower_type_expr_in_ns(
+                    self.db,
+                    te,
+                    pkg_items,
+                    &pkg_info.namespace_path,
+                    &[],
+                    &mut diags,
+                );
                 convert_tir2_ty(&tir_ty, &self.type_aliases, &self.recursive_aliases)
             })
             .unwrap_or(Ty::Null {
@@ -855,20 +869,35 @@ impl LoweringContext<'_> {
                 enclosing_class_name
                     .as_ref()
                     .and_then(|cn| {
-                        pkg_items.lookup_type(std::slice::from_ref(cn)).map(|def| {
-                            let tir_ty = baml_compiler2_tir::ty::Ty::Class(
-                                baml_compiler2_tir::lower_type_expr::qualify_def(self.db, def, cn),
-                                baml_compiler2_tir::ty::TyAttr::default(),
-                            );
-                            convert_tir2_ty(&tir_ty, &self.type_aliases, &self.recursive_aliases)
-                        })
+                        pkg_items
+                            .lookup_type(&pkg_info.namespace_path, cn)
+                            .map(|def| {
+                                let tir_ty = baml_compiler2_tir::ty::Ty::Class(
+                                    baml_compiler2_tir::lower_type_expr::qualify_def(
+                                        self.db, def, cn,
+                                    ),
+                                    baml_compiler2_tir::ty::TyAttr::default(),
+                                );
+                                convert_tir2_ty(
+                                    &tir_ty,
+                                    &self.type_aliases,
+                                    &self.recursive_aliases,
+                                )
+                            })
                     })
                     .unwrap_or(Ty::Null {
                         attr: TyAttr::default(),
                     })
             } else {
                 let mut diags = Vec::new();
-                let tir_ty = lower_type_expr(self.db, param_te, pkg_items, &[], &mut diags);
+                let tir_ty = lower_type_expr_in_ns(
+                    self.db,
+                    param_te,
+                    pkg_items,
+                    &pkg_info.namespace_path,
+                    &[],
+                    &mut diags,
+                );
                 convert_tir2_ty(&tir_ty, &self.type_aliases, &self.recursive_aliases)
             };
             let local = self

baml_language/crates/baml_compiler2_ppir/src/expand.rs

@@ -19,12 +19,18 @@ pub enum SymbolKind {
 }
 
 pub fn classify_type(package_items: &PackageItems<'_>, path: &[Name]) -> Option<SymbolKind> {
-    package_items.lookup_type(path).and_then(|def| match def {
-        Definition::Class(_) => Some(SymbolKind::Class),
-        Definition::Enum(_) => Some(SymbolKind::Enum),
-        Definition::TypeAlias(_) => Some(SymbolKind::TypeAlias),
-        _ => None,
-    })
+    if path.is_empty() {
+        return None;
+    }
+    let item = path.last().unwrap();
+    package_items
+        .lookup_type(&path[..path.len() - 1], item)
+        .and_then(|def| match def {
+            Definition::Class(_) => Some(SymbolKind::Class),
+            Definition::Enum(_) => Some(SymbolKind::Enum),
+            Definition::TypeAlias(_) => Some(SymbolKind::TypeAlias),
+            _ => None,
+        })
 }
 
 /// Classify a type, falling back to `root.*` prefix handling, bare-name
@@ -62,61 +68,112 @@ fn classify_type_cross_pkg(path: &[Name], ctx: &ExpandCtx<'_>) -> Option<SymbolK
 // ── Namespace-aware key resolution ──────────────────────────────────────────
 
 /// Resolve a path within a single package to its qualified key
-/// `[package_name, ...ns_path, item_name]`.
+/// `[package_name, ...namespace, item_name]`.
 fn resolve_in_package(
-    path: &[Name],
+    namespace: &[Name],
+    item: &Name,
     pkg_name: &Name,
     pkg_items: &PackageItems<'_>,
 ) -> Option<Vec<Name>> {
-    if path.is_empty() {
-        return None;
-    }
-    for split in (0..path.len()).rev() {
-        let ns_path = &path[..split];
-        let item_name = &path[split];
-        if let Some(ns) = pkg_items.namespaces.get(ns_path) {
-            if ns.types.contains_key(item_name) {
-                let mut key = vec![pkg_name.clone()];
-                key.extend_from_slice(ns_path);
-                key.push(item_name.clone());
-                return Some(key);
-            }
-        }
-    }
-    None
+    pkg_items.lookup_type(namespace, item).map(|_| {
+        let mut key = vec![pkg_name.clone()];
+        key.extend_from_slice(namespace);
+        key.push(item.clone());
+        key
+    })
 }
 
 /// Resolve a PPIR type path to its qualified key `[package, ...ns, name]`.
 /// Handles direct lookup, `root.*` prefix, bare names in non-root namespaces,
 /// and cross-package references.
 fn resolve_qualified_key(path: &[Name], ctx: &ExpandCtx<'_>) -> Option<Vec<Name>> {
+    if path.is_empty() {
+        return None;
+    }
+    let item = path.last().unwrap();
     // 1. Direct lookup in current package
-    if let Some(key) = resolve_in_package(path, ctx.package_name, ctx.package_items) {
+    let ns = &path[..path.len() - 1];
+    if let Some(key) = resolve_in_package(ns, item, ctx.package_name, ctx.package_items) {
         return Some(key);
     }
     // 2. Handle `root.*` prefix
     if path.len() >= 2 && path[0].as_str() == "root" {
-        if let Some(key) = resolve_in_package(&path[1..], ctx.package_name, ctx.package_items) {
+        let after_root = &path[1..];
+        let root_item = after_root.last().unwrap();
+        let root_ns = &after_root[..after_root.len() - 1];
+        if let Some(key) =
+            resolve_in_package(root_ns, root_item, ctx.package_name, ctx.package_items)
+        {
             return Some(key);
         }
     }
     // 3. Bare name in current (non-root) namespace
     if path.len() == 1 && !ctx.namespace_path.is_empty() {
-        let mut ns_qualified: Vec<Name> = ctx.namespace_path.to_vec();
-        ns_qualified.push(path[0].clone());
-        if let Some(key) = resolve_in_package(&ns_qualified, ctx.package_name, ctx.package_items) {
+        if let Some(key) = resolve_in_package(
+            ctx.namespace_path,
+            item,
+            ctx.package_name,
+            ctx.package_items,
+        ) {
             return Some(key);
         }
     }
     // 4. Cross-package (first segment = package name)
     if path.len() >= 2 {
         if let Some(foreign_items) = ctx.all_package_items.get(&path[0]) {
-            return resolve_in_package(&path[1..], &path[0], foreign_items);
+            let after_pkg = &path[1..];
+            let pkg_item = after_pkg.last().unwrap();
+            let pkg_ns = &after_pkg[..after_pkg.len() - 1];
+            return resolve_in_package(pkg_ns, pkg_item, &path[0], foreign_items);
         }
     }
     None
 }
 
+/// When a type alias in one namespace resolves to a type in a different
+/// namespace, the resulting `Named` path (and paths inside unions/lists/etc.)
+/// must be qualified so that `lower_type_expr_in_ns` can find them from the
+/// caller's namespace. Prepends `root.` to single-segment `Named` paths when
+/// the alias namespace differs from the caller namespace.
+fn requalify_for_caller(ty: PpirTy, alias_ns: &[Name], caller_ns: &[Name]) -> PpirTy {
+    if alias_ns == caller_ns {
+        return ty;
+    }
+    match ty {
+        PpirTy::Named { path, attrs } if path.len() == 1 && path[0].as_str() != "root" => {
+            let mut qualified = Vec::with_capacity(alias_ns.len() + 2);
+            qualified.push(SmolStr::from("root"));
+            qualified.extend(alias_ns.iter().cloned());
+            qualified.extend(path);
+            PpirTy::Named {
+                path: qualified,
+                attrs,
+            }
+        }
+        PpirTy::Union { variants, attrs } => PpirTy::Union {
+            variants: variants
+                .into_iter()
+                .map(|v| requalify_for_caller(v, alias_ns, caller_ns))
+                .collect(),
+            attrs,
+        },
+        PpirTy::List { inner, attrs } => PpirTy::List {
+            inner: Box::new(requalify_for_caller(*inner, alias_ns, caller_ns)),
+            attrs,
+        },
+        PpirTy::Optional { inner, attrs } => PpirTy::Optional {
+            inner: Box::new(requalify_for_caller(*inner, alias_ns, caller_ns)),
+            attrs,
+        },
+        PpirTy::Map { key, value, attrs } => PpirTy::Map {
+            key: Box::new(requalify_for_caller(*key, alias_ns, caller_ns)),
+            value: Box::new(requalify_for_caller(*value, alias_ns, caller_ns)),
+            attrs,
+        },
+        other => other,
+    }
+}
+
 // ── Context ──────────────────────────────────────────────────────────────────
 
 /// Shared context threaded through all stream-expansion functions.
@@ -379,11 +436,32 @@ fn stream_expand_inner(ty: &PpirTy, ctx: &ExpandCtx<'_>, depth: u32) -> (PpirTy,
                         if depth < MAX_ALIAS_DEPTH {
                             if let Some(key) = resolve_qualified_key(path, ctx) {
                                 if let Some(body) = ctx.alias_bodies.get(&key) {
-                                    // Set merged attrs on the resolved body and recurse
+                                    // The alias body's paths are relative to the alias
+                                    // definition's namespace (key = [pkg, ...ns, name]).
+                                    // Recurse with the alias's namespace so that
+                                    // classify_type / resolve_qualified_key resolve
+                                    // the body's bare names correctly.
+                                    let alias_ns = key[1..key.len() - 1].to_vec();
+                                    let alias_ctx = ExpandCtx {
+                                        namespace_path: &alias_ns,
+                                        ..*ctx
+                                    };
                                     let mut resolved = body.clone();
                                     resolved.attrs_mut().stream_must_exist = must_exist;
                                     resolved.attrs_mut().stream_done = done;
-                                    return stream_expand_inner(&resolved, ctx, depth + 1);
+                                    let (result_ty, sap) =
+                                        stream_expand_inner(&resolved, &alias_ctx, depth + 1);
+                                    // The result's Named paths are relative to alias_ns.
+                                    // If the caller is in a different namespace, qualify
+                                    // them so lower_type_expr_in_ns can resolve them.
+                                    return (
+                                        requalify_for_caller(
+                                            result_ty,
+                                            &alias_ns,
+                                            ctx.namespace_path,
+                                        ),
+                                        sap,
+                                    );
                                 }
                             }
                         }