Merge branch 'main' into oac.lazy

src/ir/possible-contents.cpp

@@ -503,12 +503,30 @@ struct CollectedFuncInfo {
   std::unordered_map<Expression*, Expression*> childParents;
 };
 
+// Does a walk while maintaining a map of names of branch targets to those
+// expressions, so they can be found by their name.
+// TODO: can this replace ControlFlowWalker in other places?
+template<typename SubType, typename VisitorType = Visitor<SubType>>
+struct BreakTargetWalker : public PostWalker<SubType, VisitorType> {
+  std::unordered_map<Name, Expression*> breakTargets;
+
+  Expression* findBreakTarget(Name name) { return breakTargets[name]; }
+
+  static void scan(SubType* self, Expression** currp) {
+    auto* curr = *currp;
+    BranchUtils::operateOnScopeNameDefs(
+      curr, [&](Name name) { self->breakTargets[name] = curr; });
+
+    PostWalker<SubType, VisitorType>::scan(self, currp);
+  }
+};
+
 // Walk the wasm and find all the links we need to care about, and the locations
 // and roots related to them. This builds up a CollectedFuncInfo data structure.
 // After all InfoCollectors run, those data structures will be merged and the
 // main flow will begin.
 struct InfoCollector
-  : public PostWalker<InfoCollector, OverriddenVisitor<InfoCollector>> {
+  : public BreakTargetWalker<InfoCollector, OverriddenVisitor<InfoCollector>> {
   CollectedFuncInfo& info;
 
   InfoCollector(CollectedFuncInfo& info) : info(info) {}
@@ -553,9 +571,6 @@ struct InfoCollector
       return;
     }
 
-    // Values sent to breaks to this block must be received here.
-    handleBreakTarget(curr);
-
     // The final item in the block can flow a value to here as well.
     receiveChildValue(curr->list.back(), curr);
   }
@@ -1151,8 +1166,7 @@ struct InfoCollector
         for (Index i = 0; i < params.size(); i++) {
           if (isRelevant(params[i])) {
             info.links.push_back(
-              {TagLocation{tag, i},
-               BreakTargetLocation{getFunction(), target, i}});
+              {TagLocation{tag, i}, getBreakTargetLocation(target, i)});
           }
         }
 
@@ -1164,7 +1178,7 @@ struct InfoCollector
         addRoot(location,
                 PossibleContents::fromType(Type(HeapType::exn, NonNullable)));
         info.links.push_back(
-          {location, BreakTargetLocation{getFunction(), target, exnrefIndex}});
+          {location, getBreakTargetLocation(target, exnrefIndex)});
       }
     }
   }
@@ -1279,6 +1293,13 @@ struct InfoCollector
 
   // Helpers
 
+  // Returns the location of a break target by the name (e.g. returns the
+  // location of a block, if the name is the name of a block). Also receives the
+  // index in a tuple, if this is part of a tuple value.
+  Location getBreakTargetLocation(Name target, Index i) {
+    return ExpressionLocation{findBreakTarget(target), i};
+  }
+
   // Handles the value sent in a break instruction. Does not handle anything
   // else like the condition etc.
   void handleBreakValue(Expression* curr) {
@@ -1288,26 +1309,13 @@ struct InfoCollector
           for (Index i = 0; i < value->type.size(); i++) {
             // Breaks send the contents of the break value to the branch target
             // that the break goes to.
-            info.links.push_back(
-              {ExpressionLocation{value, i},
-               BreakTargetLocation{getFunction(), target, i}});
+            info.links.push_back({ExpressionLocation{value, i},
+                                  getBreakTargetLocation(target, i)});
           }
         }
       });
   }
 
-  // Handles receiving values from breaks at the target (as in a block).
-  void handleBreakTarget(Expression* curr) {
-    if (isRelevant(curr->type)) {
-      BranchUtils::operateOnScopeNameDefs(curr, [&](Name target) {
-        for (Index i = 0; i < curr->type.size(); i++) {
-          info.links.push_back({BreakTargetLocation{getFunction(), target, i},
-                                ExpressionLocation{curr, i}});
-        }
-      });
-    }
-  }
-
   // Connect a child's value to the parent, that is, all content in the child is
   // now considered possible in the parent as well.
   void receiveChildValue(Expression* child, Expression* parent) {
@@ -2400,16 +2408,16 @@ bool Flower::updateContents(LocationIndex locationIndex,
     }
   }
 
-  // After filtering we should always have more precise information than "many"
-  // - in the worst case, we can have the type declared in the wasm.
-  assert(!contents.isMany());
-
 #if defined(POSSIBLE_CONTENTS_DEBUG) && POSSIBLE_CONTENTS_DEBUG >= 2
   std::cout << "  updateContents has something new\n";
   contents.dump(std::cout, &wasm);
   std::cout << '\n';
 #endif
 
+  // After filtering we should always have more precise information than "many"
+  // - in the worst case, we can have the type declared in the wasm.
+  assert(!contents.isMany());
+
   // Add a work item if there isn't already.
   workQueue.insert(locationIndex);
 
@@ -2896,9 +2904,6 @@ void Flower::dump(Location location) {
               << '\n';
   } else if (auto* loc = std::get_if<GlobalLocation>(&location)) {
     std::cout << "  globalloc " << loc->name << '\n';
-  } else if (auto* loc = std::get_if<BreakTargetLocation>(&location)) {
-    std::cout << "  branchloc " << loc->func->name << " : " << loc->target
-              << " tupleIndex " << loc->tupleIndex << '\n';
   } else if (std::get_if<SignatureParamLocation>(&location)) {
     std::cout << "  sigparamloc " << '\n';
   } else if (std::get_if<SignatureResultLocation>(&location)) {

src/ir/possible-contents.h

@@ -402,21 +402,6 @@ struct ResultLocation {
   }
 };
 
-// The location of a break target in a function, identified by its name.
-struct BreakTargetLocation {
-  Function* func;
-  Name target;
-  // As in ExpressionLocation, the index inside the tuple, or 0 if not a tuple.
-  // That is, if the branch target has a tuple type, then each branch to that
-  // location sends a tuple, and we'll have a separate BreakTargetLocation for
-  // each, indexed by the index in the tuple that the branch sends.
-  Index tupleIndex;
-  bool operator==(const BreakTargetLocation& other) const {
-    return func == other.func && target == other.target &&
-           tupleIndex == other.tupleIndex;
-  }
-};
-
 // The location of a global in the module.
 struct GlobalLocation {
   Name name;
@@ -523,7 +508,6 @@ using Location = std::variant<ExpressionLocation,
                               ParamLocation,
                               LocalLocation,
                               ResultLocation,
-                              BreakTargetLocation,
                               GlobalLocation,
                               SignatureParamLocation,
                               SignatureResultLocation,
@@ -577,13 +561,6 @@ template<> struct hash<wasm::ResultLocation> {
   }
 };
 
-template<> struct hash<wasm::BreakTargetLocation> {
-  size_t operator()(const wasm::BreakTargetLocation& loc) const {
-    return std::hash<std::tuple<size_t, wasm::Name, wasm::Index>>{}(
-      {size_t(loc.func), loc.target, loc.tupleIndex});
-  }
-};
-
 template<> struct hash<wasm::GlobalLocation> {
   size_t operator()(const wasm::GlobalLocation& loc) const {
     return std::hash<wasm::Name>{}(loc.name);

src/passes/OptimizeAddedConstants.cpp

@@ -32,13 +32,35 @@
 
 #include <ir/local-graph.h>
 #include <ir/local-utils.h>
-#include <ir/parents.h>
 #include <pass.h>
 #include <wasm-builder.h>
 #include <wasm.h>
 
 namespace wasm {
 
+namespace {
+
+// Similar to Parents from parents.h, but we only care about gets, so it is much
+// more efficient to just collect their parents.
+struct GetParents {
+  GetParents(Expression* expr) { inner.walk(expr); }
+
+  Expression* getParent(LocalGet* curr) const {
+    auto iter = inner.parentMap.find(curr);
+    assert(iter != inner.parentMap.end());
+    return iter->second;
+  }
+
+private:
+  struct Inner : public ExpressionStackWalker<Inner> {
+    void visitLocalGet(LocalGet* curr) { parentMap[curr] = getParent(); }
+
+    std::unordered_map<Expression*, Expression*> parentMap;
+  } inner;
+};
+
+} // anonymous namespace
+
 template<typename P, typename T> class MemoryAccessOptimizer {
 public:
   MemoryAccessOptimizer(P* parent,
@@ -354,7 +376,7 @@ struct OptimizeAddedConstants
     //  g(a, offset=10)
     // but if x has other uses, then avoid doing so - we'll be doing that add
     // anyhow, so the load/store offset trick won't actually help.
-    Parents parents(getFunction()->body);
+    GetParents parents(getFunction()->body);
     for (auto& [location, _] : localGraph->getLocations()) {
       if (auto* set = location->dynCast<LocalSet>()) {
         if (auto* add = set->value->dynCast<Binary>()) {