Handle transposed loops.

It is possible that a loop breaks out a few levels while the continuing
path after the loop breaks out of more levels.

This is impossible to represent normally.

We have to transpose the loop such that the breaking construct is the
merge point, and the continuing path becomes the breaking construct
instead.

This obliterates the CFG after retargeting branches like this, so we
have to recompute the CFG before doing anything else.

Ridiculously complicated ... :(

2 files changed, 297 insertions, 161 deletions

diff --git a/cfg_structurizer.cpp b/cfg_structurizer.cpp
index 85e650c..f4c4bb3 100644
--- a/cfg_structurizer.cpp
+++ b/cfg_structurizer.cpp
@@ -380,6 +380,12 @@ bool CFGStructurizer::run()
 		log_cfg_graphviz(graphviz_split.c_str());
 	}
 
+	while (rewrite_transposed_loops())
+	{
+		auto graphviz_split = graphviz_path + ".transpose-loop-rewrite";
+		log_cfg_graphviz(graphviz_split.c_str());
+	}
+
 	//LOGI("=== Structurize pass ===\n");
 	structurize(0);
 	update_structured_loop_merge_targets();
@@ -3087,127 +3093,293 @@ CFGNode *CFGStructurizer::find_common_post_dominator_with_ignored_break(Vector<C
 	return candidates.front();
 }
 
-void CFGStructurizer::find_loops()
+bool CFGStructurizer::rewrite_transposed_loops()
 {
-	for (auto index = forward_post_visit_order.size(); index; index--)
+	bool did_rewrite = false;
+
+	for (auto index = forward_post_visit_order.size(); index && !did_rewrite; index--)
 	{
 		// Visit in reverse order so we resolve outer loops first,
 		// this lets us detect ladder-breaking loops.
 		auto *node = forward_post_visit_order[index - 1];
 
-		if (node->freeze_structured_analysis)
+		if (node->freeze_structured_analysis && node->merge == MergeType::Loop)
+			continue;
+		if (!node->has_pred_back_edges())
+			continue;
+
+		auto result = analyze_loop(node);
+		auto merge_result = analyze_loop_merge(node, result);
+
+		auto *merge = merge_result.merge;
+		auto *dominated_merge = merge_result.dominated_merge;
+
+		// We might have a horribly complex scenario where a loop breaks, but it breaks to an outer scope
+		// which is not consistent with the merge block, i.e. we need structured control flow to resolve properly
+		// before we can break. This is ... problematic.
+		CFGNode *impossible_merge_target = nullptr;
+		if (merge_result.merge && merge_result.dominated_merge && !result.non_dominated_exit.empty())
 		{
-			// If we have a pre-created dummy loop for ladding breaking,
-			// just propagate the header information and be done with it.
-			if (node->merge == MergeType::Loop)
+			auto *common_break_target = find_common_post_dominator(result.non_dominated_exit);
+			if (common_break_target && common_break_target != merge &&
+			    common_break_target->reaches_domination_frontier_before_merge(merge))
 			{
-				node->loop_merge_block->headers.push_back(node);
-				continue;
+				impossible_merge_target = common_break_target;
 			}
 		}
 
-		if (!node->has_pred_back_edges())
+		if (!impossible_merge_target)
 			continue;
 
-		// There are back-edges here, this must be a loop header.
-		node->merge = MergeType::Loop;
+		// Rewrite the control flow from the inside out through a transposition.
+		// The common break target will become the merge block instead.
+		// The continue will break out to the transposed merge instead.
+		// In the ladder, we will enter a breaking path which branches out to loop_ladder.
 
-		// Now, we need to figure out which blocks belong in the loop construct.
-		// The way to figure out a natural loop is any block which is dominated by loop header
-		// and control flow passes to one of the back edges.
+		auto *ladder_break = pool.create_node();
+		ladder_break->name = node->name + ".transposed-merge.break";
+		ladder_break->ir.terminator.type = Terminator::Type::Branch;
+		ladder_break->ir.terminator.direct_block = impossible_merge_target;
+		ladder_break->immediate_post_dominator = impossible_merge_target;
+		ladder_break->forward_post_visit_order = impossible_merge_target->forward_post_visit_order;
+		ladder_break->backward_post_visit_order = impossible_merge_target->backward_post_visit_order;
 
-		// Unfortunately, it can be ambiguous which block is the merge block for a loop.
-		// Ideally, there is a unique block which is the loop exit block, but if there are multiple breaks
-		// there are multiple blocks which are not part of the loop construct.
+		auto *ladder_selection = pool.create_node();
+		ladder_selection->name = node->name + ".transposed-merge";
+		ladder_selection->forward_post_visit_order = impossible_merge_target->forward_post_visit_order;
+		ladder_selection->backward_post_visit_order = impossible_merge_target->backward_post_visit_order;
+		ladder_selection->immediate_post_dominator = merge;
+		ladder_break->immediate_dominator = ladder_selection;
 
-		LoopBacktracer tracer;
-		auto *pred = node->pred_back_edge;
+		auto ladder_preds = dominated_merge->pred;
 
-		// Back-trace from here.
-		// The CFG is reducible, so node must dominate pred.
-		// Since node dominates pred, there is no pred chain we can follow without
-		// eventually hitting node, and we'll stop traversal there.
+		ladder_selection->add_branch(ladder_break);
+		ladder_selection->add_branch(dominated_merge);
+		node->traverse_dominated_blocks_and_rewrite_branch(impossible_merge_target, ladder_selection);
+		ladder_selection->recompute_immediate_dominator();
 
-		// All nodes which are touched during this traversal must be part of the loop construct.
-		tracer.trace_to_parent(node, pred);
+		ladder_break->add_branch(impossible_merge_target);
 
-		LoopMergeTracer merge_tracer(tracer);
-		merge_tracer.trace_from_parent(node);
+		// Branches from these blocks should be rewritten to target transposed-merge.
+		for (auto *ladder_pred : ladder_preds)
+			ladder_pred->retarget_branch(dominated_merge, ladder_selection);
 
-		Vector<CFGNode *> direct_exits;
-		Vector<CFGNode *> inner_direct_exits;
-		Vector<CFGNode *> dominated_exit;
-		Vector<CFGNode *> inner_dominated_exit;
-		Vector<CFGNode *> non_dominated_exit;
+		ladder_selection->ir.terminator.type = Terminator::Type::Condition;
+		ladder_selection->ir.terminator.true_block = dominated_merge;
+		ladder_selection->ir.terminator.false_block = ladder_break;
+		ladder_selection->ir.terminator.conditional_id = module.allocate_id();
 
-		for (auto *loop_exit : merge_tracer.loop_exits)
+		PHI phi;
+		phi.id = ladder_selection->ir.terminator.conditional_id;
+		phi.type_id = module.get_builder().makeBoolType();
+		module.get_builder().addName(phi.id, (String("transposed_selector_") + node->name).c_str());
+		for (auto *ladder_pred : ladder_selection->pred)
 		{
-			auto exit_type = get_loop_exit_type(*node, *loop_exit);
-			switch (exit_type)
-			{
-			case LoopExitType::Exit:
-				direct_exits.push_back(loop_exit);
-				break;
+			IncomingValue incoming = {};
+			incoming.block = ladder_pred;
+			bool branches_to_dominated_merge =
+					std::find(ladder_preds.begin(), ladder_preds.end(), ladder_pred) != ladder_preds.end();
+			incoming.id = module.get_builder().makeBoolConstant(branches_to_dominated_merge);
+			phi.incoming.push_back(incoming);
+		}
+		ladder_selection->ir.phi.push_back(std::move(phi));
 
-			case LoopExitType::InnerLoopExit:
-				// It's not an exit for us, but the inner loop.
-				inner_direct_exits.push_back(loop_exit);
-				break;
+		// We have obliterated the existing control flow through transposition,
+		// and any domination or post-domination analysis will break.
+		// Re-traverse the CFG and try again.
+		// Continue until we have eliminated all impossible loops (should be extremely rare).
+		did_rewrite = true;
+	}
 
-			case LoopExitType::Merge:
-				dominated_exit.push_back(loop_exit);
-				break;
+	if (did_rewrite)
+		recompute_cfg();
+	return did_rewrite;
+}
 
-			case LoopExitType::InnerLoopMerge:
-				inner_dominated_exit.push_back(loop_exit);
-				break;
+CFGStructurizer::LoopAnalysis CFGStructurizer::analyze_loop(CFGNode *node) const
+{
+	LoopAnalysis result;
 
-			case LoopExitType::InnerLoopFalsePositive:
-				// In this case, the inner loop can only exit at the loop header,
-				// and thus post-dominance analysis will always fail.
-				// Ignore this case as it's a false exit.
-				break;
+	// Now, we need to figure out which blocks belong in the loop construct.
+	// The way to figure out a natural loop is any block which is dominated by loop header
+	// and control flow passes to one of the back edges.
 
-			case LoopExitType::Escape:
-				non_dominated_exit.push_back(loop_exit);
-				break;
-			}
+	// Unfortunately, it can be ambiguous which block is the merge block for a loop.
+	// Ideally, there is a unique block which is the loop exit block, but if there are multiple breaks
+	// there are multiple blocks which are not part of the loop construct.
+
+	LoopBacktracer tracer;
+	auto *pred = node->pred_back_edge;
+
+	// Back-trace from here.
+	// The CFG is reducible, so node must dominate pred.
+	// Since node dominates pred, there is no pred chain we can follow without
+	// eventually hitting node, and we'll stop traversal there.
+
+	// All nodes which are touched during this traversal must be part of the loop construct.
+	tracer.trace_to_parent(node, pred);
+
+	LoopMergeTracer merge_tracer(tracer);
+	merge_tracer.trace_from_parent(node);
+
+	for (auto *loop_exit : merge_tracer.loop_exits)
+	{
+		auto exit_type = get_loop_exit_type(*node, *loop_exit);
+		switch (exit_type)
+		{
+		case LoopExitType::Exit:
+			result.direct_exits.push_back(loop_exit);
+			break;
+
+		case LoopExitType::InnerLoopExit:
+			// It's not an exit for us, but the inner loop.
+			result.inner_direct_exits.push_back(loop_exit);
+			break;
+
+		case LoopExitType::Merge:
+			result.dominated_exit.push_back(loop_exit);
+			break;
+
+		case LoopExitType::InnerLoopMerge:
+			result.inner_dominated_exit.push_back(loop_exit);
+			break;
+
+		case LoopExitType::InnerLoopFalsePositive:
+			// In this case, the inner loop can only exit at the loop header,
+			// and thus post-dominance analysis will always fail.
+			// Ignore this case as it's a false exit.
+			break;
+
+		case LoopExitType::Escape:
+			result.non_dominated_exit.push_back(loop_exit);
+			break;
 		}
+	}
 
-		// If the only merge candidates we have are inner dominated, treat them as true dominated exits.
-		if (dominated_exit.empty() && !inner_dominated_exit.empty())
-			std::swap(dominated_exit, inner_dominated_exit);
+	// If the only merge candidates we have are inner dominated, treat them as true dominated exits.
+	if (result.dominated_exit.empty() && !result.inner_dominated_exit.empty())
+		std::swap(result.dominated_exit, result.inner_dominated_exit);
 
-		// If there are no direct exists, treat inner direct exists as direct exits.
-		if (direct_exits.empty())
-			direct_exits = std::move(inner_direct_exits);
+	// If there are no direct exists, treat inner direct exists as direct exits.
+	if (result.direct_exits.empty())
+		result.direct_exits = std::move(result.inner_direct_exits);
 
-		// A direct exit can be considered a dominated exit if there are no better candidates.
-		if (dominated_exit.empty() && !direct_exits.empty())
-			std::swap(dominated_exit, direct_exits);
+	// A direct exit can be considered a dominated exit if there are no better candidates.
+	if (result.dominated_exit.empty() && !result.direct_exits.empty())
+		std::swap(result.dominated_exit, result.direct_exits);
 
-		// If we only have one direct exit, consider it our merge block.
-		// Pick either Merge or Escape.
-		if (direct_exits.size() == 1 && dominated_exit.empty() && non_dominated_exit.empty())
+	// If we only have one direct exit, consider it our merge block.
+	// Pick either Merge or Escape.
+	if (result.direct_exits.size() == 1 && result.dominated_exit.empty() && result.non_dominated_exit.empty())
+	{
+		if (node->dominates(result.direct_exits.front()))
+			std::swap(result.dominated_exit, result.direct_exits);
+		else
+			std::swap(result.non_dominated_exit, result.direct_exits);
+	}
+
+	if (result.dominated_exit.size() >= 2)
+	{
+		// Try to see if we can reduce the number of merge blocks to just 1.
+		// This is relevant if we have various "clean" break blocks.
+		auto *post_dominator = find_common_post_dominator(result.dominated_exit);
+		if (std::find(result.dominated_exit.begin(), result.dominated_exit.end(),
+		              post_dominator) != result.dominated_exit.end())
 		{
-			if (node->dominates(direct_exits.front()))
-				std::swap(dominated_exit, direct_exits);
-			else
-				std::swap(non_dominated_exit, direct_exits);
+			result.dominated_exit.clear();
+			result.dominated_exit.push_back(post_dominator);
 		}
+	}
+
+	return result;
+}
+
+CFGStructurizer::LoopMergeAnalysis CFGStructurizer::analyze_loop_merge(CFGNode *node, const LoopAnalysis &analysis)
+{
+	// We have multiple blocks which are merge candidates. We need to figure out where execution reconvenes.
+	Vector<CFGNode *> merges;
+	merges.reserve(analysis.inner_dominated_exit.size() + analysis.dominated_exit.size() + analysis.non_dominated_exit.size());
+	merges.insert(merges.end(), analysis.inner_dominated_exit.begin(), analysis.inner_dominated_exit.end());
+	merges.insert(merges.end(), analysis.dominated_exit.begin(), analysis.dominated_exit.end());
+	merges.insert(merges.end(), analysis.non_dominated_exit.begin(), analysis.non_dominated_exit.end());
+	CFGNode *merge = CFGStructurizer::find_common_post_dominator(merges);
 
-		if (dominated_exit.size() >= 2)
+	CFGNode *dominated_merge = nullptr;
+
+	// Try to find the sensible target first.
+	// If one of our merge blocks is the successor of the continue block,
+	// this is a prime candidate for a ladder block.
+	if (node->pred_back_edge && node->pred_back_edge->succ.size() == 1 &&
+	    std::find(analysis.dominated_exit.begin(),
+	              analysis.dominated_exit.end(),
+	              node->pred_back_edge->succ.front()) != analysis.dominated_exit.end())
+	{
+		dominated_merge = node->pred_back_edge->succ.front();
+	}
+	else if (merge && !node->dominates(merge) && analysis.dominated_exit.size() > 1)
+	{
+		// Now, we might have Merge blocks which end up escaping out of the loop construct.
+		// We might have to remove candidates which end up being break blocks after all.
+		Vector<CFGNode *> non_breaking_exits;
+		non_breaking_exits.reserve(analysis.dominated_exit.size());
+		for (auto *exit : analysis.dominated_exit)
+			if (!control_flow_is_escaping(exit, merge))
+				non_breaking_exits.push_back(exit);
+
+		dominated_merge = CFGStructurizer::find_common_post_dominator(non_breaking_exits);
+	}
+	else
+	{
+		dominated_merge = CFGStructurizer::find_common_post_dominator(analysis.dominated_exit);
+	}
+
+	if (!dominated_merge)
+	{
+		LOGW("There is no candidate for ladder merging.\n");
+	}
+
+	if (dominated_merge && !node->dominates(dominated_merge))
+	{
+		LOGW("We don't dominate the merge target ...\n");
+		dominated_merge = nullptr;
+	}
+
+	LoopMergeAnalysis merge_result = {};
+	merge_result.merge = merge;
+	merge_result.dominated_merge = dominated_merge;
+	return merge_result;
+}
+
+void CFGStructurizer::find_loops()
+{
+	for (auto index = forward_post_visit_order.size(); index; index--)
+	{
+		// Visit in reverse order so we resolve outer loops first,
+		// this lets us detect ladder-breaking loops.
+		auto *node = forward_post_visit_order[index - 1];
+
+		if (node->freeze_structured_analysis)
 		{
-			// Try to see if we can reduce the number of merge blocks to just 1.
-			// This is relevant if we have various "clean" break blocks.
-			auto *post_dominator = find_common_post_dominator(dominated_exit);
-			if (std::find(dominated_exit.begin(), dominated_exit.end(), post_dominator) != dominated_exit.end())
+			// If we have a pre-created dummy loop for ladding breaking,
+			// just propagate the header information and be done with it.
+			if (node->merge == MergeType::Loop)
 			{
-				dominated_exit.clear();
-				dominated_exit.push_back(post_dominator);
+				node->loop_merge_block->headers.push_back(node);
+				continue;
 			}
 		}
 
+		if (!node->has_pred_back_edges())
+			continue;
+
+		// There are back-edges here, this must be a loop header.
+		node->merge = MergeType::Loop;
+
+		auto result = analyze_loop(node);
+		auto &dominated_exit = result.dominated_exit;
+		auto &inner_dominated_exit = result.inner_dominated_exit;
+		auto &non_dominated_exit = result.non_dominated_exit;
+
 		if (dominated_exit.empty() && inner_dominated_exit.empty() && non_dominated_exit.empty())
 		{
 			// There can be zero loop exits, i.e. infinite loop. This means we have no merge block.
@@ -3255,99 +3427,44 @@ void CFGStructurizer::find_loops()
 		}
 		else
 		{
-			// We have multiple blocks which are merge candidates. We need to figure out where execution reconvenes.
-			Vector<CFGNode *> merges;
-			merges.reserve(inner_dominated_exit.size() + dominated_exit.size() + non_dominated_exit.size());
-			merges.insert(merges.end(), inner_dominated_exit.begin(), inner_dominated_exit.end());
-			merges.insert(merges.end(), dominated_exit.begin(), dominated_exit.end());
-			merges.insert(merges.end(), non_dominated_exit.begin(), non_dominated_exit.end());
-			CFGNode *merge = CFGStructurizer::find_common_post_dominator(merges);
-
-			CFGNode *dominated_merge = nullptr;
-
-			// Try to find the sensible target first.
-			// If one of our merge blocks is the successor of the continue block,
-			// this is a prime candidate for a ladder block.
-			if (node->pred_back_edge && node->pred_back_edge->succ.size() == 1 &&
-			    std::find(dominated_exit.begin(), dominated_exit.end(), node->pred_back_edge->succ.front()) != dominated_exit.end())
-			{
-				dominated_merge = node->pred_back_edge->succ.front();
-			}
-			else if (merge && !node->dominates(merge) && dominated_exit.size() > 1)
-			{
-				// Now, we might have Merge blocks which end up escaping out of the loop construct.
-				// We might have to remove candidates which end up being break blocks after all.
-				Vector<CFGNode *> non_breaking_exits;
-				non_breaking_exits.reserve(dominated_exit.size());
-				for (auto *exit : dominated_exit)
-					if (!control_flow_is_escaping(exit, merge))
-						non_breaking_exits.push_back(exit);
-
-				dominated_merge = CFGStructurizer::find_common_post_dominator(non_breaking_exits);
-			}
-			else
-			{
-				dominated_merge = CFGStructurizer::find_common_post_dominator(dominated_exit);
-			}
-
-			if (!dominated_merge)
-			{
-				LOGW("There is no candidate for ladder merging.\n");
-			}
-
-			if (dominated_merge && !node->dominates(dominated_merge))
-			{
-				LOGW("We don't dominate the merge target ...\n");
-				dominated_merge = nullptr;
-			}
+			auto merge_result = analyze_loop_merge(node, result);
+			auto *merge = merge_result.merge;
+			auto *dominated_merge = merge_result.dominated_merge;
 
 			if (!merge)
 			{
 				LOGW("Failed to find a common merge point ...\n");
 			}
-			else
+			else if (node->can_loop_merge_to(merge))
 			{
+				// Clean merge.
+				// This is a unique merge block. There can be no other merge candidate.
+				//LOGI("Loop with simple multi-exit merge: %p (%s) -> %p (%s)\n", static_cast<const void *>(node),
+				//     node->name.c_str(), static_cast<const void *>(node->loop_merge_block),
+				//     node->loop_merge_block->name.c_str());
+
 				node->loop_merge_block = merge;
 				const_cast<CFGNode *>(node->loop_merge_block)->add_unique_header(node);
+			}
+			else
+			{
+				// Single-escape merge.
+				// It is unique, but we need workarounds later.
+				//LOGI("Loop with ladder multi-exit merge: %p (%s) -> %p (%s)\n", static_cast<const void *>(node),
+				//     node->name.c_str(), static_cast<const void *>(node->loop_merge_block),
+				//     node->loop_merge_block->name.c_str());
 
-				if (node->can_loop_merge_to(merge))
+				if (dominated_merge)
 				{
-					// Clean merge.
-					// This is a unique merge block. There can be no other merge candidate.
-					//LOGI("Loop with simple multi-exit merge: %p (%s) -> %p (%s)\n", static_cast<const void *>(node),
-					//     node->name.c_str(), static_cast<const void *>(node->loop_merge_block),
-					//     node->loop_merge_block->name.c_str());
+					//LOGI("    Ladder block: %p (%s)\n", static_cast<const void *>(dominated_merge),
+					//     dominated_merge->name.c_str());
 				}
-				else
-				{
-					// Single-escape merge.
-					// It is unique, but we need workarounds later.
-					//LOGI("Loop with ladder multi-exit merge: %p (%s) -> %p (%s)\n", static_cast<const void *>(node),
-					//     node->name.c_str(), static_cast<const void *>(node->loop_merge_block),
-					//     node->loop_merge_block->name.c_str());
 
-					if (dominated_merge)
-					{
-						//LOGI("    Ladder block: %p (%s)\n", static_cast<const void *>(dominated_merge),
-						//     dominated_merge->name.c_str());
-					}
-
-					// We will use this block as a ladder.
-					node->loop_ladder_block = dominated_merge;
-				}
-			}
+				// We will use this block as a ladder.
+				node->loop_ladder_block = dominated_merge;
+				node->loop_merge_block = merge;
 
-			if (node->loop_merge_block && node->loop_ladder_block && !non_dominated_exit.empty())
-			{
-				// We might have a horribly complex scenario where a loop breaks, but it breaks to an outer scope
-				// which is not consistent with the merge block.
-				// When we split merge scopes, we need to specially consider any header that dominates this common break target.
-				auto *common_break_target = find_common_post_dominator(non_dominated_exit);
-				if (common_break_target && common_break_target != node->loop_merge_block &&
-				    common_break_target->reaches_domination_frontier_before_merge(node->loop_merge_block))
-				{
-					LOGE("FIXME: Impossible loop merge detected.\n");
-				}
+				const_cast<CFGNode *>(node->loop_merge_block)->add_unique_header(node);
 			}
 		}
 	}
diff --git a/cfg_structurizer.hpp b/cfg_structurizer.hpp
index 62cc048..a54a92d 100644
--- a/cfg_structurizer.hpp
+++ b/cfg_structurizer.hpp
@@ -72,6 +72,25 @@ private:
 	bool query_reachability(const CFGNode &from, const CFGNode &to) const;
 	void structurize(unsigned pass);
 	void find_loops();
+	bool rewrite_transposed_loops();
+
+	struct LoopAnalysis
+	{
+		Vector<CFGNode *> direct_exits;
+		Vector<CFGNode *> inner_direct_exits;
+		Vector<CFGNode *> dominated_exit;
+		Vector<CFGNode *> inner_dominated_exit;
+		Vector<CFGNode *> non_dominated_exit;
+	};
+	LoopAnalysis analyze_loop(CFGNode *node) const;
+
+	struct LoopMergeAnalysis
+	{
+		CFGNode *merge;
+		CFGNode *dominated_merge;
+	};
+	LoopMergeAnalysis analyze_loop_merge(CFGNode *node, const LoopAnalysis &analysis);
+
 	void split_merge_scopes();
 	void eliminate_degenerate_blocks();
 	static bool ladder_chain_has_phi_dependencies(const CFGNode *chain, const CFGNode *incoming);