Implement simple do-while loop functionality

src/flow_graph.py

@@ -915,6 +915,9 @@ class NaturalLoop:
     nodes: Set[Node] = attr.ib(factory=set)
     backedges: Set[Node] = attr.ib(factory=set)
 
+    def is_self_loop(self) -> bool:
+        return len(self.nodes) == 1
+
 
 def build_graph_from_block(
     block: Block, blocks: List[Block], nodes: List[Node], asm_data: AsmData

src/if_statements.py

@@ -151,11 +151,34 @@ def format(self, fmt: Formatter) -> str:
         return "\n".join(lines)
 
 
+@attr.s
+class DoWhileLoop:
+    body: "Body" = attr.ib()
+    condition: Optional[Condition] = attr.ib(default=None)
+
+    def should_write(self) -> bool:
+        return True
+
+    def format(self, fmt: Formatter) -> str:
+        space = fmt.indent("")
+        after_do = f"\n{space}" if fmt.coding_style.newline_after_if else " "
+        cond = format_expr(self.condition, fmt) if self.condition else ""
+        with fmt.indented():
+            return "\n".join(
+                [
+                    f"{space}do{after_do}{{",
+                    self.body.format(fmt),
+                    f"{space}}} while ({cond});",
+                ]
+            )
+
+
 Statement = Union[
     SimpleStatement,
     IfElseStatement,
     LabelStatement,
     SwitchStatement,
+    DoWhileLoop,
 ]
 
 
@@ -200,6 +223,9 @@ def add_if_else(self, if_else: IfElseStatement) -> None:
 
         self.statements.append(if_else)
 
+    def add_do_while_loop(self, do_while_loop: DoWhileLoop) -> None:
+        self.statements.append(do_while_loop)
+
     def add_switch(self, switch: SwitchStatement) -> None:
         self.add_statement(switch)
 
@@ -265,7 +291,7 @@ def label_for_node(context: Context, node: Node) -> str:
         return f"block_{node.block.index}"
 
 
-def emit_node(context: Context, node: Node, body: Body) -> bool:
+def emit_node(context: Context, node: Node, body: Body, secretly: bool = False) -> bool:
     """
     Try to emit a node for the first time, together with a label for it.
     The label is only printed if something jumps to it, e.g. a loop.
@@ -291,7 +317,8 @@ def emit_node(context: Context, node: Node, body: Body) -> bool:
             )
     else:
         body.add_statement(LabelStatement(context, node))
-        context.emitted_nodes.add(node)
+        if not secretly:
+            context.emitted_nodes.add(node)
 
     body.add_node(node, comment_empty=True)
     if isinstance(node, ReturnNode):
@@ -601,7 +628,54 @@ def build_switch_between(
     return SwitchStatement(jump, body)
 
 
-def build_flowgraph_between(context: Context, start: Node, end: Node) -> Body:
+def detect_loop(
+    context: Context, start: ConditionalNode, end: Node
+) -> Optional[DoWhileLoop]:
+    # We detect edges that are accompanied by their reverse as loops.
+    # breakpoint()
+    imm_pdom: Node
+    if start.loop and start.loop.is_self_loop():
+        imm_pdom = start
+    else:
+        # imm_pdom = immediate_postdominator(context, start, end)
+        assert start.immediate_postdominator
+        imm_pdom = start.immediate_postdominator
+
+    if not (
+        isinstance(imm_pdom, ConditionalNode) and imm_pdom.conditional_edge is start
+    ):
+        return None
+
+    loop_body = Body(False, [])
+    emit_node(
+        context,
+        start,
+        loop_body,
+        secretly=(bool(start.loop and not start.loop.is_self_loop())),
+    )
+
+    if not start.loop or not start.loop.is_self_loop():
+        # There are more nodes to emit, "between" the start node
+        # and the loop edge that it connects to:
+        loop_body = build_flowgraph_between(
+            context,
+            start,
+            imm_pdom,
+            skip_loop_detection=True,
+        )
+        emit_node(context, imm_pdom, loop_body)
+
+    assert imm_pdom.block.block_info
+    assert imm_pdom.block.block_info.branch_condition
+    return DoWhileLoop(
+        loop_body,
+        imm_pdom.block.block_info.branch_condition,
+    )
+
+
+def build_flowgraph_between(
+    context: Context, start: Node, end: Node, skip_loop_detection: bool = False
+) -> Body:
     """
     Output a section of a flow graph that has already been translated to our
     symbolic AST. All nodes between start and end, including start but NOT end,
@@ -617,6 +691,22 @@ def build_flowgraph_between(context: Context, start: Node, end: Node) -> Body:
     while curr_start != end:
         assert not isinstance(curr_start, TerminalNode)
 
+        if not skip_loop_detection and isinstance(curr_start, ConditionalNode):
+            # breakpoint()
+            do_while_loop = detect_loop(context, curr_start, end)
+            if do_while_loop:
+                # breakpoint()
+                body.add_do_while_loop(do_while_loop)
+                # Move past the loop:
+                if curr_start.loop and curr_start.loop.is_self_loop():
+                    curr_start = curr_start.fallthrough_edge
+                else:
+                    # imm_pdom = immediate_postdominator(context, curr_start, end)
+                    imm_pdom = curr_start.immediate_postdominator
+                    assert isinstance(imm_pdom, ConditionalNode)
+                    curr_start = imm_pdom.fallthrough_edge
+                continue
+
         # Write the current node, or a goto, to the body
         if not emit_node(context, curr_start, body):
             # If the node was already witten, emit_node will use a goto

tests/end_to_end/andor_assignment/irix-g-out.c

@@ -22,24 +22,20 @@ s32 test(s32 arg0, s32 arg1, s32 arg2, s32 arg3) {
         sp20 = func_00400090(temp_t6);
         if ((sp20 != 0) && (arg3 != 0)) {
             if (sp1C < 5) {
-loop_13:
-                sp1C += 1;
-                sp1C *= 2;
-                if (sp1C < 5) {
-                    goto loop_13;
-                }
+                do {
+                    sp1C += 1;
+                    sp1C *= 2;
+                } while ((sp1C < 5) != 0);
             }
             sp1C += 5;
         }
     }
     if ((sp24 != 0) && (sp20 != 0) && (temp_t9 = sp24 + sp20, sp24 = temp_t9, sp20 = func_00400090(temp_t9), (sp20 != 0)) && (arg3 != 0)) {
         if (sp1C < 5) {
-loop_20:
-            sp1C += 1;
-            sp1C *= 2;
-            if (sp1C < 5) {
-                goto loop_20;
-            }
+            do {
+                sp1C += 1;
+                sp1C *= 2;
+            } while ((sp1C < 5) != 0);
         }
         sp1C += 5;
     } else {

tests/end_to_end/andor_assignment/irix-o2-noandor-out.c

@@ -81,13 +81,11 @@ s32 test(s32 arg0, s32 arg1, s32 arg2, s32 arg3) {
                     phi_v1_2 = temp_v1;
                     phi_v1_6 = temp_v1;
                     if (temp_v1 < 5) {
-loop_12:
-                        temp_t5 = (phi_v1_2 + 1) * 2;
-                        phi_v1_2 = temp_t5;
-                        phi_v1_6 = temp_t5;
-                        if (temp_t5 < 5) {
-                            goto loop_12;
-                        }
+                        do {
+                            temp_t5 = (phi_v1_2 + 1) * 2;
+                            phi_v1_2 = temp_t5;
+                            phi_v1_6 = temp_t5;
+                        } while ((temp_t5 < 5) != 0);
                     }
                     phi_t1_2 = temp_v0_2;
                     phi_v1_3 = phi_v1_6 + 5;
@@ -105,13 +103,11 @@ s32 test(s32 arg0, s32 arg1, s32 arg2, s32 arg3) {
                     phi_v1_4 = phi_v1_3;
                     phi_v1_7 = phi_v1_3;
                     if (phi_v1_3 < 5) {
-loop_19:
-                        temp_t9 = (phi_v1_4 + 1) * 2;
-                        phi_v1_4 = temp_t9;
-                        phi_v1_7 = temp_t9;
-                        if (temp_t9 < 5) {
-                            goto loop_19;
-                        }
+                        do {
+                            temp_t9 = (phi_v1_4 + 1) * 2;
+                            phi_v1_4 = temp_t9;
+                            phi_v1_7 = temp_t9;
+                        } while ((temp_t9 < 5) != 0);
                     }
                     phi_v1_5 = phi_v1_7 + 5;
                 } else {

tests/end_to_end/andor_assignment/irix-o2-out.c

@@ -64,13 +64,11 @@ s32 test(s32 arg0, s32 arg1, s32 arg2, s32 arg3) {
                     phi_v1_2 = temp_v1;
                     phi_v1_6 = temp_v1;
                     if (temp_v1 < 5) {
-loop_12:
-                        temp_t3 = (phi_v1_2 + 1) * 2;
-                        phi_v1_2 = temp_t3;
-                        phi_v1_6 = temp_t3;
-                        if (temp_t3 < 5) {
-                            goto loop_12;
-                        }
+                        do {
+                            temp_t3 = (phi_v1_2 + 1) * 2;
+                            phi_v1_2 = temp_t3;
+                            phi_v1_6 = temp_t3;
+                        } while ((temp_t3 < 5) != 0);
                     }
                     phi_s0_2 = temp_s0_2;
                     phi_t0_2 = temp_v0_2;
@@ -83,13 +81,11 @@ s32 test(s32 arg0, s32 arg1, s32 arg2, s32 arg3) {
         phi_v1_4 = phi_v1_3;
         phi_v1_7 = phi_v1_3;
         if (phi_v1_3 < 5) {
-loop_19:
-            temp_t5 = (phi_v1_4 + 1) * 2;
-            phi_v1_4 = temp_t5;
-            phi_v1_7 = temp_t5;
-            if (temp_t5 < 5) {
-                goto loop_19;
-            }
+            do {
+                temp_t5 = (phi_v1_4 + 1) * 2;
+                phi_v1_4 = temp_t5;
+                phi_v1_7 = temp_t5;
+            } while ((temp_t5 < 5) != 0);
         }
         phi_v1_5 = phi_v1_7 + 5;
     } else {

tests/end_to_end/loop/irix-g-out.c

@@ -4,12 +4,10 @@ void test(s32 arg0, s32 arg1) {
 
     sp4 = 0;
     if (arg1 > 0) {
-loop_1:
-        *(arg0 + sp4) = (u8)0;
-        temp_t9 = sp4 + 1;
-        sp4 = temp_t9;
-        if (temp_t9 < arg1) {
-            goto loop_1;
-        }
+        do {
+            *(arg0 + sp4) = (u8)0;
+            temp_t9 = sp4 + 1;
+            sp4 = temp_t9;
+        } while ((temp_t9 < arg1) != 0);
     }
 }

tests/end_to_end/loop/irix-o2-out.c

@@ -17,33 +17,29 @@ s32 test(s8 *arg0, s32 arg1) {
         if (temp_a3 != 0) {
             phi_v1 = arg0;
             phi_v0 = 0;
-loop_3:
-            temp_v0 = phi_v0 + 1;
-            *phi_v1 = (u8)0;
-            phi_v1 += 1;
-            phi_v0 = temp_v0;
-            if (temp_a3 != temp_v0) {
-                goto loop_3;
-            }
+            do {
+                temp_v0 = phi_v0 + 1;
+                *phi_v1 = (u8)0;
+                phi_v1 += 1;
+                phi_v0 = temp_v0;
+            } while (temp_a3 != temp_v0);
             phi_return = temp_v0;
             phi_v0_3 = temp_v0;
             if (temp_v0 != arg1) {
 block_5:
                 phi_v1_2 = arg0 + phi_v0_3;
                 phi_v0_2 = phi_v0_3;
-loop_6:
-                temp_v0_2 = phi_v0_2 + 4;
-                phi_v1_2->unk1 = (u8)0;
-                phi_v1_2->unk2 = (u8)0;
-                phi_v1_2->unk3 = (u8)0;
-                temp_v1 = phi_v1_2 + 4;
-                temp_v1->unk-4 = (u8)0;
-                phi_v1_2 = temp_v1;
-                phi_v0_2 = temp_v0_2;
-                phi_return = temp_v0_2;
-                if (temp_v0_2 != arg1) {
-                    goto loop_6;
-                }
+                do {
+                    temp_v0_2 = phi_v0_2 + 4;
+                    phi_v1_2->unk1 = (u8)0;
+                    phi_v1_2->unk2 = (u8)0;
+                    phi_v1_2->unk3 = (u8)0;
+                    temp_v1 = phi_v1_2 + 4;
+                    temp_v1->unk-4 = (u8)0;
+                    phi_v1_2 = temp_v1;
+                    phi_v0_2 = temp_v0_2;
+                    phi_return = temp_v0_2;
+                } while (temp_v0_2 != arg1);
             }
         } else {
             goto block_5;

tests/end_to_end/loop_nested/irix-g-out.c

@@ -6,20 +6,16 @@ s32 test(s32 arg0) {
     spC = 0;
     sp8 = 0;
     if (spC < arg0) {
-loop_1:
-        sp4 = 0;
-        if (sp4 < arg0) {
-loop_2:
-            sp8 += spC * sp4;
-            sp4 += 1;
+        do {
+            sp4 = 0;
             if (sp4 < arg0) {
-                goto loop_2;
+                do {
+                    sp8 += spC * sp4;
+                    sp4 += 1;
+                } while ((sp4 < arg0) != 0);
             }
-        }
-        spC += 1;
-        if (spC < arg0) {
-            goto loop_1;
-        }
+            spC += 1;
+        } while ((spC < arg0) != 0);
     }
     return sp8;
 }