Repository: marimo-team/marimo
Pull request: fix: per-block reference resolution
For each instance of a referenced name, we need to track the blocks it appears in;
previously, references on new blocks (scopes) were evicting
references on other scopes. These newer references could be
resolved but the other ones weren't, leading to us drop references.

Fixes #4742 

Problem statement:
[ast] missing ref with nested scopes
### Describe the bug

The following cell should have `x` as a ref, but doesn't.

```python
def f():
    print(x)
    def g():
        def h():
            x
        # oops, this will evict x as a ref
        x = 0
```

### Environment

<details>

```
0.13.2
```

</details>


### Code to reproduce

_No response_

Hints from issue discussion:

Commit message: fix: reference resolution at block granularity

For each name, we need to store its references in each block;
previously, references on new blocks (scopes) were evicting
references on other scopes. These newer references could be
resolved but the other ones weren't, leading to us missing references.

Patch to review:
diff --git a/marimo/_ast/visitor.py b/marimo/_ast/visitor.py
index 5ba929f4d17..b04b81fbaf5 100644
--- a/marimo/_ast/visitor.py
+++ b/marimo/_ast/visitor.py
@@ -124,6 +124,8 @@ class RefData:
 
     # Whether the ref was deleted
     deleted: bool
+    # Block in which this ref was referenced
+    block: Block
     # Ancestors of the block in which this ref was used
     parent_blocks: list[Block]
 
@@ -157,7 +159,10 @@ def __init__(
         self.ref_stack: list[set[Name]] = [set()]
         self.obscured_scope_stack: list[ObscuredScope] = []
         # Mapping from referenced names to their metadata
-        self._refs: dict[Name, RefData] = {}
+        # Each referenced name may be defined multiple blocks;
+        # ie, self._refs maps each name to RefData for each block it was
+        # defined in
+        self._refs: dict[Name, list[RefData]] = {}
         # Function (node, name, block stack) -> None
         self._on_def = on_def if on_def is not None else lambda *_: None
         # Function (node) -> None
@@ -189,7 +194,11 @@ def refs(self) -> set[Name]:
     @property
     def deleted_refs(self) -> set[Name]:
         """Referenced names that were deleted with `del`."""
-        return set(name for name in self._refs if self._refs[name].deleted)
+        return set(
+            name
+            for name in self._refs
+            if any(ref.deleted for ref in self._refs[name])
+        )
 
     def _if_local_then_mangle(
         self, name: str, ignore_scope: bool = False
@@ -242,25 +251,45 @@ def _add_ref(
         self, node: NamedNode | None, name: Name, deleted: bool
     ) -> None:
         """Register a referenced name."""
-        self._refs[name] = RefData(
-            deleted=deleted,
-            parent_blocks=self.block_stack[:-1],
-        )
+        if name not in self._refs:
+            self._refs[name] = []
+
+        # Register the ref if it doesn't already exist
+        current_block = self.block_stack[-1]
+        parents = self.block_stack[:-1]
+        if all(ref.block != current_block for ref in self._refs[name]):
+            # The reference does not yet exist in the current block, so
+            # we add it.
+            self._refs[name].append(
+                RefData(
+                    deleted=deleted,
+                    parent_blocks=parents,
+                    block=current_block,
+                )
+            )
+
         self.ref_stack[-1].add(name)
         if node is not None:
             self._on_ref(node)
 
-    def _remove_ref(self, name: Name) -> None:
-        """Remove a referenced name."""
-        del self._refs[name]
+    def _remove_ref(self, name: Name, block: Block) -> None:
+        """Remove references of name with block among its parents."""
+        refs = []
+        for ref in self._refs[name]:
+            if block not in ref.parent_blocks:
+                refs.append(ref)
+        self._refs[name] = refs
+        if not self._refs[name]:
+            del self._refs[name]
 
     def _define_in_block(
         self, name: Name, variable_data: VariableData, block_idx: int
     ) -> None:
         """Define a name in a given block."""
 
-        self.block_stack[block_idx].defs.add(name)
-        self.block_stack[block_idx].variable_data[name].append(variable_data)
+        block = self.block_stack[block_idx]
+        block.defs.add(name)
+        block.variable_data[name].append(variable_data)
         # If `name` is added to the top-level block, it is also evicted from
         # any captured refs (if present) --- this handles cases where a name is
         # encountered and captured before it is declared, such as in
@@ -270,12 +299,11 @@ def _define_in_block(
         #   print(x)
         # x = 0
         # ```
-        if (
-            name in self._refs
-            and self.block_stack[block_idx] in self._refs[name].parent_blocks
+        if name in self._refs and any(
+            block in ref.parent_blocks for ref in self._refs[name]
         ):
             # `name` was used as a capture, not a reference
-            self._remove_ref(name)
+            self._remove_ref(name, block)
 
     def _define(
         self, node: NamedNode | None, name: Name, variable_data: VariableData
diff --git a/tests/_ast/test_visitor.py b/tests/_ast/test_visitor.py
index fb6900d4989..308e3fddbf6 100644
--- a/tests/_ast/test_visitor.py
+++ b/tests/_ast/test_visitor.py
@@ -954,6 +954,24 @@ def test_private_ref_requirement_caught() -> None:
     }
 
 
+def test_outer_ref_not_resolved_by_inner_resolution() -> None:
+    code = cleandoc(
+        """
+        def f():
+            x
+            def g():
+                def h():
+                    x
+                x = 0
+        """
+    )
+    v = visitor.ScopedVisitor()
+    mod = ast.parse(code)
+    v.visit(mod)
+    assert v.defs == {"f"}
+    assert v.refs == {"x"}
+
+
 HAS_DEPS = DependencyManager.duckdb.has()
 
 


Code review comments:
diff_hunk: @@ -189,7 +194,11 @@ def refs(self) -> set[Name]:
     @property
     def deleted_refs(self) -> set[Name]:
         """Referenced names that were deleted with `del`."""
-        return set(name for name in self._refs if self._refs[name].deleted)
+        return set(
+            name
+            for name in self._refs
+            if any(ref.deleted for ref in self._refs[name])
path: marimo/_ast/visitor.py
text: @user1:
Hmm, I think maybe only the top block?

```python
x = 0
def fn():
    if False:
        z = x
        del x
x
```

@user1:
I think only possible with global keyword

@author:
I think looking at all the blocks is right, because the `_refs` data structure keeps one `RefData` per name per block.

In your example, `x` is not added as a ref, so there's no issue.

I did however find a related bug, and just pushed up a fix and a couple of unit tests.