Fix inlining name collision (#4206)

2 files changed, 39 insertions, 0 deletions

diff --git a/src/ir/branch-utils.h b/src/ir/branch-utils.h
index f2fa43f18..57bad2fb9 100644
--- a/src/ir/branch-utils.h
+++ b/src/ir/branch-utils.h
@@ -238,6 +238,28 @@ inline NameSet getBranchTargets(Expression* ast) {
   return scanner.targets;
 }
 
+// Check if an expression defines a particular name as a branch target anywhere
+// inside it.
+inline bool hasBranchTarget(Expression* ast, Name target) {
+  struct Scanner
+    : public PostWalker<Scanner, UnifiedExpressionVisitor<Scanner>> {
+    Name target;
+    bool has = false;
+
+    void visitExpression(Expression* curr) {
+      operateOnScopeNameDefs(curr, [&](Name& name) {
+        if (name == target) {
+          has = true;
+        }
+      });
+    }
+  };
+  Scanner scanner;
+  scanner.target = target;
+  scanner.walk(ast);
+  return scanner.has;
+}
+
 // Get the name of the branch target that is defined in the expression, or an
 // empty name if there is none.
 inline Name getDefinedName(Expression* curr) {
diff --git a/src/passes/Inlining.cpp b/src/passes/Inlining.cpp
index 7467d2559..0360974f0 100644
--- a/src/passes/Inlining.cpp
+++ b/src/passes/Inlining.cpp
@@ -299,6 +299,23 @@ doInlining(Module* module, Function* into, const InliningAction& action) {
   Builder builder(*module);
   auto* block = builder.makeBlock();
   block->name = Name(std::string("__inlined_func$") + from->name.str);
+  // In the unlikely event that the function already has a branch target with
+  // this name, fix that up, as otherwise we can get unexpected capture of our
+  // branches, that is, we could end up with this:
+  //
+  //  (block $X             ;; a new block we add as the target of returns
+  //    (from's contents
+  //      (block $X         ;; a block in from's contents with a colliding name
+  //        (br $X          ;; a new br we just added that replaces a return
+  //
+  // Here the br wants to go to the very outermost block, to represent a
+  // return from the inlined function's code, but it ends up captured by an
+  // internal block.
+  if (BranchUtils::hasBranchTarget(from->body, block->name)) {
+    auto existingNames = BranchUtils::getBranchTargets(from->body);
+    block->name = Names::getValidName(
+      block->name, [&](Name test) { return !existingNames.count(test); });
+  }
   if (call->isReturn) {
     if (retType.isConcrete()) {
       *action.callSite = builder.makeReturn(block);