diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp index e53cd4f880d..20a8597ac11 100644 --- a/src/passes/CodeFolding.cpp +++ b/src/passes/CodeFolding.cpp @@ -63,6 +63,7 @@ #include "ir/effects.h" #include "ir/eh-utils.h" #include "ir/find_all.h" +#include "ir/iteration.h" #include "ir/label-utils.h" #include "ir/utils.h" #include "pass.h" @@ -299,6 +300,8 @@ struct CodeFolding returnTails.clear(); unoptimizables.clear(); modifieds.clear(); + exitingBranchCache.clear(); + exitingBranchCachePopulated = false; if (needEHFixups) { EHUtils::handleBlockNestedPops(func, *getModule()); } @@ -306,6 +309,69 @@ struct CodeFolding } private: + // Cache of expressions that have branches exiting to targets defined + // outside them. Populated lazily on first access via hasExitingBranches(). + std::unordered_set exitingBranchCache; + bool exitingBranchCachePopulated = false; + + bool hasExitingBranches(Expression* expr) { + if (!exitingBranchCachePopulated) { + populateExitingBranchCache(getFunction()->body); + exitingBranchCachePopulated = true; + } + return exitingBranchCache.count(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) { + struct CachePopulator + : public PostWalker> { + std::unordered_set& 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> targetSets; + + CachePopulator(std::unordered_set& cache) : cache(cache) {} + + void visitExpression(Expression* curr) { + std::unordered_set 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(); + if (hasExiting) { + cache.insert(curr); + targetSets[curr] = std::move(targets); + } + } + }; + CachePopulator populator(exitingBranchCache); + populator.walk(root); + } + // check if we can move a list of items out of another item. we can't do so // if one of the items has a branch to something inside outOf that is not // inside that item @@ -637,9 +703,7 @@ struct CodeFolding // TODO: this should not be a problem in // *non*-terminating tails, but // double-verify that - if (EffectAnalyzer( - getPassOptions(), *getModule(), newItem) - .hasExternalBreakTargets()) { + if (hasExitingBranches(newItem)) { return true; } return false;