Fix LSP navigation for class fields, enum variants, and methods (#3298)

[Artifacts](https://cloud.codelayer.cloud/tasks/019d31f9-344d-731b-bc97-b0610fd44836/artifacts)
|
[Task](https://cloud.codelayer.cloud/deep/tasks/019d31f9-344d-731b-bc97-b0610fd44836)

## Problem

"Go to Definition" and "Find All References" didn't work for class
fields, enum variants, constructor fields, or methods. Clicking
`p.name`, `Status.Active`, `Person { name: ... }`, or `s.Celebrate()`
did nothing. This made navigating BAML code painful — you had to
manually search for field definitions instead of just ctrl-clicking.

Also: when you misspelled a field like `s.feelin`, the red squiggly
underlined the whole `s.feelin` expression instead of just `feelin`.

## What changed for users

- **Goto-definition** works for field access, enum variants, constructor
fields, and method calls
- **Find All References** on a class field definition finds all usage
sites across the package
- **Document outline** shows real spans for fields/variants (not
zero-width placeholders)
- **Diagnostic squigglies** for unresolved members now highlight just
the member name
- **LSP no longer crashes** on certain edge cases in MIR lowering
(panics replaced with graceful error emission)

## How it works

### The core idea

The type inference layer (TIR) already resolves every `x.foo` during
type checking — it knows whether `foo` is a field, variant, method, or
free function. But that information was thrown away. We now store it as
`MemberResolution` entries and use them in the LSP to navigate.

### Preserving span data (AST + HIR)

Field and variant name spans were dropped during HIR lowering. Two new
source maps fix this:

- **`ItemTreeSourceMap`** (`item_tree.rs`, `builder.rs`,
`semantic_index.rs`, `lib.rs`) — stores class field spans, enum variant
spans, and function name spans. Follows the existing source-map pattern
so it doesn't pollute Salsa cache comparisons.
- **`AstSourceMap::field_access_member_spans`** (`ast.rs`,
`lower_expr_body.rs`, `lower_cst.rs`) — stores the text range of just
the member name token (after the dot) for each `FieldAccess` expression.
Used for narrowing diagnostic spans.

### MemberResolution (TIR)

Renamed `MethodResolution` → `MemberResolution` and added `Field` and
`Variant` variants alongside the existing `Method` and `Free`
(`inference.rs`). The TIR builder (`builder.rs`) now inserts these
during field lookup, variant validation, and method resolution. Added
`resolve_class_loc`/`resolve_enum_loc` helpers.

### Narrower diagnostic spans (TIR)

Added `DiagnosticLocation::ExprMember` (`infer_context.rs`) — at render
time it looks up the member-name span instead of the full expression
span. All 9 `UnresolvedMember` diagnostic sites in `builder.rs` now use
this, so `s.feelin` only underlines `feelin`.

### MIR safety

The new `MemberResolution::Field`/`Variant` variants flow into MIR
lowering. Instead of crashing on unexpected variants (`unreachable!`,
`panic!`), the MIR now emits runtime panic instructions and continues
(`lower.rs`). This keeps the LSP alive.

### LSP goto-definition (`definition.rs`)

When `resolve_name_at` returns `Unknown` (which happens for all dotted
access), we now fall back to `resolve_member_at`, which:
- If cursor is inside a class/enum scope: navigates to the field/variant
definition directly
- If cursor is inside a function: finds the enclosing function, gets its
`ScopeInference`, and searches for the expression at the cursor offset.
Then dispatches based on `MemberResolution` type (field → field def
span, variant → variant def span, method → function name span) or checks
for constructor field keys.

### LSP find-references (`usages.rs`)

When on a class field definition, `find_field_definition_usages` scans
all files in the package. For each function body, it checks
`ScopeInference` for `MemberResolution::Field` matches and also scans
`Object` constructor field keys. Uses a text pre-filter (field name
appears in source) to skip files that can't possibly match.

### LSP outline (`outline.rs`)

Replaced zero-width placeholder spans with real field/variant spans from
`ItemTreeSourceMap`.

### Tests

- Flipped 5 goto-def tests and 1 find-refs test from negative to
positive assertions (`definition_at_tests.rs`, `usages_at_tests.rs`)
- Added 3 new diagnostic span tests: `unresolved_field`,
`unresolved_field_chained_access`,
`unresolved_field_span_should_narrow_to_member` (`inference.rs`)
- Updated 4 inline snapshot byte offsets for narrower spans
(`phase3a.rs`)
- Updated 4 file snapshots (`control_flow` and `format_checks` TIR +
diagnostics)

### Deleted old LSP v1 crates

- **`crates/baml_lsp_actions/`** — 15 source files, fully superseded by
`baml_lsp2_actions`
- **`crates/baml_lsp_actions_tests/`** — 7 source files + 281 test
fixtures
- Removed from `Cargo.toml` workspace members, updated `Cargo.lock`

## Deviations from plan

All 5 planned phases were implemented. Notable additions beyond the
plan:

- **Diagnostic span narrowing**
(`AstSourceMap::field_access_member_spans` +
`DiagnosticLocation::ExprMember`) was not in the plan — discovered the
wide-span issue during implementation
- **`function_name_spans`** in `ItemTreeSourceMap` was needed because
methods aren't top-level symbols, so the planned
`utils::definition_span` approach didn't work
- **MIR panic safety** — plan used `unreachable!` for the new variants,
we made it graceful instead since it runs in the LSP
- **`resolve_enum_loc`** needed more logic than planned for
cross-package namespace lookups

## How to verify

### Manual
- [ ] Ctrl-click `p.name` → jumps to field definition
- [ ] Ctrl-click `Status.Active` → jumps to variant definition
- [ ] Ctrl-click `name` in `Person { name: "..." }` → jumps to field
definition
- [ ] Ctrl-click `s.Celebrate()` → jumps to method definition
- [ ] Right-click field definition → Find All References → finds all
usages
- [ ] Type `s.feelin` → squiggly only on `feelin`

### Automated
```bash
cargo test -p baml_lsp2_actions  # 26 tests
cargo test -p baml_tests --lib   # 889 tests
```

## Changelog

Fix LSP goto-definition and find-references for class fields, enum
variants, constructor fields, and methods. Narrow unresolved-member
diagnostic spans to just the member name. Remove panics from MIR
lowering for LSP stability. Delete old LSP v1 crates.

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: hellovai

baml_language/Cargo.lock

@@ -49,7 +49,6 @@ sys_types = { path = "crates/sys_types" }
 sys_native = { path = "crates/sys_native" }
 baml_db = { path = "crates/baml_db" }
 baml_fmt = { path = "crates/baml_fmt" }
-baml_lsp_actions = { path = "crates/baml_lsp_actions" }
 baml_lsp2_actions = { path = "crates/baml_lsp2_actions" }
 baml_lsp_server = { path = "crates/baml_lsp_server" }
 baml_lsp_types = { path = "crates/baml_lsp_types" }

baml_language/Cargo.toml

@@ -6,6 +6,8 @@
 //! CST `Option` handling in one layer so everything downstream gets clean
 //! typed data and can be constructed directly in tests without parsing.
 
+use std::collections::HashMap;
+
 use baml_base::Name;
 use la_arena::{Arena, Idx};
 use text_size::TextRange;
@@ -225,6 +227,8 @@ pub struct AstSourceMap {
     pub match_arm_spans: Arena<TextRange>,
     pub type_annotation_spans: Arena<TextRange>,
     pub catch_arm_spans: Arena<TextRange>,
+    /// For `FieldAccess` expressions, the span of just the member name (after the dot).
+    pub field_access_member_spans: HashMap<ExprId, TextRange>,
 }
 
 impl AstSourceMap {
@@ -236,6 +240,7 @@ impl AstSourceMap {
             match_arm_spans: Arena::new(),
             type_annotation_spans: Arena::new(),
             catch_arm_spans: Arena::new(),
+            field_access_member_spans: HashMap::new(),
         }
     }
 
@@ -262,6 +267,15 @@ impl AstSourceMap {
             .unwrap_or_default()
     }
 
+    /// Look up the member-name span for a `FieldAccess` expression.
+    /// Returns the full expression span as fallback if no member span was recorded.
+    pub fn field_access_member_span(&self, id: ExprId) -> TextRange {
+        self.field_access_member_spans
+            .get(&id)
+            .copied()
+            .unwrap_or_else(|| self.expr_span(id))
+    }
+
     /// Look up the source span of a pattern by its `PatId`.
     pub fn pattern_span(&self, id: PatId) -> TextRange {
         let raw: u32 = id.into_raw().into_u32();

baml_language/crates/baml_compiler2_ast/src/ast.rs

@@ -348,6 +348,7 @@ pub(crate) fn synthesize_llm_builtin_call(
         match_arm_spans: Arena::new(),
         type_annotation_spans: Arena::new(),
         catch_arm_spans: Arena::new(),
+        field_access_member_spans: std::collections::HashMap::new(),
     };
 
     (body, source_map)
@@ -617,6 +618,7 @@ fn synthesize_retry_policy_let(node: &SyntaxNode) -> Option<Item> {
         match_arm_spans: la_arena::Arena::new(),
         type_annotation_spans: la_arena::Arena::new(),
         catch_arm_spans: la_arena::Arena::new(),
+        field_access_member_spans: std::collections::HashMap::new(),
     };
 
     Some(Item::Let(LetDef {
@@ -831,6 +833,7 @@ fn synthesize_client_let(
         match_arm_spans: la_arena::Arena::new(),
         type_annotation_spans: la_arena::Arena::new(),
         catch_arm_spans: la_arena::Arena::new(),
+        field_access_member_spans: std::collections::HashMap::new(),
     };
 
     Item::Let(LetDef {
@@ -1046,6 +1049,7 @@ fn synthesize_client_new_companion(
         match_arm_spans: la_arena::Arena::new(),
         type_annotation_spans: la_arena::Arena::new(),
         catch_arm_spans: la_arena::Arena::new(),
+        field_access_member_spans: std::collections::HashMap::new(),
     };
 
     let func_name = format!("{client_name}$new");

baml_language/crates/baml_compiler2_ast/src/lower_cst.rs

@@ -1319,12 +1319,12 @@ impl LoweringContext {
 
     fn lower_path_expr(&mut self, node: &SyntaxNode) -> ExprId {
         // PATH_EXPR contains WORD (or keyword-as-ident) tokens joined by DOTs.
-        let mut segments = Vec::new();
+        let mut segments: Vec<(Name, TextRange)> = Vec::new();
 
         for elem in node.children_with_tokens() {
             if let rowan::NodeOrToken::Token(token) = elem {
                 if is_ident_token(token.kind()) {
-                    segments.push(Name::new(token.text()));
+                    segments.push((Name::new(token.text()), token.text_range()));
                 }
             }
         }
@@ -1335,7 +1335,7 @@ impl LoweringContext {
 
         // Check if single segment is a literal keyword
         if segments.len() == 1 {
-            match segments[0].as_str() {
+            match segments[0].0.as_str() {
                 "true" => {
                     return self.alloc_expr(Expr::Literal(Literal::Bool(true)), node.text_range());
                 }
@@ -1351,22 +1351,27 @@ impl LoweringContext {
         //   Color.Red  → FieldAccess { base: Path(["Color"]), field: "Red" }
         //   a.b.c      → FieldAccess { base: FieldAccess { base: Path(["a"]), field: "b" }, field: "c" }
         // After this, Path is always single-segment (a bare identifier).
-        let mut base = self.alloc_expr(Expr::Path(vec![segments[0].clone()]), node.text_range());
-        for seg in &segments[1..] {
-            base = self.alloc_expr(
+        let mut base = self.alloc_expr(Expr::Path(vec![segments[0].0.clone()]), node.text_range());
+        for (seg, seg_range) in &segments[1..] {
+            let id = self.alloc_expr(
                 Expr::FieldAccess {
                     base,
                     field: seg.clone(),
                 },
                 node.text_range(),
             );
+            self.source_map
+                .field_access_member_spans
+                .insert(id, *seg_range);
+            base = id;
         }
         base
     }
 
     fn lower_field_access_expr(&mut self, node: &SyntaxNode) -> ExprId {
         let mut base = None;
         let mut field = None;
+        let mut field_range = None;
 
         for elem in node.children_with_tokens() {
             match elem {
@@ -1378,6 +1383,7 @@ impl LoweringContext {
                 rowan::NodeOrToken::Token(token) => {
                     if is_ident_token(token.kind()) && base.is_some() {
                         field = Some(Name::new(token.text()));
+                        field_range = Some(token.text_range());
                     }
                 }
             }
@@ -1386,7 +1392,11 @@ impl LoweringContext {
         let base = base.unwrap_or_else(|| self.alloc_expr(Expr::Missing, node.text_range()));
         let field = field.unwrap_or_else(|| Name::new("_"));
 
-        self.alloc_expr(Expr::FieldAccess { base, field }, node.text_range())
+        let id = self.alloc_expr(Expr::FieldAccess { base, field }, node.text_range());
+        if let Some(range) = field_range {
+            self.source_map.field_access_member_spans.insert(id, range);
+        }
+        id
     }
 
     fn lower_env_access_expr(&mut self, node: &SyntaxNode) -> ExprId {

baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs

@@ -44,6 +44,7 @@ pub struct SemanticIndexBuilder<'db> {
     expr_scopes: Vec<(ast::ExprId, FileScopeId)>,
 
     item_tree: ItemTree,
+    item_tree_source_map: crate::item_tree::ItemTreeSourceMap,
     type_contributions: Vec<(Name, Contribution<'db>)>,
     value_contributions: Vec<(Name, Contribution<'db>)>,
     diagnostics: Vec<Hir2Diagnostic>,
@@ -60,6 +61,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             class_depth: 0,
             expr_scopes: Vec::new(),
             item_tree: ItemTree::new(),
+            item_tree_source_map: crate::item_tree::ItemTreeSourceMap::default(),
             type_contributions: Vec::new(),
             value_contributions: Vec::new(),
             diagnostics: Vec::new(),
@@ -129,6 +131,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             scope_bindings: self.scope_bindings,
             scope_ids,
             item_tree: Arc::new(self.item_tree),
+            item_tree_source_map: Arc::new(self.item_tree_source_map),
             symbol_contributions: Arc::new(FileSymbolContributions {
                 types: self.type_contributions,
                 values: self.value_contributions,
@@ -349,6 +352,7 @@ impl<'db> SemanticIndexBuilder<'db> {
 
     fn lower_function(&mut self, f: &ast::FunctionDef) -> LocalItemId<FunctionMarker> {
         let local_id = self.item_tree.alloc_function(f);
+        ItemTree::collect_function_span(&mut self.item_tree_source_map, local_id, f);
         let loc = FunctionLoc::new(self.db, self.file, local_id);
 
         // Only contribute as a top-level symbol if not inside a class.
@@ -382,6 +386,7 @@ impl<'db> SemanticIndexBuilder<'db> {
 
     fn lower_class(&mut self, c: &ast::ClassDef) {
         let local_id = self.item_tree.alloc_class(c);
+        ItemTree::collect_class_spans(&mut self.item_tree_source_map, local_id, c);
         let loc = ClassLoc::new(self.db, self.file, local_id);
         self.type_contributions.push((
             c.name.clone(),
@@ -428,6 +433,7 @@ impl<'db> SemanticIndexBuilder<'db> {
 
     fn lower_enum(&mut self, e: &ast::EnumDef) {
         let local_id = self.item_tree.alloc_enum(e);
+        ItemTree::collect_enum_spans(&mut self.item_tree_source_map, local_id, e);
         let loc = EnumLoc::new(self.db, self.file, local_id);
         self.type_contributions.push((
             e.name.clone(),

baml_language/crates/baml_compiler2_hir/src/builder.rs

@@ -178,6 +178,24 @@ pub struct Let {
     pub name_span: TextRange,
 }
 
+// ── ItemTreeSourceMap ─────────────────────────────────────────────────────────
+
+/// Parallel source map for `ItemTree` — stores name spans that are
+/// deliberately excluded from the semantic `ItemTree` to avoid polluting
+/// Salsa's early-cutoff comparisons with position data.
+///
+/// Follows the same body/signature source-map pattern used by
+/// `function_body` / `function_body_source_map`.
+#[derive(Debug, Clone, Default)]
+pub struct ItemTreeSourceMap {
+    /// `name_span` for each class's fields, parallel to `Class::fields`.
+    pub class_field_spans: FxHashMap<LocalItemId<ClassMarker>, Vec<TextRange>>,
+    /// `name_span` for each enum's variants, parallel to `Enum::variants`.
+    pub enum_variant_spans: FxHashMap<LocalItemId<EnumMarker>, Vec<TextRange>>,
+    /// `name_span` for each function.
+    pub function_name_spans: FxHashMap<LocalItemId<FunctionMarker>, TextRange>,
+}
+
 // ── ItemTree ─────────────────────────────────────────────────────────────────
 
 /// Position-independent item storage for a single file.
@@ -315,6 +333,37 @@ impl ItemTree {
         id
     }
 
+    /// Populate source map spans for a class that was allocated via `alloc_class`.
+    pub fn collect_class_spans(
+        source_map: &mut ItemTreeSourceMap,
+        id: LocalItemId<ClassMarker>,
+        class_def: &ast::ClassDef,
+    ) {
+        let spans: Vec<TextRange> = class_def.fields.iter().map(|f| f.name_span).collect();
+        source_map.class_field_spans.insert(id, spans);
+    }
+
+    /// Populate source map name span for a function.
+    pub fn collect_function_span(
+        source_map: &mut ItemTreeSourceMap,
+        id: LocalItemId<FunctionMarker>,
+        func_def: &ast::FunctionDef,
+    ) {
+        source_map
+            .function_name_spans
+            .insert(id, func_def.name_span);
+    }
+
+    /// Populate source map spans for an enum that was allocated via `alloc_enum`.
+    pub fn collect_enum_spans(
+        source_map: &mut ItemTreeSourceMap,
+        id: LocalItemId<EnumMarker>,
+        enum_def: &ast::EnumDef,
+    ) {
+        let spans: Vec<TextRange> = enum_def.variants.iter().map(|v| v.name_span).collect();
+        source_map.enum_variant_spans.insert(id, spans);
+    }
+
     pub fn alloc_type_alias(&mut self, ta: &ast::TypeAliasDef) -> LocalItemId<TypeAliasMarker> {
         let id = self.alloc_id(ItemKind::TypeAlias, &ta.name);
         self.type_aliases.insert(

baml_language/crates/baml_compiler2_hir/src/item_tree.rs

@@ -33,7 +33,7 @@ pub use builder::SemanticIndexBuilder;
 
 use crate::{
     contributions::FileSymbolContributions,
-    item_tree::ItemTree,
+    item_tree::{ItemTree, ItemTreeSourceMap},
     semantic_index::{FileSemanticIndex, ScopeBindings},
 };
 
@@ -128,6 +128,14 @@ pub fn file_item_tree(db: &dyn Db, file: SourceFile) -> Arc<ItemTree> {
     Arc::clone(&index.item_tree)
 }
 
+/// Returns the item tree source map for a file (clones the Arc — O(1)).
+///
+/// Not tracked — the source map is cached via `file_semantic_index`.
+pub fn file_item_tree_source_map(db: &dyn Db, file: SourceFile) -> Arc<ItemTreeSourceMap> {
+    let index = file_semantic_index(db, file);
+    Arc::clone(&index.item_tree_source_map)
+}
+
 /// Returns the `ScopeBindings` for a given scope.
 ///
 /// Not tracked — callers use the pre-interned `ScopeId` to look up bindings.

baml_language/crates/baml_compiler2_hir/src/lib.rs

@@ -13,7 +13,7 @@ use text_size::{TextRange, TextSize};
 use crate::{
     contributions::FileSymbolContributions,
     diagnostic::Hir2Diagnostic,
-    item_tree::ItemTree,
+    item_tree::{ItemTree, ItemTreeSourceMap},
     scope::{FileScopeId, Scope, ScopeId},
 };
 
@@ -100,6 +100,9 @@ pub struct FileSemanticIndex<'db> {
     /// Per-file item tree — maps `LocalItemId` to item data.
     pub item_tree: Arc<ItemTree>,
 
+    /// Source map for item tree — field/variant name spans.
+    pub item_tree_source_map: Arc<ItemTreeSourceMap>,
+
     /// Names this file contributes to its package namespace.
     pub symbol_contributions: Arc<FileSymbolContributions<'db>>,