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

[wyt]

Summary

Previously, the transfer function void transfer(const Stmt *, ...) overriden by users is restricted to apply only on CFGStmts 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 CFGInitializers where CXXCtorInitializer AST nodes are contained.

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

104

Instead of adding virtual function, please continue following the CRTP pattern. Add the expected function declarations in the comment above the class.

112–118

116

Ditto.

143

166

"PostVisitCFG"?

217

PTAL at the intended usage of getAs() here: llvm-project/clang-tools-extra/clang-tidy/utils/ExprSequence.cpp

That is, it should be used like dyn_cast in if statements.

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

264–265

265

"CfgStmt" is violating the style guide casing rules.

275–280

Ditto.

330

Should be using castAs().

334

castAs()

444

Please add the newline.

[wyt]

Address comments.

[wyt]

Add FIXME.

[wyt]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

104

Instead of adding virtual function, please continue following the CRTP pattern. Add the expected function declarations in the comment above the class.

As discussed, moving to a CRTP pattern causes the code to break as they have not been implemented by users yet. Added a FIXME to do so after users have been migrated to implement transferCFGElement.

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

147

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

167

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

81

Please use triple slash for doc comments.

86

[wyt]

Address comments. Unable to rename check/runDataflowOnCFG to check/runDataflow currently due to ambiguous use.

[wyt]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

147

Cannot be renamed currently, as the use between the function that applies to CFGElement and CFGStmt is ambiguous. Added a FIXME to rename after we have updated the implementation for users and remove the function that acts on CFGStmts.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

86

Same as above, using the same name for the function that applies to CFGElement and CFGStmt is ambiguous.

[sgatev]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

103

Any reason not to call this one transfer too?

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

157

Let's start struct and member comments with ///.

184

312

314

328–340

The level of abstraction for built-in and user code is different in this function which makes it hard to follow. Let's move this in a builtinTransfer(const CFGElement &, TypeErasedDataflowAnalysisState &) function that dispatches to one of the other built-in transfer functions. This way transferCFGBlock won't mix implementation details with the high level algorithm.

341–347

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

164–170

[wyt]

Address comments.

[wyt]

Rebase.

[thakis]

I don't know how you commited this, but the commit message is missing all the (great!) details that are here on the review. See here: https://github.com/llvm/llvm-project/commit/4b815eb4fde0202434c6 or here: https://reviews.llvm.org/rG4b815eb4fde0202434c63

Please update your commit workflow so that the description on phab makes it into the commit going forward :)

[wyt]

@thakis
Thanks for pointing that out, will revert and recommit - since it also ran into some build failures.

[wyt]

Mark override for virtual transfer function.

[wyt]

Replace use of virtual functions with sfinea and CRTP to call transfer functions defined by the user.

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

public:
 
   bool isEqualTypeErased(const TypeErasedLattice &E1,
                          const TypeErasedLattice &E2) final {
     const Lattice &L1 = llvm::any_cast<const Lattice &>(E1.Value);
     const Lattice &L2 = llvm::any_cast<const Lattice &>(E2.Value);
     return L1 == L2;
   }
 
-  void transferTypeErased(const Stmt *Stmt, TypeErasedLattice &E,
+  /// Deprecated. Use the more general `transferCFGElement` function.
+  ///
+  /// Transfer function for statements in the code being analysed.
+  virtual void transfer(const Stmt *Stmt, Lattice &L, Environment &Env) {}
+
+  /// Transfer function for elements in the control flow graph built from the
+  /// code being analysed.
+  virtual void transferCFGElement(const CFGElement *Element, Lattice &L,
+                                  Environment &Env) {}
+
+  // FIXME: Use CRTP pattern and remove virtual transfer functions after users
+  // have been updated to implement transferCFGElement.
+  // (e.g. static_cast<Derived*>(this)->transferCFGElement(Element, L, Env))
+  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);
+    transferCFGElement(Element, L, Env);
+
+    // FIXME: Remove after users have been updated to implement
+    // transferCFGElement.
+    if (auto Stmt = Element->getAs<CFGStmt>()) {
+      transfer(Stmt->getStmt(), L, Env);
+    }
   }
 
 private:
   ASTContext &Context;
 };
 
 // Model of the program at a given program point.
 template <typename LatticeT> struct DataflowAnalysisState {
   // Model of a program property.
   LatticeT Lattice;
 
   // Model of the state of the program (store and heap).
   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<
+    std::function<void(const CFGElement &, 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) {
+        PostVisitCFG = nullptr) {
+  std::function<void(const CFGElement &,
+                     const TypeErasedDataflowAnalysisState &)>
+      PostVisitCFGClosure = nullptr;
+  if (PostVisitCFG) {
+    PostVisitCFGClosure = [&PostVisitCFG](
+                              const CFGElement &Element,
       auto *Lattice =
           llvm::any_cast<typename AnalysisT::Lattice>(&State.Lattice.Value);
-      PostVisitStmt(Stmt, DataflowAnalysisState<typename AnalysisT::Lattice>{
+      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();
 
   std::vector<
       llvm::Optional<DataflowAnalysisState<typename AnalysisT::Lattice>>>
       BlockStates;
   BlockStates.reserve(TypeErasedBlockStates->size());
 
   llvm::transform(std::move(*TypeErasedBlockStates),
                   std::back_inserter(BlockStates), [](auto &OptState) {
                     return std::move(OptState).map([](auto &&State) {
                       return DataflowAnalysisState<typename AnalysisT::Lattice>{
                           llvm::any_cast<typename AnalysisT::Lattice>(
                               std::move(State.Lattice.Value)),
                           std::move(State.Env)};
                     });
                   });
   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.
 class DataflowModel : public Environment::ValueModel {
 public:
   /// Return value indicates whether the model processed the `Stmt`.
   virtual bool transfer(const Stmt *Stmt, Environment &Env) = 0;
 };
 
 } // namespace dataflow
 } // namespace clang
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DATAFLOWANALYSIS_H

clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h

public:
   virtual LatticeJoinEffect joinTypeErased(TypeErasedLattice &,
                                            const TypeErasedLattice &) = 0;
 
   /// Returns true if and only if the two given type-erased lattice elements are
   /// equal.
   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
   /// `Environment`) are to be applied, returns the options to be passed to
   /// them. Otherwise returns empty.
   llvm::Optional<TransferOptions> builtinTransferOptions() const {
     return Options.BuiltinTransferOpts;
   }
 };
 
 /// Type-erased model of the program at a given program point.
 struct TypeErasedDataflowAnalysisState {
   /// Type-erased model of a program property.
   TypeErasedLattice Lattice;
 
   /// Model of the state of the program (store and heap).
   Environment Env;
 
   TypeErasedDataflowAnalysisState(TypeErasedLattice Lattice, Environment Env)
       : 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:
 ///
 ///   All predecessors of `Block` except those with loop back edges must have
 ///   already been transferred. States in `BlockStates` that are set to
 ///   `llvm::None` represent basic blocks that are not evaluated yet.
 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 = 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
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_TYPEERASEDDATAFLOWANALYSIS_H

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

private:
   }
 
   const StmtToEnvMap &StmtToEnv;
   Environment &Env;
   int BlockSuccIdx;
   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 `TC.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()) {
     // This handles a special case where the code that produced the CFG includes
     // a conditional operator with a branch that constructs a temporary and
     // calls a destructor annotated as noreturn. The CFG models this as follows:
     //
     // B1 (contains the condition of the conditional operator) - succs: B2, B3
if (Block.getTerminator().isTemporaryDtorsBranch()) {
     // noreturn block and its state should not be used as input for the state
     // of the block that has a temporary destructor terminator (B4 above). This
     // holds regardless of which branch of the ternary operator calls the
     // noreturn destructor. However, it doesn't cases where a nested ternary
     // operator includes a branch that contains a noreturn destructor call.
     //
     // 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) {
     // Skip if the `Block` is unreachable or control flow cannot get past it.
     if (!Pred || Pred->hasNoReturnElement())
       continue;
 
     // 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)
             .Visit(PredTerminatorStmt);
       }
     }
 
     if (MaybeState) {
       Analysis.joinTypeErased(MaybeState->Lattice, PredState.Lattice);
       MaybeState->Env.join(PredState.Env, Analysis);
     } else {
       MaybeState = std::move(PredState);
     }
   }
   if (!MaybeState) {
     // 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 builtinTransfer(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 builtinTransfer(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);
   }
 }
 
-TypeErasedDataflowAnalysisState transferBlock(
-    const ControlFlowContext &CFCtx,
-    llvm::ArrayRef<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates,
-    const CFGBlock &Block, const Environment &InitEnv,
-    TypeErasedDataflowAnalysis &Analysis,
-    std::function<void(const CFGStmt &,
+/// Transfers `State` by evaluating each element in the `Block` based on the
+/// `TC.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 &)>
-        HandleTransferredStmt) {
-  TypeErasedDataflowAnalysisState State =
-      computeBlockInputState(CFCtx, BlockStates, Block, InitEnv, Analysis);
-  for (const CFGElement &Element : Block) {
+                     PostVisitCFG = nullptr) {
+  auto State = computeBlockInputState(Block, AC);
+  for (const auto &Element : Block) {
+    // Built-in analysis
+    if (AC.Analysis.builtinTransferOptions()) {
       switch (Element.getKind()) {
-    case CFGElement::Statement:
-      transferCFGStmt(CFCtx, BlockStates, *Element.getAs<CFGStmt>(), Analysis,
-                      State, HandleTransferredStmt);
+      case CFGElement::Statement: {
+        builtinTransfer(Element.castAs<CFGStmt>(), State, AC);
         break;
-    case CFGElement::Initializer:
-      if (Analysis.builtinTransferOptions())
-        transferCFGInitializer(*Element.getAs<CFGInitializer>(), State);
+      }
+      case CFGElement::Initializer: {
+        builtinTransfer(Element.castAs<CFGInitializer>(), State);
         break;
+      }
       default:
         // FIXME: Evaluate other kinds of `CFGElement`.
         break;
       }
     }
+    // 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 CFGElement &,
+                       const TypeErasedDataflowAnalysisState &)>
+        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);
 
   std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates(
       CFCtx.getCFG().size(), llvm::None);
 
   // The entry basic block doesn't contain statements so it can be skipped.
   const CFGBlock &Entry = CFCtx.getCFG().getEntry();
   BlockStates[Entry.getBlockID()] = {Analysis.typeErasedInitialElement(),
                                      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.
   // FIXME: Consider making the maximum number of iterations configurable.
   // FIXME: Consider restricting the number of backedges followed, rather than
   // iterations.
   // FIXME: Set up statistics (see llvm/ADT/Statistic.h) to count average number
   // of iterations, number of functions that time out, etc.
   static constexpr uint32_t MaxAverageVisitsPerBlock = 4;
   static constexpr uint32_t AbsoluteMaxIterations = 1 << 16;
   const uint32_t RelativeMaxIterations =
       MaxAverageVisitsPerBlock * BlockStates.size();
   const uint32_t MaxIterations =
       std::min(RelativeMaxIterations, AbsoluteMaxIterations);
   uint32_t Iterations = 0;
   while (const CFGBlock *Block = Worklist.dequeue()) {
     if (++Iterations > MaxIterations) {
       return llvm::createStringError(std::errc::timed_out,
                                      "maximum number of iterations reached");
     }
 
     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,
                                    NewBlockState.Lattice) &&
         OldBlockState->Env.equivalentTo(NewBlockState.Env, Analysis)) {
       // The state of `Block` didn't change after transfer so there's no need to
       // revisit its successors.
       continue;
     }
 
     BlockStates[Block->getBlockID()] = std::move(NewBlockState);
 
     // Do not add unreachable successor blocks to `Worklist`.
     if (Block->hasNoReturnElement())
       continue;
 
     Worklist.enqueueSuccessors(Block);
   }
   // 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);
     }
   }
 
   return BlockStates;
 }
 
 } // namespace dataflow
 } // namespace clang

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

struct AnalysisData {
   ASTContext &ASTCtx;
   const ControlFlowContext &CFCtx;
   const Environment &Env;
   TypeErasedDataflowAnalysis &Analysis;
   llvm::DenseMap<const clang::Stmt *, std::string> &Annotations;
   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);
   auto Unit = tooling::buildASTFromCodeWithArgs(
       AnnotatedCode.code(), Args, "input.cc", "clang-dataflow-test",
       std::make_shared<PCHContainerOperations>(),
       tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
   auto &Context = Unit->getASTContext();
llvm::Error checkDataflowOnCFG(
   auto CFCtx = ControlFlowContext::build(F, F->getBody(), &F->getASTContext());
   if (!CFCtx)
     return CFCtx.takeError();
 
   DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>());
   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,
+      PostVisitCFG(Context, Element, State);
     };
   }
 
   llvm::Expected<llvm::DenseMap<const clang::Stmt *, std::string>>
       StmtToAnnotations = buildStatementToAnnotationMapping(F, AnnotatedCode);
   if (!StmtToAnnotations)
     return StmtToAnnotations.takeError();
   auto &Annotations = *StmtToAnnotations;
 
   llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
       MaybeBlockStates = runTypeErasedDataflowAnalysis(*CFCtx, Analysis, Env,
-                                                       PostVisitStmtClosure);
+                                                       PostVisitCFGClosure);
   if (!MaybeBlockStates)
     return MaybeBlockStates.takeError();
   auto &BlockStates = *MaybeBlockStates;
 
   AnalysisData AnalysisData{Context,  *CFCtx,      Env,
                             Analysis, Annotations, BlockStates};
   VerifyResults(AnalysisData);
   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>
 llvm::Error checkDataflow(
     llvm::StringRef Code,
     ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     std::function<void(
         llvm::ArrayRef<std::pair<
             std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
         ASTContext &)>
         VerifyResults,
     ArrayRef<std::string> Args,
     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);
           return;
         }
 
         auto &Annotations = AnalysisData.Annotations;
 
         // Compute a map from statement annotations to the state computed for
         // the program point immediately after the annotated statement.
         std::vector<std::pair<std::string, StateT>> Results;
         for (const CFGBlock *Block : AnalysisData.CFCtx.getCFG()) {
           // Skip blocks that were not evaluated.
           if (!AnalysisData.BlockStates[Block->getBlockID()])
             continue;
 
           transferBlock(
               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>(
                     &State.Lattice.Value);
                 Results.emplace_back(It->second, StateT{*Lattice, State.Env});
               });
         }
         VerifyResults(Results, AnalysisData.ASTCtx);