Detect and reroll unrolled do-while loops

src/flow_graph.py

@@ -785,11 +785,19 @@ def name(self) -> str:
     def children(self) -> List["Node"]:
         ...
 
+    @abc.abstractmethod
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        ...
+
 
 @attr.s(eq=False)
 class BasicNode(BaseNode):
     successor: "Node" = attr.ib()
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        if self.successor is replace_this:
+            self.successor = with_this
+
     def children(self) -> List["Node"]:
         # TODO: Should we also include the fallthrough node if `emit_goto` is True?
         return [self.successor]
@@ -810,6 +818,12 @@ class ConditionalNode(BaseNode):
     conditional_edge: "Node" = attr.ib()
     fallthrough_edge: "Node" = attr.ib()
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        if self.conditional_edge is replace_this:
+            self.conditional_edge = with_this
+        if self.fallthrough_edge is replace_this:
+            self.fallthrough_edge = with_this
+
     def children(self) -> List["Node"]:
         if self.conditional_edge == self.fallthrough_edge:
             return [self.conditional_edge]
@@ -837,6 +851,9 @@ class ReturnNode(BaseNode):
     def children(self) -> List["Node"]:
         return [self.terminal]
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        return None
+
     def name(self) -> str:
         name = super().name()
         return name if self.is_real() else f"{name}.{self.index}"
@@ -858,6 +875,11 @@ def __str__(self) -> str:
 class SwitchNode(BaseNode):
     cases: List["Node"] = attr.ib()
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        for i, case in enumerate(self.cases):
+            if case is replace_this:
+                self.cases[i] = with_this
+
     def children(self) -> List["Node"]:
         # Deduplicate nodes in `self.cases`
         seen = set()
@@ -891,6 +913,9 @@ def terminal() -> "TerminalNode":
     def children(self) -> List["Node"]:
         return []
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        return None
+
     def __str__(self) -> str:
         return "return"
 
@@ -1322,7 +1347,10 @@ def is_reducible(self) -> bool:
             for s in n.children():
                 if s.loop is not None and n in s.loop.backedges:
                     continue
-                incoming_edges[s].remove(n)
+                try:
+                    incoming_edges[s].remove(n)
+                except KeyError:
+                    break
                 if not incoming_edges[s]:
                     queue.add(s)
 

src/loop_rerolling.py

@@ -0,0 +1,138 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+from .flow_graph import (
+    BasicNode,
+    ConditionalNode,
+    FlowGraph,
+    Node,
+    compute_relations,
+)
+from .parse_instruction import Instruction
+
+
+def replace_node_references(
+    flow_graph: FlowGraph, replace_this: Node, with_this: Node
+) -> None:
+    for node_to_modify in flow_graph.nodes:
+        node_to_modify.replace_any_children(replace_this, with_this)
+
+
+def remove_node(flow_graph: FlowGraph, to_delete: Node, new_child: Node) -> None:
+    flow_graph.nodes.remove(to_delete)
+    replace_node_references(flow_graph, to_delete, new_child)
+
+
+def replace_node(flow_graph: FlowGraph, replace_this: Node, with_this: Node) -> None:
+    replacement_index = flow_graph.nodes.index(replace_this)
+    flow_graph.nodes[replacement_index] = with_this
+    replace_node_references(flow_graph, replace_this, with_this)
+
+
+PatternGraph = Dict[int, Union[int, Tuple[int, int]]]
+
+IDO_O2_SIMPLE_LOOP: PatternGraph = {
+    0: (1, 7),
+    1: (2, 5),
+    2: 3,
+    3: (4, 3),
+    4: (5, 7),
+    5: 6,
+    6: (7, 6),
+}
+
+
+def detect_pattern(
+    pattern: PatternGraph, flow_graph: FlowGraph, start: Node
+) -> Optional[Tuple[Node, ...]]:
+    indices = [node.block.index for node in flow_graph.nodes]
+    assert sorted(indices) == indices, "FlowGraphs should be sorted"
+
+    offset = start.block.index
+    for label in pattern.keys():
+        try:
+            node = flow_graph.nodes[label + offset]
+            target = pattern[label]
+            if isinstance(target, int):
+                if not (
+                    isinstance(node, BasicNode)
+                    and node.successor is flow_graph.nodes[offset + target]
+                ):
+                    return None
+            else:
+                (fallthrough, conditional) = target
+                if not (
+                    isinstance(node, ConditionalNode)
+                    and node.conditional_edge is flow_graph.nodes[offset + conditional]
+                    and node.fallthrough_edge is flow_graph.nodes[offset + fallthrough]
+                ):
+                    return None
+        except IndexError:
+            return None
+
+    # TODO: Check that the subgraph is self-contained except for the entry
+    # and exit nodes, which themselves should be somehow "designated" using an enum.
+
+    all_nodes_in_pattern = (
+        {offset + label for label in pattern.keys()}
+        | {offset + label[0] for label in pattern.values() if isinstance(label, tuple)}
+        | {offset + label[1] for label in pattern.values() if isinstance(label, tuple)}
+    )
+    return tuple(
+        node
+        for i, node in enumerate(flow_graph.nodes)
+        if node is not start and i in all_nodes_in_pattern
+    )
+
+
+def remove_and_replace_nodes(flow_graph: FlowGraph, nodes: Tuple[Node, ...]) -> None:
+    (node_1, node_2, node_3, node_4, node_5, node_6, node_7) = nodes
+    new_node_1 = BasicNode(node_1.block, node_1.emit_goto, node_2)
+    replace_node(flow_graph, node_1, new_node_1)
+    remove_node(flow_graph, node_4, node_7)
+    remove_node(flow_graph, node_5, node_7)
+    remove_node(flow_graph, node_6, node_7)  # TODO: assert didn't execute anything?.
+
+
+def reroll_loop(flow_graph: FlowGraph, start: ConditionalNode) -> bool:
+    nodes = detect_pattern(IDO_O2_SIMPLE_LOOP, flow_graph, start)
+    if nodes is None:
+        return False
+    (node_1, node_2, node_3, node_4, node_5, node_6, node_7) = nodes
+
+    def modify_node_1_instructions(instructions: List[Instruction]) -> bool:
+        # First, we check that the node has the instructions we
+        # think it has.
+        branches = [instr for instr in instructions if instr.is_branch_instruction()]
+        if len(branches) != 1:
+            return False
+        andi_instrs = [instr for instr in instructions if instr.mnemonic == "andi"]
+        if len(andi_instrs) != 1:
+            return False
+        # We are now free to modify the instructions, as we have verified
+        # that this node fits the criteria.
+        instructions.remove(branches[0])
+        andi = andi_instrs[0]
+        move = Instruction.derived("move", [andi.args[0], andi.args[1]], andi)
+        instructions[instructions.index(andi)] = move
+        return True
+
+    if not modify_node_1_instructions(node_1.block.instructions):
+        return False
+
+    remove_and_replace_nodes(flow_graph, nodes)
+
+    return True
+
+
+def reroll_loops(flow_graph: FlowGraph) -> FlowGraph:
+    changed: bool = True
+    while changed:
+        changed = False
+        for node in flow_graph.nodes:
+            if not isinstance(node, ConditionalNode):
+                continue
+            changed = reroll_loop(flow_graph, node)
+            if changed:
+                compute_relations(flow_graph.nodes)
+                break
+    return flow_graph

src/main.py

@@ -6,7 +6,8 @@
 
 from .c_types import TypeMap, build_typemap, dump_typemap
 from .error import DecompFailure
-from .flow_graph import visualize_flowgraph
+from .flow_graph import FlowGraph, build_flowgraph, visualize_flowgraph
+from .loop_rerolling import reroll_loops
 from .if_statements import get_function_text
 from .options import CodingStyle, Options
 from .parse_file import MIPSFile, parse_file
@@ -94,7 +95,12 @@ def run(options: Options) -> int:
     function_infos: List[Union[FunctionInfo, Exception]] = []
     for function in functions:
         try:
-            info = translate_to_ast(function, options, global_info)
+            flowgraph: FlowGraph = build_flowgraph(function, mips_file.asm_data)
+
+            if options.loop_rerolling:
+                flowgraph = reroll_loops(flowgraph)
+
+            info = translate_to_ast(function, flowgraph, options, global_info)
             function_infos.append(info)
         except Exception as e:
             # Store the exception for later, to preserve the order in the output
@@ -174,6 +180,12 @@ def parse_flags(flags: List[str]) -> Options:
         help="disable detection of &&/||",
         action="store_false",
     )
+    parser.add_argument(
+        "--no-reroll",
+        dest="loop_rerolling",
+        help="disable detection and fixing of unrolled loops",
+        action="store_false",
+    )
     parser.add_argument(
         "--no-casts",
         dest="skip_casts",
@@ -304,6 +316,7 @@ def parse_flags(flags: List[str]) -> Options:
         void=args.void,
         ifs=args.ifs,
         andor_detection=args.andor_detection,
+        loop_rerolling=args.loop_rerolling,
         skip_casts=args.skip_casts,
         reg_vars=reg_vars,
         goto_patterns=args.goto_patterns,

src/options.py

@@ -19,6 +19,7 @@ class Options:
     void: bool = attr.ib()
     ifs: bool = attr.ib()
     andor_detection: bool = attr.ib()
+    loop_rerolling: bool = attr.ib()
     skip_casts: bool = attr.ib()
     reg_vars: List[str] = attr.ib()
     goto_patterns: List[str] = attr.ib()

src/translate.py

@@ -17,7 +17,6 @@
     ReturnNode,
     SwitchNode,
     TerminalNode,
-    build_flowgraph,
 )
 from .options import Formatter, Options
 from .parse_file import AsmData, AsmDataEntry
@@ -3877,6 +3876,7 @@ class FunctionInfo:
 
 def translate_to_ast(
     function: Function,
+    flow_graph: FlowGraph,
     options: Options,
     global_info: GlobalInfo,
 ) -> FunctionInfo:
@@ -3886,7 +3886,6 @@ def translate_to_ast(
     branch condition.
     """
     # Initialize info about the function.
-    flow_graph: FlowGraph = build_flowgraph(function, global_info.asm_data)
     stack_info = get_stack_info(function, global_info, flow_graph)
     typemap = global_info.typemap
 

tests/end_to_end/loop/irix-o2-no-reroll-flags.txt

@@ -0,0 +1 @@
+--no-reroll

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

@@ -0,0 +1,49 @@
+s32 test(s8 *arg0, s32 arg1) {
+    s32 temp_a3;
+    s32 temp_v0;
+    s32 temp_v0_2;
+    void *temp_v1;
+    s8 *phi_v1;
+    s32 phi_v0;
+    void *phi_v1_2;
+    s32 phi_v0_2;
+    s32 phi_return;
+    s32 phi_v0_3;
+
+    phi_return = 0;
+    if (arg1 > 0) {
+        temp_a3 = arg1 & 3;
+        phi_v0_3 = 0;
+        if (temp_a3 != 0) {
+            phi_v1 = arg0;
+            phi_v0 = 0;
+            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;
+                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;
+        }
+    }
+    return phi_return;
+}

tests/end_to_end/loop/irix-o2-no-reroll.s

@@ -0,0 +1,34 @@
+.set noat      # allow manual use of $at
+.set noreorder # don't insert nops after branches
+
+
+glabel test
+/* 000090 00400090 18A00012 */  blez  $a1, .L004000DC
+/* 000094 00400094 00001025 */   move  $v0, $zero
+/* 000098 00400098 30A70003 */  andi  $a3, $a1, 3
+/* 00009C 0040009C 10E00007 */  beqz  $a3, .L004000BC
+/* 0000A0 004000A0 00E03025 */   move  $a2, $a3
+/* 0000A4 004000A4 00801821 */  move  $v1, $a0
+.L004000A8:
+/* 0000A8 004000A8 24420001 */  addiu $v0, $v0, 1
+/* 0000AC 004000AC A0600000 */  sb    $zero, ($v1)
+/* 0000B0 004000B0 14C2FFFD */  bne   $a2, $v0, .L004000A8
+/* 0000B4 004000B4 24630001 */   addiu $v1, $v1, 1
+/* 0000B8 004000B8 10450008 */  beq   $v0, $a1, .L004000DC
+.L004000BC:
+/* 0000BC 004000BC 00821821 */   addu  $v1, $a0, $v0
+.L004000C0:
+/* 0000C0 004000C0 24420004 */  addiu $v0, $v0, 4
+/* 0000C4 004000C4 A0600001 */  sb    $zero, 1($v1)
+/* 0000C8 004000C8 A0600002 */  sb    $zero, 2($v1)
+/* 0000CC 004000CC A0600003 */  sb    $zero, 3($v1)
+/* 0000D0 004000D0 24630004 */  addiu $v1, $v1, 4
+/* 0000D4 004000D4 1445FFFA */  bne   $v0, $a1, .L004000C0
+/* 0000D8 004000D8 A060FFFC */   sb    $zero, -4($v1)
+.L004000DC:
+/* 0000DC 004000DC 03E00008 */  jr    $ra
+/* 0000E0 004000E0 00000000 */   nop
+
+/* 0000E4 004000E4 00000000 */  nop
+/* 0000E8 004000E8 00000000 */  nop
+/* 0000EC 004000EC 00000000 */  nop