fix(lambda): use scope graph env for edit capture checks

lang/lambda/edits/pkg.generated.mbti

@@ -10,8 +10,6 @@ import {
 }
 
 // Values
-pub fn collect_lam_env(@core.NodeId, Map[@core.NodeId, @core.ProjNode[@ast.Term]]) -> @hashset.HashSet[String]
-
 pub fn compute_text_edit(TreeEditOp, EditContext[@ast.Term]) -> Result[(Array[@core.SpanEdit], @core.FocusHint)?, String]
 
 pub fn find_binding_for_init(@core.NodeId, @proj.FlatProj) -> (@core.NodeId, Int)?

lang/lambda/edits/scope.mbt

@@ -40,6 +40,187 @@ fn collect_var_usages(
   }
 }
 
+///|
+fn name_captured_at_node(
+  g : @scope.ScopeGraph,
+  node_id : NodeId,
+  name : String,
+) -> Bool {
+  @scope.enclosing_env(g, node_id).contains(name)
+}
+
+///|
+fn free_names_captured_at_use_sites(
+  g : @scope.ScopeGraph,
+  node : ProjNode[@ast.Term],
+  free_names : @immut/hashset.HashSet[String],
+) -> Bool {
+  let mut captured = false
+  free_names.each(fn(v) {
+    let var_ids : Array[NodeId] = []
+    collect_var_usages(node, v, var_ids)
+    if var_ids.any(fn(vid) { name_captured_at_node(g, vid, v) }) {
+      captured = true
+    }
+  })
+  captured
+}
+
+///|
+fn def_cutoff_at_node(
+  flat_proj : FlatProj,
+  source_map : SourceMap,
+  node_id : NodeId,
+) -> Int {
+  guard source_map.get_range(node_id) is Some(r) else {
+    return flat_proj.defs.length()
+  }
+  flat_proj.defs
+  .search_by(fn(def) {
+    source_map.get_range(def.1.id()) is Some(dr) &&
+    r.start >= dr.start &&
+    r.start < dr.end
+  })
+  .unwrap_or(flat_proj.defs.length())
+}
+
+///|
+fn scope_id_for_node(
+  g : @scope.ScopeGraph,
+  node_id : NodeId,
+) -> @scope.ScopeId? {
+  match g.refs.iter().find_first(fn(r) { r.node_id == node_id }) {
+    Some(r) => Some(r.scope)
+    None =>
+      g.decls
+      .iter()
+      .find_first(fn(d) { d.node_id == node_id })
+      .map(fn(d) { d.scope })
+  }
+}
+
+///|
+fn root_scope_id(g : @scope.ScopeGraph) -> @scope.ScopeId? {
+  g.scopes.iter().find_first(fn(s) { s.parent is None }).map(fn(s) { s.id })
+}
+
+///|
+fn declaration_id_for_name_from_scope(
+  g : @scope.ScopeGraph,
+  start_scope : @scope.ScopeId,
+  cutoff : Int,
+  name : String,
+) -> @scope.DeclId? {
+  let mut cur : @scope.ScopeId? = Some(start_scope)
+  while cur is Some(sid) {
+    let @scope.ScopeId(si) = sid
+    let scope = g.scopes[si]
+    let mut hit : @scope.DeclId? = None
+    for did in scope.decl_ids {
+      let @scope.DeclId(di) = did
+      let decl = g.decls[di]
+      if decl.name == name {
+        match decl.kind {
+          @scope.DeclKind::ModuleDef(def_index~) =>
+            if def_index < cutoff {
+              hit = Some(did)
+            }
+          @scope.DeclKind::LamParam(_) => hit = Some(did)
+        }
+      }
+    }
+    if hit is Some(_) {
+      return hit
+    }
+    cur = scope.parent
+  }
+  None
+}
+
+///|
+fn declaration_id_for_name_at_node(
+  g : @scope.ScopeGraph,
+  flat_proj : FlatProj,
+  source_map : SourceMap,
+  node_id : NodeId,
+  name : String,
+) -> @scope.DeclId? {
+  guard scope_id_for_node(g, node_id) is Some(start_scope) else { return None }
+  declaration_id_for_name_from_scope(
+    g,
+    start_scope,
+    def_cutoff_at_node(flat_proj, source_map, node_id),
+    name,
+  )
+}
+
+///|
+fn declaration_id_for_name_at_module_end(
+  g : @scope.ScopeGraph,
+  flat_proj : FlatProj,
+  name : String,
+) -> @scope.DeclId? {
+  guard root_scope_id(g) is Some(root_scope) else { return None }
+  declaration_id_for_name_from_scope(
+    g,
+    root_scope,
+    flat_proj.defs.length(),
+    name,
+  )
+}
+
+///|
+fn free_names_would_rebind_at_node(
+  g : @scope.ScopeGraph,
+  flat_proj : FlatProj,
+  source_map : SourceMap,
+  source_node : ProjNode[@ast.Term],
+  target_node_id : NodeId,
+  free_names : @immut/hashset.HashSet[String],
+) -> Bool {
+  let mut rebound = false
+  free_names.each(fn(v) {
+    let target_decl = declaration_id_for_name_at_node(
+      g, flat_proj, source_map, target_node_id, v,
+    )
+    if free_name_would_rebind_to(g, source_node, v, target_decl) {
+      rebound = true
+    }
+  })
+  rebound
+}
+
+///|
+fn free_names_would_rebind_at_module_end(
+  g : @scope.ScopeGraph,
+  flat_proj : FlatProj,
+  source_node : ProjNode[@ast.Term],
+  free_names : @immut/hashset.HashSet[String],
+) -> Bool {
+  let mut rebound = false
+  free_names.each(fn(v) {
+    let target_decl = declaration_id_for_name_at_module_end(g, flat_proj, v)
+    if free_name_would_rebind_to(g, source_node, v, target_decl) {
+      rebound = true
+    }
+  })
+  rebound
+}
+
+///|
+fn free_name_would_rebind_to(
+  g : @scope.ScopeGraph,
+  source_node : ProjNode[@ast.Term],
+  name : String,
+  target_decl : @scope.DeclId?,
+) -> Bool {
+  let var_ids : Array[NodeId] = []
+  collect_var_usages(source_node, name, var_ids)
+  var_ids.any(fn(vid) {
+    @scope.declaration(g, vid).map(fn(d) { d.id }) != target_decl
+  })
+}
+
 ///|
 /// Look up the FlatProj binding NodeId for a given init expression NodeId.
 /// The UI selects init expression ProjNodes (children of Module); this function
@@ -55,27 +236,3 @@ pub fn find_binding_for_init(
   }
   None
 }
-
-///|
-/// Collect Lam parameter names in scope at a node position (walking up the tree).
-pub fn collect_lam_env(
-  node_id : NodeId,
-  registry : Map[NodeId, ProjNode[@ast.Term]],
-) -> @immut/hashset.HashSet[String] {
-  let mut env : @immut/hashset.HashSet[String] = @immut/hashset.new()
-  fn walk_up(nid : NodeId) {
-    for pid, pnode in registry {
-      for child in pnode.children {
-        if child.id() == nid {
-          if pnode.kind is @ast.Term::Lam(param, _) {
-            env = env.add(param)
-          }
-          walk_up(pid)
-          return
-        }
-      }
-    }
-  }
-  walk_up(node_id)
-  env
-}

lang/lambda/edits/text_edit_refactor.mbt

@@ -11,33 +11,14 @@ fn compute_extract_to_let(
   guard source_map.get_range(node_id) is Some(range) else {
     return Err("Range not found")
   }
-  // Free-var validation: check for lambda-captured variables
-  let mut module_env : @immut/hashset.HashSet[String] = @immut/hashset.new()
-  // Find which def contains this node (or if it's in the body)
-  let mut in_def_index : Int = flat_proj.defs.length() // body sentinel
-  for i, def in flat_proj.defs {
-    if source_map.get_range(def.1.id()) is Some(r) &&
-      range.start >= r.start &&
-      range.end <= r.end {
-      in_def_index = i
-      break
-    }
-  }
-  // Add def names in scope at the extraction site
-  let scope_limit = if in_def_index < flat_proj.defs.length() {
-    in_def_index
+  let g = @scope.build(flat_proj, registry, source_map)
+  let fv = free_vars(node.kind, @immut/hashset.new())
+  let would_capture = if flat_proj.defs.is_empty() {
+    free_names_captured_at_use_sites(g, node, fv)
   } else {
-    flat_proj.defs.length()
+    free_names_would_rebind_at_module_end(g, flat_proj, node, fv)
   }
-  for i = 0; i < scope_limit; i = i + 1 {
-    module_env = module_env.add(flat_proj.defs[i].0)
-  }
-  // Check for lambda-captured vars
-  let lam_env = collect_lam_env(node_id, registry)
-  let fv = free_vars(node.kind, module_env)
-  let mut captured_var : String? = None
-  fv.each(fn(v) { if lam_env.contains(v) { captured_var = Some(v) } })
-  if captured_var is Some(_) {
+  if would_capture {
     return Err("Cannot extract: expression captures lambda-bound variables")
   }
   // Build edits
@@ -150,16 +131,16 @@ fn compute_inline_definition(
     DeclKind::ModuleDef(def_index~) => {
       let def = flat_proj.defs[def_index]
       let init_text = @ast.print_term(def.1.kind)
-      // Check for capture: would any free var in init be captured by a lambda at the usage site?
+      // Check for capture: would any free var in init be captured by an enclosing scope at the usage site?
       let init_fv = free_vars(def.1.kind, @immut/hashset.new())
-      let usage_lam_env = collect_lam_env(node_id, registry)
-      let mut would_capture = false
-      init_fv.each(fn(v) {
-        if usage_lam_env.contains(v) {
-          would_capture = true
-        }
-      })
-      if would_capture {
+      if free_names_would_rebind_at_node(
+          g,
+          flat_proj,
+          source_map,
+          def.1,
+          node_id,
+          init_fv,
+        ) {
         return Err(
           "Cannot inline: would capture variables bound by enclosing lambda",
         )
@@ -257,13 +238,17 @@ fn compute_inline_all_usages(
     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?
+  // Check for capture: would any free var in init be captured by an enclosing scope at any reference site?
   let init_fv = free_vars(def.1.kind, @immut/hashset.new())
   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 ref_lam_env.contains(v) { would_capture = true } })
-    if would_capture {
+    if free_names_would_rebind_at_node(
+        g,
+        flat_proj,
+        source_map,
+        def.1,
+        rid,
+        init_fv,
+      ) {
       return Err(
         "Cannot inline: would capture variables bound by enclosing lambda",
       )

lang/lambda/edits/text_edit_rename.mbt

@@ -187,10 +187,9 @@ fn rename_module_binding(
     Ok(ids) => ids
     Err(e) => return Err(e)
   }
-  // Guard: check if new_name would be captured by an enclosing lambda at any reference site
+  // Guard: check if new_name would be captured by an enclosing scope 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) {
+    if new_name != old_name && name_captured_at_node(g, ref_id, new_name) {
       return Err(
         "Cannot rename: '" +
         new_name +