perf: cache repeated tree walks to avoid O(N^2) in optimizeTerminatingTails

Author: Changqing-JING

src/passes/CodeFolding.cpp

@@ -140,6 +140,12 @@ struct CodeFolding
   // walk to avoid O(N^2) repeated tree traversals on nested blocks.
   std::unordered_map<Expression*, bool> exitingBranchCache_;
 
+  // Cached data for the function body, computed once per
+  // optimizeTerminatingTails call to avoid repeated O(N) tree walks.
+  BranchUtils::NameSet bodyBranchTargets_;
+  bool bodyHasTry_ = false;
+  bool bodyHasTryTable_ = false;
+
   // walking
 
   void visitExpression(Expression* curr) {
@@ -368,6 +374,27 @@ struct CodeFolding
   // inside that item
   bool canMove(const std::vector<Expression*>& items, Expression* outOf) {
     auto allTargets = BranchUtils::getBranchTargets(outOf);
+    bool hasTry = false;
+    bool hasTryTable = false;
+    if (getModule()->features.hasExceptionHandling()) {
+      hasTry = FindAll<Try>(outOf).has();
+      hasTryTable = FindAll<TryTable>(outOf).has();
+    }
+    return canMoveImpl(items, allTargets, hasTry, hasTryTable);
+  }
+
+  // Like canMove, but uses precomputed branch targets and Try/TryTable
+  // presence. This avoids repeated O(N) tree walks when outOf is the
+  // function body and canMove is called multiple times.
+  bool canMoveWithCachedBodyInfo(const std::vector<Expression*>& items) {
+    return canMoveImpl(
+      items, bodyBranchTargets_, bodyHasTry_, bodyHasTryTable_);
+  }
+
+  bool canMoveImpl(const std::vector<Expression*>& items,
+                   const BranchUtils::NameSet& allTargets,
+                   bool hasTry,
+                   bool hasTryTable) {
     for (auto* item : items) {
       auto exiting = BranchUtils::getExitingBranches(item);
       std::vector<Name> intersection;
@@ -398,9 +425,7 @@ struct CodeFolding
         // conservative approximation because there can be cases that
         // 'try'/'try_table' is within the expression that may throw so it is
         // safe to take the expression out.
-        // TODO: optimize this check to avoid two FindAlls.
-        if (effects.throws() &&
-            (FindAll<Try>(outOf).has() || FindAll<TryTable>(outOf).has())) {
+        if (effects.throws() && (hasTry || hasTryTable)) {
           return false;
         }
       }
@@ -610,6 +635,13 @@ struct CodeFolding
     // exitingBranchCache_ lookups are O(1).
     if (num == 0) {
       populateExitingBranchCache(getFunction()->body);
+      // Cache branch targets and Try/TryTable presence for the function body
+      // so that canMove doesn't re-walk the entire body on every call.
+      bodyBranchTargets_ = BranchUtils::getBranchTargets(getFunction()->body);
+      if (getModule()->features.hasExceptionHandling()) {
+        bodyHasTry_ = FindAll<Try>(getFunction()->body).has();
+        bodyHasTryTable_ = FindAll<TryTable>(getFunction()->body).has();
+      }
     }
     // remove things that are untoward and cannot be optimized
     tails.erase(
@@ -672,8 +704,7 @@ struct CodeFolding
       cost += WORTH_ADDING_BLOCK_TO_REMOVE_THIS_MUCH;
       // if we cannot merge to the end, then we definitely need 2 blocks,
       // and a branch
-      // TODO: efficiency, entire body
-      if (!canMove(items, getFunction()->body)) {
+      if (!canMoveWithCachedBodyInfo(items)) {
         cost += 1 + WORTH_ADDING_BLOCK_TO_REMOVE_THIS_MUCH;
         // TODO: to do this, we need to maintain a map of element=>parent,
         //       so that we can insert the new blocks in the right place