[clang][dataflow] Refactor `TestingSupport.h`

[wyt]

Summary

• Add AnalysisInputs struct as the parameters for checkDataflow, and renamed AnalysisData struct to AnalysisOutputs which contains the data structures generated from a dataflow analysis run.

• Remove compulsory binding from statement to annotations. Instead, checkDataflow in the most general form takes a VerifyResults callback which takes as input an AnalysisOutputs struct. This struct contains the data structures generated by the analysis that can then be tested. We then introduce two overloads/wrappers of checkDataflow for different mechanisms of testing - one which exposes annotation line numbers and is not restricted to statements, and the other which exposes states computed after annotated statements. In the future, we should look at retrieving the analysis states for constructs other than statements.

Depends On D131616

[wyt]

Small fixes.

[gribozavr2]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

60

AnalysisInputs?

82–93

AnalysisOutputs?

125

Could we capture the lattice elements in the PostVisitCFG callback? It already receives the lattice element, we can easily copy it when the CFGElement is the Stmt of interest.

That should allow us to remove this function completely!

[wyt]

Address comments: rename AnalysisArguments to AnalysisInputs and AnalysisData to AnalysisOutputs.

[wyt]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

125

Coming in next patch.

[wyt]

Propagate change from parent patch.

[wyt]

Rebase.

[sgatev]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

60

Could you please indicate which members are mandatory and which are optional. Perhaps comments could start either with /// Mandatory. ... or /// Optional. ....

244

Should requirements include the full list of the requirements of checkDataflow below?

294–311

Let's clearly indicate that it's deprecated.

338–355

Let's clearly indicate that it's deprecated.

[wyt]

Address comments.

[wyt]

Fix typo and indentation.

[gribozavr2]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

125

(if possible)

139

(if possible)

175

254–255

264

[wyt]

Address comments: add const qualifiers where applicable.

[wyt]

Remove use of designated initializers which are only allowed on C++20. Instead, introduce a constructor that sets required fields, and withFieldName methods that sets optional fields when constructing an AnalysisInputs struct.

[wyt]

This is a re-update of a previous update which uploaded the wrong diff.
Remove use of designated initializers which are only allowed on C++20. Instead, introduce a constructor that sets required fields, and withFieldName methods that sets optional fields when constructing an AnalysisInputs struct.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

 std::ostream &operator<<(std::ostream &OS,
                          const DataflowAnalysisState<Lattice> &S) {
   // FIXME: add printing support for the environment.
   return OS << "{lattice=" << S.Lattice << ", environment=...}";
 }
 
 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);
+/// Arguments for building the dataflow analysis.
+template <typename AnalysisT> struct AnalysisArguments {
+  /// Input code that is analyzed.
+  llvm::StringRef Code;
+  /// The body of the function which matches this matcher is analyzed.
+  ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher;
+  /// 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;
+  /// 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;
+
+  /// Options for building the AST context.
+  ArrayRef<std::string> ASTBuildArgs = {};
+  /// Options for building the AST context.
+  const tooling::FileContentMappings &ASTBuildVirtualMappedFiles = {};
+};
 
+/// Contains data structures required and produced by a dataflow analysis run.
 struct AnalysisData {
+  /// 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;
 };
 
-// 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 `AD.Code`.
+llvm::DenseMap<unsigned, std::string>
+getAnnotationLinesAndContent(AnalysisData &AD);
+
+// 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 `AD.Code`.
+template <typename AnalysisT>
+llvm::Expected<std::vector<
+    std::pair<std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>>
+getAnnotationStates(AnalysisData &AD) {
+  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(AD.Target, AD.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 : AD.CFCtx.getCFG()) {
+    // Skip blocks that were not evaluated.
+    if (!AD.BlockStates[Block->getBlockID()])
+      continue;
+
+    transferBlock(
+        AD.CFCtx, AD.BlockStates, *Block, AD.InitEnv, AD.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 `AA.MakeAnalysis` and `AA.PostVisitCFG` on the
+/// body of the function that matches `AA.TargetFuncMatcher` in `Code`.
+/// Given the analysis data, `VerifyResults` checks that the results from the
+/// analysis are correct.
 ///
 /// Requirements:
 ///
 ///  `AnalysisT` contains a type `Lattice`.
 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(AnalysisArguments<AnalysisT> AA,
+              std::function<llvm::Error(AnalysisData &)> VerifyResults) {
+  // Build AST context from code.
+  llvm::Annotations AnnotatedCode(AA.Code);
   auto Unit = tooling::buildASTFromCodeWithArgs(
-      AnnotatedCode.code(), Args, "input.cc", "clang-dataflow-test",
+      AnnotatedCode.code(), AA.ASTBuildArgs, "input.cc", "clang-dataflow-test",
       std::make_shared<PCHContainerOperations>(),
-      tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
+      tooling::getClangStripDependencyFileAdjuster(),
+      AA.ASTBuildVirtualMappedFiles);
   auto &Context = Unit->getASTContext();
 
   if (Context.getDiagnostics().getClient()->getNumErrors() != 0) {
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Source file has syntax or type errors, "
                                       "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(),
+                                               AA.TargetFuncMatcher)
+  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 = AA.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 (AA.PostVisitCFG) {
+    PostVisitCFGClosure =
+        [&AA, &Context](const CFGElement &Element,
+          AA.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 data.
+  AnalysisData AD{AnnotatedCode, Context, Target,     CFCtx,
+                  Analysis,      InitEnv, BlockStates};
+  return VerifyResults(AD);
 }
 
+/// Runs dataflow specified from `AA.MakeAnalysis` and `AA.PostVisitCFG` on the
+/// body of the function that matches `AA.TargetFuncMatcher` in `AA.Code`. Given
+/// the annotation line numbers and analysis data, `VerifyResults` checks that
+/// the results from the analysis are correct.
+///
+/// Requirements:
+///
+///  `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(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);
-          }
-        };
+    AnalysisArguments<AnalysisT> AA,
+    std::function<void(llvm::DenseMap<unsigned, std::string> &, AnalysisData &)>
+        VerifyResults) {
+  return checkDataflow<AnalysisT>(
+      std::move(AA), [&VerifyResults](AnalysisData &AD) -> llvm::Error {
+        auto AnnotationLinesAndContent = getAnnotationLinesAndContent(AD);
+        VerifyResults(AnnotationLinesAndContent, AD);
+        return llvm::Error::success();
+      });
+
+/// Runs dataflow specified from `AA.MakeAnalysis` and `AA.PostVisitCFG` on the
+/// body of the function that matches `AA.TargetFuncMatcher` in `AA.Code`. Given
+/// the state computed at each annotated statement and analysis data,
+/// `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.
+///
+///   There can be at most one annotation attached per statement.
+///
+///   Annotations must not be repeated.
+template <typename AnalysisT>
+llvm::Error checkDataflow(
+    AnalysisArguments<AnalysisT> AA,
+    std::function<void(
+        llvm::ArrayRef<std::pair<
+            std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
+        AnalysisData &)>
+        VerifyResults) {
+  return checkDataflow<AnalysisT>(
+      std::move(AA), [&VerifyResults](AnalysisData &AD) -> llvm::Error {
+        auto MaybeAnnotationStates = getAnnotationStates<AnalysisT>(AD);
+        if (!MaybeAnnotationStates) {
+          return MaybeAnnotationStates.takeError();
         }
-  return checkDataflowOnCFG(Code, TargetFuncMatcher, MakeAnalysis, PostVisitCFG,
-                            VerifyResults, Args, VirtualMappedFiles);
+        auto &AnnotationStates = *MaybeAnnotationStates;
+        VerifyResults(AnnotationStates, AD);
+        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`.
+// FIXME: Remove this function after usage has been updated to the overload
+// which uses the `AnalysisArguments` 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`.
+///
+///   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(
         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 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);
+  return checkDataflow<AnalysisT>(
+      {
+          .Code = Code,
+          .TargetFuncMatcher = TargetFuncMatcher,
+          .MakeAnalysis = MakeAnalysis,
+          .ASTBuildArgs = Args,
+          .ASTBuildVirtualMappedFiles = VirtualMappedFiles,
       },
-      Args, VirtualMappedFiles);
+      [&VerifyResults](
+          llvm::ArrayRef<std::pair<
+              std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>
+              AnnotationStates,
+          AnalysisData &AD) { VerifyResults(AnnotationStates, AD.ASTCtx); });
 }
 
-// Runs dataflow on the body of the function named `target_fun` in code snippet
-// `code`.
+// FIXME: Remove this function after usage has been updated to the overload
+// which uses the `AnalysisArguments` 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.
+///
+///   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,
     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 = {}) {
-  return checkDataflow(Code, ast_matchers::hasName(TargetFun),
-                       std::move(MakeAnalysis), std::move(VerifyResults), Args,
+  return checkDataflow<AnalysisT>(Code, ast_matchers::hasName(TargetFun),
+                                  MakeAnalysis, VerifyResults, Args,
                                   VirtualMappedFiles);
 }
 
 /// Returns the `ValueDecl` for the given identifier.
 ///
 /// Requirements:
 ///
 ///  `Name` must be unique in `ASTCtx`.
 const ValueDecl *findValueDecl(ASTContext &ASTCtx, llvm::StringRef Name);

clang/unittests/Analysis/FlowSensitive/TestingSupport.cpp

while (NextToken && SourceManager.getFileOffset(NextToken->getLocation()) <
 
     NextToken = Lexer::findNextToken(NextToken->getEndLoc(), SourceManager,
                                      LangOptions);
   }
 
   return true;
 }
 
+llvm::DenseMap<unsigned, std::string>
+test::getAnnotationLinesAndContent(AnalysisData &AD) {
+  llvm::DenseMap<unsigned, std::string> LineNumberToContent;
+  auto Code = AD.Code.code();
+  auto Annotations = AD.Code.ranges();
+  auto &SM = AD.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) {
   llvm::DenseMap<const Stmt *, std::string> Result;
 
   auto StmtMatcher =
       findAll(stmt(unless(anyOf(hasParent(expr()), hasParent(returnStmt()))))
                   .bind("stmt"));

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

Headers.emplace_back("unchecked_optional_access_test.h", R"(
       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,
+        /*AA:AnalysisArguments=*/{
+            .Code = SourceCode,
+            .TargetFuncMatcher = FuncMatcher,
+            .MakeAnalysis =
                 [Options](ASTContext &Ctx, Environment &) {
                   return UncheckedOptionalAccessModel(Ctx, Options);
                 },
-        [&Diagnostics, Diagnoser = UncheckedOptionalAccessDiagnoser(Options)](
-            ASTContext &Ctx, const CFGStmt &Stmt,
+            .PostVisitCFG =
+                [&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);
+                      Diagnoser.diagnose(Ctx, Stmt->getStmt(), State.Env);
                   llvm::move(StmtDiagnostics, std::back_inserter(Diagnostics));
                 },
-        [&Diagnostics](AnalysisData AnalysisData) {
-          auto &SrcMgr = AnalysisData.ASTCtx.getSourceManager();
-
+            .ASTBuildArgs = {"-fsyntax-only", "-std=c++17",
+                             "-Wno-undefined-inline"},
+            .ASTBuildVirtualMappedFiles = FileContents,
+        },
+        /*VerifyResults=*/[&Diagnostics](llvm::DenseMap<unsigned, std::string>
+                                             &Annotations,
+                                         AnalysisData &AD) {
           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 = AD.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));
   }
 };
 
 INSTANTIATE_TEST_SUITE_P(
     UncheckedOptionalUseTestInst, UncheckedOptionalAccessTest,
     ::testing::Values(OptionalTypeIdentifier{"std", "optional"},