diff --git a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
--- a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
+++ b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
@@ -56,47 +56,160 @@
 
 namespace test {
 
-// Returns assertions based on annotations that are present after statements in
-// `AnnotatedCode`.
-llvm::Expected<llvm::DenseMap<const Stmt *, std::string>>
-buildStatementToAnnotationMapping(const FunctionDecl *Func,
-                                  llvm::Annotations AnnotatedCode);
-
-struct AnalysisData {
+/// Contains data structures required and produced by a dataflow analysis run.
+struct AnalysisOutputs {
+  /// Input code that is analyzed. Points within the code may be marked with
+  /// annotations to facilitate testing.
+  ///
+  /// Example:
+  /// void target(int *x) {
+  ///   *x; // [[p]]
+  /// }
+  /// From the annotation `p`, the line number and analysis state immediately
+  /// after the statement `*x` can be retrieved and verified.
+  llvm::Annotations Code;
+  /// AST context generated from `Code`.
   ASTContext &ASTCtx;
+  /// The function whose body is analyzed.
+  const FunctionDecl *Target;
+  /// Contains the control flow graph built from the body of the `Target`
+  /// function and is analyzed.
   const ControlFlowContext &CFCtx;
-  const Environment &Env;
+  /// The analysis to be run.
   TypeErasedDataflowAnalysis &Analysis;
-  llvm::DenseMap<const clang::Stmt *, std::string> &Annotations;
-  std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates;
+  /// 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;
+};
+
+/// Arguments for building the dataflow analysis.
+template <typename AnalysisT> struct AnalysisInputs {
+  /// Required fields are set in constructor.
+  AnalysisInputs(
+      llvm::StringRef CodeArg,
+      ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcherArg,
+      std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysisArg)
+      : Code(CodeArg), TargetFuncMatcher(std::move(TargetFuncMatcherArg)),
+        MakeAnalysis(std::move(MakeAnalysisArg)) {}
+
+  /// Optional fields can be set with methods of the form `withFieldName(...)`.
+  AnalysisInputs<AnalysisT> &&
+  withPostVisitCFG(std::function<void(ASTContext &, const CFGElement &,
+                                      const TypeErasedDataflowAnalysisState &)>
+                       Arg) && {
+    PostVisitCFG = std::move(Arg);
+    return std::move(*this);
+  }
+  AnalysisInputs<AnalysisT> &&withASTBuildArgs(ArrayRef<std::string> Arg) && {
+    ASTBuildArgs = std::move(Arg);
+    return std::move(*this);
+  }
+  AnalysisInputs<AnalysisT> &&
+  withASTBuildVirtualMappedFiles(tooling::FileContentMappings Arg) && {
+    ASTBuildVirtualMappedFiles = std::move(Arg);
+    return std::move(*this);
+  }
+
+  /// Required. Input code that is analyzed.
+  llvm::StringRef Code;
+  /// Required. The body of the function which matches this matcher is analyzed.
+  ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher;
+  /// Required. The analysis to be run is constructed with this function that
+  /// takes as argument the AST generated from the code being analyzed and the
+  /// initial state from which the analysis starts with.
+  std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis;
+  /// Optional. If provided, this function is applied on each CFG element after
+  /// the analysis has been run.
+  std::function<void(ASTContext &, const CFGElement &,
+                     const TypeErasedDataflowAnalysisState &)>
+      PostVisitCFG = nullptr;
+
+  /// Optional. Options for building the AST context.
+  ArrayRef<std::string> ASTBuildArgs = {};
+  /// Optional. Options for building the AST context.
+  tooling::FileContentMappings ASTBuildVirtualMappedFiles = {};
 };
 
-// FIXME: Rename to `checkDataflow` after usages of the overload that applies to
-// `CFGStmt` have been replaced.
+/// Returns assertions based on annotations that are present after statements in
+/// `AnnotatedCode`.
+llvm::Expected<llvm::DenseMap<const Stmt *, std::string>>
+buildStatementToAnnotationMapping(const FunctionDecl *Func,
+                                  llvm::Annotations AnnotatedCode);
+
+/// Returns line numbers and content of the annotations in `AO.Code`.
+llvm::DenseMap<unsigned, std::string>
+getAnnotationLinesAndContent(const AnalysisOutputs &AO);
+
+// FIXME: Return a string map instead of a vector of pairs.
 //
-/// 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.
+/// Returns the analysis states at each annotated statement in `AO.Code`.
+template <typename AnalysisT>
+llvm::Expected<std::vector<
+    std::pair<std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>>
+getAnnotationStates(const AnalysisOutputs &AO) {
+  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
+  // FIXME: Extend to annotations on non-statement constructs.
+  // Get annotated statements.
+  llvm::Expected<llvm::DenseMap<const Stmt *, std::string>>
+      MaybeStmtToAnnotations =
+          buildStatementToAnnotationMapping(AO.Target, AO.Code);
+  if (!MaybeStmtToAnnotations)
+    return MaybeStmtToAnnotations.takeError();
+  auto &StmtToAnnotations = *MaybeStmtToAnnotations;
+
+  // 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 : AO.CFCtx.getCFG()) {
+    // Skip blocks that were not evaluated.
+    if (!AO.BlockStates[Block->getBlockID()])
+      continue;
+
+    transferBlock(
+        AO.CFCtx, AO.BlockStates, *Block, AO.InitEnv, AO.Analysis,
+        [&Results,
+         &StmtToAnnotations](const clang::CFGElement &Element,
+                             const TypeErasedDataflowAnalysisState &State) {
+          auto Stmt = Element.getAs<CFGStmt>();
+          if (!Stmt)
+            return;
+          auto It = StmtToAnnotations.find(Stmt->getStmt());
+          if (It == StmtToAnnotations.end())
+            return;
+          auto *Lattice =
+              llvm::any_cast<typename AnalysisT::Lattice>(&State.Lattice.Value);
+          Results.emplace_back(It->second, StateT{*Lattice, State.Env});
+        });
+  }
+
+  return Results;
+}
+
+/// Runs dataflow specified from `AI.MakeAnalysis` and `AI.PostVisitCFG` on the
+/// body of the function that matches `AI.TargetFuncMatcher` in `AI.Code`.
+/// Given the analysis outputs, `VerifyResults` checks that the results from the
+/// analysis are correct.
 ///
 /// Requirements:
 ///
-///  `AnalysisT` contains a type `Lattice`.
+///   `AnalysisT` contains a type `Lattice`.
+///
+///   `Code`, `TargetFuncMatcher` and `MakeAnalysis` must be provided in `AI`.
+///
+///   `VerifyResults` must be provided.
 template <typename AnalysisT>
-llvm::Error checkDataflowOnCFG(
-    llvm::StringRef Code,
-    ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
-    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
-    std::function<void(ASTContext &, const CFGElement &,
-                       const TypeErasedDataflowAnalysisState &)>
-        PostVisitCFG,
-    std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
-    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  llvm::Annotations AnnotatedCode(Code);
+llvm::Error checkDataflow(
+    AnalysisInputs<AnalysisT> AI,
+    std::function<llvm::Error(const AnalysisOutputs &)> VerifyResults) {
+  // Build AST context from code.
+  llvm::Annotations AnnotatedCode(AI.Code);
   auto Unit = tooling::buildASTFromCodeWithArgs(
-      AnnotatedCode.code(), Args, "input.cc", "clang-dataflow-test",
+      AnnotatedCode.code(), AI.ASTBuildArgs, "input.cc", "clang-dataflow-test",
       std::make_shared<PCHContainerOperations>(),
-      tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
+      tooling::getClangStripDependencyFileAdjuster(),
+      AI.ASTBuildVirtualMappedFiles);
   auto &Context = Unit->getASTContext();
 
   if (Context.getDiagnostics().getClient()->getNumErrors() != 0) {
@@ -105,82 +218,141 @@
                                       "they were printed to the test log");
   }
 
-  const FunctionDecl *F = ast_matchers::selectFirst<FunctionDecl>(
-      "target",
-      ast_matchers::match(ast_matchers::functionDecl(
-                              ast_matchers::isDefinition(), TargetFuncMatcher)
-                              .bind("target"),
-                          Context));
-  if (F == nullptr)
+  // Get AST node of target function.
+  const FunctionDecl *Target = ast_matchers::selectFirst<FunctionDecl>(
+      "target", ast_matchers::match(
+                    ast_matchers::functionDecl(ast_matchers::isDefinition(),
+                                               AI.TargetFuncMatcher)
+                        .bind("target"),
+                    Context));
+  if (Target == nullptr)
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Could not find target function.");
 
-  auto CFCtx = ControlFlowContext::build(F, F->getBody(), &F->getASTContext());
-  if (!CFCtx)
-    return CFCtx.takeError();
+  // Build control flow graph from body of target function.
+  auto MaybeCFCtx =
+      ControlFlowContext::build(Target, *Target->getBody(), Context);
+  if (!MaybeCFCtx)
+    return MaybeCFCtx.takeError();
+  auto &CFCtx = *MaybeCFCtx;
 
+  // Initialize states and run dataflow analysis.
   DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>());
-  Environment Env(DACtx, *F);
-  auto Analysis = MakeAnalysis(Context, Env);
-
+  Environment InitEnv(DACtx, *Target);
+  auto Analysis = AI.MakeAnalysis(Context, InitEnv);
   std::function<void(const CFGElement &,
                      const TypeErasedDataflowAnalysisState &)>
       PostVisitCFGClosure = nullptr;
-  if (PostVisitCFG) {
-    PostVisitCFGClosure = [&PostVisitCFG, &Context](
-                              const CFGElement &Element,
-                              const TypeErasedDataflowAnalysisState &State) {
-      PostVisitCFG(Context, Element, State);
-    };
+  if (AI.PostVisitCFG) {
+    PostVisitCFGClosure =
+        [&AI, &Context](const CFGElement &Element,
+                        const TypeErasedDataflowAnalysisState &State) {
+          AI.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;
-
+  // If successful, the run returns a mapping from block IDs to the
+  // post-analysis states for the CFG blocks that have been evaluated.
   llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
-      MaybeBlockStates = runTypeErasedDataflowAnalysis(*CFCtx, Analysis, Env,
+      MaybeBlockStates = runTypeErasedDataflowAnalysis(CFCtx, Analysis, InitEnv,
                                                        PostVisitCFGClosure);
   if (!MaybeBlockStates)
     return MaybeBlockStates.takeError();
   auto &BlockStates = *MaybeBlockStates;
 
-  AnalysisData AnalysisData{Context,  *CFCtx,      Env,
-                            Analysis, Annotations, BlockStates};
-  VerifyResults(AnalysisData);
-  return llvm::Error::success();
+  // Verify dataflow analysis outputs.
+  AnalysisOutputs AO{AnnotatedCode, Context, Target,     CFCtx,
+                     Analysis,      InitEnv, BlockStates};
+  return VerifyResults(AO);
 }
 
+/// Runs dataflow specified from `AI.MakeAnalysis` and `AI.PostVisitCFG` on the
+/// body of the function that matches `AI.TargetFuncMatcher` in `AI.Code`. Given
+/// the annotation line numbers and analysis outputs, `VerifyResults` checks
+/// that the results from the analysis are correct.
+///
+/// Requirements:
+///
+///   `AnalysisT` contains a type `Lattice`.
+///
+///   `Code`, `TargetFuncMatcher` and `MakeAnalysis` must be provided in `AI`.
+///
+///   `VerifyResults` must be provided.
+template <typename AnalysisT>
+llvm::Error
+checkDataflow(AnalysisInputs<AnalysisT> AI,
+              std::function<void(const llvm::DenseMap<unsigned, std::string> &,
+                                 const AnalysisOutputs &)>
+                  VerifyResults) {
+  return checkDataflow<AnalysisT>(
+      std::move(AI),
+      [&VerifyResults](const AnalysisOutputs &AO) -> llvm::Error {
+        auto AnnotationLinesAndContent = getAnnotationLinesAndContent(AO);
+        VerifyResults(AnnotationLinesAndContent, AO);
+        return llvm::Error::success();
+      });
+}
+
+/// Runs dataflow specified from `AI.MakeAnalysis` and `AI.PostVisitCFG` on the
+/// body of the function that matches `AI.TargetFuncMatcher` in `AI.Code`. Given
+/// the state computed at each annotated statement and analysis outputs,
+/// `VerifyResults` checks that the results from the analysis are correct.
+///
+/// Requirements:
+///
+///   `AnalysisT` contains a type `Lattice`.
+///
+///   `Code`, `TargetFuncMatcher` and `MakeAnalysis` must be provided in `AI`.
+///
+///   `VerifyResults` must be provided.
+///
+///   Any annotations appearing in `Code` must come after a statement.
+///
+///   There can be at most one annotation attached per statement.
+///
+///   Annotations must not be repeated.
 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);
+    AnalysisInputs<AnalysisT> AI,
+    std::function<void(
+        llvm::ArrayRef<std::pair<
+            std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
+        const AnalysisOutputs &)>
+        VerifyResults) {
+  return checkDataflow<AnalysisT>(
+      std::move(AI),
+      [&VerifyResults](const AnalysisOutputs &AO) -> llvm::Error {
+        auto MaybeAnnotationStates = getAnnotationStates<AnalysisT>(AO);
+        if (!MaybeAnnotationStates) {
+          return MaybeAnnotationStates.takeError();
+        }
+        auto &AnnotationStates = *MaybeAnnotationStates;
+        VerifyResults(AnnotationStates, AO);
+        return llvm::Error::success();
+      });
 }
 
-// Runs dataflow on the body of the function that matches `TargetFuncMatcher` in
-// code snippet `Code`. Requires: `AnalysisT` contains a type `Lattice`.
+// Deprecated.
+// FIXME: Remove this function after usage has been updated to the overload
+// which uses the `AnalysisInputs` struct.
+//
+/// Runs dataflow specified from `MakeAnalysis` on the body of the function that
+/// matches `TargetFuncMatcher` in `Code`. Given the state computed at each
+/// annotated statement, `VerifyResults` checks that the results from the
+/// analysis are correct.
+///
+/// Requirements:
+///
+///   `AnalysisT` contains a type `Lattice`.
+///
+///   `Code`, `TargetFuncMatcher`, `MakeAnalysis` and `VerifyResults` must be
+///   provided.
+///
+///   Any annotations appearing in `Code` must come after a statement.
+///
+///   There can be at most one annotation attached per statement.
+///
+///   Annotations must not be repeated.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code,
@@ -193,52 +365,40 @@
         VerifyResults,
     ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
-
-  return checkDataflowOnCFG(
-      Code, std::move(TargetFuncMatcher), std::move(MakeAnalysis),
-      /*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::CFGElement &Element,
-                             const TypeErasedDataflowAnalysisState &State) {
-                // 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);
-      },
-      Args, VirtualMappedFiles);
+  return checkDataflow<AnalysisT>(
+      AnalysisInputs<AnalysisT>(Code, std::move(TargetFuncMatcher),
+                                std::move(MakeAnalysis))
+          .withASTBuildArgs(std::move(Args))
+          .withASTBuildVirtualMappedFiles(std::move(VirtualMappedFiles)),
+      [&VerifyResults](
+          llvm::ArrayRef<std::pair<
+              std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>
+              AnnotationStates,
+          const AnalysisOutputs &AO) {
+        VerifyResults(AnnotationStates, AO.ASTCtx);
+      });
 }
 
-// Runs dataflow on the body of the function named `target_fun` in code snippet
-// `code`.
+// Deprecated.
+// FIXME: Remove this function after usage has been updated to the overload
+// which uses the `AnalysisInputs` struct.
+//
+/// Runs dataflow specified from `MakeAnalysis` on the body of the function
+/// named `TargetFun` in `Code`. Given the state computed at each annotated
+/// statement, `VerifyResults` checks that the results from the analysis are
+/// correct.
+///
+/// Requirements:
+///
+///   `AnalysisT` contains a type `Lattice`.
+///
+///   Any annotations appearing in `Code` must come after a statement.
+///
+///   `Code`, `TargetFun`, `MakeAnalysis` and `VerifyResults` must be provided.
+///
+///   There can be at most one annotation attached per statement.
+///
+///   Annotations must not be repeated.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code, llvm::StringRef TargetFun,
@@ -250,16 +410,16 @@
         VerifyResults,
     ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  return checkDataflow(Code, ast_matchers::hasName(TargetFun),
-                       std::move(MakeAnalysis), std::move(VerifyResults), Args,
-                       VirtualMappedFiles);
+  return checkDataflow<AnalysisT>(
+      Code, ast_matchers::hasName(TargetFun), std::move(MakeAnalysis),
+      std::move(VerifyResults), Args, VirtualMappedFiles);
 }
 
 /// Returns the `ValueDecl` for the given identifier.
 ///
 /// Requirements:
 ///
-///  `Name` must be unique in `ASTCtx`.
+///   `Name` must be unique in `ASTCtx`.
 const ValueDecl *findValueDecl(ASTContext &ASTCtx, llvm::StringRef Name);
 
 /// Creates and owns constraints which are boolean values.
diff --git a/clang/unittests/Analysis/FlowSensitive/TestingSupport.cpp b/clang/unittests/Analysis/FlowSensitive/TestingSupport.cpp
--- a/clang/unittests/Analysis/FlowSensitive/TestingSupport.cpp
+++ b/clang/unittests/Analysis/FlowSensitive/TestingSupport.cpp
@@ -52,6 +52,22 @@
   return true;
 }
 
+llvm::DenseMap<unsigned, std::string>
+test::getAnnotationLinesAndContent(const AnalysisOutputs &AO) {
+  llvm::DenseMap<unsigned, std::string> LineNumberToContent;
+  auto Code = AO.Code.code();
+  auto Annotations = AO.Code.ranges();
+  auto &SM = AO.ASTCtx.getSourceManager();
+  for (auto &AnnotationRange : Annotations) {
+    auto LineNumber =
+        SM.getPresumedLineNumber(SM.getLocForStartOfFile(SM.getMainFileID())
+                                     .getLocWithOffset(AnnotationRange.Begin));
+    auto Content = Code.slice(AnnotationRange.Begin, AnnotationRange.End).str();
+    LineNumberToContent[LineNumber] = Content;
+  }
+  return LineNumberToContent;
+}
+
 llvm::Expected<llvm::DenseMap<const Stmt *, std::string>>
 test::buildStatementToAnnotationMapping(const FunctionDecl *Func,
                                         llvm::Annotations AnnotatedCode) {
diff --git a/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp b/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
--- a/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
+++ b/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
@@ -1234,49 +1234,48 @@
       template <typename T>
       T Make();
     )");
-    const tooling::FileContentMappings FileContents(Headers.begin(),
-                                                    Headers.end());
     UncheckedOptionalAccessModelOptions Options{
         /*IgnoreSmartPointerDereference=*/true};
     std::vector<SourceLocation> Diagnostics;
     llvm::Error Error = checkDataflow<UncheckedOptionalAccessModel>(
-        SourceCode, FuncMatcher,
-        [Options](ASTContext &Ctx, Environment &) {
-          return UncheckedOptionalAccessModel(Ctx, Options);
-        },
-        [&Diagnostics, Diagnoser = UncheckedOptionalAccessDiagnoser(Options)](
-            ASTContext &Ctx, const CFGStmt &Stmt,
-            const TypeErasedDataflowAnalysisState &State) mutable {
-          auto StmtDiagnostics =
-              Diagnoser.diagnose(Ctx, Stmt.getStmt(), State.Env);
-          llvm::move(StmtDiagnostics, std::back_inserter(Diagnostics));
-        },
-        [&Diagnostics](AnalysisData AnalysisData) {
-          auto &SrcMgr = AnalysisData.ASTCtx.getSourceManager();
-
+        AnalysisInputs<UncheckedOptionalAccessModel>(
+            SourceCode, std::move(FuncMatcher),
+            [Options](ASTContext &Ctx, Environment &) {
+              return UncheckedOptionalAccessModel(Ctx, Options);
+            })
+            .withPostVisitCFG(
+                [&Diagnostics,
+                 Diagnoser = UncheckedOptionalAccessDiagnoser(Options)](
+                    ASTContext &Ctx, const CFGElement &Elt,
+                    const TypeErasedDataflowAnalysisState &State) mutable {
+                  auto Stmt = Elt.getAs<CFGStmt>();
+                  if (!Stmt) {
+                    return;
+                  }
+                  auto StmtDiagnostics =
+                      Diagnoser.diagnose(Ctx, Stmt->getStmt(), State.Env);
+                  llvm::move(StmtDiagnostics, std::back_inserter(Diagnostics));
+                })
+            .withASTBuildArgs(
+                {"-fsyntax-only", "-std=c++17", "-Wno-undefined-inline"})
+            .withASTBuildVirtualMappedFiles(
+                tooling::FileContentMappings(Headers.begin(), Headers.end())),
+        /*VerifyResults=*/[&Diagnostics](
+                              const llvm::DenseMap<unsigned, std::string>
+                                  &Annotations,
+                              const AnalysisOutputs &AO) {
           llvm::DenseSet<unsigned> AnnotationLines;
-          for (const auto &Pair : AnalysisData.Annotations) {
-            auto *Stmt = Pair.getFirst();
-            AnnotationLines.insert(
-                SrcMgr.getPresumedLineNumber(Stmt->getBeginLoc()));
-            // We add both the begin and end locations, so that if the
-            // statement spans multiple lines then the test will fail.
-            //
-            // FIXME: Going forward, we should change this to instead just
-            // get the single line number from the annotation itself, rather
-            // than looking at the statement it's attached to.
-            AnnotationLines.insert(
-                SrcMgr.getPresumedLineNumber(Stmt->getEndLoc()));
+          for (const auto &[Line, _] : Annotations) {
+            AnnotationLines.insert(Line);
           }
-
+          auto &SrcMgr = AO.ASTCtx.getSourceManager();
           llvm::DenseSet<unsigned> DiagnosticLines;
           for (SourceLocation &Loc : Diagnostics) {
             DiagnosticLines.insert(SrcMgr.getPresumedLineNumber(Loc));
           }
 
           EXPECT_THAT(DiagnosticLines, ContainerEq(AnnotationLines));
-        },
-        {"-fsyntax-only", "-std=c++17", "-Wno-undefined-inline"}, FileContents);
+        });
     if (Error)
       FAIL() << llvm::toString(std::move(Error));
   }