Detect and reroll unrolled do-while loops

src/flow_graph.py

@@ -778,18 +778,34 @@ class BaseNode(abc.ABC):
     # i.e. there is an invariant `(node.loop is None) or (node.loop.head is node)`
     loop: Optional["NaturalLoop"] = attr.ib(init=False, default=None)
 
+    def to_basic_node(self, successor: "Node") -> "BasicNode":
+        new_node = BasicNode(self.block, self.emit_goto, successor)
+        new_node.parents = self.parents
+        new_node.dominators = self.dominators
+        new_node.immediate_dominator = self.immediate_dominator
+        new_node.immediately_dominates = self.immediately_dominates
+        return new_node
+
     def name(self) -> str:
         return str(self.block.index)
 
     @abc.abstractmethod
     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 +826,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 +859,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 +883,15 @@ def __str__(self) -> str:
 class SwitchNode(BaseNode):
     cases: List["Node"] = attr.ib()
 
+    def replace_any_children(self, replace_this: "Node", with_this: "Node") -> None:
+        new_cases: List["Node"] = []
+        for case in self.cases:
+            if case is replace_this:
+                new_cases.append(with_this)
+            else:
+                new_cases.append(case)
+        self.cases = new_cases
+
     def children(self) -> List["Node"]:
         # Deduplicate nodes in `self.cases`
         seen = set()
@@ -891,6 +925,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 +1359,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,116 @@
+from typing import List
+
+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)
+
+
+def reroll_loop(flow_graph: FlowGraph, start: ConditionalNode) -> bool:
+    node_1 = start.fallthrough_edge
+    node_7 = start.conditional_edge
+
+    if not isinstance(node_1, ConditionalNode):
+        return False
+    node_2 = node_1.fallthrough_edge
+    node_5 = node_1.conditional_edge
+
+    if not isinstance(node_2, BasicNode):
+        return False
+    node_3 = node_2.successor
+
+    if not (
+        isinstance(node_3, ConditionalNode)
+        and node_3.loop
+        and node_3.conditional_edge is node_3
+    ):
+        return False
+    node_4 = node_3.fallthrough_edge
+
+    if not (
+        isinstance(node_4, ConditionalNode)
+        and node_4.fallthrough_edge is node_5
+        and node_4.conditional_edge is node_7
+    ):
+        return False
+
+    if not isinstance(node_5, BasicNode):
+        return False
+    node_6 = node_5.successor
+
+    if not (
+        isinstance(node_6, ConditionalNode)
+        and node_6.loop
+        and node_6.conditional_edge is node_6
+        and node_6.fallthrough_edge is node_7
+    ):
+        return False
+
+    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
+
+    new_node_1 = node_1.to_basic_node(
+        successor=node_2  # node_2 doesn't know it's a parent yet
+    )
+    replace_node(flow_graph, node_1, new_node_1)  # now it does
+    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?.
+
+    return True
+
+
+def reroll_loops(flow_graph: FlowGraph) -> FlowGraph:
+    # TODO: What if knocking out nodes reveals another set of nodes
+    # that look identical? We will incorrectly be merging two
+    # adjacent for-loops.
+    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:
+                break
+    compute_relations(flow_graph.nodes)
+    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

@@ -4,6 +4,7 @@
 import sys
 import traceback
 from contextlib import contextmanager
+from copy import copy
 from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Union
 
 import attr
@@ -2981,7 +2982,7 @@ def regs_clobbered_until_dominator(
     if node.immediate_dominator is None:
         return set()
     seen = {node.immediate_dominator}
-    stack = node.parents[:]
+    stack = copy(node.parents)
     clobbered = set()
     while stack:
         n = stack.pop()
@@ -3003,7 +3004,7 @@ def reg_always_set(
     if node.immediate_dominator is None:
         return False
     seen = {node.immediate_dominator}
-    stack = node.parents[:]
+    stack = copy(node.parents)
     while stack:
         n = stack.pop()
         if n == node.immediate_dominator and not dom_set:
@@ -3032,11 +3033,17 @@ def assign_phis(used_phis: List[PhiExpr], stack_info: StackInfo) -> None:
     while i < len(used_phis):
         phi = used_phis[i]
         assert phi.num_usages > 0
-        assert len(phi.node.parents) >= 2
+        # assert len(phi.node.parents) >= 2
+        if len(phi.node.parents) < 2:
+            i += 1
+            continue
         exprs = []
         for node in phi.node.parents:
             block_info = node.block.block_info
-            assert isinstance(block_info, BlockInfo)
+            # assert isinstance(block_info, BlockInfo)
+            if not block_info:
+                i += 1
+                continue
             exprs.append(block_info.final_register_states[phi.reg])
 
         first_uw = early_unwrap(exprs[0])
@@ -3094,7 +3101,9 @@ def compute_has_custom_return(nodes: List[Node]) -> None:
                 continue
             for p in n.parents:
                 block_info2 = p.block.block_info
-                assert isinstance(block_info2, BlockInfo)
+                # assert isinstance(block_info2, BlockInfo)
+                if not block_info2:
+                    continue
                 if block_info2.has_custom_return:
                     block_info.has_custom_return = True
                     changed = True
@@ -3877,6 +3886,7 @@ class FunctionInfo:
 
 def translate_to_ast(
     function: Function,
+    flow_graph: FlowGraph,
     options: Options,
     global_info: GlobalInfo,
 ) -> FunctionInfo:
@@ -3886,7 +3896,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;
+}