Implement simple do-while loop functionality

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: Condition = attr.ib()
+
+ 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)
+ 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)
 
@@ -601,11 +627,47 @@ 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: Node, end: Node) -> Optional[DoWhileLoop]:
+ assert start.loop
+
+ # Find the the condition for the do-while, if it exists
+ condition: Optional[Condition] = None
+ for node in start.loop.backedges:
+ if (
+ node in start.postdominators
+ and isinstance(node, ConditionalNode)
+ and node.fallthrough_edge == end
+ ):
+ assert node.block.block_info
+ assert node.block.block_info.branch_condition
+ condition = node.block.block_info.branch_condition
+ new_end = node
+ break
+ if not condition:
+ return None
+
+ loop_body = build_flowgraph_between(
+ context,
+ start,
+ new_end,
+ skip_loop_detection=True,
+ )
+ emit_node(context, new_end, loop_body)
+
+ return DoWhileLoop(loop_body, 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,
- will be printed out using if-else statements and block info
+ will be printed out using if-else statements and block info.
+
+ `skip_loop_detection` is used to prevent infinite recursion, since (in the
+ case of loops) this function can be recursively called by itself (via
+ `detect_loop`) with the same `start` argument.
  """
  curr_start: Node = start
  body = Body(print_node_comment=context.options.debug)
@@ -617,6 +679,27 @@ 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 curr_start.loop
+ and not curr_start in context.emitted_nodes
+ ):
+ # Find the immediate postdominator to the whole loop,
+ # i.e. the first node outside the loop body
+ imm_pdom: Node = curr_start
+ while imm_pdom in curr_start.loop.nodes:
+ assert imm_pdom.immediate_postdominator is not None
+ imm_pdom = imm_pdom.immediate_postdominator
+
+ # Construct the do-while loop
+ do_while_loop = detect_loop(context, curr_start, imm_pdom)
+ if do_while_loop:
+ body.add_do_while_loop(do_while_loop)
+
+ # Move on.
+ curr_start = imm_pdom
+ 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;
 }