Make br_if with a value's type match the spec

CHANGELOG.md

@@ -21,6 +21,15 @@ Current Trunk
  - Add a new `BinaryenModuleReadWithFeatures` function to the C API that allows
    to configure which features to enable in the parser.
 
+ - Fix the result type of `br_if` to match the spec. This means that `br_if` now
+   returns a less precise type than before. In the IR it means that the type of
+   `br_if` with a value depends on the block it targets, which means it must be
+   provided when creating it, and updated if altered.
+   - As a result the C API `BinaryenBreak` now has a final parameter to allow
+     setting the type, which is a `Type*`. If NULL then the type will be
+     inferred in common cases (MVP types, or anything but a `br_if` with a
+     value).
+
 v117
 ----
 

src/binaryen-c.cpp

@@ -1120,10 +1120,21 @@ BinaryenExpressionRef BinaryenLoop(BinaryenModuleRef module,
 BinaryenExpressionRef BinaryenBreak(BinaryenModuleRef module,
                                     const char* name,
                                     BinaryenExpressionRef condition,
-                                    BinaryenExpressionRef value) {
+                                    BinaryenExpressionRef value,
+                                    BinaryenType* type) {
+  std::optional<Type> optType;
+  if (type) {
+    optType = Type(*type);
+  } else {
+    // As a convenience for users, if type is not provided but we can infer the
+    // type - which is the case for MVP types - then infer it.
+    if (condition && value && !value->type.containsRef()) {
+      optType = value->type;
+    }
+  }
   return static_cast<Expression*>(
     Builder(*(Module*)module)
-      .makeBreak(name, (Expression*)value, (Expression*)condition));
+      .makeBreak(name, (Expression*)value, (Expression*)condition, optType));
 }
 BinaryenExpressionRef BinaryenSwitch(BinaryenModuleRef module,
                                      const char** names,

src/binaryen-c.h

@@ -740,12 +740,16 @@ BINARYEN_API BinaryenExpressionRef BinaryenIf(BinaryenModuleRef module,
 BINARYEN_API BinaryenExpressionRef BinaryenLoop(BinaryenModuleRef module,
                                                 const char* in,
                                                 BinaryenExpressionRef body);
-// Break: value and condition can be NULL
+// Break: value and condition can be NULL. type can be NULL, but must be
+//        provided for a `br_if` with a value (as the type depends on the block
+//        that is targeted). The type can also be NULL if the value type is an
+//        MVP type.
 BINARYEN_API BinaryenExpressionRef
 BinaryenBreak(BinaryenModuleRef module,
               const char* name,
               BinaryenExpressionRef condition,
-              BinaryenExpressionRef value);
+              BinaryenExpressionRef value,
+              BinaryenType* type);
 // Switch: value can be NULL
 BINARYEN_API BinaryenExpressionRef
 BinaryenSwitch(BinaryenModuleRef module,
@@ -1265,6 +1269,9 @@ BINARYEN_API void BinaryenBreakSetCondition(BinaryenExpressionRef expr,
 BINARYEN_API BinaryenExpressionRef
 BinaryenBreakGetValue(BinaryenExpressionRef expr);
 // Sets the value expression, if any, of a `br` or `br_if` expression.
+// Note that for a `br_if` with a value, if you change the type of the block
+// that this targets then you also need to change the type of the `br_if`, which
+// you can do using BinaryenExpressionSetType().
 BINARYEN_API void BinaryenBreakSetValue(BinaryenExpressionRef expr,
                                         BinaryenExpressionRef valueExpr);
 

src/ir/ReFinalize.cpp

@@ -24,6 +24,51 @@ static Type getValueType(Expression* value) {
   return value ? value->type : Type::none;
 }
 
+void ReFinalize::doWalkFunction(Function* func) {
+  using Super =
+    WalkerPass<PostWalker<ReFinalize, OverriddenVisitor<ReFinalize>>>;
+
+  while (1) {
+    // Walk normally.
+    Super::doWalkFunction(func);
+
+    // After that walk we may have br_ifs in need of refinalization. Update them
+    // and refinalize again, as they may enable further improvements. This is in
+    // theory very slow, but in practice one or two cycles suffices and we can't
+    // try to be frugal here as must propagate all the possible improvements, or
+    // else we'd end up with a situation where ReFinalize^2 != ReFinalize. (If
+    // this ends up being slow in practice we'd want to build a custom graph IR
+    // for this operation here.)
+    //
+    // Note that this loop must surely terminate because the lattice of types is
+    // finite.
+    auto updatedBr = false;
+
+    for (auto& [_, info] : blockBrInfoMap) {
+      auto* block = info.block;
+      assert(block->type.isConcrete());
+      auto blockType = block->type;
+      for (auto* br : info.brs) {
+        // We would have ignored an unreachable or MVP br_if before; only
+        // references must be refined.
+        assert(br->type.containsRef());
+        if (br->type != blockType) {
+          br->type = blockType;
+          updatedBr = true;
+        }
+      }
+    }
+
+    if (!updatedBr) {
+      return;
+    }
+
+    // Clear all state and loop/walk again.
+    breakTypes.clear();
+    blockBrInfoMap.clear();
+  }
+}
+
 void ReFinalize::visitBlock(Block* curr) {
   if (curr->list.size() == 0) {
     curr->type = Type::none;
@@ -38,6 +83,12 @@ void ReFinalize::visitBlock(Block* curr) {
       auto& types = iter->second;
       types.insert(curr->list.back()->type);
       curr->type = Type::getLeastUpperBound(types);
+
+      // If we have br_ifs then we must note ourselves there.
+      auto iter = blockBrInfoMap.find(curr->name);
+      if (iter != blockBrInfoMap.end()) {
+        iter->second.block = curr;
+      }
       return;
     }
   }
@@ -65,6 +116,10 @@ void ReFinalize::visitBreak(Break* curr) {
   } else {
     updateBreakValueType(curr->name, valueType);
   }
+  // Note relevant br_ifs for type updating later.
+  if (curr->condition && curr->value && curr->type.containsRef()) {
+    blockBrInfoMap[curr->name].brs.push_back(curr);
+  }
 }
 void ReFinalize::visitSwitch(Switch* curr) {
   curr->finalize();

src/ir/utils.h

@@ -129,10 +129,25 @@ struct ReFinalize
 
   ReFinalize() { name = "refinalize"; }
 
-  // block finalization is O(bad) if we do each block by itself, so do it in
+  // Block finalization is O(bad) if we do each block by itself, so do it in
   // bulk, tracking break value types so we just do a linear pass.
   std::unordered_map<Name, std::unordered_set<Type>> breakTypes;
 
+  // Track br_ifs with values as we go, as those must be updated if their block
+  // type changes. We must also track block names to blocks for this, so we
+  // track all that using the block name as the key.
+  struct BlockBrInfo {
+    // The block and its relevant breaks.
+    Block* block = nullptr;
+    std::vector<Break*> brs;
+
+    BlockBrInfo() {}
+    BlockBrInfo(Block* block) : block(block) {}
+  };
+  std::unordered_map<Name, BlockBrInfo> blockBrInfoMap;
+
+  void doWalkFunction(Function* func);
+
 #define DELEGATE(CLASS_TO_VISIT)                                               \
   void visit##CLASS_TO_VISIT(CLASS_TO_VISIT* curr);
 

src/passes/Heap2Local.cpp

@@ -343,8 +343,13 @@ struct Heap2LocalOptimizer {
       }
 
       // Breaks that our allocation flows through may change type, as we now
-      // have a nullable type there.
-      curr->finalize();
+      // have a nullable type there. This is simple to fix as we know this is a
+      // break with a value, and the only possible change there is to become
+      // nullable.
+      assert(curr->value);
+      if (curr->type.isNonNullable()) {
+        curr->type = Type(curr->type.getHeapType(), Nullable);
+      }
     }
 
     void visitStructNew(StructNew* curr) {

src/passes/RemoveUnusedBrs.cpp

@@ -140,6 +140,29 @@ static bool tooCostlyToRunUnconditionally(const PassOptions& passOptions,
   return tooCostlyToRunUnconditionally(passOptions, max);
 }
 
+// Utility to update a br_if type after a change. This receives the br and a
+// mapping of block names to their types, which it uses if it needs to (in
+// particular, we do not need the mapping to figure out the type of a br_if
+// when the value has MVP type, so the mapping can only contain things relevant
+// to GC). This also calls finalize as needed, so that after calling it the
+// br_if type is fully updated.
+using BlockTypeMap = std::unordered_map<Name, Type>;
+static void updateBrIfType(Break* br, const BlockTypeMap& blockTypeMap) {
+  assert(br->condition);
+  assert(br->condition);
+  if (br->value) {
+    if (br->value->type.containsRef()) {
+      auto iter = blockTypeMap.find(br->name);
+      assert(iter != blockTypeMap.end());
+      br->type = iter->second;
+    } else {
+      // A simple type we can infer from the value.
+      br->type = br->value->type;
+    }
+  }
+  br->finalize();
+}
+
 struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
   bool isFunctionParallel() override { return true; }
 
@@ -251,11 +274,14 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           break;
         }
       }
+      // Otherwise it is ok for properly-typed values to flow out.
+      self->stopUnrefinedValueFlow(curr->type);
     } else if (curr->is<Nop>()) {
       // ignore (could be result of a previous cycle)
       self->stopValueFlow();
     } else if (curr->is<Loop>()) {
-      // do nothing - it's ok for values to flow out
+      // As with a block, it is ok for properly-typed values to flow out.
+      self->stopUnrefinedValueFlow(curr->type);
     } else if (auto* sw = curr->dynCast<Switch>()) {
       self->stopFlow();
       self->optimizeSwitch(sw);
@@ -282,6 +308,52 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
 
   void stopValueFlow() { removeValueFlow(flows); }
 
+  // Stop a value from flowing that is not as refined as a given type. For
+  // example:
+  //
+  //  (block (result $sub)
+  //    ..maybe other branches to the block..
+  //    (return (super))    ;; the function returns super; the block is sub
+  //  )
+  //
+  //  If we let the value flow through, we'd be forced to unrefine the block:
+  //
+  //  (block (result $super)
+  //    ..maybe other branches to the block..
+  //    (super)
+  //  )
+  //
+  // It is best to avoid unrefining like that, so block the flow of such values
+  // here.
+  void stopUnrefinedValueFlow(Type type) {
+    if (type == Type::unreachable) {
+      // We are flowing through something currently unreachable. Us flowing
+      // through it may turn it reachable, which is fine, and there are no
+      // subtyping risks here, so nothing to check.
+      return;
+    }
+    if (!type.containsRef()) {
+      // Subtyping cannot be an issue here.
+      return;
+    }
+    flows.erase(std::remove_if(flows.begin(),
+                               flows.end(),
+                               [&](Expression** currp) {
+                                 auto* curr = *currp;
+                                 Expression* value;
+                                 if (auto* ret = curr->dynCast<Return>()) {
+                                   value = ret->value;
+                                 } else {
+                                   value = curr->cast<Break>()->value;
+                                 }
+                                 // Remove if there is a value, and it is
+                                 // unrefined.
+                                 return value &&
+                                        !Type::isSubType(value->type, type);
+                               }),
+                flows.end());
+  }
+
   static void clear(RemoveUnusedBrs* self, Expression** currp) {
     self->flows.clear();
   }
@@ -417,7 +489,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             br->condition =
               builder.makeSelect(br->condition, curr->condition, zero);
           }
-          br->finalize();
+          updateBrIfType(br, blockTypes);
           replaceCurrent(Builder(*getModule()).dropIfConcretelyTyped(br));
           anotherCycle = true;
         }
@@ -460,9 +532,9 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
   static void scan(RemoveUnusedBrs* self, Expression** currp) {
     self->pushTask(visitAny, currp);
 
-    auto* iff = (*currp)->dynCast<If>();
+    auto* curr = *currp;
 
-    if (iff) {
+    if (auto* iff = curr->dynCast<If>()) {
       if (iff->condition->type == Type::unreachable) {
         // avoid trying to optimize this, we never reach it anyhow
         return;
@@ -478,11 +550,23 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
       self->pushTask(scan, &iff->ifTrue);
       self->pushTask(clear, currp); // clear all flow after the condition
       self->pushTask(scan, &iff->condition);
-    } else {
-      super::scan(self, currp);
+      return;
+    }
+
+    if (auto* block = curr->dynCast<Block>()) {
+      if (block->name.is() && block->type.containsRef()) {
+        // This block has a type that may be needed for updating br_ifs, so
+        // note it.
+        self->blockTypes[block->name] = block->type;
+      }
     }
+
+    super::scan(self, currp);
   }
 
+  // We track the types of blocks that have relevant types for updateBrIfType.
+  std::unordered_map<Name, Type> blockTypes;
+
   // optimizes a loop. returns true if we made changes
   bool optimizeLoop(Loop* loop) {
     // if a loop ends in
@@ -1066,7 +1150,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             // we are an if-else where the ifTrue is a break without a
             // condition, so we can do this
             ifTrueBreak->condition = iff->condition;
-            ifTrueBreak->finalize();
+            updateBrIfType(ifTrueBreak, blockTypes);
             list[i] = Builder(*getModule()).dropIfConcretelyTyped(ifTrueBreak);
             ExpressionManipulator::spliceIntoBlock(curr, i + 1, iff->ifFalse);
             continue;
@@ -1080,7 +1164,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
                                 *getModule())) {
             ifFalseBreak->condition =
               Builder(*getModule()).makeUnary(EqZInt32, iff->condition);
-            ifFalseBreak->finalize();
+            updateBrIfType(ifFalseBreak, blockTypes);
             list[i] = Builder(*getModule()).dropIfConcretelyTyped(ifFalseBreak);
             ExpressionManipulator::spliceIntoBlock(curr, i + 1, iff->ifTrue);
             continue;
@@ -1676,6 +1760,21 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
           start = end;
         }
       }
+
+      // Note types of blocks. This parallels the logic in the parent class,
+      // but we do not want to just reuse the data structure there: things may
+      // have changed since then.
+      static void scan(FinalOptimizer* self, Expression** currp) {
+        if (auto* block = (*currp)->dynCast<Block>()) {
+          if (block->name.is() && block->type.containsRef()) {
+            self->blockTypes[block->name] = block->type;
+          }
+        }
+
+        PostWalker<FinalOptimizer>::scan(self, currp);
+      }
+
+      std::unordered_map<Name, Type> blockTypes;
     };
     FinalOptimizer finalOptimizer(getPassOptions());
     finalOptimizer.setModule(getModule());

src/passes/SimplifyLocals.cpp

@@ -608,11 +608,14 @@ struct SimplifyLocals
       auto* set = (*breakLocalSetPointer)->template cast<LocalSet>();
       if (br->condition) {
         br->value = set;
-        set->makeTee(this->getFunction()->getLocalType(set->index));
+        auto localType = this->getFunction()->getLocalType(set->index);
+        set->makeTee(localType);
         *breakLocalSetPointer =
           this->getModule()->allocator.template alloc<Nop>();
-        // in addition, as this is a conditional br that now has a value, it now
-        // returns a value, so it must be dropped
+        // In addition, as this is a conditional br that now has a value, it now
+        // returns a value, so we must update the type for that and also drop
+        // it.
+        br->type = localType;
         br->finalize();
         *brp = Builder(*this->getModule()).makeDrop(br);
       } else {