[clang][dataflow] Extend transfer functions for other `CFGElement`s

Previously, the transfer function `void transfer(const Stmt *, ...)` overriden by users is restricted to apply only on `CFGStmt`s and its contained `Stmt`.

By using a transfer function (`void transfer(const CFGElement *, ...)`) that takes a `CFGElement` as input, this patch extends user-defined analysis to all kinds of `CFGElement`. For example, users can now handle `CFGInitializer`s where `CXXCtorInitializer` AST nodes are contained.

Reviewed By: gribozavr2, sgatev

Differential Revision: https://reviews.llvm.org/D131614

5 files changed, 265 insertions, 121 deletions

diff --git a/clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h b/clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h
index 4d1f5248f211..4e084d57ba01 100644
--- a/clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h
+++ b/clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h
@@ -15,11 +15,13 @@
 #define LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DATAFLOWANALYSIS_H
 
 #include <iterator>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Stmt.h"
+#include "clang/Analysis/CFG.h"
 #include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #include "clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h"
@@ -31,6 +33,17 @@
 namespace clang {
 namespace dataflow {
 
+template <typename AnalysisT, typename LatticeT, typename InputT,
+          typename = std::void_t<>>
+struct HasTransferFor : std::false_type {};
+
+template <typename AnalysisT, typename LatticeT, typename InputT>
+struct HasTransferFor<
+    AnalysisT, LatticeT, InputT,
+    std::void_t<decltype(std::declval<AnalysisT>().transfer(
+        std::declval<const InputT *>(), std::declval<LatticeT &>(),
+        std::declval<Environment &>()))>> : std::true_type {};
+
 /// Base class template for dataflow analyses built on a single lattice type.
 ///
 /// Requirements:
@@ -39,8 +52,9 @@ namespace dataflow {
 ///  must provide the following public members:
 ///   * `LatticeT initialElement()` - returns a lattice element that models the
 ///     initial state of a basic block;
-///   * `void transfer(const Stmt *, LatticeT &, Environment &)` - applies the
-///     analysis transfer function for a given statement and lattice element.
+///   * `void transfer(const CFGElement *, LatticeT &, Environment &)` - applies
+///     the analysis transfer function for a given CFG element and lattice
+///     element.
 ///
 ///  `Derived` can optionally override the following members:
 ///   * `bool merge(QualType, const Value &, const Value &, Value &,
@@ -93,10 +107,20 @@ public:
     return L1 == L2;
   }
 
-  void transferTypeErased(const Stmt *Stmt, TypeErasedLattice &E,
+  void transferTypeErased(const CFGElement *Element, TypeErasedLattice &E,
                           Environment &Env) final {
     Lattice &L = llvm::any_cast<Lattice &>(E.Value);
-    static_cast<Derived *>(this)->transfer(Stmt, L, Env);
+    if constexpr (HasTransferFor<Derived, LatticeT, CFGElement>::value) {
+      static_cast<Derived *>(this)->transfer(Element, L, Env);
+    }
+
+    // FIXME: Remove after users have been updated to implement `transfer` on
+    // `CFGElement`.
+    if constexpr (HasTransferFor<Derived, LatticeT, Stmt>::value) {
+      if (auto Stmt = Element->getAs<CFGStmt>()) {
+        static_cast<Derived *>(this)->transfer(Stmt->getStmt(), L, Env);
+      }
+    }
   }
 
 private:
@@ -112,37 +136,41 @@ template <typename LatticeT> struct DataflowAnalysisState {
   Environment Env;
 };
 
+// FIXME: Rename to `runDataflowAnalysis` after usages of the overload that
+// applies to `CFGStmt` have been replaced.
+//
 /// Performs dataflow analysis and returns a mapping from basic block IDs to
 /// dataflow analysis states that model the respective basic blocks. The
 /// returned vector, if any, will have the same size as the number of CFG
 /// blocks, with indices corresponding to basic block IDs. Returns an error if
 /// the dataflow analysis cannot be performed successfully. Otherwise, calls
-/// `PostVisitStmt` on each statement with the final analysis results at that
+/// `PostVisitCFG` on each CFG element with the final analysis results at that
 /// program point.
 template <typename AnalysisT>
 llvm::Expected<std::vector<
     llvm::Optional<DataflowAnalysisState<typename AnalysisT::Lattice>>>>
-runDataflowAnalysis(
+runDataflowAnalysisOnCFG(
     const ControlFlowContext &CFCtx, AnalysisT &Analysis,
     const Environment &InitEnv,
-    std::function<void(const CFGStmt &, const DataflowAnalysisState<
-                                            typename AnalysisT::Lattice> &)>
-        PostVisitStmt = nullptr) {
-  std::function<void(const CFGStmt &, const TypeErasedDataflowAnalysisState &)>
-      PostVisitStmtClosure = nullptr;
-  if (PostVisitStmt != nullptr) {
-    PostVisitStmtClosure = [&PostVisitStmt](
-                               const CFGStmt &Stmt,
-                               const TypeErasedDataflowAnalysisState &State) {
+    std::function<void(const CFGElement &, const DataflowAnalysisState<
+                                               typename AnalysisT::Lattice> &)>
+        PostVisitCFG = nullptr) {
+  std::function<void(const CFGElement &,
+                     const TypeErasedDataflowAnalysisState &)>
+      PostVisitCFGClosure = nullptr;
+  if (PostVisitCFG) {
+    PostVisitCFGClosure = [&PostVisitCFG](
+                              const CFGElement &Element,
+                              const TypeErasedDataflowAnalysisState &State) {
       auto *Lattice =
           llvm::any_cast<typename AnalysisT::Lattice>(&State.Lattice.Value);
-      PostVisitStmt(Stmt, DataflowAnalysisState<typename AnalysisT::Lattice>{
-                              *Lattice, State.Env});
+      PostVisitCFG(Element, DataflowAnalysisState<typename AnalysisT::Lattice>{
+                                *Lattice, State.Env});
     };
   }
 
   auto TypeErasedBlockStates = runTypeErasedDataflowAnalysis(
-      CFCtx, Analysis, InitEnv, PostVisitStmtClosure);
+      CFCtx, Analysis, InitEnv, PostVisitCFGClosure);
   if (!TypeErasedBlockStates)
     return TypeErasedBlockStates.takeError();
 
@@ -163,6 +191,41 @@ runDataflowAnalysis(
   return BlockStates;
 }
 
+/// Deprecated. Use `runDataflowAnalysisOnCFG` instead.
+///
+/// Performs dataflow analysis and returns a mapping from basic block IDs to
+/// dataflow analysis states that model the respective basic blocks. The
+/// returned vector, if any, will have the same size as the number of CFG
+/// blocks, with indices corresponding to basic block IDs. Returns an error if
+/// the dataflow analysis cannot be performed successfully. Otherwise, calls
+/// `PostVisitStmt` on each statement with the final analysis results at that
+/// program point.
+template <typename AnalysisT>
+llvm::Expected<std::vector<
+    llvm::Optional<DataflowAnalysisState<typename AnalysisT::Lattice>>>>
+runDataflowAnalysis(
+    const ControlFlowContext &CFCtx, AnalysisT &Analysis,
+    const Environment &InitEnv,
+    std::function<void(const CFGStmt &, const DataflowAnalysisState<
+                                            typename AnalysisT::Lattice> &)>
+        PostVisitStmt = nullptr) {
+  std::function<void(
+      const CFGElement &,
+      const DataflowAnalysisState<typename AnalysisT::Lattice> &)>
+      PostVisitCFG = nullptr;
+  if (PostVisitStmt) {
+    PostVisitCFG =
+        [&PostVisitStmt](
+            const CFGElement &Element,
+            const DataflowAnalysisState<typename AnalysisT::Lattice> &State) {
+          if (auto Stmt = Element.getAs<CFGStmt>()) {
+            PostVisitStmt(*Stmt, State);
+          }
+        };
+  }
+  return runDataflowAnalysisOnCFG(CFCtx, Analysis, InitEnv, PostVisitCFG);
+}
+
 /// Abstract base class for dataflow "models": reusable analysis components that
 /// model a particular aspect of program semantics in the `Environment`. For
 /// example, a model may capture a type and its related functions.
diff --git a/clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h b/clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h
index 58acda7e6389..c7dccc36a6a2 100644
--- a/clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h
+++ b/clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h
@@ -79,9 +79,9 @@ public:
   virtual bool isEqualTypeErased(const TypeErasedLattice &,
                                  const TypeErasedLattice &) = 0;
 
-  /// Applies the analysis transfer function for a given statement and
-  /// type-erased lattice element.
-  virtual void transferTypeErased(const Stmt *, TypeErasedLattice &,
+  /// Applies the analysis transfer function for a given control flow graph
+  /// element and type-erased lattice element.
+  virtual void transferTypeErased(const CFGElement *, TypeErasedLattice &,
                                   Environment &) = 0;
 
   /// If the built-in transfer functions (which model the heap and stack in the
@@ -104,10 +104,10 @@ struct TypeErasedDataflowAnalysisState {
       : Lattice(std::move(Lattice)), Env(std::move(Env)) {}
 };
 
-/// Transfers the state of a basic block by evaluating each of its statements in
+/// Transfers the state of a basic block by evaluating each of its elements in
 /// the context of `Analysis` and the states of its predecessors that are
-/// available in `BlockStates`. `HandleTransferredStmt` (if provided) will be
-/// applied to each statement in the block, after it is evaluated.
+/// available in `BlockStates`. `PostVisitCFG` (if provided) will be applied to
+/// each element in the block, after it is evaluated.
 ///
 /// Requirements:
 ///
@@ -119,23 +119,23 @@ TypeErasedDataflowAnalysisState transferBlock(
     llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates,
     const CFGBlock &Block, const Environment &InitEnv,
     TypeErasedDataflowAnalysis &Analysis,
-    std::function<void(const CFGStmt &,
+    std::function<void(const CFGElement &,
                        const TypeErasedDataflowAnalysisState &)>
-        HandleTransferredStmt = nullptr);
+        PostVisitCFG = nullptr);
 
 /// Performs dataflow analysis and returns a mapping from basic block IDs to
 /// dataflow analysis states that model the respective basic blocks. Indices of
 /// the returned vector correspond to basic block IDs. Returns an error if the
 /// dataflow analysis cannot be performed successfully. Otherwise, calls
-/// `PostVisitStmt` on each statement with the final analysis results at that
+/// `PostVisitCFG` on each CFG element with the final analysis results at that
 /// program point.
 llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
 runTypeErasedDataflowAnalysis(
     const ControlFlowContext &CFCtx, TypeErasedDataflowAnalysis &Analysis,
     const Environment &InitEnv,
-    std::function<void(const CFGStmt &,
+    std::function<void(const CFGElement &,
                        const TypeErasedDataflowAnalysisState &)>
-        PostVisitStmt = nullptr);
+        PostVisitCFG = nullptr);
 
 } // namespace dataflow
 } // namespace clang
diff --git a/clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp b/clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp
index dc2ecef771b6..b85a72a8a7b7 100644
--- a/clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp
+++ b/clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp
@@ -154,19 +154,37 @@ private:
   TransferOptions TransferOpts;
 };
 
+/// Holds data structures required for running dataflow analysis.
+struct AnalysisContext {
+  AnalysisContext(
+      const ControlFlowContext &CFCtx, TypeErasedDataflowAnalysis &Analysis,
+      const Environment &InitEnv,
+      llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>>
+          BlockStates)
+      : CFCtx(CFCtx), Analysis(Analysis), InitEnv(InitEnv),
+        BlockStates(BlockStates) {}
+
+  /// Contains the CFG being analyzed.
+  const ControlFlowContext &CFCtx;
+  /// The analysis to be run.
+  TypeErasedDataflowAnalysis &Analysis;
+  /// Initial state to start the analysis.
+  const Environment &InitEnv;
+  /// Stores the state of a CFG block if it has been evaluated by the analysis.
+  /// The indices correspond to the block IDs.
+  llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates;
+};
+
 /// Computes the input state for a given basic block by joining the output
 /// states of its predecessors.
 ///
 /// Requirements:
 ///
 ///   All predecessors of `Block` except those with loop back edges must have
-///   already been transferred. States in `BlockStates` that are set to
+///   already been transferred. States in `AC.BlockStates` that are set to
 ///   `llvm::None` represent basic blocks that are not evaluated yet.
-static TypeErasedDataflowAnalysisState computeBlockInputState(
-    const ControlFlowContext &CFCtx,
-    llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates,
-    const CFGBlock &Block, const Environment &InitEnv,
-    TypeErasedDataflowAnalysis &Analysis) {
+static TypeErasedDataflowAnalysisState
+computeBlockInputState(const CFGBlock &Block, AnalysisContext &AC) {
   llvm::DenseSet<const CFGBlock *> Preds;
   Preds.insert(Block.pred_begin(), Block.pred_end());
   if (Block.getTerminator().isTemporaryDtorsBranch()) {
@@ -193,13 +211,16 @@ static TypeErasedDataflowAnalysisState computeBlockInputState(
     //
     // See `NoreturnDestructorTest` for concrete examples.
     if (Block.succ_begin()->getReachableBlock()->hasNoReturnElement()) {
-      auto StmtBlock = CFCtx.getStmtToBlock().find(Block.getTerminatorStmt());
-      assert(StmtBlock != CFCtx.getStmtToBlock().end());
+      auto &StmtToBlock = AC.CFCtx.getStmtToBlock();
+      auto StmtBlock = StmtToBlock.find(Block.getTerminatorStmt());
+      assert(StmtBlock != StmtToBlock.end());
       Preds.erase(StmtBlock->getSecond());
     }
   }
 
   llvm::Optional<TypeErasedDataflowAnalysisState> MaybeState;
+
+  auto &Analysis = AC.Analysis;
   auto BuiltinTransferOpts = Analysis.builtinTransferOptions();
 
   for (const CFGBlock *Pred : Preds) {
@@ -210,14 +231,14 @@ static TypeErasedDataflowAnalysisState computeBlockInputState(
     // Skip if `Pred` was not evaluated yet. This could happen if `Pred` has a
     // loop back edge to `Block`.
     const llvm::Optional<TypeErasedDataflowAnalysisState> &MaybePredState =
-        BlockStates[Pred->getBlockID()];
+        AC.BlockStates[Pred->getBlockID()];
     if (!MaybePredState)
       continue;
 
     TypeErasedDataflowAnalysisState PredState = MaybePredState.value();
     if (BuiltinTransferOpts) {
       if (const Stmt *PredTerminatorStmt = Pred->getTerminatorStmt()) {
-        const StmtToEnvMapImpl StmtToEnv(CFCtx, BlockStates);
+        const StmtToEnvMapImpl StmtToEnv(AC.CFCtx, AC.BlockStates);
         TerminatorVisitor(StmtToEnv, PredState.Env,
                           blockIndexInPredecessor(*Pred, Block),
                           *BuiltinTransferOpts)
@@ -236,107 +257,125 @@ static TypeErasedDataflowAnalysisState computeBlockInputState(
     // FIXME: Consider passing `Block` to `Analysis.typeErasedInitialElement()`
     // to enable building analyses like computation of dominators that
     // initialize the state of each basic block differently.
-    MaybeState.emplace(Analysis.typeErasedInitialElement(), InitEnv);
+    MaybeState.emplace(Analysis.typeErasedInitialElement(), AC.InitEnv);
   }
   return *MaybeState;
 }
 
-/// Transfers `State` by evaluating `CfgStmt` in the context of `Analysis`.
-/// `HandleTransferredStmt` (if provided) will be applied to `CfgStmt`, after it
-/// is evaluated.
-static void transferCFGStmt(
-    const ControlFlowContext &CFCtx,
-    llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates,
-    const CFGStmt &CfgStmt, TypeErasedDataflowAnalysis &Analysis,
-    TypeErasedDataflowAnalysisState &State,
-    std::function<void(const CFGStmt &,
-                       const TypeErasedDataflowAnalysisState &)>
-        HandleTransferredStmt) {
-  const Stmt *S = CfgStmt.getStmt();
+/// Built-in transfer function for `CFGStmt`.
+void builtinTransferStatement(const CFGStmt &Elt,
+                              TypeErasedDataflowAnalysisState &InputState,
+                              AnalysisContext &AC) {
+  const Stmt *S = Elt.getStmt();
   assert(S != nullptr);
-
-  auto BuiltinTransferOpts = Analysis.builtinTransferOptions();
-  if (BuiltinTransferOpts)
-    transfer(StmtToEnvMapImpl(CFCtx, BlockStates), *S, State.Env,
-             *BuiltinTransferOpts);
-  Analysis.transferTypeErased(S, State.Lattice, State.Env);
-
-  if (HandleTransferredStmt != nullptr)
-    HandleTransferredStmt(CfgStmt, State);
+  transfer(StmtToEnvMapImpl(AC.CFCtx, AC.BlockStates), *S, InputState.Env,
+           *AC.Analysis.builtinTransferOptions());
 }
 
-/// Transfers `State` by evaluating `CfgInit`.
-static void transferCFGInitializer(const CFGInitializer &CfgInit,
-                                   TypeErasedDataflowAnalysisState &State) {
-  const auto &ThisLoc = *cast<AggregateStorageLocation>(
-      State.Env.getThisPointeeStorageLocation());
+/// Built-in transfer function for `CFGInitializer`.
+void builtinTransferInitializer(const CFGInitializer &Elt,
+                                TypeErasedDataflowAnalysisState &InputState) {
+  const CXXCtorInitializer *Init = Elt.getInitializer();
+  assert(Init != nullptr);
 
-  const CXXCtorInitializer *Initializer = CfgInit.getInitializer();
-  assert(Initializer != nullptr);
+  auto &Env = InputState.Env;
+  const auto &ThisLoc =
+      *cast<AggregateStorageLocation>(Env.getThisPointeeStorageLocation());
 
-  const FieldDecl *Member = Initializer->getMember();
+  const FieldDecl *Member = Init->getMember();
   if (Member == nullptr)
     // Not a field initializer.
     return;
 
-  auto *InitStmt = Initializer->getInit();
+  auto *InitStmt = Init->getInit();
   assert(InitStmt != nullptr);
 
-  auto *InitStmtLoc =
-      State.Env.getStorageLocation(*InitStmt, SkipPast::Reference);
+  auto *InitStmtLoc = Env.getStorageLocation(*InitStmt, SkipPast::Reference);
   if (InitStmtLoc == nullptr)
     return;
 
-  auto *InitStmtVal = State.Env.getValue(*InitStmtLoc);
+  auto *InitStmtVal = Env.getValue(*InitStmtLoc);
   if (InitStmtVal == nullptr)
     return;
 
   if (Member->getType()->isReferenceType()) {
     auto &MemberLoc = ThisLoc.getChild(*Member);
-    State.Env.setValue(MemberLoc,
-                       State.Env.takeOwnership(
-                           std::make_unique<ReferenceValue>(*InitStmtLoc)));
+    Env.setValue(MemberLoc, Env.takeOwnership(std::make_unique<ReferenceValue>(
+                                *InitStmtLoc)));
   } else {
     auto &MemberLoc = ThisLoc.getChild(*Member);
-    State.Env.setValue(MemberLoc, *InitStmtVal);
+    Env.setValue(MemberLoc, *InitStmtVal);
+  }
+}
+
+void builtinTransfer(const CFGElement &Elt,
+                     TypeErasedDataflowAnalysisState &State,
+                     AnalysisContext &AC) {
+  switch (Elt.getKind()) {
+  case CFGElement::Statement: {
+    builtinTransferStatement(Elt.castAs<CFGStmt>(), State, AC);
+    break;
+  }
+  case CFGElement::Initializer: {
+    builtinTransferInitializer(Elt.castAs<CFGInitializer>(), State);
+    break;
+  }
+  default:
+    // FIXME: Evaluate other kinds of `CFGElement`.
+    break;
   }
 }
 
+/// Transfers `State` by evaluating each element in the `Block` based on the
+/// `AC.Analysis` specified.
+///
+/// Built-in transfer functions (if the option for `ApplyBuiltinTransfer` is set
+/// by the analysis) will be applied to the element before evaluation by the
+/// user-specified analysis.
+/// `PostVisitCFG` (if provided) will be applied to the element after evaluation
+/// by the user-specified analysis.
+TypeErasedDataflowAnalysisState
+transferCFGBlock(const CFGBlock &Block, AnalysisContext &AC,
+                 std::function<void(const CFGElement &,
+                                    const TypeErasedDataflowAnalysisState &)>
+                     PostVisitCFG = nullptr) {
+  auto State = computeBlockInputState(Block, AC);
+  for (const auto &Element : Block) {
+    // Built-in analysis
+    if (AC.Analysis.builtinTransferOptions()) {
+      builtinTransfer(Element, State, AC);
+    }
+
+    // User-provided analysis
+    AC.Analysis.transferTypeErased(&Element, State.Lattice, State.Env);
+
+    // Post processing
+    if (PostVisitCFG) {
+      PostVisitCFG(Element, State);
+    }
+  }
+  return State;
+}
+
 TypeErasedDataflowAnalysisState transferBlock(
     const ControlFlowContext &CFCtx,
     llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates,
     const CFGBlock &Block, const Environment &InitEnv,
     TypeErasedDataflowAnalysis &Analysis,
-    std::function<void(const CFGStmt &,
+    std::function<void(const CFGElement &,
                        const TypeErasedDataflowAnalysisState &)>
-        HandleTransferredStmt) {
-  TypeErasedDataflowAnalysisState State =
-      computeBlockInputState(CFCtx, BlockStates, Block, InitEnv, Analysis);
-  for (const CFGElement &Element : Block) {
-    switch (Element.getKind()) {
-    case CFGElement::Statement:
-      transferCFGStmt(CFCtx, BlockStates, *Element.getAs<CFGStmt>(), Analysis,
-                      State, HandleTransferredStmt);
-      break;
-    case CFGElement::Initializer:
-      if (Analysis.builtinTransferOptions())
-        transferCFGInitializer(*Element.getAs<CFGInitializer>(), State);
-      break;
-    default:
-      // FIXME: Evaluate other kinds of `CFGElement`.
-      break;
-    }
-  }
-  return State;
+        PostVisitCFG) {
+  AnalysisContext AC(CFCtx, Analysis, InitEnv, BlockStates);
+  return transferCFGBlock(Block, AC, PostVisitCFG);
 }
 
 llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
 runTypeErasedDataflowAnalysis(
     const ControlFlowContext &CFCtx, TypeErasedDataflowAnalysis &Analysis,
     const Environment &InitEnv,
-    std::function<void(const CFGStmt &,
+    std::function<void(const CFGElement &,
                        const TypeErasedDataflowAnalysisState &)>
-        PostVisitStmt) {
+        PostVisitCFG) {
   PostOrderCFGView POV(&CFCtx.getCFG());
   ForwardDataflowWorklist Worklist(CFCtx.getCFG(), &POV);
 
@@ -349,6 +388,8 @@ runTypeErasedDataflowAnalysis(
                                      InitEnv};
   Worklist.enqueueSuccessors(&Entry);
 
+  AnalysisContext AC(CFCtx, Analysis, InitEnv, BlockStates);
+
   // Bugs in lattices and transfer functions can prevent the analysis from
   // converging. To limit the damage (infinite loops) that these bugs can cause,
   // limit the number of iterations.
@@ -373,7 +414,7 @@ runTypeErasedDataflowAnalysis(
     const llvm::Optional<TypeErasedDataflowAnalysisState> &OldBlockState =
         BlockStates[Block->getBlockID()];
     TypeErasedDataflowAnalysisState NewBlockState =
-        transferBlock(CFCtx, BlockStates, *Block, InitEnv, Analysis);
+        transferCFGBlock(*Block, AC);
 
     if (OldBlockState &&
         Analysis.isEqualTypeErased(OldBlockState.value().Lattice,
@@ -395,14 +436,12 @@ runTypeErasedDataflowAnalysis(
   // FIXME: Consider evaluating unreachable basic blocks (those that have a
   // state set to `llvm::None` at this point) to also analyze dead code.
 
-  if (PostVisitStmt) {
+  if (PostVisitCFG) {
     for (const CFGBlock *Block : CFCtx.getCFG()) {
       // Skip blocks that were not evaluated.
       if (!BlockStates[Block->getBlockID()])
         continue;
-
-      transferBlock(CFCtx, BlockStates, *Block, InitEnv, Analysis,
-                    PostVisitStmt);
+      transferCFGBlock(*Block, AC, PostVisitCFG);
     }
   }
 
diff --git a/clang/unittests/Analysis/FlowSensitive/SingleVarConstantPropagationTest.cpp b/clang/unittests/Analysis/FlowSensitive/SingleVarConstantPropagationTest.cpp
index 3c95ead0ab8a..664950065844 100644
--- a/clang/unittests/Analysis/FlowSensitive/SingleVarConstantPropagationTest.cpp
+++ b/clang/unittests/Analysis/FlowSensitive/SingleVarConstantPropagationTest.cpp
@@ -194,8 +194,8 @@ void RunDataflow(llvm::StringRef Code, Matcher Expectations) {
           },
           [&Expectations](
               llvm::ArrayRef<std::pair<
-                  std::string, DataflowAnalysisState<
-                                    ConstantPropagationAnalysis::Lattice>>>
+                  std::string,
+                  DataflowAnalysisState<ConstantPropagationAnalysis::Lattice>>>
                   Results,
               ASTContext &) { EXPECT_THAT(Results, Expectations); },
           {"-fsyntax-only", "-std=c++17"}),
diff --git a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
index f09d5f9db21a..cdcf61ee2257 100644
--- a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
+++ b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
@@ -71,14 +71,25 @@ struct AnalysisData {
   std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates;
 };
 
+// FIXME: Rename to `checkDataflow` after usages of the overload that applies to
+// `CFGStmt` have been replaced.
+//
+/// Runs dataflow analysis (specified from `MakeAnalysis`) and the
+/// `PostVisitCFG` function (if provided) on the body of the function that
+/// matches `TargetFuncMatcher` in code snippet `Code`. `VerifyResults` checks
+/// that the results from the analysis are correct.
+///
+/// Requirements:
+///
+///  `AnalysisT` contains a type `Lattice`.
 template <typename AnalysisT>
-llvm::Error checkDataflow(
+llvm::Error checkDataflowOnCFG(
     llvm::StringRef Code,
     ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
-    std::function<void(ASTContext &, const CFGStmt &,
+    std::function<void(ASTContext &, const CFGElement &,
                        const TypeErasedDataflowAnalysisState &)>
-        PostVisitStmt,
+        PostVisitCFG,
     std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
   llvm::Annotations AnnotatedCode(Code);
@@ -112,13 +123,14 @@ llvm::Error checkDataflow(
   Environment Env(DACtx, *F);
   auto Analysis = MakeAnalysis(Context, Env);
 
-  std::function<void(const CFGStmt &, const TypeErasedDataflowAnalysisState &)>
-      PostVisitStmtClosure = nullptr;
-  if (PostVisitStmt != nullptr) {
-    PostVisitStmtClosure = [&PostVisitStmt, &Context](
-                               const CFGStmt &Stmt,
-                               const TypeErasedDataflowAnalysisState &State) {
-      PostVisitStmt(Context, Stmt, State);
+  std::function<void(const CFGElement &,
+                     const TypeErasedDataflowAnalysisState &)>
+      PostVisitCFGClosure = nullptr;
+  if (PostVisitCFG) {
+    PostVisitCFGClosure = [&PostVisitCFG, &Context](
+                              const CFGElement &Element,
+                              const TypeErasedDataflowAnalysisState &State) {
+      PostVisitCFG(Context, Element, State);
     };
   }
 
@@ -130,7 +142,7 @@ llvm::Error checkDataflow(
 
   llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
       MaybeBlockStates = runTypeErasedDataflowAnalysis(*CFCtx, Analysis, Env,
-                                                       PostVisitStmtClosure);
+                                                       PostVisitCFGClosure);
   if (!MaybeBlockStates)
     return MaybeBlockStates.takeError();
   auto &BlockStates = *MaybeBlockStates;
@@ -141,6 +153,32 @@ llvm::Error checkDataflow(
   return llvm::Error::success();
 }
 
+template <typename AnalysisT>
+llvm::Error checkDataflow(
+    llvm::StringRef Code,
+    ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
+    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
+    std::function<void(ASTContext &, const CFGStmt &,
+                       const TypeErasedDataflowAnalysisState &)>
+        PostVisitStmt,
+    std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
+    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
+  std::function<void(ASTContext & Context, const CFGElement &,
+                     const TypeErasedDataflowAnalysisState &)>
+      PostVisitCFG = nullptr;
+  if (PostVisitStmt) {
+    PostVisitCFG =
+        [&PostVisitStmt](ASTContext &Context, const CFGElement &Element,
+                         const TypeErasedDataflowAnalysisState &State) {
+          if (auto Stmt = Element.getAs<CFGStmt>()) {
+            PostVisitStmt(Context, *Stmt, State);
+          }
+        };
+  }
+  return checkDataflowOnCFG(Code, TargetFuncMatcher, MakeAnalysis, PostVisitCFG,
+                            VerifyResults, Args, VirtualMappedFiles);
+}
+
 // Runs dataflow on the body of the function that matches `TargetFuncMatcher` in
 // code snippet `Code`. Requires: `AnalysisT` contains a type `Lattice`.
 template <typename AnalysisT>
@@ -157,9 +195,9 @@ llvm::Error checkDataflow(
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
   using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
 
-  return checkDataflow(
+  return checkDataflowOnCFG(
       Code, std::move(TargetFuncMatcher), std::move(MakeAnalysis),
-      /*PostVisitStmt=*/nullptr,
+      /*PostVisitCFG=*/nullptr,
       [&VerifyResults](AnalysisData AnalysisData) {
         if (AnalysisData.BlockStates.empty()) {
           VerifyResults({}, AnalysisData.ASTCtx);
@@ -180,9 +218,13 @@ llvm::Error checkDataflow(
               AnalysisData.CFCtx, AnalysisData.BlockStates, *Block,
               AnalysisData.Env, AnalysisData.Analysis,
               [&Results,
-               &Annotations](const clang::CFGStmt &Stmt,
+               &Annotations](const clang::CFGElement &Element,
                              const TypeErasedDataflowAnalysisState &State) {
-                auto It = Annotations.find(Stmt.getStmt());
+                // FIXME: Extend testing annotations to non statement constructs
+                auto Stmt = Element.getAs<CFGStmt>();
+                if (!Stmt)
+                  return;
+                auto It = Annotations.find(Stmt->getStmt());
                 if (It == Annotations.end())
                   return;
                 auto *Lattice = llvm::any_cast<typename AnalysisT::Lattice>(