Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 63 additions & 3 deletions src/passes/CodeFolding.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -135,6 +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_;

// walking

void visitExpression(Expression* curr) {
Expand Down Expand Up @@ -299,13 +304,65 @@ struct CodeFolding
returnTails.clear();
unoptimizables.clear();
modifieds.clear();
exitingBranchCache_.clear();
if (needEHFixups) {
EHUtils::handleBlockNestedPops(func, *getModule());
}
}
}

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);
}
}
};
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
Expand Down Expand Up @@ -549,6 +606,11 @@ 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) {
Copy link
Copy Markdown
Member

@kripken kripken Apr 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. We are called more than once with num == 0, so I think this is doing more work than needed? (there are three calls to this, two with num == 0 as the default value)
  2. We may also not end up needing the cache at all, if other issues stop us earlier.
  3. We will also only need the cache for some expressions, not the entire function.

To fix those issues, how about making new line 702 call a function that checks for external break targets. That function would lazily populate a cache internally, that is, given a specific expression it would compute it and cache results for that expression and all children (avoiding walking children already in the cache).

populateExitingBranchCache(getFunction()->body);
}
// remove things that are untoward and cannot be optimized
tails.erase(
std::remove_if(tails.begin(),
Expand Down Expand Up @@ -637,9 +699,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 (exitingBranchCache_[newItem]) {
return true;
}
return false;
Expand Down
Loading