Fix stability of Definition entities

Author: lionel-

crates/oak_db/src/definition.rs

@@ -2,9 +2,12 @@ use aether_syntax::RBinaryExpression;
 use aether_syntax::RSyntaxKind;
 use biome_rowan::AstNode;
 use biome_rowan::TextRange;
+use oak_semantic::semantic_index::DefinitionId;
 use oak_semantic::semantic_index::DefinitionKind;
 use oak_semantic::semantic_index::ScopeId;
+use rustc_hash::FxHashMap;
 
+use crate::Db;
 use crate::File;
 use crate::Name;
 
@@ -86,3 +89,70 @@ fn is_right_assignment(node: &RBinaryExpression) -> bool {
         )
     })
 }
+
+#[salsa::tracked]
+impl<'db> File {
+    /// Look up the `Definition` entity for the binding at `(scope, def_id)` in
+    /// this file. The entity has already been minted in [`File::definitions`]
+    /// and this is a plain lookup. The salsa id stays stable across edits that
+    /// renumber `def_id`.
+    pub(crate) fn definition(
+        self,
+        db: &'db dyn Db,
+        scope: ScopeId,
+        def_id: DefinitionId,
+    ) -> Option<Definition<'db>> {
+        self.definitions(db)
+            .by_site
+            .get(&DefinitionSite { scope, def_id })
+            .copied()
+    }
+
+    /// Mint every `Definition` in this file, keyed by its `(scope, def_id)`
+    /// site. This is the single `Definition::new` call site: every resolution
+    /// path looks entities up here rather than minting its own, so the same
+    /// binding has one salsa id no matter how it's reached.
+    ///
+    /// Identity is `(file, scope, name)` (see [`Definition`]). `def_id` is only
+    /// the lookup key, never part of identity, so inserting a binding earlier
+    /// in a scope doesn't churn the ids of the others.
+    #[salsa::tracked(returns(ref))]
+    fn definitions(self, db: &'db dyn Db) -> FileDefinitions<'db> {
+        let index = self.semantic_index(db);
+        let mut by_site = FxHashMap::default();
+
+        for scope in index.scope_ids() {
+            let symbols = index.symbols(scope);
+            for (def_id, def) in index.definitions(scope).iter() {
+                let name = Name::new(db, symbols.symbol(def.symbol()).name());
+                let definition = Definition::new(db, self, scope, name, def.kind().clone());
+                by_site.insert(DefinitionSite { scope, def_id }, definition);
+            }
+        }
+
+        FileDefinitions { by_site }
+    }
+}
+
+/// Every `Definition` in a file, keyed by its definition site.
+///
+/// A `salsa::Update` struct rather than a bare map so the salsa-tracked
+/// `File::definitions` query can store it: salsa needs the `Update` impl to
+/// handle the `Definition<'db>` lifetime that the map values carry.
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
+struct FileDefinitions<'db> {
+    by_site: FxHashMap<DefinitionSite, Definition<'db>>,
+}
+
+/// Map key for [`FileDefinitions`]: a binding's `(scope, def_id)` site.
+///
+/// A named struct rather than a `(ScopeId, DefinitionId)` tuple so it can
+/// derive `salsa::Update` (the salsa map key must implement it). The derive
+/// falls back to a plain replace for the `'static` index fields, which don't
+/// implement `Update` themselves (`oak_semantic` has no salsa dependency).
+/// Mirrors ty's `DefinitionNodeKey`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, salsa::Update)]
+struct DefinitionSite {
+    scope: ScopeId,
+    def_id: DefinitionId,
+}

crates/oak_db/src/file_exports.rs

@@ -61,7 +61,7 @@ impl File {
     #[salsa::tracked(returns(ref), cycle_result = exports_cycle_result)]
     pub fn exports(self, db: &dyn Db) -> FileExports {
         let mut entries: HashMap<String, ExportEntry> = HashMap::new();
-        for (name, def) in self.semantic_index(db).exports() {
+        for (name, (_def_id, def)) in self.semantic_index(db).exports() {
             let entry = match def.kind() {
                 DefinitionKind::Import {
                     file: target_url,

crates/oak_db/src/file_resolve.rs

@@ -2,7 +2,6 @@ use std::collections::HashSet;
 use std::rc::Rc;
 
 use biome_rowan::TextSize;
-use oak_semantic::DefinitionId;
 use oak_semantic::ScopeId;
 
 use crate::Db;
@@ -93,8 +92,9 @@ impl<'db> File {
     ///    visible subset of attaches / collation predecessors.
     ///
     /// Not `#[salsa::tracked]` because keying on `(self, offset)` would
-    /// balloon the cache. `Definition` creation delegates to the tracked
-    /// [`File::resolve_export()`] and [`File::intern_definition()`].
+    /// balloon the cache. The `Definition` entities it returns are minted by
+    /// the tracked [`File::definitions()`], so they stay cached and
+    /// identity-stable even though this lookup isn't.
     pub fn resolve_at(self, db: &'db dyn Db, offset: TextSize) -> Option<Definition<'db>> {
         let index = self.semantic_index(db);
         let (use_scope, _use_id, use_site) = index.use_at(offset)?;
@@ -111,12 +111,7 @@ impl<'db> File {
         let file_scope = ScopeId::from(0);
         if let Some((binding_scope, def_id, _def)) = index.resolve(&name, use_scope) {
             if binding_scope != file_scope {
-                return Some(self.intern_definition(
-                    db,
-                    binding_scope,
-                    def_id,
-                    Name::new(db, name.as_str()),
-                ));
+                return self.definition(db, binding_scope, def_id);
             }
         }
 
@@ -173,23 +168,13 @@ impl<'db> File {
 
             match current_file.exports(db).get(current_name.as_ref())? {
                 ExportEntry::Local => {
-                    // Fetch exports again, this time through the semantic index
-                    // to get the volatile `kind` field that the firewall query
-                    // `File::exports()` doesn't expose.
-                    let kind = current_file
-                        .semantic_index(db)
-                        .exports()
-                        .get(current_name.as_ref())
-                        .map(|def| def.kind().clone())?;
-
-                    let file_scope = ScopeId::from(0);
-                    return Some(Definition::new(
-                        db,
-                        current_file,
-                        file_scope,
-                        Name::new(db, current_name.as_ref()),
-                        kind,
-                    ));
+                    // Look up the file-scope binding through the semantic index
+                    // to recover its `def_id`, then fetch the interned
+                    // definition. `exports()` returns the last-wins
+                    // definitions, so this is the right binding for the name.
+                    let index = current_file.semantic_index(db);
+                    let def_id = index.exports().get(current_name.as_ref())?.0;
+                    return current_file.definition(db, ScopeId::from(0), def_id);
                 },
 
                 ExportEntry::Import { file, name } => {
@@ -199,22 +184,4 @@ impl<'db> File {
             }
         }
     }
-
-    /// Intern the salsa-tracked `Definition` entity for a binding identified
-    /// by `(scope, def_id)` in this file's semantic index. Wraps
-    /// `Definition::new` in a tracked context, which is required to construct
-    /// tracked structs.
-    #[salsa::tracked]
-    fn intern_definition(
-        self,
-        db: &'db dyn Db,
-        scope: ScopeId,
-        def_id: DefinitionId,
-        name: Name<'db>,
-    ) -> Definition<'db> {
-        let kind = self.semantic_index(db).definitions(scope)[def_id]
-            .kind()
-            .clone();
-        Definition::new(db, self, scope, name, kind)
-    }
 }

crates/oak_db/src/tests/file_resolve.rs

@@ -271,6 +271,138 @@ fn test_definition_id_stable_across_body_edits() {
     assert_ne!(range1, range2);
 }
 
+#[test]
+fn test_definition_id_stable_across_def_id_renumber_local_path() {
+    // The function-scope local path looks up its `Definition` from the single
+    // mint site `File::definitions`, whose identity is `(file, scope, name)`.
+    // `def_id` is only the lookup key, never part of identity, so prepending
+    // an unrelated binding inside the function (which renumbers x's `def_id`)
+    // leaves x's salsa id unchanged. This is the same stability the export
+    // path has in `test_definition_id_stable_across_body_edits`, now extended
+    // to the local path.
+    use biome_rowan::TextSize;
+    use salsa::plumbing::AsId;
+
+    let content1 = "f <- function() {\nx <- 1\nx\n}\n";
+    let use1 = content1.find("\nx\n").expect("standalone use of x") + 1;
+
+    let mut db = TestDb::new();
+    let files = setup_workspace(&mut db, &[("w/a.R", content1)]);
+    let file = files[0];
+
+    let id1 = file
+        .resolve_at(&db, TextSize::from(use1 as u32))
+        .expect("use of x resolves to its function-scope binding")
+        .as_id();
+
+    // Prepend an unrelated binding inside the function so x's DefinitionId
+    // shifts 0 -> 1 within the function scope.
+    let content2 = "f <- function() {\nw <- 0\nx <- 1\nx\n}\n";
+    let use2 = content2.find("\nx\n").expect("standalone use of x") + 1;
+    file.set_contents(&mut db).to(content2.to_string());
+
+    let id2 = file
+        .resolve_at(&db, TextSize::from(use2 as u32))
+        .expect("use of x still resolves")
+        .as_id();
+
+    assert_eq!(id1, id2);
+}
+
+#[test]
+fn test_definitions_mints_distinct_entities_for_same_name_redefinition() {
+    // Two file-scope `x` bindings share the `(file, scope, name)` id-fields.
+    // The single mint site must create two distinct salsa entities rather than
+    // collide or panic; salsa disambiguates same-id-field tracked structs by
+    // creation order. Resolving `x` forces the mint of both (via `definitions`)
+    // and must land on the last definition (offset 7).
+    let mut db = TestDb::new();
+    let files = setup_workspace(&mut db, &[("w/a.R", "x <- 1\nx <- 2\n")]);
+    let file = files[0];
+
+    let def = file.resolve(&db, name(&db, "x")).expect("x should resolve");
+    assert_eq!(def.name(&db).text(&db).as_str(), "x");
+    let range = def.name_range(&db).expect("local binding has a name range");
+    assert_eq!(usize::from(range.start()), 7);
+}
+
+#[test]
+fn test_position_shift_keeps_id_and_does_not_invalidate_identity_consumers() {
+    // A pure position shift (prepend a comment, no binding added or removed)
+    // moves the binding's AstPtr but leaves `(file, scope, name)` and its
+    // ordinal unchanged, so the salsa id is stable. A downstream query that
+    // reads only identity therefore stays cached across the rebuild; only
+    // consumers of `kind` (the moved AstPtr) would re-run.
+    use salsa::plumbing::AsId;
+
+    use crate::Db;
+    use crate::Definition;
+
+    #[salsa::tracked]
+    fn name_len<'db>(db: &'db dyn Db, def: Definition<'db>) -> usize {
+        def.name(db).text(db).len()
+    }
+
+    let mut db = TestDb::new();
+    let files = setup_workspace(&mut db, &[("w/a.R", "x <- 1\n")]);
+    let file = files[0];
+
+    let (id1, range1) = {
+        let def = file.resolve(&db, name(&db, "x")).expect("x resolves");
+        let _ = name_len(&db, def);
+        (def.as_id(), def.name_range(&db))
+    };
+    assert_eq!(db.executions("name_len"), 1);
+
+    // Pure position shift: x moves down a line, no binding added or removed.
+    file.set_contents(&mut db)
+        .to("# comment\nx <- 1\n".to_string());
+
+    let (id2, range2) = {
+        let def = file.resolve(&db, name(&db, "x")).expect("x still resolves");
+        let _ = name_len(&db, def);
+        (def.as_id(), def.name_range(&db))
+    };
+
+    // Same entity, and the name range moved (it really was a position shift).
+    assert_eq!(id1, id2);
+    assert_ne!(range1, range2);
+    // The identity-only consumer was not re-executed by the position shift.
+    assert_eq!(db.executions("name_len"), 1);
+}
+
+#[test]
+fn test_same_name_sibling_insertion_churns_later_definition_id() {
+    // TRACKING TEST for a known boundary, not a guarantee to preserve.
+    //
+    // Identity is `(file, scope, name)` plus salsa's creation-order
+    // disambiguator among same-name siblings. Inserting *another* `x` earlier
+    // in the scope shifts the ordinals of the later `x` definitions, so their
+    // salsa ids churn even though their position-stability would otherwise
+    // hold. This matches ty's `push_additional_definition` ordering. The test
+    // exists to notice if salsa's disambiguation ever changes.
+    use salsa::plumbing::AsId;
+
+    let mut db = TestDb::new();
+    let files = setup_workspace(&mut db, &[("w/a.R", "x <- 1\nx <- 2\n")]);
+    let file = files[0];
+
+    // `resolve` is last-wins, so it returns the final `x` (`x <- 2`), ordinal 1.
+    let id1 = file.resolve(&db, name(&db, "x")).expect("x resolves").as_id();
+
+    // Insert another `x` at the top. The final `x` is still last-wins, but its
+    // ordinal among same-name siblings shifts from 1 to 2.
+    file.set_contents(&mut db)
+        .to("x <- 0\nx <- 1\nx <- 2\n".to_string());
+
+    let id2 = file
+        .resolve(&db, name(&db, "x"))
+        .expect("x still resolves")
+        .as_id();
+
+    assert_ne!(id1, id2);
+}
+
 #[test]
 fn test_resolve_unbound_name_in_package_does_not_cycle() {
     // Without exports-only sibling chase, A's `resolve` would walk into

crates/oak_semantic/src/semantic_index.rs

@@ -126,14 +126,14 @@ impl SemanticIndex {
     /// Top-level definitions exported by this file (definitions in the file scope).
     /// Includes `Import`-kind forwarding definitions from `source()` calls.
     /// Last definition of each name wins (later assignments shadow earlier ones).
-    pub fn exports(&self) -> FxHashMap<&str, &Definition> {
+    pub fn exports(&self) -> FxHashMap<&str, (DefinitionId, &Definition)> {
         let file_scope = ScopeId::from(0);
         let symbols = &self.symbol_tables[file_scope];
 
         let mut exports = FxHashMap::default();
-        for (_id, def) in self.definitions[file_scope].iter() {
+        for (id, def) in self.definitions[file_scope].iter() {
             let name = symbols.symbol(def.symbol()).name();
-            exports.insert(name, def);
+            exports.insert(name, (id, def));
         }
 
         exports

crates/oak_semantic/tests/integration/builder.rs

@@ -1512,7 +1512,7 @@ fn test_file_exports_last_def_wins() {
     let exports = index.exports();
     assert_eq!(exports.len(), 2);
     // The range should be the second `foo` (offset 9..12)
-    let def = exports.get("foo").unwrap();
+    let (_def_id, def) = exports.get("foo").unwrap();
     assert_eq!(def.range().start(), biome_rowan::TextSize::from(9));
 }