Fix review

Author: Changqing-JING

src/passes/CodeFolding.cpp

@@ -300,76 +300,96 @@ struct CodeFolding
       returnTails.clear();
       unoptimizables.clear();
       modifieds.clear();
-      exitingBranchCache.clear();
-      exitingBranchCachePopulated = false;
+      hasExitingBranchCache.clear();
       if (needEHFixups) {
         EHUtils::handleBlockNestedPops(func, *getModule());
       }
     }
   }
 
 private:
-  // Cache of expressions that have branches exiting to targets defined
-  // outside them. Populated lazily on first access via hasExitingBranches().
-  std::unordered_set<Expression*> exitingBranchCache;
-  bool exitingBranchCachePopulated = false;
+  // Cache of hasExitingBranches() bool results, populated on demand.
+  std::unordered_map<Expression*, bool> hasExitingBranchCache;
 
   bool hasExitingBranches(Expression* expr) {
-    if (!exitingBranchCachePopulated) {
-      populateExitingBranchCache(getFunction()->body);
-      exitingBranchCachePopulated = true;
+    auto it = hasExitingBranchCache.find(expr);
+    if (it != hasExitingBranchCache.end()) {
+      return it->second;
     }
-    return exitingBranchCache.count(expr);
+    return populateExitingBranchCache(expr);
   }
 
-  // Pre-populate the exiting branch cache for all sub-expressions of root
-  // in a single O(N) bottom-up walk. After this, exitingBranchCache
-  // lookups are O(1).
-  void populateExitingBranchCache(Expression* root) {
+  // Walk |root| bottom-up computing exiting branches. Name sets are kept
+  // transiently (moved from children, erased after merge). Only the root's
+  // bool result is persisted to keep the cache small. Already-cached
+  // subtrees are skipped via scan().
+  bool populateExitingBranchCache(Expression* root) {
     struct CachePopulator
       : public PostWalker<CachePopulator,
                           UnifiedExpressionVisitor<CachePopulator>> {
-      std::unordered_set<Expression*>& cache;
-      // Track unresolved branch targets at each node. We propagate children's
-      // targets upward: add uses, remove defs. If any remain, the expression
-      // has exiting branches.
-      std::unordered_map<Expression*, std::unordered_set<Name>> targetSets;
-
-      CachePopulator(std::unordered_set<Expression*>& cache) : cache(cache) {}
+      std::unordered_map<Expression*, bool>& resultCache;
+      Expression* root;
+      bool rootResult = false;
+      std::unordered_map<Expression*, std::unordered_set<Name>> nameSets;
+
+      CachePopulator(std::unordered_map<Expression*, bool>& resultCache,
+                     Expression* root)
+        : resultCache(resultCache), root(root) {}
+
+      static void scan(CachePopulator* self, Expression** currp) {
+        auto* curr = *currp;
+        if (self->resultCache.count(curr)) {
+          return;
+        }
+        PostWalker<CachePopulator,
+                   UnifiedExpressionVisitor<CachePopulator>>::scan(self, currp);
+      }
 
       void visitExpression(Expression* curr) {
         std::unordered_set<Name> targets;
-        // Merge children's target sets into ours (move to avoid copies)
+
+        bool childFromPriorWalkHasExiting = false;
         ChildIterator children(curr);
         for (auto* child : children) {
-          auto it = targetSets.find(child);
-          if (it != targetSets.end()) {
+          auto it = nameSets.find(child);
+          if (it != nameSets.end()) {
             if (targets.empty()) {
               targets = std::move(it->second);
             } else {
               targets.merge(it->second);
             }
-            targetSets.erase(it);
+            nameSets.erase(it);
+          } else {
+            // Conservatively propagate prior-cached children's results.
+            auto cacheIt = resultCache.find(child);
+            if (cacheIt != resultCache.end() && cacheIt->second) {
+              childFromPriorWalkHasExiting = true;
+            }
           }
         }
-        // Add branch uses (names this expression branches to)
+
         BranchUtils::operateOnScopeNameUses(
           curr, [&](Name& name) { targets.insert(name); });
-        // Remove branch defs (names this expression defines as targets)
+
         BranchUtils::operateOnScopeNameDefs(curr, [&](Name& name) {
           if (name.is()) {
             targets.erase(name);
           }
         });
-        bool hasExiting = !targets.empty();
-        if (hasExiting) {
-          cache.insert(curr);
-          targetSets[curr] = std::move(targets);
+
+        bool hasExiting = !targets.empty() || childFromPriorWalkHasExiting;
+        if (!targets.empty()) {
+          nameSets[curr] = std::move(targets);
+        }
+        if (curr == root) {
+          resultCache[curr] = hasExiting;
+          rootResult = hasExiting;
         }
       }
     };
-    CachePopulator populator(exitingBranchCache);
+    CachePopulator populator(hasExitingBranchCache, root);
     populator.walk(root);
+    return populator.rootResult;
   }
 
   // check if we can move a list of items out of another item. we can't do so