[clang][dataflow] Add `SetupTest` parameter for `AnalysisInputs`.

[wyt]

Summary

Moves the work required for retrieving annotation states into the SetupTest and PostVisitCFG callback to avoid having to run a separate pass over the CFG after analysis has completed.

Depends on D132147

[wyt]

Propagate change from parent patch.

[wyt]

Rebase.

[gribozavr2]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

106

Since SetupTest runs before PostVisitCFG, could we declare them in order?

297

304

std::move?

315

Unused variable?

317

[wyt]

Address comments.

[sgatev]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

88

Why move this? It makes it hard to tell if there are other changes. If there are other changes, let's keep it where it is to have a clean diff and move it in a separate commit.

[wyt]

Propagate change from parent patch.

[wyt]

Propagate change from parent patch.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

88

It was moved as the new parameter declared uses AnalysisOutputs, hence AnalysisOutputs needs to come before AnalysisInputs. The move is now fixed by reordering these two struct declarations in the parent patch where they were first introduced.

[ymandel]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

222

formatting seems funny here?

225

no braces for a single-statement if. could also put the binding in the if condition: if (auto Error = ...

[wyt]

Address comment.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

222

this is returned from clang-format, seems to align the arguments.

[wyt]

Update according to change in parent patch to replace designated initialisers when constructing AnalysisInputs.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

struct AnalysisOutputs {
   /// 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> &&
+  withSetupTest(std::function<llvm::Error(AnalysisOutputs &)> Arg) && {
+    SetupTest = std::move(Arg);
+    return std::move(*this);
+  }
+  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 executed immediately before
+  /// running the dataflow analysis to allow for additional setup. All fields in
+  /// the `AnalysisOutputs` argument will be initialized except for the
+  /// `BlockStates` field which is only computed later during the analysis.
+  std::function<llvm::Error(AnalysisOutputs &)> SetupTest = nullptr;
   /// 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 = {};
 };
 
 /// 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`.
+/// Returns line numbers and content of the annotations in `AnnotatedCode`.
 llvm::DenseMap<unsigned, std::string>
-getAnnotationLinesAndContent(const AnalysisOutputs &AO);
-
-// FIXME: Return a string map instead of a vector of pairs.
-//
-/// 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;
-}
+buildLineToAnnotationMapping(SourceManager &SM,
+                             llvm::Annotations AnnotatedCode);
 
 /// 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`.
 ///
 ///   `Code`, `TargetFuncMatcher` and `MakeAnalysis` must be provided in `AI`.
 ///
 ///   `VerifyResults` must be provided.
 template <typename AnalysisT>
-llvm::Error checkDataflow(
-    AnalysisInputs<AnalysisT> AI,
-    std::function<llvm::Error(const AnalysisOutputs &)> VerifyResults) {
+llvm::Error
+checkDataflow(AnalysisInputs<AnalysisT> AI,
+              std::function<void(const AnalysisOutputs &)> VerifyResults) {
   // Build AST context from code.
   llvm::Annotations AnnotatedCode(AI.Code);
   auto Unit = tooling::buildASTFromCodeWithArgs(
       AnnotatedCode.code(), AI.ASTBuildArgs, "input.cc", "clang-dataflow-test",
       std::make_shared<PCHContainerOperations>(),
       tooling::getClangStripDependencyFileAdjuster(),
       AI.ASTBuildVirtualMappedFiles);
   auto &Context = Unit->getASTContext();
checkDataflow(AnalysisInputs<AnalysisT> AI,
 
   // 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.
+  // Initialize states for running dataflow analysis.
   DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>());
   Environment InitEnv(DACtx, *Target);
   auto Analysis = AI.MakeAnalysis(Context, InitEnv);
   std::function<void(const CFGElement &,
                      const TypeErasedDataflowAnalysisState &)>
       PostVisitCFGClosure = nullptr;
   if (AI.PostVisitCFG) {
     PostVisitCFGClosure =
         [&AI, &Context](const CFGElement &Element,
                         const TypeErasedDataflowAnalysisState &State) {
           AI.PostVisitCFG(Context, Element, State);
         };
   }
 
-  // If successful, the run returns a mapping from block IDs to the
-  // post-analysis states for the CFG blocks that have been evaluated.
+  // Additional test setup.
+  AnalysisOutputs AO{AnnotatedCode, Context, Target, CFCtx,
+                     Analysis,      InitEnv, {}};
+  if (AI.SetupTest) {
+    if (auto Error = AI.SetupTest(AO))
+      return Error;
+  }
+
+  // If successful, the dataflow analysis 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, InitEnv,
                                                        PostVisitCFGClosure);
   if (!MaybeBlockStates)
     return MaybeBlockStates.takeError();
-  auto &BlockStates = *MaybeBlockStates;
+  AO.BlockStates = *MaybeBlockStates;
 
   // Verify dataflow analysis outputs.
-  AnalysisOutputs AO{AnnotatedCode, Context, Target,     CFCtx,
-                     Analysis,      InitEnv, BlockStates};
-  return VerifyResults(AO);
+  VerifyResults(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 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);
+      std::move(AI), [&VerifyResults](const AnalysisOutputs &AO) {
+        auto AnnotationLinesAndContent =
+            buildLineToAnnotationMapping(AO.ASTCtx.getSourceManager(), AO.Code);
         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.
 ///

 template <typename AnalysisT>
 llvm::Error checkDataflow(
     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;
+  // Compute mapping from nodes of annotated statements to the content in the
+  // annotation.
+  llvm::DenseMap<const Stmt *, std::string> StmtToAnnotations;
+  auto SetupTest = [&StmtToAnnotations,
+                    PrevSetupTest = std::move(AI.SetupTest)](
+                       AnalysisOutputs &AO) -> llvm::Error {
+    auto MaybeStmtToAnnotations = buildStatementToAnnotationMapping(
+        cast<FunctionDecl>(AO.InitEnv.getDeclCtx()), AO.Code);
+    if (!MaybeStmtToAnnotations) {
+      return MaybeStmtToAnnotations.takeError();
+    }
+    StmtToAnnotations = std::move(*MaybeStmtToAnnotations);
+    return PrevSetupTest ? PrevSetupTest(AO) : llvm::Error::success();
+  };
+
+  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
+
+  // FIXME: Use a string map instead of a vector of pairs.
+  //
+  // Save the states computed for program points immediately following annotated
+  // statements. The saved states are keyed by the content of the annotation.
+  std::vector<std::pair<std::string, StateT>> AnnotationStates;
+  auto PostVisitCFG = [&StmtToAnnotations, &AnnotationStates,
+                       PrevPostVisitCFG = std::move(AI.PostVisitCFG)](
+                          ASTContext &Ctx, const CFGElement &Elt,
+                          const TypeErasedDataflowAnalysisState &State) {
+    if (PrevPostVisitCFG) {
+      PrevPostVisitCFG(Ctx, Elt, State);
+    // FIXME: Extend retrieval of state for non statement constructs.
+    auto Stmt = Elt.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);
+    AnnotationStates.emplace_back(It->second, StateT{*Lattice, State.Env});
+  };
+  return checkDataflow<AnalysisT>(
+      std::move(AI)
+          .withSetupTest(std::move(SetupTest))
+          .withPostVisitCFG(std::move(PostVisitCFG)),
+      [&VerifyResults, &AnnotationStates](const AnalysisOutputs &AO) {
         VerifyResults(AnnotationStates, AO);
-        return llvm::Error::success();
       });
 }
 
 // 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

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(const AnalysisOutputs &AO) {
+test::buildLineToAnnotationMapping(SourceManager &SM,
+                                   llvm::Annotations AnnotatedCode) {
   llvm::DenseMap<unsigned, std::string> LineNumberToContent;
-  auto Code = AO.Code.code();
-  auto Annotations = AO.Code.ranges();
-  auto &SM = AO.ASTCtx.getSourceManager();
+  auto Code = AnnotatedCode.code();
+  auto Annotations = AnnotatedCode.ranges();
   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;