From 7ce131ee6354aefcf4780186f8063dc154d9aa40 Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Mon, 2 Sep 2024 12:25:01 +0800
Subject: [PATCH] update scope exit logic

---
 vyper/codegen/context.py | 13 +++++++++++--
 vyper/codegen/stmt.py    |  3 ++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/vyper/codegen/context.py b/vyper/codegen/context.py
index 405d888505..6c3b031901 100644
--- a/vyper/codegen/context.py
+++ b/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,
diff --git a/vyper/codegen/stmt.py b/vyper/codegen/stmt.py
index 93e3442bea..ff6af4f98b 100644
--- a/vyper/codegen/stmt.py
+++ b/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