[clang][dataflow] Change the diagnoser API to receive a correctly typed lattice element

[merrymeerkat]

Summary

Previously, the diagnoser could only receive the Environment at a given program point. Now, it receives the complete dataflow state: the environment and lattice element.

This change does not contain any tests because we modify the checkDataflow function to rely on the newly introduced lattice element in PostVisitCFG, and existing tests that verify lattice elements depend on this overload of checkDataflow.

[ymandel]

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

134

nit: I think capturing by move is a bit better, since Arg is passed by copy to withPostVisitCFG: [Arg = std::move(Arg)]

[merrymeerkat]

Address review comment

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

template <typename LatticeT> struct TransferState {
   TransferState(LatticeT &Lattice, Environment &Env)
       : Lattice(Lattice), Env(Env) {}
 
   /// Current lattice element.
   LatticeT &Lattice;
   Environment &Env;
 };
 
+/// A read-only version of TransferState.
+template <typename LatticeT> struct TransferStateForDiagnostics {
+  TransferStateForDiagnostics(const LatticeT &Lattice, const Environment &Env)
+      : Lattice(Lattice), Env(Env) {}
+
+  /// Current lattice element.
+  const LatticeT &Lattice;
+  const Environment &Env;
+};
+
 template <typename T>
 using MatchSwitchMatcher = ast_matchers::internal::Matcher<T>;
 
 template <typename T, typename State, typename Result = void>
 using MatchSwitchAction = std::function<Result(
     const T *, const ast_matchers::MatchFinder::MatchResult &, State &)>;
 
 template <typename BaseT, typename State, typename Result = void>

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

 #include "clang/AST/Stmt.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/ASTMatchers/ASTMatchersInternal.h"
 #include "clang/Analysis/CFG.h"
 #include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/DataflowAnalysis.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
+#include "clang/Analysis/FlowSensitive/MatchSwitch.h"
 #include "clang/Analysis/FlowSensitive/WatchedLiteralsSolver.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Serialization/PCHContainerOperations.h"
 #include "clang/Tooling/ArgumentsAdjusters.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
AnalysisInputs(
         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 TransferStateForDiagnostics<typename AnalysisT::Lattice> &)>
+          Arg) && {
+    PostVisitCFG = std::move(Arg);
+    return std::move(*this);
+  }
+
   AnalysisInputs<AnalysisT> &&
   withPostVisitCFG(std::function<void(ASTContext &, const CFGElement &,
                                       const TypeErasedDataflowAnalysisState &)>
                        Arg) && {
-    PostVisitCFG = std::move(Arg);
+    PostVisitCFG =
+        [Arg = std::move(Arg)](ASTContext &Context, const CFGElement &Element,
+            const TransferStateForDiagnostics<typename AnalysisT::Lattice>
+                &State) {
+          Arg(Context, Element,
+              TypeErasedDataflowAnalysisState({State.Lattice}, State.Env));
+        };
     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) && {
template <typename AnalysisT> struct AnalysisInputs {
   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 &)>
+  std::function<void(
+      ASTContext &, const CFGElement &,
+      const TransferStateForDiagnostics<typename AnalysisT::Lattice> &)>
       PostVisitCFG = nullptr;
 
   /// Optional. Options for building the AST context.
   ArrayRef<std::string> ASTBuildArgs = {};
   /// Optional. Options for building the AST context.
   tooling::FileContentMappings ASTBuildVirtualMappedFiles = {};
 };
 
checkDataflow(AnalysisInputs<AnalysisT> AI,
   // 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,
+    PostVisitCFGClosure = [&AI, &Context](
+                              const CFGElement &Element,
                               const TypeErasedDataflowAnalysisState &State) {
-          AI.PostVisitCFG(Context, Element, State);
-        };
+      AI.PostVisitCFG(Context, Element,
+                      TransferStateForDiagnostics<typename AnalysisT::Lattice>(
+                          llvm::any_cast<const typename AnalysisT::Lattice &>(
+                              State.Lattice.Value),
+                          State.Env));
   }
 
   // Additional test setup.
   AnalysisOutputs AO{AnnotatedCode, Context, Target, CFCtx,
                      Analysis,      InitEnv, {}};
   if (AI.SetupTest) {
     if (auto Error = AI.SetupTest(AO))
       return Error;
auto SetupTest = [&StmtToAnnotations,
     return PrevSetupTest ? PrevSetupTest(AO) : llvm::Error::success();
   };
 
   using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
 
   // Save the states computed for program points immediately following annotated
   // statements. The saved states are keyed by the content of the annotation.
   llvm::StringMap<StateT> AnnotationStates;
-  auto PostVisitCFG = [&StmtToAnnotations, &AnnotationStates,
-                       PrevPostVisitCFG = std::move(AI.PostVisitCFG)](
-                          ASTContext &Ctx, const CFGElement &Elt,
-                          const TypeErasedDataflowAnalysisState &State) {
+  auto PostVisitCFG =
+      [&StmtToAnnotations, &AnnotationStates,
+          const TransferStateForDiagnostics<typename AnalysisT::Lattice>
+              &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);
-    auto [_, InsertSuccess] =
-        AnnotationStates.insert({It->second, StateT{*Lattice, State.Env}});
+        auto [_, InsertSuccess] = AnnotationStates.insert(
+            {It->second, StateT{State.Lattice, State.Env}});
         (void)_;
         (void)InsertSuccess;
         assert(InsertSuccess);
       };
   return checkDataflow<AnalysisT>(
       std::move(AI)
           .withSetupTest(std::move(SetupTest))
           .withPostVisitCFG(std::move(PostVisitCFG)),
       [&VerifyResults, &AnnotationStates](const AnalysisOutputs &AO) {
         VerifyResults(AnnotationStates, AO);
       });
 }

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

llvm::Error Error = checkDataflow<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 {
+                    const TransferStateForDiagnostics<NoopLattice>
+                        &State) mutable {
                   auto EltDiagnostics =
                       Diagnoser.diagnose(Ctx, &Elt, State.Env);
                   llvm::move(EltDiagnostics, std::back_inserter(Diagnostics));
                 })
             .withASTBuildArgs(
                 {"-fsyntax-only", "-std=c++17", "-Wno-undefined-inline"})
             .withASTBuildVirtualMappedFiles(
                 tooling::FileContentMappings(Headers.begin(), Headers.end())),