3 files changed, 54 insertions, 42 deletions

diff --git a/src/wasm-validator.h b/src/wasm-validator.h
index 8e94fd368..893306897 100644
--- a/src/wasm-validator.h
+++ b/src/wasm-validator.h
@@ -56,9 +56,6 @@ public:
   }
   void visitLoop(Loop *curr) {
     if (curr->in.is()) {
-      LoopChildChecker childChecker(curr->in);
-      childChecker.walk(curr->body);
-      shouldBeTrue(childChecker.valid, curr, "loop must return none");
       breakTypes.erase(curr->in);
     }
     if (curr->out.is()) {
@@ -285,20 +282,6 @@ public:
 
 private:
 
-  // the "in" label has a none type, since no one can receive its value. make sure no one breaks to it with a value.
-  struct LoopChildChecker : public PostWalker<LoopChildChecker, Visitor<LoopChildChecker>> {
-    Name in;
-    bool valid = true;
-
-    LoopChildChecker(Name in) : in(in) {}
-
-    void visitBreak(Break *curr) {
-      if (curr->name == in && curr->value) {
-        valid = false;
-      }
-    }
-  };
-
   // helpers
 
   std::ostream& fail() {
diff --git a/src/wasm.cpp b/src/wasm.cpp
index 09d9c3c16..d00ddd42f 100644
--- a/src/wasm.cpp
+++ b/src/wasm.cpp
@@ -20,53 +20,52 @@
 
 namespace wasm {
 
-struct BlockTypeSeeker : public PostWalker<BlockTypeSeeker, Visitor<BlockTypeSeeker>> {
-  Block* target; // look for this one
+struct TypeSeeker : public PostWalker<TypeSeeker, Visitor<TypeSeeker>> {
+  Expression* target; // look for this one
+  Name targetName;
   std::vector<WasmType> types;
 
-  BlockTypeSeeker(Block* target) : target(target) {}
+  TypeSeeker(Expression* target, Name targetName) : target(target), targetName(targetName) {
+    Expression* temp = target;
+    walk(temp);
+  }
 
-  void visitBreak(Break *curr) {
-    if (curr->name == target->name) {
+  void visitBreak(Break* curr) {
+    if (curr->name == targetName) {
       types.push_back(curr->value ? curr->value->type : none);
     }
   }
 
-  void visitSwitch(Switch *curr) {
+  void visitSwitch(Switch* curr) {
     for (auto name : curr->targets) {
-      if (name == target->name) types.push_back(curr->value ? curr->value->type : none);
+      if (name == targetName) types.push_back(curr->value ? curr->value->type : none);
     }
   }
 
-  void visitBlock(Block *curr) {
+  void visitBlock(Block* curr) {
     if (curr == target) {
       if (curr->list.size() > 0) {
         types.push_back(curr->list.back()->type);
       } else {
         types.push_back(none);
       }
-    } else if (curr->name == target->name) {
+    } else if (curr->name == targetName) {
       types.clear(); // ignore all breaks til now, they were captured by someone with the same name
     }
   }
-};
 
-void Block::finalize() {
-  if (!name.is()) {
-    // nothing branches here, so this is easy
-    if (list.size() > 0) {
-      type = list.back()->type;
-    } else {
-      type = unreachable;
+  void visitLoop(Loop* curr) {
+    if (curr == target) {
+      types.push_back(curr->body->type);
+    } else if (curr->in == targetName || curr->out == targetName) {
+      types.clear(); // ignore all breaks til now, they were captured by someone with the same name
     }
-    return;
   }
+};
 
-  BlockTypeSeeker seeker(this);
-  Expression* temp = this;
-  seeker.walk(temp);
-  type = unreachable;
-  for (auto other : seeker.types) {
+static WasmType mergeTypes(std::vector<WasmType>& types) {
+  WasmType type = unreachable;
+  for (auto other : types) {
     // once none, stop. it then indicates a poison value, that must not be consumed
     // and ignore unreachable
     if (type != none) {
@@ -81,6 +80,32 @@ void Block::finalize() {
       }
     }
   }
+  return type;
+}
+
+void Block::finalize() {
+  if (!name.is()) {
+    // nothing branches here, so this is easy
+    if (list.size() > 0) {
+      type = list.back()->type;
+    } else {
+      type = unreachable;
+    }
+    return;
+  }
+
+  TypeSeeker seeker(this, this->name);
+  type = mergeTypes(seeker.types);
+}
+
+void Loop::finalize() {
+  if (!out.is()) {
+    type = body->type;
+    return;
+  }
+
+  TypeSeeker seeker(this, this->out);
+  type = mergeTypes(seeker.types);
 }
 
 } // namespace wasm
diff --git a/src/wasm.h b/src/wasm.h
index fc0337700..839dcf3c2 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -974,9 +974,13 @@ public:
   Name out, in;
   Expression *body;
 
-  void finalize() {
-    type = body->type; // loop might have a type, if the body ends in something that does not break
+  // set the type of a loop if you already know it
+  void finalize(WasmType type_) {
+    type = type_;
   }
+
+  // set the type of a loop based on its contents. this scans the loop, so it is not fast
+  void finalize();
 };
 
 class Break : public SpecificExpression<Expression::BreakId> {