fix: Don't inline TEAL blocks with a single predecessor when their la…

…bel is referenced multiple times. This issue could occur when there is a default fall through block in a switch or match scenario, which is also targeted by said op, and would cause a compilation failure.
algorandfoundation · Jan 23, 2025 · d193042 · d193042
1 parent d50fd88
commit d193042
Show file tree

Hide file tree

Showing 7 changed files with 547 additions and 39 deletions.
diff --git a/examples/sizes.txt b/examples/sizes.txt
@@ -102,7 +102,7 @@
  nested_loops/Nested                            213     200       -    |        132     120       - 
  regression_tests/Issue118                      136     102      96    |         75      52      49 
  regression_tests/Issue194                       31       8       -    |         21       4       - 
- regression_tests/TealSwitchInlining             27      22    None    |         14      10    None 
+ regression_tests/TealSwitchInlining             27      22      19    |         14      10       8 
  reinterpret_cast                               137      83      76    |         70      38      35 
  scratch_slots                                   94      84       -    |         56      43       - 
  scratch_slots/2                                 94      84       -    |         56      43       - 
@@ -139,4 +139,4 @@
  unssa/UnSSA                                    420     266       -    |        237     153       - 
  voting/VotingRoundApp                         1584    1426    1415    |        725     624     625 
  with_reentrancy/WithReentrancy                 245     214       -    |        126     108       - 
- Total                                        70547   39139   36225    |      33310   18259   16976 
+ Total                                        70547   39139   36244    |      33310   18259   16984 
diff --git a/src/puya/teal/optimize/main.py b/src/puya/teal/optimize/main.py
@@ -174,13 +174,13 @@ def _remove_unreachable_blocks(teal_sub: models.TealSubroutine) -> None:
 
 
 def _inline_singly_referenced_blocks(teal_sub: models.TealSubroutine) -> None:
-    predecessors = defaultdict[str, set[str]](set)
-    predecessors[teal_sub.blocks[0].label].add("<entrypoint>")
+    predecessors = defaultdict[str, list[str]](list)
+    predecessors[teal_sub.blocks[0].label].append("<entrypoint>")
     for block in teal_sub.blocks:
         for op in block.ops:
             if isinstance(op, models.ControlOp):
                 for target_label in op.targets:
-                    predecessors[target_label].add(block.label)
+                    predecessors[target_label].append(block.label)
 
     pairs = dict[str, str]()
     inlineable = set[str]()
@@ -189,7 +189,7 @@ def _inline_singly_referenced_blocks(teal_sub: models.TealSubroutine) -> None:
             case models.Branch(target=target_label):
                 target_predecessors = predecessors[target_label]
                 if len(target_predecessors) == 1:
-                    assert target_predecessors == {block.label}
+                    assert target_predecessors == [block.label]
                     pairs[block.label] = target_label
                     inlineable.add(target_label)
 

diff --git a/test_cases/regression_tests/out_O2/TealSwitchInlining.approval.puya.map b/test_cases/regression_tests/out_O2/TealSwitchInlining.approval.puya.map
@@ -0,0 +1,57 @@
+{
+  "version": 3,
+  "sources": [
+    "../teal_switch_inlining.py"
+  ],
+  "mappings": ";AAUc;;AAAN;;;;;;AAMe;;AAAP;AAEO;;AAAP",
+  "op_pc_offset": 0,
+  "pc_events": {
+    "1": {
+      "subroutine": "test_cases.regression_tests.teal_switch_inlining.TealSwitchInlining.approval_program",
+      "params": {},
+      "block": "main",
+      "stack_in": [],
+      "op": "txn NumAppArgs",
+      "defined_out": [
+        "tmp%0#0"
+      ],
+      "stack_out": [
+        "tmp%0#0"
+      ]
+    },
+    "3": {
+      "op": "switch main_switch_case_0@1 main_switch_case_1@2",
+      "stack_out": []
+    },
+    "9": {
+      "block": "main_switch_case_0@1",
+      "stack_in": [],
+      "op": "pushint 1 // 1",
+      "defined_out": [
+        "1"
+      ],
+      "stack_out": [
+        "1"
+      ]
+    },
+    "11": {
+      "op": "return",
+      "stack_out": []
+    },
+    "12": {
+      "block": "main_switch_case_1@2",
+      "stack_in": [],
+      "op": "pushint 0 // 0",
+      "defined_out": [
+        "0"
+      ],
+      "stack_out": [
+        "0"
+      ]
+    },
+    "14": {
+      "op": "return",
+      "stack_out": []
+    }
+  }
+}
diff --git a/test_cases/regression_tests/out_O2/TealSwitchInlining.approval.teal b/test_cases/regression_tests/out_O2/TealSwitchInlining.approval.teal
@@ -0,0 +1,15 @@
+#pragma version 10
+#pragma typetrack false
+
+// test_cases.regression_tests.teal_switch_inlining.TealSwitchInlining.approval_program() -> uint64:
+main:
+    txn NumAppArgs
+    switch main_switch_case_0@1 main_switch_case_1@2
+
+main_switch_case_0@1:
+    pushint 1 // 1
+    return
+
+main_switch_case_1@2:
+    pushint 0 // 0
+    return
diff --git a/test_cases/regression_tests/out_O2/TealSwitchInlining.clear.puya.map b/test_cases/regression_tests/out_O2/TealSwitchInlining.clear.puya.map
@@ -0,0 +1,27 @@
+{
+  "version": 3,
+  "sources": [
+    "../teal_switch_inlining.py"
+  ],
+  "mappings": ";AAuBe;;AAAP",
+  "op_pc_offset": 0,
+  "pc_events": {
+    "1": {
+      "subroutine": "test_cases.regression_tests.teal_switch_inlining.TealSwitchInlining.clear_state_program",
+      "params": {},
+      "block": "main",
+      "stack_in": [],
+      "op": "pushint 1 // 1",
+      "defined_out": [
+        "1"
+      ],
+      "stack_out": [
+        "1"
+      ]
+    },
+    "3": {
+      "op": "return",
+      "stack_out": []
+    }
+  }
+}
diff --git a/test_cases/regression_tests/out_O2/TealSwitchInlining.clear.teal b/test_cases/regression_tests/out_O2/TealSwitchInlining.clear.teal
@@ -0,0 +1,7 @@
+#pragma version 10
+#pragma typetrack false
+
+// test_cases.regression_tests.teal_switch_inlining.TealSwitchInlining.clear_state_program() -> uint64:
+main:
+    pushint 1 // 1
+    return