[clang][FlowSensitive] Do a quick check and bail early for massive CFGs

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

@@ -11,6 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <cstddef>
 #include <optional>
 #include <system_error>
 #include <utility>
@@ -34,6 +35,7 @@
 #include "clang/Basic/LLVM.h"
 #include "clang/Support/Compiler.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
@@ -479,6 +481,27 @@ transferCFGBlock(const CFGBlock &Block, AnalysisContext &AC,
   return State;
 }
 
+// Returns the number of reachable blocks (would be visited if we only visit
+// each reachable block once). This is a light version of the main fixpoint loop
+// (keep in sync).
+size_t NumBlockVisitsIfVisitOnce(const CFG &CFG) {
+  PostOrderCFGView POV(&CFG);
+  ForwardDataflowWorklist Worklist(CFG, &POV);
+  llvm::BitVector VisitedBlocks(CFG.size());
+  const CFGBlock &Entry = CFG.getEntry();
+  Worklist.enqueueSuccessors(&Entry);
+  while (const CFGBlock *Block = Worklist.dequeue()) {
+    if (VisitedBlocks[Block->getBlockID()])
+      continue;
+    VisitedBlocks[Block->getBlockID()] = true;
+    // Do not add unreachable successor blocks to `Worklist`.
+    if (Block->hasNoReturnElement())
+      continue;
+    Worklist.enqueueSuccessors(Block);
+  }
+  return VisitedBlocks.count();
+}
+
 llvm::Expected<std::vector<std::optional<TypeErasedDataflowAnalysisState>>>
 runTypeErasedDataflowAnalysis(
     const AdornedCFG &ACFG, TypeErasedDataflowAnalysis &Analysis,
@@ -496,6 +519,20 @@ runTypeErasedDataflowAnalysis(
       MaybeStartingEnv ? *MaybeStartingEnv : InitEnv;
 
   const clang::CFG &CFG = ACFG.getCFG();
+
+  // Bail out if the number of reachable blocks is already beyond the maximum
+  // number of block visits to save time and avoid runnign out of memory for
+  // massive CFGs. Note: CFG.size() could have unreachable blocks,
+  // but it is a faster to compare that to MaxBlockVisits first before doing the
+  // reachable block count.
+  if (CFG.size() > static_cast<size_t>(MaxBlockVisits)) {
+    if (NumBlockVisitsIfVisitOnce(CFG) > static_cast<size_t>(MaxBlockVisits)) {
+      return llvm::createStringError(
+          std::errc::timed_out, "number of blocks in cfg will lead to "
+                                "exceeding maximum number of block visits");
+    }
+  }
+
   PostOrderCFGView POV(&CFG);
   ForwardDataflowWorklist Worklist(CFG, &POV);
   llvm::SmallDenseSet<const CFGBlock *> NonStructLoopBackedgeNodes =

clang/unittests/Analysis/FlowSensitive/TypeErasedDataflowAnalysisTest.cpp

@@ -58,7 +58,8 @@ class DataflowAnalysisTest : public Test {
   template <typename AnalysisT>
   llvm::Expected<std::vector<
       std::optional<DataflowAnalysisState<typename AnalysisT::Lattice>>>>
-  runAnalysis(llvm::StringRef Code, AnalysisT (*MakeAnalysis)(ASTContext &)) {
+  runAnalysis(llvm::StringRef Code, AnalysisT (*MakeAnalysis)(ASTContext &),
+              int32_t MaxBlockVisits = MaxBlockVisitsInAnalysis) {
     AST = tooling::buildASTFromCodeWithArgs(Code, {"-std=c++11"});
 
     auto *Func = selectFirst<FunctionDecl>(
@@ -75,7 +76,7 @@ class DataflowAnalysisTest : public Test {
         std::make_unique<WatchedLiteralsSolver>());
     Environment Env(*DACtx, *Func);
 
-    return runDataflowAnalysis(*ACFG, Analysis, Env);
+    return runDataflowAnalysis(*ACFG, Analysis, Env, {}, MaxBlockVisits);
   }
 
   /// Returns the `CFGBlock` containing `S` (and asserts that it exists).
@@ -450,6 +451,34 @@ TEST_F(DataflowAnalysisTest, NonConvergingAnalysis) {
             "maximum number of blocks processed");
 }
 
+TEST_F(DataflowAnalysisTest, MaxBlockVisitsWithUnreachableBlocks) {
+  std::string Code = R"(
+    void target() {
+      if (false) {
+        (void)0;
+      }
+    }
+  )";
+
+  // There are 4 blocks in the CFG:
+  //   entry, if-header, then-body, and exit.
+  // But, only 2 of them are reachable, if we do not count the entry block,
+  // which is how we count BlockVisits. So:
+  // - MaxBlockVisits = 2 should still succeed (being less than CFG.size() = 4,
+  //   shouldn't prevent the analysis from succeeding)
+  // - MaxBlockVisits = 1 should fail, as it is not enough to cover all of the
+  //   reachable blocks.
+  EXPECT_THAT_EXPECTED(
+      runAnalysis<NoopAnalysis>(
+          Code, [](ASTContext &C) { return NoopAnalysis(C); }, 2),
+      llvm::Succeeded());
+  EXPECT_THAT_EXPECTED(
+      runAnalysis<NoopAnalysis>(
+          Code, [](ASTContext &C) { return NoopAnalysis(C); }, 1),
+      llvm::FailedWithMessage("number of blocks in cfg will lead to exceeding "
+                              "maximum number of block visits"));
+}
+
 // Regression test for joins of bool-typed lvalue expressions. The first loop
 // results in two passes through the code that follows. Each pass results in a
 // different `StorageLocation` for the pointee of `v`. Then, the second loop