bpo-24340: Fix estimation of the code stack size. (#5076)

2 files changed, 298 insertions, 9 deletions

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index ae3043a..29aa362 100644
--- a/Lib/test/test_compile.py
+++ b/Lib/test/test_compile.py
@@ -672,7 +672,7 @@ if 1:
         compile("42", PathLike("test_compile_pathlike"), "single")
 
 
-class TestStackSize(unittest.TestCase):
+class TestExpressionStackSize(unittest.TestCase):
     # These tests check that the computed stack size for a code object
     # stays within reasonable bounds (see issue #21523 for an example
     # dysfunction).
@@ -710,5 +710,294 @@ class TestStackSize(unittest.TestCase):
         self.check_stack_size(code)
 
 
+class TestStackSizeStability(unittest.TestCase):
+    # Check that repeating certain snippets doesn't increase the stack size
+    # beyond what a single snippet requires.
+
+    def check_stack_size(self, snippet, async_=False):
+        def compile_snippet(i):
+            ns = {}
+            script = """def func():\n""" + i * snippet
+            if async_:
+                script = "async " + script
+            code = compile(script, "<script>", "exec")
+            exec(code, ns, ns)
+            return ns['func'].__code__
+
+        sizes = [compile_snippet(i).co_stacksize for i in range(2, 5)]
+        if len(set(sizes)) != 1:
+            import dis, io
+            out = io.StringIO()
+            dis.dis(compile_snippet(1), file=out)
+            self.fail("stack sizes diverge with # of consecutive snippets: "
+                      "%s\n%s\n%s" % (sizes, snippet, out.getvalue()))
+
+    def test_if(self):
+        snippet = """
+            if x:
+                a
+            """
+        self.check_stack_size(snippet)
+
+    def test_if_else(self):
+        snippet = """
+            if x:
+                a
+            elif y:
+                b
+            else:
+                c
+            """
+        self.check_stack_size(snippet)
+
+    def test_try_except_bare(self):
+        snippet = """
+            try:
+                a
+            except:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_try_except_qualified(self):
+        snippet = """
+            try:
+                a
+            except ImportError:
+                b
+            except:
+                c
+            else:
+                d
+            """
+        self.check_stack_size(snippet)
+
+    def test_try_except_as(self):
+        snippet = """
+            try:
+                a
+            except ImportError as e:
+                b
+            except:
+                c
+            else:
+                d
+            """
+        self.check_stack_size(snippet)
+
+    def test_try_finally(self):
+        snippet = """
+                try:
+                    a
+                finally:
+                    b
+            """
+        self.check_stack_size(snippet)
+
+    def test_with(self):
+        snippet = """
+            with x as y:
+                a
+            """
+        self.check_stack_size(snippet)
+
+    def test_while_else(self):
+        snippet = """
+            while x:
+                a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for(self):
+        snippet = """
+            for x in y:
+                a
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_else(self):
+        snippet = """
+            for x in y:
+                a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_break_continue(self):
+        snippet = """
+            for x in y:
+                if z:
+                    break
+                elif u:
+                    continue
+                else:
+                    a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_break_continue_inside_try_finally_block(self):
+        snippet = """
+            for x in y:
+                try:
+                    if z:
+                        break
+                    elif u:
+                        continue
+                    else:
+                        a
+                finally:
+                    f
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_break_inside_finally_block(self):
+        snippet = """
+            for x in y:
+                try:
+                    t
+                finally:
+                    if z:
+                        break
+                    else:
+                        a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_break_continue_inside_except_block(self):
+        snippet = """
+            for x in y:
+                try:
+                    t
+                except:
+                    if z:
+                        break
+                    elif u:
+                        continue
+                    else:
+                        a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_for_break_continue_inside_with_block(self):
+        snippet = """
+            for x in y:
+                with c:
+                    if z:
+                        break
+                    elif u:
+                        continue
+                    else:
+                        a
+            else:
+                b
+            """
+        self.check_stack_size(snippet)
+
+    def test_return_inside_try_finally_block(self):
+        snippet = """
+            try:
+                if z:
+                    return
+                else:
+                    a
+            finally:
+                f
+            """
+        self.check_stack_size(snippet)
+
+    def test_return_inside_finally_block(self):
+        snippet = """
+            try:
+                t
+            finally:
+                if z:
+                    return
+                else:
+                    a
+            """
+        self.check_stack_size(snippet)
+
+    def test_return_inside_except_block(self):
+        snippet = """
+            try:
+                t
+            except:
+                if z:
+                    return
+                else:
+                    a
+            """
+        self.check_stack_size(snippet)
+
+    def test_return_inside_with_block(self):
+        snippet = """
+            with c:
+                if z:
+                    return
+                else:
+                    a
+            """
+        self.check_stack_size(snippet)
+
+    def test_async_with(self):
+        snippet = """
+            async with x as y:
+                a
+            """
+        self.check_stack_size(snippet, async_=True)
+
+    def test_async_for(self):
+        snippet = """
+            async for x in y:
+                a
+            """
+        self.check_stack_size(snippet, async_=True)
+
+    def test_async_for_else(self):
+        snippet = """
+            async for x in y:
+                a
+            else:
+                b
+            """
+        self.check_stack_size(snippet, async_=True)
+
+    def test_for_break_continue_inside_async_with_block(self):
+        snippet = """
+            for x in y:
+                async with c:
+                    if z:
+                        break
+                    elif u:
+                        continue
+                    else:
+                        a
+            else:
+                b
+            """
+        self.check_stack_size(snippet, async_=True)
+
+    def test_return_inside_async_with_block(self):
+        snippet = """
+            async with c:
+                if z:
+                    return
+                else:
+                    a
+            """
+        self.check_stack_size(snippet, async_=True)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index bfbbee2..590f041 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -284,18 +284,18 @@ dis_traceback = """\
              22 POP_TOP
              24 STORE_FAST               0 (e)
              26 POP_TOP
-             28 SETUP_FINALLY           12 (to 42)
+             28 SETUP_FINALLY           10 (to 40)
 
 %3d          30 LOAD_FAST                0 (e)
              32 LOAD_ATTR                1 (__traceback__)
              34 STORE_FAST               1 (tb)
              36 POP_BLOCK
-             38 POP_EXCEPT
-             40 LOAD_CONST               0 (None)
-        >>   42 LOAD_CONST               0 (None)
-             44 STORE_FAST               0 (e)
-             46 DELETE_FAST              0 (e)
-             48 END_FINALLY
+             38 LOAD_CONST               0 (None)
+        >>   40 LOAD_CONST               0 (None)
+             42 STORE_FAST               0 (e)
+             44 DELETE_FAST              0 (e)
+             46 END_FINALLY
+             48 POP_EXCEPT
              50 JUMP_FORWARD             2 (to 54)
         >>   52 END_FINALLY
 
@@ -741,7 +741,7 @@ Filename:          (.*)
 Argument count:    0
 Kw-only arguments: 0
 Number of locals:  2
-Stack size:        17
+Stack size:        10
 Flags:             OPTIMIZED, NEWLOCALS, NOFREE, COROUTINE
 Constants:
    0: None