3 files changed, 33 insertions, 8 deletions
diff --git a/src/relooper/Relooper.cpp b/src/relooper/Relooper.cpp
index cb989e20..e91c6e74 100644
--- a/src/relooper/Relooper.cpp
+++ b/src/relooper/Relooper.cpp
@@ -893,10 +893,22 @@ void Relooper::Calculate(Block *Entry) {
         func(Loop->Next); \
       }
 
+    // Find the single block that must be hit in a shape, or NULL if there is more than one
+    Block *FollowNaturalFlow(Shape *S) {
+      SHAPE_SWITCH(S, {
+        return Simple->Inner;
+      }, {
+        return NULL;
+      }, {
+        return FollowNaturalFlow(Loop->Inner);
+      });
+    }
+
     // Remove unneeded breaks and continues.
     // A flow operation is trivially unneeded if the shape we naturally get to by normal code
     // execution is the same as the flow forces us to.
     void RemoveUnneededFlows(Shape *Root, Shape *Natural=NULL) {
+      Block *NaturalBlock = FollowNaturalFlow(Natural);
       Shape *Next = Root;
       while (Next) {
         Root = Next;
@@ -911,7 +923,7 @@ void Relooper::Calculate(Block *Entry) {
             for (BlockBranchMap::iterator iter = Simple->Inner->ProcessedBranchesOut.begin(); iter != Simple->Inner->ProcessedBranchesOut.end(); iter++) {
               Block *Target = iter->first;
               Branch *Details = iter->second;
-              if (Details->Type != Branch::Direct && Target->Parent == Natural) {
+              if (Details->Type != Branch::Direct && Target == NaturalBlock) { // note: cannot handle split blocks
                 Details->Type = Branch::Direct;
                 if (MultipleShape *Multiple = Shape::IsMultiple(Details->Ancestor)) {
                   Multiple->NeedLoop--;
@@ -1026,7 +1038,7 @@ void Relooper::SetAsmJSMode(int On) {
 #if DEBUG
 // Debugging
 
-void DebugDump(BlockSet &Blocks, const char *prefix) {
+void Debugging::Dump(BlockSet &Blocks, const char *prefix) {
   if (prefix) printf("%s ", prefix);
   for (BlockSet::iterator iter = Blocks.begin(); iter != Blocks.end(); iter++) {
     Block *Curr = *iter;
@@ -1039,6 +1051,21 @@ void DebugDump(BlockSet &Blocks, const char *prefix) {
   }
 }
 
+void Debugging::Dump(Shape *S, const char *prefix) {
+  if (prefix) printf("%s ", prefix);
+  printf(" %d ", S->Id);
+  SHAPE_SWITCH(S, {
+    printf("<< Simple with block %d\n", Simple->Inner->Id);
+  }, {
+    printf("<< Multiple\n");
+    for (BlockShapeMap::iterator iter = Multiple->InnerMap.begin(); iter != Multiple->InnerMap.end(); iter++) {
+      printf("     with entry %d\n", iter->first->Id);
+    }
+  }, {
+    printf("<< Loop\n");
+  });
+}
+
 static void PrintDebug(const char *Format, ...) {
   printf("// ");
   va_list Args;
diff --git a/src/relooper/Relooper.h b/src/relooper/Relooper.h
index 5eabe827..34b6db08 100644
--- a/src/relooper/Relooper.h
+++ b/src/relooper/Relooper.h
@@ -211,6 +211,7 @@ typedef std::map<Block*, BlockSet> BlockBlockSetMap;
 #if DEBUG
 struct Debugging {
   static void Dump(BlockSet &Blocks, const char *prefix=NULL);
+  static void Dump(Shape *S, const char *prefix=NULL);
 };
 #endif
 
diff --git a/src/relooper/test.txt b/src/relooper/test.txt
index 2ff70e66..84a8c905 100644
--- a/src/relooper/test.txt
+++ b/src/relooper/test.txt
@@ -129,12 +129,9 @@ do {
     break;
   }
 } while(0);
-do {
-  if (label == 33) {
-    // block C;
-    break;
-  }
-} while(0);
+if (label == 33) {
+  // block C;
+}
 while(1) {
   // block E
   // block F