From e865fa75b643e1a7bd09c195d47dc2779cd8ecb2 Mon Sep 17 00:00:00 2001
From: Dan Gohman <dan433584@gmail.com>
Date: Fri, 2 Oct 2015 21:11:36 +0000
Subject: [PATCH] [WebAssembly] Fix CFG stackification of nested loops.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@249187 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../WebAssembly/WebAssemblyCFGStackify.cpp    | 19 +++++++++---
 test/CodeGen/WebAssembly/cfg-stackify.ll      | 30 +++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 27698969967..077273c17d7 100644
--- a/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
@@ -194,8 +194,10 @@ static void SortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI) {
     if (MachineLoop *Loop = MLI.getLoopFor(&MBB)) {
       // Assert that loops are contiguous.
       assert(Loop->getHeader() == Loop->getTopBlock());
-      assert(Loop->getHeader() == &MBB ||
-             MLI.getLoopFor(prev(MachineFunction::iterator(&MBB))) == Loop);
+      assert((Loop->getHeader() == &MBB ||
+              Loop->contains(
+                  MLI.getLoopFor(prev(MachineFunction::iterator(&MBB))))) &&
+             "Loop isn't contiguous");
     } else {
       // Assert that non-loops have no backedge predecessors.
       for (auto Pred : MBB.predecessors())
@@ -245,9 +247,18 @@ static void PlaceMarkers(MachineFunction &MF, const MachineLoopInfo &MLI,
   for (auto &MBB : MF) {
     // Place the LOOP for loops.
     if (MachineLoop *Loop = MLI.getLoopFor(&MBB))
-      if (Loop->getHeader() == &MBB)
+      if (Loop->getHeader() == &MBB) {
+        // The operand of a LOOP is the first block after the loop. If the loop
+        // is the bottom of the function, insert a dummy block at the end.
+        MachineBasicBlock *Bottom = Loop->getBottomBlock();
+        auto Iter = next(MachineFunction::iterator(Bottom));
+        if (Iter == MF.end()) {
+          MF.push_back(MF.CreateMachineBasicBlock());
+          Iter = next(MachineFunction::iterator(Bottom));
+        }
         BuildMI(MBB, MBB.begin(), DebugLoc(), TII.get(WebAssembly::LOOP))
-            .addMBB(Loop->getBottomBlock());
+            .addMBB(Iter);
+      }
 
     // Check for forward branches and switches that need BLOCKS placed.
     for (auto &Term : MBB.terminators())
diff --git a/test/CodeGen/WebAssembly/cfg-stackify.ll b/test/CodeGen/WebAssembly/cfg-stackify.ll
index 205384e6102..8ba6bcf61b2 100644
--- a/test/CodeGen/WebAssembly/cfg-stackify.ll
+++ b/test/CodeGen/WebAssembly/cfg-stackify.ll
@@ -272,3 +272,33 @@ exit:
   store volatile i32 4, i32* %p
   ret i32 0
 }
+
+; Test that nested loops are handled.
+
+declare void @bar()
+
+define void @test3(i32 %w)  {
+entry:
+  br i1 undef, label %outer.ph, label %exit
+
+outer.ph:
+  br label %outer
+
+outer:
+  %tobool = icmp eq i32 undef, 0
+  br i1 %tobool, label %inner, label %unreachable
+
+unreachable:
+  unreachable
+
+inner:
+  %c = icmp eq i32 undef, %w
+  br i1 %c, label %if.end, label %inner
+
+exit:
+  ret void
+
+if.end:
+  call void @bar()
+  br label %outer
+}
-- 
2.34.1