correctly rename bound variables that appear as dict keys (#297)

Author: jfeser

effectful/ops/syntax.py

@@ -901,6 +901,18 @@ def defterm(__dispatch: Callable[[type], Callable[[T], Expr[T]]], value: T):
         return __dispatch(type(value))(value)
 
 
+def _map_structure_and_keys(func, structure):
+    def _map_value(value):
+        if isinstance(value, dict):
+            return {func(k): v for k, v in value.items()}
+        elif not tree.is_nested(value):
+            return func(value)
+        else:
+            return value
+
+    return tree.traverse(_map_value, structure, top_down=False)
+
+
 @_CustomSingleDispatchCallable
 def defdata(
     __dispatch: Callable[[type], Callable[..., Expr[T]]],
@@ -977,7 +989,7 @@ def _(op, *args, **kwargs):
         *{k: (v, kwarg_ctxs[k]) for k, v in kwargs.items()}.items(),
     ):
         if c:
-            v = tree.map_structure(
+            v = _map_structure_and_keys(
                 lambda a: renaming.get(a, a) if isinstance(a, Operation) else a, v
             )
             res = evaluate(

tests/test_ops_syntax.py

@@ -10,6 +10,7 @@
 from effectful.ops.syntax import (
     Scoped,
     _CustomSingleDispatchCallable,
+    _map_structure_and_keys,
     deffn,
     defop,
     defstream,
@@ -112,6 +113,13 @@ def f(x):
     assert f_op != ff_op
 
 
+def test_map_structure_and_keys():
+    s = {1: 2, 3: [4, 5, (6, {7: 8})]}
+    expected = {2: 3, 4: [5, 6, (7, {8: 9})]}
+    actual = _map_structure_and_keys(lambda x: x + 1, s)
+    assert actual == expected
+
+
 def test_scoped_collections():
     """Test that Scoped annotations work with tree-structured collections containing Operations."""
 
@@ -135,6 +143,11 @@ def let_many(
     term = let_many(bindings, body)
     free_vars = fvsof(term)
 
+    new_x = list(term.args[0].keys())[0]
+    new_y = list(term.args[0].keys())[1]
+    assert new_x == term.args[1].args[0].args[0].op and new_x != x
+    assert new_y == term.args[1].args[0].args[1].op and new_y != y
+
     assert x not in free_vars
     assert y not in free_vars
     assert z in free_vars