fix(scope): use identity-based references for module edits

lang/lambda/edits/pkg.generated.mbti

@@ -16,8 +16,6 @@ pub fn compute_text_edit(TreeEditOp, EditContext[@ast.Term]) -> Result[(Array[@c
 
 pub fn find_binding_for_init(@core.NodeId, @proj.FlatProj) -> (@core.NodeId, Int)?
 
-pub fn find_usages(String, Int, @proj.FlatProj, Map[@core.NodeId, @core.ProjNode[@ast.Term]]) -> Array[@core.NodeId]
-
 pub fn free_vars(@ast.Term, @hashset.HashSet[String]) -> @hashset.HashSet[String]
 
 pub fn get_actions_for_node(@ast.Term, NodeContext) -> Array[Action]

lang/lambda/edits/scope.mbt

@@ -1,28 +1,21 @@
 ///|
-/// Find all Var nodes referencing the given name from scope_start_index onward
-/// in the Module defs + body, respecting shadowing.
-pub fn find_usages(
-  var_name : String,
-  scope_start_index : Int,
-  flat_proj : FlatProj,
-  registry : Map[NodeId, ProjNode[@ast.Term]],
-) -> Array[NodeId] {
-  let results : Array[NodeId] = []
-  let mut shadowed = false
-  for i = scope_start_index; i < flat_proj.defs.length(); i = i + 1 {
-    let def = flat_proj.defs[i]
-    // If this def shadows the name, stop
-    if def.0 == var_name && i > scope_start_index {
-      shadowed = true
-      break
+/// Identity-based references to a module definition, delegated to the canonical
+/// scope graph query. The old name-based usage scan could over-match a later
+/// shadowing definition's body; this helper first locates the graph-local
+/// `DeclId` for the requested module def, then asks `@scope.references` for only
+/// refs resolved to that declaration.
+fn module_def_references(
+  g : @scope.ScopeGraph,
+  def_index : Int,
+) -> Result[Array[NodeId], String] {
+  for decl in g.decls {
+    match decl.kind {
+      @scope.DeclKind::ModuleDef(def_index=i) if i == def_index =>
+        return Ok(@scope.references(g, decl.id))
+      _ => ()
     }
-    collect_var_usages(def.1, var_name, results)
-  }
-  // Check body only if not shadowed
-  if !shadowed && flat_proj.final_expr is Some(body) {
-    collect_var_usages(body, var_name, results)
   }
-  results
+  Err("No scope declaration for module def index " + def_index.to_string())
 }
 
 ///|

lang/lambda/edits/scope_wbtest.mbt

@@ -14,34 +14,51 @@ fn scope_test_registry(
 }
 
 ///|
-test "find_usages: finds Var references in Module body" {
+test "module_def_references: finds Var references in Module body" {
   let text = "let x = 0\nx + x"
   let (proj, _) = parse_to_proj_node(text)
   let fp = FlatProj::from_proj_node(proj)
+  let source_map = SourceMap::from_ast(proj)
   let registry = scope_test_registry(proj)
-  let usages = find_usages("x", 0, fp, registry)
-  inspect(usages.length(), content="2")
+  let g = @scope.build(fp, registry, source_map)
+  match module_def_references(g, 0) {
+    Ok(refs) => inspect(refs.length(), content="2")
+    Err(e) => fail(e)
+  }
 }
 
 ///|
-test "find_usages: respects shadowing by inner Lam" {
+test "module_def_references: respects shadowing by inner Lam" {
   let text = "let x = 0\nλx. x"
   let (proj, _) = parse_to_proj_node(text)
   let fp = FlatProj::from_proj_node(proj)
+  let source_map = SourceMap::from_ast(proj)
   let registry = scope_test_registry(proj)
-  let usages = find_usages("x", 0, fp, registry)
-  inspect(usages.length(), content="0")
+  let g = @scope.build(fp, registry, source_map)
+  match module_def_references(g, 0) {
+    Ok(refs) => inspect(refs.length(), content="0")
+    Err(e) => fail(e)
+  }
 }
 
 ///|
-test "find_usages: shadowed by later def suppresses body" {
-  let text = "let x = 0\nlet x = 1\nx"
+test "module_def_references: later duplicate shadows body only" {
+  let text = "let x = 0\nlet x = x\nx"
   let (proj, _) = parse_to_proj_node(text)
   let fp = FlatProj::from_proj_node(proj)
+  let source_map = SourceMap::from_ast(proj)
   let registry = scope_test_registry(proj)
-  // Starting from def 0, def 1 shadows "x" — body should not be searched for def 0's usages
-  let usages = find_usages("x", 0, fp, registry)
-  inspect(usages.length(), content="0")
+  let g = @scope.build(fp, registry, source_map)
+  match (module_def_references(g, 0), module_def_references(g, 1)) {
+    (Ok(first_refs), Ok(second_refs)) => {
+      // The first def is referenced by the second def's init; the body resolves
+      // to the second def. This is the identity-based case the old name scan
+      // could not distinguish when callers started at def_index + 1.
+      inspect(first_refs.length(), content="1")
+      inspect(second_refs.length(), content="1")
+    }
+    (Err(e), _) | (_, Err(e)) => fail(e)
+  }
 }
 
 ///|

lang/lambda/edits/text_edit_refactor.mbt

@@ -164,9 +164,12 @@ fn compute_inline_definition(
           "Cannot inline: would capture variables bound by enclosing lambda",
         )
       }
-      // Check if this is the sole usage
-      let usages = find_usages(var_name, def_index + 1, flat_proj, registry)
-      if usages.length() <= 1 {
+      // Check if this is the sole resolved reference.
+      let refs = match module_def_references(g, def_index) {
+        Ok(ids) => ids
+        Err(e) => return Err(e)
+      }
+      if refs.length() <= 1 {
         // Sole usage — inline and delete the binding
         let (let_start, binding_end) = match
           get_binding_text_range(source_text, source_map, def) {
@@ -248,32 +251,36 @@ fn compute_inline_all_usages(
   let def = flat_proj.defs[def_index]
   let var_name = def.0
   let init_text = @ast.print_term(def.1.kind)
-  // Find all usages
-  let usage_ids = find_usages(var_name, def_index + 1, flat_proj, registry)
-  // Check for capture: would any free var in init be captured by a lambda at any usage site?
+  // Find all references resolved to this binding.
+  let g = @scope.build(flat_proj, registry, source_map)
+  let ref_ids = match module_def_references(g, def_index) {
+    Ok(ids) => ids
+    Err(e) => return Err(e)
+  }
+  // Check for capture: would any free var in init be captured by a lambda at any reference site?
   let init_fv = free_vars(def.1.kind, @immut/hashset.new())
-  for uid in usage_ids {
-    let usage_lam_env = collect_lam_env(uid, registry)
+  for rid in ref_ids {
+    let ref_lam_env = collect_lam_env(rid, registry)
     let mut would_capture = false
-    init_fv.each(fn(v) { if usage_lam_env.contains(v) { would_capture = true } })
+    init_fv.each(fn(v) { if ref_lam_env.contains(v) { would_capture = true } })
     if would_capture {
       return Err(
         "Cannot inline: would capture variables bound by enclosing lambda",
       )
     }
   }
-  // Build replacement edits for all usages (reverse document order)
+  // Build replacement edits for all references (reverse document order)
   let edits : Array[SpanEdit] = []
-  // Collect usage ranges, sorted by position (descending)
-  let usage_ranges : Array[(Int, Int)] = []
-  for uid in usage_ids {
-    if source_map.get_range(uid) is Some(r) {
-      usage_ranges.push((r.start, r.end))
+  // Collect reference ranges, sorted by position (descending)
+  let ref_ranges : Array[(Int, Int)] = []
+  for rid in ref_ids {
+    if source_map.get_range(rid) is Some(r) {
+      ref_ranges.push((r.start, r.end))
     }
   }
-  usage_ranges.sort_by(fn(a, b) { b.0.compare(a.0) })
-  for ur in usage_ranges {
-    edits.push({ start: ur.0, delete_len: ur.1 - ur.0, inserted: init_text })
+  ref_ranges.sort_by(fn(a, b) { b.0.compare(a.0) })
+  for rr in ref_ranges {
+    edits.push({ start: rr.0, delete_len: rr.1 - rr.0, inserted: init_text })
   }
   // Delete the binding
   let (let_start, binding_end) = match

lang/lambda/edits/text_edit_rename.mbt

@@ -181,12 +181,16 @@ fn rename_module_binding(
       )
     }
   }
-  // Find all usages
-  let usage_ids = find_usages(old_name, def_index + 1, flat_proj, registry)
-  // Guard: check if new_name would be captured by an enclosing lambda at any usage site
-  for usage_id in usage_ids {
-    let usage_lam_env = collect_lam_env(usage_id, registry)
-    if usage_lam_env.contains(new_name) {
+  // Find all references resolved to this binding.
+  let g = @scope.build(flat_proj, registry, source_map)
+  let ref_ids = match module_def_references(g, def_index) {
+    Ok(ids) => ids
+    Err(e) => return Err(e)
+  }
+  // Guard: check if new_name would be captured by an enclosing lambda at any reference site
+  for ref_id in ref_ids {
+    let ref_lam_env = collect_lam_env(ref_id, registry)
+    if ref_lam_env.contains(new_name) {
       return Err(
         "Cannot rename: '" +
         new_name +
@@ -209,16 +213,16 @@ fn rename_module_binding(
         "No token span for binding name at def index " + def_index.to_string(),
       )
   }
-  let usage_ranges : Array[(Int, Int)] = []
-  for uid in usage_ids {
-    if source_map.get_range(uid) is Some(r) {
-      usage_ranges.push((r.start, r.end))
+  let ref_ranges : Array[(Int, Int)] = []
+  for rid in ref_ids {
+    if source_map.get_range(rid) is Some(r) {
+      ref_ranges.push((r.start, r.end))
     }
   }
   // Sort descending by position
-  usage_ranges.sort_by(fn(a, b) { b.0.compare(a.0) })
-  for ur in usage_ranges {
-    edits.push({ start: ur.0, delete_len: ur.1 - ur.0, inserted: new_name })
+  ref_ranges.sort_by(fn(a, b) { b.0.compare(a.0) })
+  for rr in ref_ranges {
+    edits.push({ start: rr.0, delete_len: rr.1 - rr.0, inserted: new_name })
   }
   // Sort all edits descending
   edits.sort_by(fn(a, b) { b.start.compare(a.start) })

lang/lambda/edits/text_edit_wbtest.mbt

@@ -877,6 +877,25 @@ test "compute_text_edit: InlineAllUsages with unknown binding returns error" {
   inspect(result is Err(_), content="true")
 }
 
+///|
+test "compute_text_edit: InlineAllUsages first duplicate preserves shadowed body" {
+  let text = "let x = 0\nlet x = x\nx"
+  let (_, sm, registry, fp) = parse_with_token_spans(text)
+  let result = compute_text_edit(
+    InlineAllUsages(binding_node_id=fp.defs[0].3),
+    make_edit_ctx(text, sm, registry, fp),
+  )
+  match result {
+    Ok(Some((edits, _))) => {
+      let new_text = apply_edits(text, edits)
+      // The init Var range includes its leading space, matching existing inline
+      // whitespace behavior; the assertion pins the resolved-reference set.
+      inspect(new_text, content="let x =0\nx")
+    }
+    _ => fail("expected Some edits, got: " + @debug.to_string(result))
+  }
+}
+
 ///|
 test "compute_text_edit: InlineDefinition on lambda-bound var returns error" {
   let text = "λx. x"
@@ -990,6 +1009,25 @@ test "compute_text_edit: Rename module binding via Var" {
   }
 }
 
+///|
+test "compute_text_edit: Rename first duplicate module binding preserves shadowed body" {
+  let text = "let x = 0\nlet x = x\nx"
+  let (_, sm, registry, fp) = parse_with_token_spans(text)
+  let result = compute_text_edit(
+    Rename(node_id=fp.defs[0].3, new_name="y"),
+    make_edit_ctx(text, sm, registry, fp),
+  )
+  match result {
+    Ok(Some((edits, _))) => {
+      let new_text = apply_edits(text, edits)
+      // The init Var range includes its leading space; the body x must remain
+      // shadowed by the second def and therefore unchanged.
+      inspect(new_text, content="let y = 0\nlet x =y\nx")
+    }
+    _ => fail("expected Some edits, got: " + @debug.to_string(result))
+  }
+}
+
 ///|
 test "compute_text_edit: Rename module binding with multiple usages" {
   let text = "let x = 42\nx + x"