Repository: leonardt/ast_tools

Issue:
Regression in fault test suite
I'm not sure how important this is, but with the latest release the following example taken from fault doesn't work anymore with the SSA pass

```python

import magma as m
import fault


class MWrapperMeta(m.MagmaProtocolMeta):
    def __getitem__(cls, T):
        assert cls is MWrapper
        return type(cls)(f'MWrapper[{T}]', (cls,), {'_T_': T})

    def _to_magma_(cls):
        return cls._T_

    def _qualify_magma_(cls, d):
        return MWrapper[cls._T_.qualify(d)]

    def _flip_magma_(cls):
        return MWrapper[cls._T_.flip()]

    def _from_magma_value_(cls, value):
        return cls(value)


class MWrapper(m.MagmaProtocol, metaclass=MWrapperMeta):
    def __init__(self, val):
        if not isinstance(val, type(self)._T_):
            raise TypeError()
        self._value_ = val

    def _get_magma_value_(self):
        return self._value_

    def apply(self, f):
        return f(self._value_)


WrappedBits8 = MWrapper[m.UInt[8]]


@m.sequential2()
class Foo:
    def __call__(self, val: WrappedBits8) -> m.UInt[8]:
        return val.apply(lambda x: x + 1)


def test_proto():
    tester = fault.Tester(Foo)
    tester.circuit.val = 1
    tester.eval()
    tester.circuit.O.expect(2)
    tester.compile_and_run("verilator", flags=['-Wno-unused'])
```

It produces this error

```python
../ast_tools/ast_tools/passes/ssa.py:514: in leave_Assign
    assert not self._assigned_names, (to_module(original_node).code, self._assigned_names)
E   AssertionError: ('__0_return_0 = val.apply(lambda x: x + 1)', ['x_0'])
```

Upon investigation, it seems like the SSA pass is treating the `lambda x: x + 1` expression to contain an assignment (perhaps because the `x` inside the lambda parameter is treated as a "store").  Adding support for lambda (not sure whether it worked before or just missed this case) seems doable (we need to treat it as a store, but remove the assigned names after processing the lambda, so we use the new name for the parameters, but the correct SSA names for anything else coming from the outer scope).  Not sure if this is something we need though, so just documenting it here.

CC @cdonovick 

Comments:

Patch:
diff --git a/ast_tools/passes/ssa.py b/ast_tools/passes/ssa.py
index 01d90fa..b37a642 100644
--- a/ast_tools/passes/ssa.py
+++ b/ast_tools/passes/ssa.py
@@ -394,9 +394,7 @@ class SSATransformer(NodeTrackingTransformer):
 
             # Need to visit params to get them to be rebuilt and therfore
             # tracked to build the symbol table
-            self._skip += 1
             update_params = updated_node.params.visit(self)
-            self._skip -= 1
             assert not self._skip
             assert not self._assigned_names, self._assigned_names
             new_body = updated_node.body.visit(self)
@@ -531,6 +529,16 @@ class SSATransformer(NodeTrackingTransformer):
     def leave_Arg_keyword(self, node: cst.Arg):
         self._skip -= 1
 
+    def visit_Parameters(self, node: cst.Parameters) -> tp.Optional[bool]:
+        self._skip += 1
+        return True
+
+    def leave_Parameters(self,
+            original_node: cst.Parameters,
+            updated_node: cst.Parameters) -> cst.Parameters:
+        self._skip -= 1
+        return updated_node
+
     def leave_Name(self,
             original_node: cst.Name,
             updated_node: cst.Name) -> cst.Name:


Test patch:
diff --git a/tests/test_ssa.py b/tests/test_ssa.py
index 4c35f7e..56786c4 100644
--- a/tests/test_ssa.py
+++ b/tests/test_ssa.py
@@ -36,6 +36,7 @@ def _do_ssa(func, strict, **kwargs):
         func = dec(func)
     return func
 
+
 @pytest.mark.parametrize('strict', [True, False])
 @pytest.mark.parametrize('a', template_options)
 @pytest.mark.parametrize('b', template_options)
@@ -457,3 +458,20 @@ def test_call_in_annotations(strict, x, y):
     f1 = exec_def_in_file(tree, env)
     f2 = apply_passes([ssa(strict)])(f1)
 
+
+@pytest.mark.parametrize('strict', [True, False])
+def test_issue_79(strict):
+    class Wrapper:
+        def __init__(self, val):
+            self.val = val
+        def apply(self, f):
+            return f(self.val)
+
+    def f1(x):
+        return x.apply(lambda x: x+1)
+
+    f2 = apply_passes([ssa(strict)])(f1)
+
+    for _ in range(8):
+        x = Wrapper(random.randint(0, 1<<10))
+        assert f1(x) == f2(x)