update scope exit logic

vyper/codegen/context.py

@@ -225,9 +225,18 @@ def mark_for_deallocation(self, varname):
 
     # "mark-and-sweep", haha
     def sweep(self):
+        tmp = set()
         for varname in self._to_deallocate:
-            self.deallocate_variable(varname, self.vars[varname])
-        self._to_deallocate.clear()
+            var = self.vars[varname]
+            for s in self._scopes:
+                if s not in var.blockscopes:
+                    # defer deallocation until we hit the end of its scope
+                    tmp.add(varname)
+                    break
+            else:
+                self.deallocate_variable(varname, self.vars[varname])
+
+        self._to_deallocate = tmp
 
     def _new_variable(
         self,

vyper/codegen/stmt.py

@@ -34,13 +34,14 @@ def __init__(self, node: vy_ast.VyperNode, context: Context) -> None:
             fn = getattr(self, fn_name)
             with context.internal_memory_scope():
                 self.ir_node = fn()
-            context.sweep()
 
             assert isinstance(self.ir_node, IRnode), self.ir_node
 
         self.ir_node.annotation = self.stmt.get("node_source_code")
         self.ir_node.ast_source = self.stmt
 
+        context.sweep()
+
     def parse_Expr(self):
         return Expr(self.stmt.value, self.context, is_stmt=True).ir_node