Merge pull request #66846 from kavon/standardize-workqueue

Add a `BasicBlockWorkqueue` as a common utility.

Author: swift-ci

include/swift/SIL/BasicBlockDatastructures.h

@@ -19,6 +19,7 @@
 
 #include "swift/SIL/StackList.h"
 #include "swift/SIL/BasicBlockBits.h"
+#include <deque>
 
 namespace swift {
 
@@ -61,10 +62,12 @@ class BasicBlockSetVector {
   }
 };
 
-/// A utility for processing basic blocks in a worklist.
+/// A utility for processing basic blocks in a worklist that uses
+/// last-in-first-out order.
 ///
-/// It is basically a combination of a block vector and a block set. It can be
-/// used for typical worklist-processing algorithms.
+/// It is basically a combination of a stack list and a block set. It can be
+/// used for typical worklist-processing algorithms and is recommended in
+/// nearly all situations as its implementation is well-optimized.
 class BasicBlockWorklist {
   StackList<SILBasicBlock *> worklist;
   BasicBlockSet visited;
@@ -119,6 +122,70 @@ class BasicBlockWorklist {
   bool isVisited(SILBasicBlock *block) const { return visited.contains(block); }
 };
 
+
+/// A utility for processing basic blocks in a work-queue that uses
+/// first-in-first-out order.
+///
+/// It is basically a combination of a std::deque and a block set. It can be
+/// used for certain kinds of worklist-processing algorithms.
+///
+/// When in doubt, you should reach for a BasicBlockWorklist instead.
+class BasicBlockWorkqueue {
+  std::deque<SILBasicBlock*> workqueue;
+  BasicBlockSet visited;
+
+public:
+  /// Construct an empty workqueue.
+  BasicBlockWorkqueue(SILFunction *function)
+      : workqueue(), visited(function) {}
+
+  /// Initialize the workqueue with \p initialBlock.
+  BasicBlockWorkqueue(SILBasicBlock *initialBlock)
+      : BasicBlockWorkqueue(initialBlock->getParent()) {
+    push(initialBlock);
+  }
+
+  /// Pops the next element from the front of the workqueue or returns null,
+  /// if the workqueue is empty.
+  SILBasicBlock *pop() {
+    if (workqueue.empty())
+      return nullptr;
+    SILBasicBlock *block = workqueue.front();
+    workqueue.pop_front();
+    return block;
+  }
+
+  /// Pushes \p block onto the end of the workqueue if \p block has never been
+  /// pushed before.
+  bool pushIfNotVisited(SILBasicBlock *block) {
+    if (visited.insert(block)) {
+      workqueue.push_back(block);
+      return true;
+    }
+    return false;
+  }
+
+  /// Like `pushIfNotVisited`, but requires that \p block has never been on the
+  /// workqueue before.
+  void push(SILBasicBlock *block) {
+    assert(!visited.contains(block));
+    visited.insert(block);
+    workqueue.push_back(block);
+  }
+
+  /// Like `pop`, but marks the returned block as "unvisited". This means, that
+  /// the block can be pushed onto the queue again.
+  SILBasicBlock *popAndForget() {
+    SILBasicBlock *block = pop();
+    if (block)
+      visited.erase(block);
+    return block;
+  }
+
+  /// Returns true if \p block was visited, i.e. has been added to the workqueue
+  bool isVisited(SILBasicBlock *block) const { return visited.contains(block); }
+};
+
 } // namespace swift
 
 #endif

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -1890,25 +1890,16 @@ bool PullbackCloner::Implementation::run() {
   // ignored.
   // The original blocks in traversal order for pullback generation.
   SmallVector<SILBasicBlock *, 8> originalBlocks;
-  // The set of visited original blocks.
-  SmallDenseSet<SILBasicBlock *, 8> visitedBlocks;
+  // The workqueue used for bookkeeping during the breadth-first traversal.
+  BasicBlockWorkqueue workqueue = {originalExitBlock};
 
   // Perform BFS from the original exit block.
   {
-    std::deque<SILBasicBlock *> worklist = {};
-    worklist.push_back(originalExitBlock);
-    visitedBlocks.insert(originalExitBlock);
-    while (!worklist.empty()) {
-      auto *BB = worklist.front();
-      worklist.pop_front();
-
+    while (auto *BB = workqueue.pop()) {
       originalBlocks.push_back(BB);
 
       for (auto *nextBB : BB->getPredecessorBlocks()) {
-        if (!visitedBlocks.count(nextBB)) {
-          worklist.push_back(nextBB);
-          visitedBlocks.insert(nextBB);
-        }
+        workqueue.pushIfNotVisited(nextBB);
       }
     }
   }
@@ -2004,7 +1995,7 @@ bool PullbackCloner::Implementation::run() {
     //   pullback original block, passing adjoint values of active values.
     for (auto *succBB : origBB->getSuccessorBlocks()) {
       // Skip generating pullback block for original unreachable blocks.
-      if (!visitedBlocks.count(succBB))
+      if (!workqueue.isVisited(succBB))
         continue;
       auto *pullbackTrampolineBB = pullback.createBasicBlockBefore(pullbackBB);
       pullbackTrampolineBBMap.insert({{origBB, succBB}, pullbackTrampolineBB});

lib/SILOptimizer/Mandatory/FlowIsolation.cpp

@@ -20,6 +20,7 @@
 #include "swift/SIL/BitDataflow.h"
 #include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 
 using namespace swift;
@@ -277,34 +278,13 @@ SILInstruction *AnalysisInfo::findNonisolatedBlame(SILInstruction* startInst) {
   SILBasicBlock* firstBlk = startInst->getParent();
   assert(firstBlk->getParent() == getFunction());
 
-  // workList for breadth-first search to find one of the closest blocks.
-  std::deque<SILBasicBlock*> workList;
-  BasicBlockSet visited(getFunction());
-
-  // seed the search
-  workList.push_back(firstBlk);
-  SILBasicBlock::reverse_iterator cursor = startInst->getReverseIterator();
-
-  while (!workList.empty()) {
-    auto *block = workList.front();
-    workList.pop_front();
+  // searches the a block starting at the provided position in reverse
+  // order of instructions (i.e., from terminator to first instruction).
+  auto searchBlockForNonisolated =
+      [&](SILBasicBlock::reverse_iterator cursor) -> SILInstruction * {
+    SILBasicBlock *block = cursor->getParent();
     auto &state = flow[block];
 
-    // if this block doesn't have exiting nonisolation, then there's no
-    // way we'll find nonisolation in a predecessor.
-    assert(state.exitSet[State::Nonisolated] && "nonisolation is unreachable!");
-
-    // If this is the first time we're scanning the start block, then leave
-    // the cursor alone and do a partial scan. If the block is part of
-    // a cycle, we want to scan the block entirely on the second visit, so we
-    // do not count this as a visit.
-    if (startInst) {
-      startInst = nullptr; // make sure second visit scans entirely.
-    } else {
-      cursor = block->rbegin();
-      visited.insert(block);
-    }
-
     // does this block generate non-isolation?
     if (state.genSet[State::Nonisolated]) {
       auto &data = this->operator[](block);
@@ -321,16 +301,36 @@ SILInstruction *AnalysisInfo::findNonisolatedBlame(SILInstruction* startInst) {
       }
     }
 
-    for (auto *pred : block->getPredecessorBlocks()) {
-      // skip visited
-      if (visited.contains(pred))
-        continue;
+    return nullptr;
+  };
+
+  // whether we should visit a given predecessor block in the search.
+  auto shouldVisit = [&](SILBasicBlock *pred) {
+    // visit blocks that contribute nonisolation to successors.
+    return flow[pred].exitSet[State::Nonisolated];
+  };
 
-      // skip blocks that do not contribute nonisolation.
-      if (flow[pred].exitSet[State::Nonisolated] == false)
-        continue;
+  // first check if the nonisolated use precedes the start instruction in
+  // this same block.
+  if (auto *inst = searchBlockForNonisolated(startInst->getReverseIterator()))
+    return inst;
 
-      workList.push_back(pred);
+  // Seed a workQueue with the predecessors of this start block to
+  // begin a breadth-first search to find one of the closest predecessors.
+  BasicBlockWorkqueue workQueue(firstBlk->getFunction());
+  for (auto *pred : firstBlk->getPredecessorBlocks())
+    if (shouldVisit(pred))
+      workQueue.push(pred);
+
+  while (auto *block = workQueue.pop()) {
+    // do we have a nonisolated use here?
+    if (auto *inst = searchBlockForNonisolated(block->rbegin()))
+      return inst;
+
+    // otherwise keep looking
+    for (auto *pred : block->getPredecessorBlocks()) {
+      if (shouldVisit(pred))
+        workQueue.pushIfNotVisited(pred);
     }
   }
 

lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp

@@ -290,12 +290,8 @@ void DiagnosticEmitter::emitMissingConsumeInDiscardingContext(
     //          / \        / \
     //         d   d      .   .
     //
-    BasicBlockSet visited(destroyPoint->getFunction());
-    std::deque<SILBasicBlock *> bfsWorklist = {destroyPoint->getParent()};
-    while (auto *bb = bfsWorklist.front()) {
-      visited.insert(bb);
-      bfsWorklist.pop_front();
-
+    BasicBlockWorkqueue bfsWorklist = {destroyPoint->getParent()};
+    while (auto *bb = bfsWorklist.pop()) {
       TermInst *term = bb->getTerminator();
 
       // Looking for a block that exits the function or terminates the program.
@@ -315,8 +311,7 @@ void DiagnosticEmitter::emitMissingConsumeInDiscardingContext(
       }
 
       for (auto *nextBB : term->getSuccessorBlocks())
-        if (!visited.contains(nextBB))
-          bfsWorklist.push_back(nextBB);
+        bfsWorklist.pushIfNotVisited(nextBB);
     }
   }