Fix review

Author: Changqing-JING

src/passes/SimplifyLocals.cpp

@@ -93,19 +93,15 @@ struct SimplifyLocals
 
   // Reverse index: for each local L, tracks which sinkable keys have effects
   // that read L. This allows checkInvalidations to find conflicting sinkables
-  // in O(|current effects|) instead of O(|all sinkables|).
+  // in O(|current effects|) instead of O(|all sinkables|) when the current
+  // expression only has local effects.
   std::unordered_map<Index, std::unordered_set<Index>> localReadBySinkable_;
 
   // Reverse index: for each local L, tracks which sinkable keys have effects
   // that write L. A sinkable at key K always writes K, but may also write
   // other locals if its value contains nested local.sets.
   std::unordered_map<Index, std::unordered_set<Index>> localWrittenBySinkable_;
 
-  // Sinkable keys that have non-local ordering-relevant effects (calls,
-  // memory, control flow, etc.). These need full orderedAfter checks when
-  // the current expression also has non-local effects.
-  std::unordered_set<Index> heavySinkables_;
-
   void registerSinkable(Index key) {
     auto& effects = sinkables.at(key).effects;
     for (auto L : effects.localsRead) {
@@ -114,9 +110,6 @@ struct SimplifyLocals
     for (auto L : effects.localsWritten) {
       localWrittenBySinkable_[L].insert(key);
     }
-    if (effects.hasNonLocalOrderingEffects()) {
-      heavySinkables_.insert(key);
-    }
   }
 
   void unregisterSinkable(Index key) {
@@ -143,20 +136,17 @@ struct SimplifyLocals
         }
       }
     }
-    heavySinkables_.erase(key);
   }
 
   void clearSinkables() {
     sinkables.clear();
     localReadBySinkable_.clear();
     localWrittenBySinkable_.clear();
-    heavySinkables_.clear();
   }
 
   Sinkables takeSinkables() {
     localReadBySinkable_.clear();
     localWrittenBySinkable_.clear();
-    heavySinkables_.clear();
     return std::move(sinkables);
   }
 
@@ -384,56 +374,56 @@ struct SimplifyLocals
   }
 
   void checkInvalidations(EffectAnalyzer& effects) {
-    // Use targeted lookups instead of iterating all sinkables.
-    // We collect candidate sinkable keys that *might* conflict, then verify.
-    std::unordered_set<Index> candidates;
-
-    // Local conflicts via reverse indices.
-    // When the current expression reads local L, any sinkable that writes L
-    // has a write-read conflict.
-    for (auto L : effects.localsRead) {
-      auto it = localWrittenBySinkable_.find(L);
-      if (it != localWrittenBySinkable_.end()) {
-        candidates.insert(it->second.begin(), it->second.end());
+    // If the current expression only has local effects, we can use reverse
+    // indices to find exactly the conflicting sinkables in O(|locals|) instead
+    // of iterating all sinkables. This is the common case for local.get,
+    // local.set, const, and simple arithmetic expressions.
+    if (!effects.hasNonLocalOrderingEffects()) {
+      std::unordered_set<Index> candidates;
+      // When the current expression reads local L, any sinkable that writes L
+      // has a write-read conflict.
+      for (auto L : effects.localsRead) {
+        auto it = localWrittenBySinkable_.find(L);
+        if (it != localWrittenBySinkable_.end()) {
+          candidates.insert(it->second.begin(), it->second.end());
+        }
       }
-    }
-    // When the current expression writes local L, any sinkable that reads L
-    // (read-write conflict) or writes L (write-write conflict) is a candidate.
-    for (auto L : effects.localsWritten) {
-      auto it = localReadBySinkable_.find(L);
-      if (it != localReadBySinkable_.end()) {
-        candidates.insert(it->second.begin(), it->second.end());
+      // When the current expression writes local L, any sinkable that reads L
+      // (read-write conflict) or writes L (write-write conflict) is a
+      // candidate.
+      for (auto L : effects.localsWritten) {
+        auto it = localReadBySinkable_.find(L);
+        if (it != localReadBySinkable_.end()) {
+          candidates.insert(it->second.begin(), it->second.end());
+        }
+        auto it2 = localWrittenBySinkable_.find(L);
+        if (it2 != localWrittenBySinkable_.end()) {
+          candidates.insert(it2->second.begin(), it2->second.end());
+        }
       }
-      auto it2 = localWrittenBySinkable_.find(L);
-      if (it2 != localWrittenBySinkable_.end()) {
-        candidates.insert(it2->second.begin(), it2->second.end());
+      std::vector<Index> invalidated;
+      for (auto key : candidates) {
+        auto it = sinkables.find(key);
+        if (it != sinkables.end() && effects.orderedAfter(it->second.effects)) {
+          invalidated.push_back(key);
+        }
       }
-    }
-
-    // Non-local conflicts: if the current expression has non-local effects,
-    // check sinkables that also have non-local effects.
-    if (effects.hasNonLocalOrderingEffects()) {
-      candidates.insert(heavySinkables_.begin(), heavySinkables_.end());
-    }
-    // If current transfers control flow, all sinkables with any side effects
-    // (including local access) are invalidated. Since all sinkables access
-    // locals, this means all of them.
-    if (effects.transfersControlFlow()) {
-      for (auto& [key, _] : sinkables) {
-        candidates.insert(key);
+      for (auto key : invalidated) {
+        eraseSinkable(key);
       }
+      return;
     }
 
-    // Verify candidates with the full ordering check and invalidate.
+    // The current expression has non-local effects (memory, calls, control
+    // flow, etc.), so we must check all sinkables.
     std::vector<Index> invalidated;
-    for (auto key : candidates) {
-      auto it = sinkables.find(key);
-      if (it != sinkables.end() && effects.orderedAfter(it->second.effects)) {
-        invalidated.push_back(key);
+    for (auto& [index, info] : sinkables) {
+      if (effects.orderedAfter(info.effects)) {
+        invalidated.push_back(index);
       }
     }
-    for (auto key : invalidated) {
-      eraseSinkable(key);
+    for (auto index : invalidated) {
+      eraseSinkable(index);
     }
   }