Merge pull request #125 from HansKristian-Work/infinite-loop-fixes

Handle more cases where we need to rewrite infinite loops.

1 files changed, 27 insertions, 4 deletions

diff --git a/cfg_structurizer.cpp b/cfg_structurizer.cpp
index 41ee5f8..11cee8d 100644
--- a/cfg_structurizer.cpp
+++ b/cfg_structurizer.cpp
@@ -1034,6 +1034,21 @@ void CFGStructurizer::insert_phi()
 		for (auto *op : node->ir.operations)
 			if (op->id)
 				value_id_to_block[op->id] = node;
+
+		// If we inserted dummy branches from back-edge to rewrite infinite loops, we must prune these branches
+		// now, so we don't end up creating a wrong amount of PHI incoming values.
+		// We don't have to recompute the CFG since we don't really care about post-visit orders at this stage.
+		if (node->pred_back_edge && node->pred_back_edge->ir.terminator.type == Terminator::Type::Branch &&
+		    node->pred_back_edge->succ_back_edge == node->pred_back_edge->ir.terminator.direct_block &&
+		    node->pred_back_edge->succ.size() == 1)
+		{
+			auto *back_edge = node->pred_back_edge;
+			auto *succ = back_edge->succ.front();
+			back_edge->succ.clear();
+			auto itr = std::find(succ->pred.begin(), succ->pred.end(), back_edge);
+			assert(itr != succ->pred.end());
+			succ->pred.erase(itr);
+		}
 	}
 
 	// Resolve phi-nodes top-down since PHI nodes may depend on other PHI nodes.
@@ -1651,15 +1666,23 @@ void CFGStructurizer::backwards_visit()
 					return a->forward_post_visit_order > b->forward_post_visit_order;
 				});
 
-				bool exit_is_pure_return = false;
+				bool transpose_loop_exit = false;
 				if (exits.size() == 1)
 				{
 					auto *exit_node = exits.front();
-					exit_is_pure_return = exit_node->ir.phi.empty() && exit_node->ir.operations.empty() &&
-					                      exit_node->ir.terminator.type == Terminator::Type::Return;
+					// If this is true, we never really leave the loop, which is problematic.
+					transpose_loop_exit = exit_node->dominates_all_reachable_exits();
+
+					// Only transpose if we're the innermost header, otherwise, inner loops which try to branch
+					// to the return will be considered a multi-break which is very awkward.
+					if (transpose_loop_exit)
+					{
+						auto *innermost_header = get_innermost_loop_header_for(node, exit_node);
+						transpose_loop_exit = innermost_header == node;
+					}
 				}
 
-				if (exit_is_pure_return)
+				if (transpose_loop_exit)
 				{
 					for (auto *f : exits)
 						node->pred_back_edge->add_branch(f);