From 1d6e21299c5f58f220d992e74847a3ce69d3f971 Mon Sep 17 00:00:00 2001 From: Avery Townsend Date: Mon, 27 Jul 2026 15:39:42 -0700 Subject: [PATCH] fix(compiler): preserve local binding identity --- .../crates/baml_compiler2_hir/src/builder.rs | 99 ++-- .../crates/baml_compiler2_hir/src/lib.rs | 2 +- .../baml_compiler2_hir/src/semantic_index.rs | 104 ++-- .../crates/baml_compiler2_mir/src/lower.rs | 471 ++++++++---------- .../crates/baml_compiler2_tir/src/builder.rs | 28 +- .../baml_lsp2_actions/src/annotations.rs | 2 +- .../crates/baml_lsp2_actions/src/usages.rs | 137 +++-- .../baml_lsp2_actions/src/usages_at_tests.rs | 27 + .../ns_lexical_scoping/lexical_scoping.baml | 24 + .../prompt_tag_runtime.baml | 15 + .../snapshots/baml_src/lexical_scoping.snap | 77 +++ .../baml_src/prompt_tag_runtime.snap | 214 +++++++- 12 files changed, 788 insertions(+), 412 deletions(-) diff --git a/baml_language/crates/baml_compiler2_hir/src/builder.rs b/baml_language/crates/baml_compiler2_hir/src/builder.rs index 642f057ae76..c4bbc1d0ef5 100644 --- a/baml_language/crates/baml_compiler2_hir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_hir/src/builder.rs @@ -26,8 +26,9 @@ use crate::{ }, scope::{FileScopeId, ItemScopeOwner, Scope, ScopeId, ScopeKind}, semantic_index::{ - BindingId, DefinitionSite, FileSemanticIndex, LocalBinding, PathResolution, ScopeBindings, - SemanticIndexExtra, visible_binding_at_in_scopes, + BindingId, DefinitionSite, ExprMetadataKey, ExprMetadataScope, FileSemanticIndex, + LocalBinding, PathResolution, ScopeBindings, SemanticIndexExtra, + visible_binding_at_in_scopes, }, }; @@ -74,12 +75,13 @@ pub struct SemanticIndexBuilder<'db> { /// contribute to top-level symbols — they belong to the class scope). class_depth: u32, - /// Expression → scope mappings, sorted by `ExprId` at the end. - expr_scopes: Vec<(ast::ExprId, FileScopeId)>, + /// Expression to lexical scope mappings, sorted by arena-safe key at the end. + expr_scopes: Vec<(ExprMetadataKey, FileScopeId)>, - /// Path root resolutions for multi-segment `Path` expressions. - /// Collected during `walk_expr_body`, sorted by `ExprId` at the end. - path_resolutions: Vec<(ast::ExprId, PathResolution)>, + /// Path root resolutions, sorted by arena-safe expression key at the end. + path_resolutions: Vec<(ExprMetadataKey, PathResolution)>, + /// Arena namespace active while walking an expression body or defaults. + expr_metadata_scope_stack: Vec, /// Path-root references collected while walking source order. Unlike /// `expr_scopes`, this carries the scope and innermost lambda context at /// collection time so capture analysis does not rely on arena-local `ExprId`s. @@ -108,6 +110,7 @@ impl<'db> SemanticIndexBuilder<'db> { class_depth: 0, expr_scopes: Vec::new(), path_resolutions: Vec::new(), + expr_metadata_scope_stack: Vec::new(), path_root_references: Vec::new(), lambda_stack: Vec::new(), item_tree: crate::item_tree::builder::ItemTreeBuilder::new(), @@ -174,10 +177,10 @@ impl<'db> SemanticIndexBuilder<'db> { self.pop_scope(); // Project // Sort expr_scopes for binary search - self.expr_scopes.sort_by_key(|(id, _)| *id); + self.expr_scopes.sort_by_key(|(key, _)| *key); // Sort path_resolutions for binary search - self.path_resolutions.sort_by_key(|(id, _)| *id); + self.path_resolutions.sort_by_key(|(key, _)| *key); // Pre-intern ScopeIds for each FileScopeId let scope_ids: Vec> = (0..self.scopes.len()) @@ -264,8 +267,17 @@ impl<'db> SemanticIndexBuilder<'db> { // ── Expression recording ───────────────────────────────────────────────── /// Record that an expression belongs to the current scope. + fn current_expr_metadata_key(&self, expr_id: ast::ExprId) -> ExprMetadataKey { + let scope = *self + .expr_metadata_scope_stack + .last() + .expect("expression walked without an arena namespace"); + ExprMetadataKey::new(scope, expr_id) + } + fn record_expr_scope(&mut self, expr_id: ast::ExprId) { - self.expr_scopes.push((expr_id, self.current_scope_id())); + let key = self.current_expr_metadata_key(expr_id); + self.expr_scopes.push((key, self.current_scope_id())); } /// Build a dotted scope path from the current scope stack, e.g. `Foo.Bar`. @@ -327,9 +339,30 @@ impl<'db> SemanticIndexBuilder<'db> { /// Walk an `ExprBody` arena in source order, recording expression ownership /// and local bindings in the lexical scope that owns each expression. fn walk_expr_body(&mut self, body: &ast::ExprBody, source_map: &ast::AstSourceMap) { + let metadata_scope = ExprMetadataScope::Body(self.current_scope_id()); + self.expr_metadata_scope_stack.push(metadata_scope); if let Some(root_expr) = body.root_expr { self.walk_expr(root_expr, body, source_map, false); } + let popped = self.expr_metadata_scope_stack.pop(); + debug_assert_eq!(popped, Some(metadata_scope)); + } + + fn walk_parameter_defaults(&mut self, function: &ast::FunctionDef) { + let metadata_scope = ExprMetadataScope::ParameterDefault(self.current_scope_id()); + self.expr_metadata_scope_stack.push(metadata_scope); + for param in &function.params { + if let Some(default) = param.default { + self.walk_expr( + default.expr(), + &function.defaults.exprs, + &function.defaults.source_map, + true, + ); + } + } + let popped = self.expr_metadata_scope_stack.pop(); + debug_assert_eq!(popped, Some(metadata_scope)); } /// Walk an expression, recording its `FileScopeId` and (for `Block`s) @@ -684,9 +717,7 @@ impl<'db> SemanticIndexBuilder<'db> { let use_scope = self.current_scope_id(); let use_offset = source_map.expr_span(expr_id).start(); self.record_path_root_reference(root, use_scope, use_offset); - if segments.len() >= 2 { - self.classify_path_expr(expr_id, segments, use_scope, use_offset); - } + self.resolve_path_expr(expr_id, root, use_scope, use_offset); } } ast::Expr::GenericApply { base, .. } => { @@ -1094,16 +1125,7 @@ impl<'db> SemanticIndexBuilder<'db> { } self.emit_duplicate_param_diagnostics(&func_def.params); self.lambda_stack.push(scope_id); - for param in &func_def.params { - if let Some(default) = param.default { - self.walk_expr( - default.expr(), - &func_def.defaults.exprs, - &func_def.defaults.source_map, - true, - ); - } - } + self.walk_parameter_defaults(func_def); if let Some(ast::FunctionBodyDef::Expr(lambda_body, lambda_source_map)) = &func_def.body { self.walk_expr_body(lambda_body, lambda_source_map); self.analyze_lambda_captures(scope_id, lambda_body, lambda_source_map); @@ -1170,26 +1192,18 @@ impl<'db> SemanticIndexBuilder<'db> { false } - fn classify_path_expr( + fn resolve_path_expr( &mut self, expr_id: ast::ExprId, - segments: &[Name], + root: &Name, use_scope: FileScopeId, use_offset: TextSize, ) { - if segments.len() < 2 { - return; - } - let root = &segments[0]; - let resolution = if self + let resolution = self .visible_binding_at(use_scope, use_offset, root) - .is_some() - { - PathResolution::Local { name: root.clone() } - } else { - PathResolution::Unknown - }; - self.path_resolutions.push((expr_id, resolution)); + .map_or(PathResolution::Unknown, PathResolution::Local); + let key = self.current_expr_metadata_key(expr_id); + self.path_resolutions.push((key, resolution)); } fn record_path_root_reference( @@ -1252,16 +1266,7 @@ impl<'db> SemanticIndexBuilder<'db> { .push((param.name.clone(), idx)); } self.emit_duplicate_param_diagnostics(&f.params); - for param in &f.params { - if let Some(default) = param.default { - self.walk_expr( - default.expr(), - &f.defaults.exprs, - &f.defaults.source_map, - true, - ); - } - } + self.walk_parameter_defaults(f); if let Some(ast::FunctionBodyDef::Expr(ref body, ref source_map)) = f.body { self.walk_expr_body(body, source_map); diff --git a/baml_language/crates/baml_compiler2_hir/src/lib.rs b/baml_language/crates/baml_compiler2_hir/src/lib.rs index 99b3b782bde..f21dad5b3f7 100644 --- a/baml_language/crates/baml_compiler2_hir/src/lib.rs +++ b/baml_language/crates/baml_compiler2_hir/src/lib.rs @@ -31,7 +31,7 @@ use std::sync::Arc; use baml_base::SourceFile; pub use builder::SemanticIndexBuilder; -pub use semantic_index::PathResolution; +pub use semantic_index::{ExprMetadataKey, ExprMetadataScope, PathResolution}; use crate::{ contributions::FileSymbolContributions, 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 8fb907ac3db..5250fcb08a6 100644 --- a/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs +++ b/baml_language/crates/baml_compiler2_hir/src/semantic_index.rs @@ -10,27 +10,40 @@ use baml_base::Name; use baml_compiler2_ast::{ExprId, LoweringDiagnostic, PatId, StmtId}; use text_size::{TextRange, TextSize}; -// ── PathResolution ──────────────────────────────────────────────────────────── +// Expression identity and path resolution -/// Scope-level resolution of a multi-segment `Path` expression's root segment. +/// Namespace for arena-local expression IDs. /// -/// Produced during HIR scope building for `Path` nodes with ≥ 2 segments. -/// Only distinguishes whether the root is a locally-declared variable/param -/// (in which case segments[1..] are field accesses on the local) or not -/// (in which case the path is a package-qualified name, to be resolved by TIR). +/// A function or lambda body and its parameter-default expressions use +/// separate `ExprBody` arenas, and every nested lambda starts another pair of +/// arenas. Pairing the arena owner with `ExprId` makes expression identity +/// unique within a file. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum ExprMetadataScope { + Body(FileScopeId), + ParameterDefault(FileScopeId), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ExprMetadataKey { + pub scope: ExprMetadataScope, + pub expr: ExprId, +} + +impl ExprMetadataKey { + pub fn new(scope: ExprMetadataScope, expr: ExprId) -> Self { + Self { scope, expr } + } +} + +/// HIR resolution of a value path's root segment. /// -/// Note: Package-item resolution requires cross-file `package_items` queries -/// that are not available during HIR building (would create a circular -/// dependency: `file_semantic_index` → `package_items` → `file_semantic_index`). -/// Full resolution (namespace vs package-item vs unknown) is done in TIR. -#[derive(Debug, Clone, PartialEq, Eq)] +/// Local roots are resolved to declaration identity while the lexical scope +/// stack is active. Non-local roots are completed by TIR because package-item +/// resolution would introduce a HIR query cycle. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PathResolution { - /// Root segment is a local variable, parameter, or capture in the current scope. - /// Segments[1..] are field/method accesses on the local, resolved by TIR type - /// checking. - Local { name: Name }, - /// Root segment is not a known local. May be a package name, namespace - /// shorthand, or truly unresolved — TIR determines which. + Local(BindingId), Unknown, } @@ -209,14 +222,14 @@ pub struct FileSemanticIndex<'db> { /// Scope tree — arena of `Scope` nodes indexed by `FileScopeId`. pub scopes: Vec, - /// Expression → owning scope mapping. + /// Expression to owning lexical scope mapping. /// /// Every `ExprId` in the file's `ExprBody` arenas is mapped to the /// `FileScopeId` of its innermost containing scope. Built during the /// `SemanticIndexBuilder` walk. /// - /// Sorted by `ExprId` for binary search (more compact than `HashMap`). - pub expr_scopes: Vec<(ExprId, FileScopeId)>, + /// Sorted by arena-safe expression key for binary search. + pub expr_scopes: Vec<(ExprMetadataKey, FileScopeId)>, /// Per-scope local bindings, indexed by `FileScopeId`. /// Parallel to `scopes` — `scope_bindings[i]` holds bindings for `scopes[i]`. @@ -253,15 +266,9 @@ pub struct FileSemanticIndex<'db> { /// following Ty's `Option>` pattern. pub extra: Option>, - /// Root-segment resolution for multi-segment `Path` expressions. - /// - /// For each `Path` with ≥ 2 segments, records whether the root is a local - /// variable/parameter (`PathResolution::Local`) or not - /// (`PathResolution::Unknown`). Sorted by `ExprId` for binary-search lookup. - /// - /// Package-item resolution (namespace, split point) is deferred to TIR - /// since it requires cross-file `package_items` queries. - pub path_resolutions: Vec<(ExprId, PathResolution)>, + /// Root resolution for every value `Path` expression, sorted by its + /// arena-safe expression key for binary-search lookup. + pub path_resolutions: Vec<(ExprMetadataKey, PathResolution)>, /// Environment variable references (`env.X`) found in this file's expression bodies. pub env_var_refs: Vec, @@ -311,6 +318,32 @@ impl FileSemanticIndex<'_> { }) } + /// Find a `Lambda` scope with `span` inside `owner_scope`. + /// + /// Restricting the search to the owner's DFS subtree disambiguates + /// synthesized companion functions that share source ranges. + pub fn lambda_scope_for_within( + &self, + owner_scope: FileScopeId, + span: text_size::TextRange, + ) -> Option { + let descendants = self + .scopes + .get(owner_scope.index() as usize)? + .descendants + .clone(); + let start = descendants.start.index() as usize; + let end = descendants.end.index() as usize; + self.scopes[start..end] + .iter() + .enumerate() + .find(|(_, scope)| matches!(scope.kind, ScopeKind::Lambda) && scope.range == span) + .map(|(offset, _)| { + #[allow(clippy::cast_possible_truncation)] + FileScopeId::new((start + offset) as u32) + }) + } + /// Find the innermost scope containing `offset`. /// /// Scopes are in DFS pre-order. We walk in reverse (deepest first) @@ -364,14 +397,21 @@ impl FileSemanticIndex<'_> { FileScopeId::ROOT } - /// Look up which scope owns an expression. - pub fn expression_scope(&self, expr_id: ExprId) -> Option { + /// Look up which lexical scope owns an expression. + pub fn expression_scope(&self, key: ExprMetadataKey) -> Option { self.expr_scopes - .binary_search_by_key(&expr_id, |(id, _)| *id) + .binary_search_by_key(&key, |(candidate, _)| *candidate) .ok() .map(|idx| self.expr_scopes[idx].1) } + pub fn path_resolution(&self, key: ExprMetadataKey) -> Option { + self.path_resolutions + .binary_search_by_key(&key, |(candidate, _)| *candidate) + .ok() + .map(|idx| self.path_resolutions[idx].1) + } + /// Walk ancestor scopes from `scope_id` upward to the root. pub fn ancestor_scopes(&self, scope_id: FileScopeId) -> Vec { let mut ancestors = vec![scope_id]; diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index a83cf2dc2f0..02a9a566de2 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -1572,7 +1572,10 @@ use baml_compiler2_hir::{ loc::{FunctionLoc, LetLoc}, package::{PackageId, package_items}, scope::FileScopeId, - semantic_index::{BindingId, DefinitionSite}, + semantic_index::{ + BindingId, DefinitionSite, ExprMetadataKey, ExprMetadataScope as MetadataScope, + PathResolution, + }, }; use baml_compiler2_ppir::file_semantic_index; use baml_compiler2_tir::{ @@ -1927,13 +1930,6 @@ struct InterfaceFieldCandidate { field_idx: usize, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -enum MetadataScope { - Body(FileScopeId), - ParameterDefault(FileScopeId), -} - -type ExprMetadataKey = (MetadataScope, AstExprId); type PatMetadataKey = (MetadataScope, AstPatId); struct LoweringContext<'db> { @@ -3115,38 +3111,62 @@ impl<'db> LoweringContext<'db> { self.any_pattern_binding_is_captured(pattern, DefinitionSite::CatchBinding(pattern)) } - fn binding_id_for_name_at(&self, expr_id: AstExprId, name: &Name) -> Option { + fn path_resolution(&self, expr_id: AstExprId) -> Option { let index = file_semantic_index(self.db, self.file); - let (scope_id, offset) = if let Some(source_map) = self.source_map.as_ref() { - let offset = source_map.expr_span(expr_id).start(); - ( - index.scope_at_offset(offset, self.scope_func_name.as_ref()), - offset, - ) - } else { - // The source-map-less branch is only valid for **synthesized** - // expressions emitted by the lowering itself (e.g. for-loop index - // increments, capture forwarding, init function bodies). The - // fallback uses `current_scope` and the scope's end offset, which - // is correct for synthesized refs at the end of the current scope - // but would silently pick the post-shadow binding for a - // user-written name lowered without a source map. - // - // If you find yourself adding a user-visible expression that - // hits this path: the right fix is to thread a `BindingId` - // through to the call site, not to widen this fallback. - let scope_id = self.current_scope; - let offset = index.scopes[scope_id.index() as usize].range.end(); - (scope_id, offset) + index.path_resolution(self.expr_metadata_key(expr_id)) + } + + fn hir_binding_id_for_path(&self, expr_id: AstExprId) -> Option { + match self.path_resolution(expr_id)? { + PathResolution::Local(binding_id) => Some(binding_id), + PathResolution::Unknown => None, + } + } + + fn binding_id_for_path(&self, expr_id: AstExprId, name: &Name) -> Option { + let hir_binding = self.hir_binding_id_for_path(expr_id); + let Some(&tagged_binding) = self.tagged_body_param_bindings.get(name) else { + return hir_binding; }; - index.visible_binding_at(scope_id, offset, name) + let Some(hir_binding) = hir_binding else { + return Some(tagged_binding); + }; + + // Tagged-template body parameters are synthetic MIR bindings. Prefer a + // HIR binding only when it is inside the tagged body scope, where it + // represents a real lexical shadow. + let index = file_semantic_index(self.db, self.file); + if index + .ancestor_scopes(hir_binding.scope) + .contains(&tagged_binding.scope) + { + Some(hir_binding) + } else { + Some(tagged_binding) + } + } + + fn local_for_path(&self, expr_id: AstExprId, name: &Name) -> Option { + let binding_id = self.binding_id_for_path(expr_id, name)?; + self.binding_locals.get(&binding_id).copied() } - fn capture_index_for_name_at(&self, expr_id: AstExprId, name: &Name) -> Option { - let binding_id = self.binding_id_for_name_at(expr_id, name)?; - self.capture_indices + fn place_for_path(&mut self, expr_id: AstExprId, name: &Name) -> Option { + let binding_id = self.binding_id_for_path(expr_id, name)?; + if let Some(&local) = self.binding_locals.get(&binding_id) { + return Some(Place::Local(local)); + } + if let Some(capture) = self + .capture_indices .as_ref() .and_then(|captures| captures.get(&binding_id).copied()) + { + return Some(Place::Capture(capture)); + } + if self.tagged_body_param_bindings.get(name) == Some(&binding_id) { + return Some(Place::Capture(self.ensure_transitive_capture(binding_id))); + } + None } /// Return the current lambda's capture index for `binding_id`, allocating a @@ -3235,7 +3255,7 @@ impl<'db> LoweringContext<'db> { /// Get the `baml_type::RuntimeTy` for an expression by looking up in the aggregated map /// and converting from TIR `Ty`. Uses `current_metadata_scope` as the arena namespace. fn expr_metadata_key(&self, expr_id: AstExprId) -> ExprMetadataKey { - (self.current_metadata_scope, expr_id) + ExprMetadataKey::new(self.current_metadata_scope, expr_id) } fn pat_metadata_key(&self, pat_id: AstPatId) -> PatMetadataKey { @@ -3261,10 +3281,10 @@ impl<'db> LoweringContext<'db> { } fn tir_expr_type(&self, key: ExprMetadataKey) -> Option<&'db Tir2Ty> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.expression_type(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.expression_type(key.expr), MetadataScope::ParameterDefault(fsi) => { - self.inference_for(fsi)?.default_expression_type(key.1) + self.inference_for(fsi)?.default_expression_type(key.expr) } } } @@ -3282,30 +3302,30 @@ impl<'db> LoweringContext<'db> { &self, key: ExprMetadataKey, ) -> Option<&'db baml_compiler2_tir::inference::MemberResolution<'db>> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.resolution(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.resolution(key.expr), MetadataScope::ParameterDefault(fsi) => { - self.inference_for(fsi)?.default_resolution(key.1) + self.inference_for(fsi)?.default_resolution(key.expr) } } } fn tir_is_exhaustive_match(&self, key: ExprMetadataKey) -> bool { - match key.0 { + match key.scope { MetadataScope::Body(fsi) => self .inference_for(fsi) - .is_some_and(|inf| inf.is_exhaustive_match(key.1)), + .is_some_and(|inf| inf.is_exhaustive_match(key.expr)), MetadataScope::ParameterDefault(fsi) => self .inference_for(fsi) - .is_some_and(|inf| inf.default_is_exhaustive_match(key.1)), + .is_some_and(|inf| inf.default_is_exhaustive_match(key.expr)), } } fn tir_path_root_type(&self, key: ExprMetadataKey) -> Option<&'db Tir2Ty> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.path_root_type(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.path_root_type(key.expr), MetadataScope::ParameterDefault(fsi) => { - self.inference_for(fsi)?.default_path_root_type(key.1) + self.inference_for(fsi)?.default_path_root_type(key.expr) } } } @@ -3323,11 +3343,11 @@ impl<'db> LoweringContext<'db> { &self, key: ExprMetadataKey, ) -> Option<&'db [baml_compiler2_tir::inference::MemberResolution<'db>]> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.path_member_resolution(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.path_member_resolution(key.expr), MetadataScope::ParameterDefault(fsi) => self .inference_for(fsi)? - .default_path_member_resolution(key.1), + .default_path_member_resolution(key.expr), } } @@ -3335,10 +3355,10 @@ impl<'db> LoweringContext<'db> { &self, key: ExprMetadataKey, ) -> Option<&'db baml_compiler2_tir::inference::CallPlan> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.call_plan(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.call_plan(key.expr), MetadataScope::ParameterDefault(fsi) => { - self.inference_for(fsi)?.default_call_plan(key.1) + self.inference_for(fsi)?.default_call_plan(key.expr) } } } @@ -3347,10 +3367,10 @@ impl<'db> LoweringContext<'db> { &self, key: ExprMetadataKey, ) -> Option<&'db baml_compiler2_tir::inference::FunctionCoercion> { - match key.0 { - MetadataScope::Body(fsi) => self.inference_for(fsi)?.function_coercion(key.1), + match key.scope { + MetadataScope::Body(fsi) => self.inference_for(fsi)?.function_coercion(key.expr), MetadataScope::ParameterDefault(fsi) => { - self.inference_for(fsi)?.default_function_coercion(key.1) + self.inference_for(fsi)?.default_function_coercion(key.expr) } } } @@ -3545,7 +3565,7 @@ impl<'db> LoweringContext<'db> { expr_id: AstExprId, name: &Name, ) -> Option { - let binding_id = self.binding_id_for_name_at(expr_id, name)?; + let binding_id = self.binding_id_for_path(expr_id, name)?; self.source_param_interface_view_for_binding(name, binding_id) } @@ -5131,31 +5151,9 @@ impl LoweringContext<'_> { // function's captures. Disambiguate by preferring the lambda scope // nested within the function currently being lowered; fall back to // the first range match. - let current_descendants = index - .scopes - .get(self.current_scope.index() as usize) - .map(|s| s.descendants.clone()); - let is_in_current_fn = |id: FileScopeId| { - current_descendants - .as_ref() - .is_some_and(|d| id.index() >= d.start.index() && id.index() < d.end.index()) - }; - let mut first = None; - let mut scoped = None; - for (i, scope) in index.scopes.iter().enumerate() { - if scope.kind == baml_compiler2_hir::scope::ScopeKind::Lambda && scope.range == span - { - let id = FileScopeId::new(i as u32); - if first.is_none() { - first = Some(id); - } - if is_in_current_fn(id) { - scoped = Some(id); - break; - } - } - } - let found = scoped.or(first); + let found = index + .lambda_scope_for_within(self.current_scope, span) + .or_else(|| index.lambda_scope_for(span)); debug_assert!( found.is_some(), "no HIR Lambda scope for tagged template at {span:?}" @@ -6035,11 +6033,11 @@ impl LoweringContext<'_> { /// Whether `segments` is rooted at the BEP-044 `default` receiver keyword /// and that keyword is not shadowed by a local of the same name. See /// [`baml_compiler2_ast::DEFAULT_RECEIVER_KEYWORD`]. - fn is_default_receiver_root(&self, segments: &[Name]) -> bool { + fn is_default_receiver_root(&self, expr_id: AstExprId, segments: &[Name]) -> bool { segments .first() .is_some_and(|s| s.as_str() == baml_compiler2_ast::DEFAULT_RECEIVER_KEYWORD) - && !self.locals.contains_key(&segments[0]) + && self.binding_id_for_path(expr_id, &segments[0]).is_none() } fn lower_literal(lit: &AstLiteral) -> Constant { @@ -6087,16 +6085,10 @@ impl<'db> LoweringContext<'db> { if let Some(item) = resolution_to_item_ref(self.db, &resolution) { let receiver_segments = &segments[..segments.len() - 1]; let receiver_op = if receiver_segments.len() == 1 { - if let Some(&recv_local) = self.locals.get(&receiver_segments[0]) { - Operand::Copy(Place::Local(recv_local)) - } else if let Some(cap_idx) = - self.capture_index_for_name_at(expr_id, &receiver_segments[0]) - { - // Receiver is a captured variable — use capture slot. - Operand::Copy(Place::Capture(cap_idx)) - } else { - Operand::Constant(Constant::Null) - } + self.place_for_path(expr_id, &segments[0]).map_or_else( + || Operand::Constant(Constant::Null), + Operand::Copy, + ) } else { // Multi-segment receiver (e.g. `cfg.encoder`): lower as field chain. let recv_ty = self.expr_ty(expr_id); @@ -6126,7 +6118,7 @@ impl<'db> LoweringContext<'db> { Some( MemberResolution::InterfaceVirtualMethod { .. } | MemberResolution::InterfaceConcreteMethod { .. }, - ) if self.locals.contains_key(&segments[0]) => {} + ) if self.binding_id_for_path(expr_id, &segments[0]).is_some() => {} Some( MemberResolution::UnboundMethod { .. } | MemberResolution::Free { .. } @@ -6172,16 +6164,10 @@ impl<'db> LoweringContext<'db> { if let Some(item) = resolution_to_item_ref(self.db, &resolution) { let receiver_segments = &segments[..segments.len() - 1]; let receiver_op = if receiver_segments.len() == 1 { - if let Some(&recv_local) = self.locals.get(&receiver_segments[0]) { - Operand::Copy(Place::Local(recv_local)) - } else if let Some(cap_idx) = - self.capture_index_for_name_at(expr_id, &receiver_segments[0]) - { - // Receiver is a captured variable — use capture slot. - Operand::Copy(Place::Capture(cap_idx)) - } else { - Operand::Constant(Constant::Null) - } + self.place_for_path(expr_id, &segments[0]).map_or_else( + || Operand::Constant(Constant::Null), + Operand::Copy, + ) } else { let recv_ty = self.expr_ty(expr_id); let recv_local = self.builder.temp(recv_ty); @@ -6207,7 +6193,7 @@ impl<'db> LoweringContext<'db> { // below captures the receiver and binds its impl at runtime. MemberResolution::InterfaceVirtualMethod { .. } | MemberResolution::InterfaceConcreteMethod { .. } - if self.locals.contains_key(&segments[0]) => {} + if self.binding_id_for_path(expr_id, &segments[0]).is_some() => {} MemberResolution::UnboundMethod { .. } | MemberResolution::Free { .. } | MemberResolution::InterfaceVirtualMethod { .. } @@ -6247,7 +6233,7 @@ impl<'db> LoweringContext<'db> { // `resolutions` above and returns before reaching here, so the // receiver is always `segments[..len-1]`. if segments.len() >= 2 - && let Some(&recv_root_local) = self.locals.get(&segments[0]) + && let Some(recv_root_local) = self.local_for_path(expr_id, &segments[0]) { let method_name = segments.last().unwrap().clone(); let recv_seg_idx = if segments.len() == 2 { @@ -6278,15 +6264,14 @@ impl<'db> LoweringContext<'db> { return; } } - if self.locals.contains_key(&segments[0]) - || self - .capture_index_for_name_at(expr_id, &segments[0]) - .is_some() + if self + .binding_id_for_path(expr_id, &segments[0]) + .is_some() // BEP-044 wf3 #4: `default.` as a value — the field-chain // lowerer maps the `default` root to `self`-viewed-as-interface. // (The `default.method(...)` call form is intercepted earlier in // `lower_call`, so this only catches the value/field form.) - || self.is_default_receiver_root(segments) + || self.is_default_receiver_root(expr_id, segments) { self.lower_multi_segment_path_as_field_chain(expr_id, segments, dest); return; @@ -6323,6 +6308,15 @@ impl<'db> LoweringContext<'db> { return; } + if let Some(place) = self.place_for_path(expr_id, name) { + self.builder.assign(dest, Rvalue::Use(Operand::Copy(place))); + return; + } + if self.binding_id_for_path(expr_id, name).is_some() { + self.emit_panic_call(&format!("unresolved local: {name}"), expr_id); + return; + } + let span_start = self .source_map .as_ref() @@ -6337,22 +6331,6 @@ impl<'db> LoweringContext<'db> { self.scope_func_name.as_ref(), ); match resolved { - ResolvedName::Local { - name: local_name, .. - } => { - if let Some(&local) = self.locals.get(&local_name) { - self.builder - .assign(dest, Rvalue::Use(Operand::Copy(Place::Local(local)))); - } else if let Some(cap_idx) = self.capture_index_for_name_at(expr_id, &local_name) { - // This variable is captured from an enclosing scope. - // Emit a LoadCapture via Place::Capture. - self.builder - .assign(dest, Rvalue::Use(Operand::Copy(Place::Capture(cap_idx)))); - } else { - let msg = format!("unresolved local: {local_name}"); - self.emit_panic_call(&msg, expr_id); - } - } ResolvedName::Item(def) => { self.lower_item_ref(expr_id, def, dest); } @@ -6363,27 +6341,7 @@ impl<'db> LoweringContext<'db> { Rvalue::Use(Operand::Constant(Constant::Function(item))), ); } - ResolvedName::Unknown if self.tagged_body_param_bindings.contains_key(name) => { - // A tagged-template body-lambda parameter (BEP-049 §10 / M4e.1): - // a MIR-only local that `build_tagged_body_closure` injects. It has - // no HIR binding (the tag can't be resolved during the HIR walk), - // so `resolve_name_at_in_scope` returns `Unknown`. - if let Some(&local) = self.locals.get(name) { - // The reference sits directly in the body closure: a plain local. - self.builder - .assign(dest, Rvalue::Use(Operand::Copy(Place::Local(local)))); - } else { - // Referenced from a nested lambda inside the interpolations — - // the param lives in an enclosing frame. HIR can't list it as a - // capture (no binding), so capture it transitively by its stored - // synthetic BindingId, the same way grandparent locals thread up. - let binding_id = self.tagged_body_param_bindings[name]; - let cap_idx = self.ensure_transitive_capture(binding_id); - self.builder - .assign(dest, Rvalue::Use(Operand::Copy(Place::Capture(cap_idx)))); - } - } - ResolvedName::Unknown => { + ResolvedName::Local { .. } | ResolvedName::Unknown => { if self .tir_expr_type(self.expr_metadata_key(expr_id)) .is_some() @@ -6412,55 +6370,57 @@ impl<'db> LoweringContext<'db> { segments: &[Name], dest: Place, ) { - let (mut current_place, mut current_ty) = - if let Some(&root_local) = self.locals.get(&segments[0]) { - let place = Place::Local(root_local); - let ty = if let Some(tir_root) = self.path_root_ty(expr_id) { - // If TIR inferred a more specific type for the root local, - // update the MIR local's declared type so the emitter can - // resolve field names for display (e.g. `load_field .index`). - if matches!( - self.builder.local_ty(root_local), - RuntimeTy::BuiltinUnknown { .. } - ) && !matches!( - tir_root, - RuntimeTy::BuiltinUnknown { .. } | RuntimeTy::Void { .. } - ) { - self.builder.local_decl_mut(root_local).ty = tir_root.clone(); + let root_place = self.place_for_path(expr_id, &segments[0]); + let (mut current_place, mut current_ty) = if let Some(place) = root_place { + let ty = match place { + Place::Local(root_local) => { + if let Some(tir_root) = self.path_root_ty(expr_id) { + // If TIR inferred a more specific type for the root local, + // update the MIR local's declared type so the emitter can + // resolve field names for display (e.g. `load_field .index`). + if matches!( + self.builder.local_ty(root_local), + RuntimeTy::BuiltinUnknown { .. } + ) && !matches!( + tir_root, + RuntimeTy::BuiltinUnknown { .. } | RuntimeTy::Void { .. } + ) { + self.builder.local_decl_mut(root_local).ty = tir_root.clone(); + } + tir_root + } else { + self.builder.local_ty(root_local) } - tir_root - } else { - self.builder.local_ty(root_local) - }; - (place, ty) - } else if let Some(cap_idx) = self.capture_index_for_name_at(expr_id, &segments[0]) { - let place = Place::Capture(cap_idx); - let ty = self - .path_root_ty(expr_id) - .unwrap_or_else(|| RuntimeTy::BuiltinUnknown { - attr: TyAttr::default(), - }); - (place, ty) - } else if self.is_default_receiver_root(segments) - && let Some(&self_local) = self.locals.get(&Name::new("self")) - { - // BEP-044 wf3 #4: `default.` denotes the enclosing `self` - // viewed at the declaring interface. TIR typed the root as - // `RuntimeTy::Interface`, so reuse that and let the interface-prefix - // routing below resolve the field view (same path as - // `self.as.field`). Without this the `default` root is not a - // local → null → `string + null` VM crash. - let place = Place::Local(self_local); - let ty = self - .path_root_ty(expr_id) - .unwrap_or_else(|| self.builder.local_ty(self_local)); - (place, ty) - } else { - // Root not found as a local or capture — emit null. - self.builder - .assign(dest, Rvalue::Use(Operand::Constant(Constant::Null))); - return; + } + Place::Capture(_) => { + self.path_root_ty(expr_id) + .unwrap_or_else(|| RuntimeTy::BuiltinUnknown { + attr: TyAttr::default(), + }) + } + _ => unreachable!("path roots are locals or captures"), }; + (place, ty) + } else if self.is_default_receiver_root(expr_id, segments) + && let Some(&self_local) = self.locals.get(&Name::new("self")) + { + // BEP-044 wf3 #4: `default.` denotes the enclosing `self` + // viewed at the declaring interface. TIR typed the root as + // `RuntimeTy::Interface`, so reuse that and let the interface-prefix + // routing below resolve the field view (same path as + // `self.as.field`). Without this the `default` root is not a + // local -> null -> `string + null` VM crash. + let place = Place::Local(self_local); + let ty = self + .path_root_ty(expr_id) + .unwrap_or_else(|| self.builder.local_ty(self_local)); + (place, ty) + } else { + // Root not found as a local or capture; emit null. + self.builder + .assign(dest, Rvalue::Use(Operand::Constant(Constant::Null))); + return; + }; let mut skip_next_segment = false; for (offset, seg) in segments[1..].iter().enumerate() { @@ -7595,15 +7555,10 @@ impl<'db> LoweringContext<'db> { // (Can't reuse `lower_path_receiver_to_local`: it assumes a local // root and `expr_ty(callee)` would ICE on the Unknown callee.) let recv_op = if receiver_segments.len() == 1 { - if let Some(&recv_local) = self.locals.get(&receiver_segments[0]) { - Operand::Copy(Place::Local(recv_local)) - } else if let Some(cap_idx) = - self.capture_index_for_name_at(callee, &receiver_segments[0]) - { - Operand::Copy(Place::Capture(cap_idx)) - } else { + let Some(place) = self.place_for_path(callee, &receiver_segments[0]) else { return false; - } + }; + Operand::Copy(place) } else { let recv_ty = self .tir_path_segment_type(( @@ -7737,15 +7692,10 @@ impl<'db> LoweringContext<'db> { AstExpr::Path(segments) => { let receiver_segments = &segments[..segments.len() - 1]; let recv_op = if receiver_segments.len() == 1 { - if let Some(&recv_local) = self.locals.get(&receiver_segments[0]) { - Operand::Copy(Place::Local(recv_local)) - } else if let Some(cap_idx) = - self.capture_index_for_name_at(callee, &receiver_segments[0]) - { - Operand::Copy(Place::Capture(cap_idx)) - } else { + let Some(place) = self.place_for_path(callee, &receiver_segments[0]) else { return false; - } + }; + Operand::Copy(place) } else { let recv_ty = self .tir_path_segment_type(( @@ -7855,14 +7805,12 @@ impl<'db> LoweringContext<'db> { let static_receiver = match &callee_expr { AstExpr::MemberAccess { base, .. } => match &self.body.exprs[*base] { AstExpr::Path(segs) if !segs.is_empty() => { - !self.locals.contains_key(&segs[0]) - && self.capture_index_for_name_at(*base, &segs[0]).is_none() + self.binding_id_for_path(*base, &segs[0]).is_none() } _ => false, }, AstExpr::Path(segs) if segs.len() >= 2 => { - !self.locals.contains_key(&segs[0]) - && self.capture_index_for_name_at(callee, &segs[0]).is_none() + self.binding_id_for_path(callee, &segs[0]).is_none() } _ => false, }; @@ -7996,7 +7944,7 @@ impl<'db> LoweringContext<'db> { // the override is being deliberately bypassed. if let AstExpr::Path(segments) = &callee_expr && segments.len() == 2 - && self.is_default_receiver_root(segments) + && self.is_default_receiver_root(callee, segments) && let Some(target) = self.implements_block_iface_target() && matches!( target.type_refs[target.target].kind, @@ -8175,7 +8123,7 @@ impl<'db> LoweringContext<'db> { // The segment just before the method name may be a real field // access (`r.a.b.c.d.e.speak()`) whose static type is an interface. if segments.len() >= 2 - && let Some(&recv_root_local) = self.locals.get(&segments[0]) + && let Some(recv_root_local) = self.local_for_path(callee, &segments[0]) { let method_name = segments.last().unwrap().clone(); let prefix_idx = segments.len() - 2; @@ -8212,7 +8160,7 @@ impl<'db> LoweringContext<'db> { // so the dispatch below routes it to a virtual call. .or_else(|| { if segments.len() == 2 { - self.binding_id_for_name_at(callee, &segments[0]) + self.binding_id_for_path(callee, &segments[0]) .and_then(|bid| { self.source_param_tir_ty_for_binding(&segments[0], bid) }) @@ -8374,10 +8322,7 @@ impl<'db> LoweringContext<'db> { // TIR types (`Interface`, `Class`) but are not runtime values. let base_is_value = match &self.body.exprs[*base] { AstExpr::Path(segments) if !segments.is_empty() => { - self.locals.contains_key(&segments[0]) - || self - .capture_index_for_name_at(*base, &segments[0]) - .is_some() + self.binding_id_for_path(*base, &segments[0]).is_some() } _ => self .tir_expr_type(self.expr_metadata_key(*base)) @@ -8525,15 +8470,8 @@ impl<'db> LoweringContext<'db> { } else { let receiver_op = if receiver_segments.len() == 1 { // Simple local variable receiver (e.g. `self`). - if let Some(&recv_local) = self.locals.get(&receiver_segments[0]) { - Operand::Copy(Place::Local(recv_local)) - } else if let Some(cap_idx) = - self.capture_index_for_name_at(callee, &receiver_segments[0]) - { - Operand::Copy(Place::Capture(cap_idx)) - } else { - Operand::Constant(Constant::Null) - } + self.place_for_path(callee, &receiver_segments[0]) + .map_or_else(|| Operand::Constant(Constant::Null), Operand::Copy) } else { // Multi-segment receiver (e.g. `user.profile.items`): lower as field chain. let recv_ty = self.expr_ty(callee); // approximation; actual type not critical here @@ -8565,13 +8503,7 @@ impl<'db> LoweringContext<'db> { Some(item) => Operand::Constant(Constant::Function(item)), None => self.lower_to_operand(callee), }; - let first_seg = &segments[0]; - let receiver_op = if let Some(&receiver_local) = self.locals.get(first_seg) { - Some(Operand::Copy(Place::Local(receiver_local))) - } else { - self.capture_index_for_name_at(callee, first_seg) - .map(|cap_idx| Operand::Copy(Place::Capture(cap_idx))) - }; + let receiver_op = self.place_for_path(callee, &segments[0]).map(Operand::Copy); if let Some(receiver_op) = receiver_op { let prefix_idx = segments.len() - 2; receiver_path_tir_ty = self @@ -10511,7 +10443,7 @@ impl<'db> LoweringContext<'db> { let expr = &self.body.exprs[expr_id]; if let AstExpr::Path(segments) = expr { if segments.len() == 1 { - if let Some(&local) = self.locals.get(&segments[0]) { + if let Some(local) = self.local_for_path(expr_id, &segments[0]) { return Some(local); } } @@ -12031,12 +11963,8 @@ impl LoweringContext<'_> { let expr = self.body.exprs[expr_id].clone(); match &expr { AstExpr::Path(segments) if segments.len() == 1 => { - if let Some(&local) = self.locals.get(&segments[0]) { - Place::Local(local) - } else if let Some(cap_idx) = self.capture_index_for_name_at(expr_id, &segments[0]) - { - // Assignment to a captured variable in a closure body. - Place::Capture(cap_idx) + if let Some(place) = self.place_for_path(expr_id, &segments[0]) { + place } else { // Unresolved single-segment assignment target. This is // only reachable for programs TIR already rejected (an @@ -12062,32 +11990,31 @@ impl LoweringContext<'_> { AstExpr::Path(segments) if segments.len() >= 2 => { // Multi-segment path lvalue: `a.b` or `a.b.c`. // Chain field projections from the root local or capture. - let (mut current_place, mut current_ty) = if let Some(&l) = - self.locals.get(&segments[0]) - { - let ty = self - .path_root_ty(expr_id) - .unwrap_or_else(|| self.builder.local_ty(l)); - (Place::Local(l), ty) - } else if let Some(cap_idx) = self.capture_index_for_name_at(expr_id, &segments[0]) - { - let ty = - self.path_root_ty(expr_id) - .unwrap_or_else(|| RuntimeTy::BuiltinUnknown { - attr: TyAttr::default(), - }); - (Place::Capture(cap_idx), ty) - } else { - let tmp = self.builder.temp(RuntimeTy::Null { - attr: TyAttr::default(), - }); - ( - Place::Local(tmp), - RuntimeTy::Null { + let (mut current_place, mut current_ty) = + if let Some(place) = self.place_for_path(expr_id, &segments[0]) { + let ty = match place { + Place::Local(local) => self + .path_root_ty(expr_id) + .unwrap_or_else(|| self.builder.local_ty(local)), + Place::Capture(_) => self.path_root_ty(expr_id).unwrap_or_else(|| { + RuntimeTy::BuiltinUnknown { + attr: TyAttr::default(), + } + }), + _ => unreachable!("path roots are locals or captures"), + }; + (place, ty) + } else { + let tmp = self.builder.temp(RuntimeTy::Null { attr: TyAttr::default(), - }, - ) - }; + }); + ( + Place::Local(tmp), + RuntimeTy::Null { + attr: TyAttr::default(), + }, + ) + }; for (offset, seg) in segments[1..].iter().enumerate() { let seg_idx = offset + 1; @@ -12886,7 +12813,7 @@ impl LoweringContext<'_> { self.builder.set_current_block(bb_body); let saved_locals = self.locals.clone(); - self.bind_pattern(scrutinee, part); + self.bind_pattern_inner(scrutinee, part, arm.pattern, part, false); if let Some(guard) = arm.guard { let guard_op = self.lower_to_operand(guard); let bb_guarded = self.builder.create_block(); diff --git a/baml_language/crates/baml_compiler2_tir/src/builder.rs b/baml_language/crates/baml_compiler2_tir/src/builder.rs index 340cc7cb557..ec768e73d28 100644 --- a/baml_language/crates/baml_compiler2_tir/src/builder.rs +++ b/baml_language/crates/baml_compiler2_tir/src/builder.rs @@ -24,6 +24,7 @@ use baml_compiler2_hir::{ loc::{ClassLoc, FunctionLoc}, package::{PackageId, PackageItems}, scope::{FileScopeId, ScopeId}, + semantic_index::{ExprMetadataKey, ExprMetadataScope, PathResolution}, }; // The trait must be in scope so the builder (which implements it) can call the // defaulted type-algebra methods on itself — `self.is_subtype(a, b)`. @@ -817,6 +818,8 @@ pub struct TypeInferenceBuilder<'db> { /// The scope being analyzed (kept for future use). #[allow(dead_code)] scope: ScopeId<'db>, + /// Arena namespace for the expressions currently being inferred. + expr_metadata_scope: ExprMetadataScope, /// Declared return type for the function (used to check return statements). declared_return_ty: Option, /// Resolved type alias map: alias qualified name → expanded Ty. @@ -1157,6 +1160,7 @@ impl<'db> TypeInferenceBuilder<'db> { let package_items = &res_ctx.own_items; let pkg_info = baml_compiler2_hir::file_package::file_package(db, scope.file(db)); let ns_context = pkg_info.namespace_path; + let expr_metadata_scope = ExprMetadataScope::Body(scope.file_scope_id(db)); Self { context, expressions: FxHashMap::default(), @@ -1171,6 +1175,7 @@ impl<'db> TypeInferenceBuilder<'db> { package_items, package_id, scope, + expr_metadata_scope, declared_return_ty: None, aliases, normalized_overlap_aliases: std::cell::OnceCell::new(), @@ -1452,15 +1457,11 @@ impl<'db> TypeInferenceBuilder<'db> { if !self.locals.contains_key(name) { return false; } - let Some(source_map) = self.body_source_map.as_ref() else { - return false; - }; let db = self.context.db(); let file = self.context.scope().file(db); let index = baml_compiler2_ppir::file_semantic_index(db, file); - let offset = source_map.expr_span(subject).start(); - let use_scope = index.scope_at_offset(offset, None); - let Some(binding_id) = index.visible_binding_at(use_scope, offset, name) else { + let key = ExprMetadataKey::new(self.expr_metadata_scope, subject); + let Some(PathResolution::Local(binding_id)) = index.path_resolution(key) else { return false; }; index @@ -2884,6 +2885,11 @@ impl<'db> TypeInferenceBuilder<'db> { let saved_call_type_instantiations = std::mem::take(&mut self.call_type_instantiations); let saved_function_coercions = std::mem::take(&mut self.function_coercions); let saved_lambda_effective_throws = std::mem::take(&mut self.lambda_effective_throws); + let saved_expr_metadata_scope = self.expr_metadata_scope; + let owner_scope = match saved_expr_metadata_scope { + ExprMetadataScope::Body(scope) | ExprMetadataScope::ParameterDefault(scope) => scope, + }; + self.expr_metadata_scope = ExprMetadataScope::ParameterDefault(owner_scope); let defaults = ¶meter_defaults.defaults; let saved_body_source_map = self.body_source_map.replace(defaults.source_map.clone()); let saved_locals = self.locals.clone(); @@ -3006,6 +3012,7 @@ impl<'db> TypeInferenceBuilder<'db> { self.call_type_instantiations = saved_call_type_instantiations; self.function_coercions = saved_function_coercions; self.lambda_effective_throws = saved_lambda_effective_throws; + self.expr_metadata_scope = saved_expr_metadata_scope; // Resolve all default-checking diagnostics against the defaults source // map before restoring the function body's map (otherwise their // defaults-arena `ExprId`s render at the wrong offsets). @@ -14919,6 +14926,7 @@ impl<'db> TypeInferenceBuilder<'db> { let saved_call_plans = std::mem::take(&mut self.call_plans); let saved_call_type_instantiations = std::mem::take(&mut self.call_type_instantiations); let saved_function_coercions = std::mem::take(&mut self.function_coercions); + let saved_expr_metadata_scope = self.expr_metadata_scope; // BEP-042: a lambda body is its own control-flow region — `return` // targets the lambda, and the parent's defer/loop nesting must not leak // in. Reset the counters for the body and restore them afterwards. @@ -14970,7 +14978,7 @@ impl<'db> TypeInferenceBuilder<'db> { // // Also captures the lambda's `FileScopeId` for use as a position-independent // key in `nested_lambda_types` (avoids TextRange in Salsa-cached output). - let lambda_file_scope_id = { + let (lambda_file_scope_id, lambda_metadata_scope) = { let db = self.context.db(); let file = self.context.scope().file(db); let index = baml_compiler2_ppir::file_semantic_index(db, file); @@ -15013,8 +15021,11 @@ impl<'db> TypeInferenceBuilder<'db> { self.seed_capture(capture_name, ty); } } - key_fsi + (key_fsi, seed_fsi) }; + if let Some(scope) = lambda_metadata_scope { + self.expr_metadata_scope = ExprMetadataScope::Body(scope); + } // Set return type context for return statement checking inside lambda if let Some(ret) = expected_ret { @@ -15161,6 +15172,7 @@ impl<'db> TypeInferenceBuilder<'db> { self.implements_block_interface = saved_implements_block_interface; self.loop_depth = saved_loop_depth; self.defer_loop_floors = saved_defer_loop_floors; + self.expr_metadata_scope = saved_expr_metadata_scope; // Freeze this lambda's diagnostic spans against its own source map // before restoring the enclosing map (otherwise they render at `0..0`). self.context diff --git a/baml_language/crates/baml_lsp2_actions/src/annotations.rs b/baml_language/crates/baml_lsp2_actions/src/annotations.rs index c8246c919bb..1ed7261491c 100644 --- a/baml_language/crates/baml_lsp2_actions/src/annotations.rs +++ b/baml_language/crates/baml_lsp2_actions/src/annotations.rs @@ -311,7 +311,7 @@ fn process_body( Expr::Lambda(func_def) => { if let Some(FunctionBodyDef::Expr(lbody, lsmap)) = &func_def.body { let lambda_scope = index - .lambda_scope_for(source_map.expr_span(expr_id)) + .lambda_scope_for_within(owner_scope, source_map.expr_span(expr_id)) .unwrap_or(owner_scope); process_body(db, file, index, lambda_scope, lbody, lsmap, out); } diff --git a/baml_language/crates/baml_lsp2_actions/src/usages.rs b/baml_language/crates/baml_lsp2_actions/src/usages.rs index 00b2d3b9d1f..82182b9210b 100644 --- a/baml_language/crates/baml_lsp2_actions/src/usages.rs +++ b/baml_language/crates/baml_lsp2_actions/src/usages.rs @@ -8,9 +8,9 @@ //! CSTs for `WORD` tokens that match the target name, confirm each via //! `resolve_name_at`, and collect as `Location`s. //! -//! - **Locals** (let bindings, parameters): search only within the enclosing -//! function's `ExprBody` for `Expr::Path` nodes that use the same name and -//! resolve to the same local binding. +//! - **Locals** (let bindings, parameters): search the enclosing function and +//! nested lambda arenas for `Expr::Path` nodes that resolve to the same local +//! binding. //! //! ## Optimization //! @@ -21,7 +21,13 @@ use baml_base::{Name, SourceFile}; use baml_compiler_syntax::SyntaxKind; use baml_compiler2_ast::{Expr, ExprBody}; -use baml_compiler2_hir::{body::FunctionBody, scope::ScopeKind, semantic_index::BindingId}; +use baml_compiler2_hir::{ + body::FunctionBody, + scope::{FileScopeId, ScopeKind}, + semantic_index::{ + BindingId, ExprMetadataKey, ExprMetadataScope, FileSemanticIndex, PathResolution, + }, +}; use baml_compiler2_tir::resolve::{ResolvedName, resolve_name_at}; use rowan::NodeOrToken; use text_size::{TextRange, TextSize}; @@ -150,12 +156,11 @@ fn same_item_definition(a: &ResolvedName<'_>, b: &ResolvedName<'_>) -> bool { // ── local usages ────────────────────────────────────────────────────────────── -/// Search for references to a local variable within the enclosing function's -/// expression body. +/// Search for references to a local variable within the enclosing function. /// /// We walk the `ExprBody` (span-free) and use the source map for positions. -/// For each `Expr::Path([name])` that resolves to the same local, we emit a -/// `Location` using the expression's span from the source map. +/// For each path whose root resolves to the same local, we emit a `Location` +/// using the root segment's span from the source map. fn find_local_usages( db: &dyn Db, file: SourceFile, @@ -200,59 +205,95 @@ fn find_local_usages( return Vec::new(); }; - let name = Name::new(name_text); - let mut results = Vec::new(); - - collect_local_path_usages( - db, + let mut collector = LocalUsageCollector { file, - expr_body, - &name, + index, + name: Name::new(name_text), target_binding, + results: Vec::new(), + }; + collector.collect( + enclosing_func_scope, + expr_body, + ExprMetadataScope::Body(enclosing_func_scope), &source_map, - &mut results, ); - results + let defaults = baml_compiler2_ppir::function_parameter_defaults(db, func_loc); + collector.collect( + enclosing_func_scope, + &defaults.defaults.exprs, + ExprMetadataScope::ParameterDefault(enclosing_func_scope), + &defaults.defaults.source_map, + ); + + collector.results } -/// Walk an `ExprBody` and collect `Expr::Path([name])` occurrences that -/// resolve to the same local as `target_resolved`. -fn collect_local_path_usages( - db: &dyn Db, +struct LocalUsageCollector<'index, 'db> { file: SourceFile, - expr_body: &ExprBody, - name: &Name, + index: &'index FileSemanticIndex<'db>, + name: Name, target_binding: BindingId, - source_map: &baml_compiler2_ast::AstSourceMap, - results: &mut Vec, -) { - let index = baml_compiler2_hir::file_semantic_index(db, file); - - for (expr_id, expr) in expr_body.exprs.iter() { - let Expr::Path(segments) = expr else { - continue; - }; + results: Vec, +} - // Only single-segment paths can refer to locals. - if segments.len() != 1 || &segments[0] != name { - continue; - } +impl LocalUsageCollector<'_, '_> { + /// Walk one expression arena and recurse into nested lambda arenas. + fn collect( + &mut self, + owner_scope: FileScopeId, + expr_body: &ExprBody, + metadata_scope: ExprMetadataScope, + source_map: &baml_compiler2_ast::AstSourceMap, + ) { + for (expr_id, expr) in expr_body.exprs.iter() { + match expr { + Expr::Path(segments) if segments.first() == Some(&self.name) => { + let segment_range = source_map.path_segment_span(expr_id, 0); + let range = + TextRange::at(segment_range.start(), TextSize::of(segments[0].as_str())); + if range.is_empty() { + continue; + } - // Get the span of this expression from the source map. - let range = source_map.expr_span(expr_id); - if range.is_empty() { - continue; - } + let key = ExprMetadataKey::new(metadata_scope, expr_id); + if self.index.path_resolution(key) + == Some(PathResolution::Local(self.target_binding)) + { + self.results.push(Location { + file: self.file, + range, + }); + } + } + Expr::Lambda(func_def) => { + let span = source_map.expr_span(expr_id); + let Some(lambda_scope) = self.index.lambda_scope_for_within(owner_scope, span) + else { + continue; + }; - // Confirm that this usage resolves to the exact same visible binding. - let use_offset = range.start(); - let Some(use_scope) = index.expression_scope(expr_id) else { - continue; - }; + self.collect( + lambda_scope, + &func_def.defaults.exprs, + ExprMetadataScope::ParameterDefault(lambda_scope), + &func_def.defaults.source_map, + ); - if index.visible_binding_at(use_scope, use_offset, name) == Some(target_binding) { - results.push(Location { file, range }); + if let Some(baml_compiler2_ast::FunctionBodyDef::Expr(body, body_source_map)) = + &func_def.body + { + self.collect( + lambda_scope, + body, + ExprMetadataScope::Body(lambda_scope), + body_source_map, + ); + } + } + _ => {} + } } } } diff --git a/baml_language/crates/baml_lsp2_actions/src/usages_at_tests.rs b/baml_language/crates/baml_lsp2_actions/src/usages_at_tests.rs index 2cc7b7e8a76..8f015aeefac 100644 --- a/baml_language/crates/baml_lsp2_actions/src/usages_at_tests.rs +++ b/baml_language/crates/baml_lsp2_actions/src/usages_at_tests.rs @@ -54,6 +54,33 @@ function Process(<[CURSOR]input: string) -> string { ); } + #[test] + fn test_find_refs_lambda_parameter_in_nested_lambda() { + let test = CursorTest::new( + r#" +function Test() -> int { + let f = (<[CURSOR]x: int) -> int { + let y = x + let g = () -> int { x + y } + x + g() + } + f(1) +} +"#, + ); + + let usages = test.find_all_usages(); + assert_eq!( + usages.len(), + 3, + "Should find lambda-parameter usages across nested lambda arenas, found: {:?}", + usages + .iter() + .map(|l| test.format_location_with_name(l)) + .collect::>() + ); + } + #[test] fn test_find_refs_function() { let test = CursorTest::new( diff --git a/baml_language/crates/baml_tests/baml_src/ns_lexical_scoping/lexical_scoping.baml b/baml_language/crates/baml_tests/baml_src/ns_lexical_scoping/lexical_scoping.baml index e5f99b07d1b..db816392942 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_lexical_scoping/lexical_scoping.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_lexical_scoping/lexical_scoping.baml @@ -295,3 +295,27 @@ function multi_clause_catch_main() -> int { test "multi_clause_catch_uses_clause_local_binding" { assert.is_true(baml.deep_equals(multi_clause_catch_main(), 7)) } + +function closure_parameter_in_for_iterable() -> int { + let sum = 0 + for (let n in [1, 2, 3, 4].filter((x: int) -> bool { x > 2 })) { + sum += n + } + sum +} + +test "closure_parameter_resolves_in_for_iterable" { + assert.equal(closure_parameter_in_for_iterable(), 7) +} + +function closure_parameter_in_while_condition() -> int { + let i = 0 + while (((x: int) -> bool { x < 3 })(i)) { + i += 1 + } + i +} + +test "closure_parameter_resolves_in_while_condition" { + assert.equal(closure_parameter_in_while_condition(), 3) +} diff --git a/baml_language/crates/baml_tests/baml_src/ns_prompt_tag_runtime/prompt_tag_runtime.baml b/baml_language/crates/baml_tests/baml_src/ns_prompt_tag_runtime/prompt_tag_runtime.baml index d10fb78b559..e55ca118710 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_prompt_tag_runtime/prompt_tag_runtime.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_prompt_tag_runtime/prompt_tag_runtime.baml @@ -147,6 +147,21 @@ test "prompt_interpolates_ctx_output_format" { assert.contains(pt_ctx_output_format_fn(), "age") } +function PtRenderPerson() -> Person { + client "openai/gpt-4o-mini" + prompt `Make a person.${ctx.output_format}` +} + +function pt_render_prompt_companion_output_format_fn() -> string { + PtRenderPerson$render_prompt().text() +} + +test "render_prompt_companion_interpolates_ctx_output_format" { + assert.contains(pt_render_prompt_companion_output_format_fn(), "Make a person.") + assert.contains(pt_render_prompt_companion_output_format_fn(), "name") + assert.contains(pt_render_prompt_companion_output_format_fn(), "age") +} + // BEP-049 M5b.2: `${ctx.output_format_with(prefix=..., ...)}` re-renders the // return type's schema with caller options. `Context._output_format` carries // the prebuilt schema handle; a non-default `prefix` must appear in the diff --git a/baml_language/crates/baml_tests/snapshots/baml_src/lexical_scoping.snap b/baml_language/crates/baml_tests/snapshots/baml_src/lexical_scoping.snap index 743bdb42f02..f512fbab88a 100644 --- a/baml_language/crates/baml_tests/snapshots/baml_src/lexical_scoping.snap +++ b/baml_language/crates/baml_tests/snapshots/baml_src/lexical_scoping.snap @@ -121,6 +121,20 @@ function lexical_scoping.$init_test_ns_lexical_scoping_lexical_scoping(registry: load_const null call testing.TestCollector.register_test_at pop 1 + load_var registry + load_const "root.lexical_scoping" + load_const "closure_parameter_resolves_in_for_iterable" + make_closure ., 0 + load_const null + call testing.TestCollector.register_test_at + pop 1 + load_var registry + load_const "root.lexical_scoping" + load_const "closure_parameter_resolves_in_while_condition" + make_closure ., 0 + load_const null + call testing.TestCollector.register_test_at + pop 1 load_const null return } @@ -257,6 +271,69 @@ function lexical_scoping.catch_base_uses_outer_lambda() -> int { return } +function lexical_scoping.closure_parameter_in_for_iterable() -> int { + load_const 0 + store_var sum + load_type int + load_type never + load_const 1 + load_const 2 + load_const 3 + load_const 4 + load_type int + alloc_array 4 + make_closure ., 0 + call baml.Array.filter + load_type baml.iter.Iterable + load_const "iter" + virtual_call nargs=1 ntypeargs=0 + store_var _7 + + L0: + load_var _7 + load_type baml.iter.Iterator + load_const "next" + virtual_call nargs=1 ntypeargs=0 + store_var _8 + load_var _8 + is_type baml.iter.Done + pop_jump_if_false L1 + jump L2 + + L1: + load_var2 1 3 + add_int + store_var sum + jump L0 + + L2: + load_var sum + return +} + +function lexical_scoping.closure_parameter_in_while_condition() -> int { + load_const 0 + store_var i + + L0: + load_var i + make_closure ., 0 + call_indirect + pop_jump_if_false L1 + jump L2 + + L1: + load_var i + return + + L2: + load_var i + load_const 1 + add_int + store_var i + jump L0 +} + function lexical_scoping.declared_type_restored_outer() -> string { load_const "outer" return diff --git a/baml_language/crates/baml_tests/snapshots/baml_src/prompt_tag_runtime.snap b/baml_language/crates/baml_tests/snapshots/baml_src/prompt_tag_runtime.snap index 4b2c1afdd92..acd06f27ea3 100644 --- a/baml_language/crates/baml_tests/snapshots/baml_src/prompt_tag_runtime.snap +++ b/baml_language/crates/baml_tests/snapshots/baml_src/prompt_tag_runtime.snap @@ -46,25 +46,32 @@ function prompt_tag_runtime.$init_test_ns_prompt_tag_runtime_prompt_tag_runtime( pop 1 load_var registry load_const "root.prompt_tag_runtime" - load_const "prompt_interpolates_ctx_output_format_with" + load_const "render_prompt_companion_interpolates_ctx_output_format" make_closure ., 0 load_const null call testing.TestCollector.register_test_at pop 1 load_var registry load_const "root.prompt_tag_runtime" - load_const "prompt_output_format_with_omits_leading_optional_arg" + load_const "prompt_interpolates_ctx_output_format_with" make_closure ., 0 load_const null call testing.TestCollector.register_test_at pop 1 load_var registry load_const "root.prompt_tag_runtime" - load_const "prompt_output_format_with_can_render_null_as_omit" + load_const "prompt_output_format_with_omits_leading_optional_arg" make_closure ., 0 load_const null call testing.TestCollector.register_test_at pop 1 + load_var registry + load_const "root.prompt_tag_runtime" + load_const "prompt_output_format_with_can_render_null_as_omit" + make_closure ., 0 + load_const null + call testing.TestCollector.register_test_at + pop 1 load_const null return } @@ -79,6 +86,200 @@ function prompt_tag_runtime.Labeled.baml.ToString.to_string(self: prompt_tag_run return } +function prompt_tag_runtime.PtRenderPerson(client: baml.llm.Client) -> prompt_tag_runtime.Person { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + make_closure ., 0 + call baml.llm.prompt + store_var _5 + load_type prompt_tag_runtime.Person + load_var client + load_const "PtRenderPerson" + load_type string + load_type unknown + alloc_map 0 + load_var _5 + call baml.llm.call_llm_function + return +} + +function prompt_tag_runtime.PtRenderPerson$build_request(client: baml.llm.Client) -> baml.http.Request { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + make_closure ., 0 + call baml.llm.prompt + store_var _5 + load_var client + load_const "PtRenderPerson" + load_type string + load_type unknown + alloc_map 0 + load_var _5 + call baml.llm.build_request + return +} + +function prompt_tag_runtime.PtRenderPerson$build_request_stream(client: baml.llm.Client) -> baml.http.Request { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + make_closure ., 0 + call baml.llm.prompt + store_var _5 + load_var client + load_const "PtRenderPerson" + load_type string + load_type unknown + alloc_map 0 + load_var _5 + call baml.llm.build_request_stream + return +} + +function prompt_tag_runtime.PtRenderPerson$parse(json: string, client: baml.llm.Client) -> prompt_tag_runtime.Person { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + load_type prompt_tag_runtime.Person$stream | null + load_type prompt_tag_runtime.Person + load_var json + call baml.llm.parse + return +} + +function prompt_tag_runtime.PtRenderPerson$parse_stream(sse: baml.http.SseStream, client: baml.llm.Client) -> baml.llm.Stream { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + load_type prompt_tag_runtime.Person$stream | null + load_type prompt_tag_runtime.Person + load_var2 2 1 + call baml.llm.Client.__make_stream + return +} + +function prompt_tag_runtime.PtRenderPerson$render_prompt(client: baml.llm.Client) -> baml.llm.PromptAst { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + make_closure ., 0 + call baml.llm.prompt + store_var _5 + load_var client + load_const "PtRenderPerson" + load_type string + load_type unknown + alloc_map 0 + load_var _5 + call baml.llm.render_prompt + return +} + +function prompt_tag_runtime.PtRenderPerson$stream(client: baml.llm.Client) -> baml.llm.Stream { + load_var client + load_const + cmp_op == + pop_jump_if_false L0 + load_const "openai/gpt-4o-mini" + load_const baml.llm.ClientType.Primitive + alloc_variant baml.llm.ClientType + load_type baml.llm.Client + alloc_array 0 + load_const null + load_const 0 + init_instance baml.llm.Client .name, .client_type, .sub_clients, .retry, .counter + store_var client + + L0: + make_closure ., 0 + call baml.llm.prompt + store_var _5 + load_type prompt_tag_runtime.Person$stream | null + load_type prompt_tag_runtime.Person + load_var client + load_const "PtRenderPerson" + load_type string + load_type unknown + alloc_map 0 + load_var _5 + call baml.llm.stream_llm_function + return +} + function prompt_tag_runtime.pt_class_array_interp_fn() -> string { load_var ?1 make_cell @@ -279,6 +480,13 @@ function prompt_tag_runtime.pt_output_format_with_prefix_fn() -> string { return } +function prompt_tag_runtime.pt_render_prompt_companion_output_format_fn() -> string { + load_const + call user.prompt_tag_runtime.PtRenderPerson$render_prompt + sys_op baml.llm.PromptAst.text + return +} + function prompt_tag_runtime.pt_role_construction_fn() -> baml.llm.Role { load_const "system" load_type string