Author: Changqing-JING

src/passes/CodeFolding.cpp

@@ -136,9 +136,10 @@ struct CodeFolding
     modifieds; // modified code should not be processed
                // again, wait for next pass
 
-  // Cache for hasExitingBranches results. Populated by a single bottom-up
-  // walk to avoid O(N^2) repeated tree traversals on nested blocks.
-  std::unordered_map<Expression*, bool> exitingBranchCache_;
+  // Lazy cache for exiting branch checks. Key absent = not yet computed;
+  // key present with empty set = no exiting branches; non-empty set = has
+  // exiting branches (and stores the unresolved target names for propagation).
+  std::unordered_map<Expression*, std::unordered_set<Name>> exitingBranchCache_;
 
   // walking
 
@@ -312,55 +313,46 @@ struct CodeFolding
   }
 
 private:
-  // 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) {
-    struct CachePopulator
-      : public PostWalker<CachePopulator,
-                          UnifiedExpressionVisitor<CachePopulator>> {
-      std::unordered_map<Expression*, bool>& 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_map<Expression*, bool>& cache)
-        : cache(cache) {}
-
-      void visitExpression(Expression* curr) {
-        std::unordered_set<Name> targets;
-        // Merge children's target sets into ours (move to avoid copies)
-        ChildIterator children(curr);
-        for (auto* child : children) {
-          auto it = targetSets.find(child);
-          if (it != targetSets.end()) {
-            if (targets.empty()) {
-              targets = std::move(it->second);
-            } else {
-              targets.merge(it->second);
-            }
-            targetSets.erase(it);
-          }
-        }
-        // 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();
-        cache[curr] = hasExiting;
-        if (hasExiting) {
-          targetSets[curr] = std::move(targets);
-        }
+  // Lazily check if an expression has branches that exit to targets defined
+  // outside it. On the first call for a given expression, this walks the
+  // expression and all its unvisited children, caching results so that
+  // subsequent queries are O(1).
+  bool hasExitingBranches(Expression* expr) {
+    auto it = exitingBranchCache_.find(expr);
+    if (it != exitingBranchCache_.end()) {
+      return !it->second.empty();
+    }
+    computeExitingBranches(expr);
+    return !exitingBranchCache_[expr].empty();
+  }
+
+  // Recursively compute and cache exiting-branch info for expr and all
+  // unvisited children.
+  void computeExitingBranches(Expression* expr) {
+    if (exitingBranchCache_.count(expr)) {
+      return;
+    }
+    // Compute children first (bottom-up)
+    ChildIterator children(expr);
+    for (auto* child : children) {
+      computeExitingBranches(child);
+    }
+    // Merge children's unresolved targets
+    std::unordered_set<Name> targets;
+    for (auto* child : children) {
+      auto& childTargets = exitingBranchCache_[child];
+      targets.insert(childTargets.begin(), childTargets.end());
+    }
+    // Add branch uses (names this expression branches to)
+    BranchUtils::operateOnScopeNameUses(
+      expr, [&](Name& name) { targets.insert(name); });
+    // Remove branch defs (names this expression defines as targets)
+    BranchUtils::operateOnScopeNameDefs(expr, [&](Name& name) {
+      if (name.is()) {
+        targets.erase(name);
       }
-    };
-    CachePopulator populator(exitingBranchCache_);
-    populator.walk(root);
+    });
+    exitingBranchCache_[expr] = std::move(targets);
   }
 
   // check if we can move a list of items out of another item. we can't do so
@@ -606,11 +598,6 @@ struct CodeFolding
     if (tails.size() < 2) {
       return false;
     }
-    // Pre-populate the cache once at the top level so all subsequent
-    // exitingBranchCache_ lookups are O(1).
-    if (num == 0) {
-      populateExitingBranchCache(getFunction()->body);
-    }
     // remove things that are untoward and cannot be optimized
     tails.erase(
       std::remove_if(tails.begin(),
@@ -699,7 +686,7 @@ struct CodeFolding
                                 // TODO: this should not be a problem in
                                 //       *non*-terminating tails, but
                                 //       double-verify that
-                                if (exitingBranchCache_[newItem]) {
+                                if (hasExitingBranches(newItem)) {
                                   return true;
                                 }
                                 return false;