Update nullability verification framework to adapt to the new diagnoser API in the clang dataflow framework

The new API (PHAB_REVIEW=https://reviews.llvm.org/D139868) allows diagnosers to take in `State` objects instead of `Environment` objects as argument. A `State` contains both an `Environment` and a `Lattice` element. This change will allow the pointer nullability diagnosis to use information from the dataflow lattice, which can be helpful for testing and debugging purposes. For example, the diagnoser can now access the nullability vectors computed for a given expression, which are stored in a map in the lattice.

PiperOrigin-RevId: 495350578
diff --git a/nullability_verification/BUILD b/nullability_verification/BUILD
index 7e4c3d2..d15de47 100644
--- a/nullability_verification/BUILD
+++ b/nullability_verification/BUILD
@@ -38,6 +38,7 @@
     hdrs = ["pointer_nullability_diagnosis.h"],
     deps = [
         ":pointer_nullability",
+        ":pointer_nullability_lattice",
         ":pointer_nullability_matchers",
         "@llvm-project//clang:analysis",
         "@llvm-project//clang:ast",
diff --git a/nullability_verification/pointer_nullability_diagnosis.cc b/nullability_verification/pointer_nullability_diagnosis.cc
index 6559beb..32ccad6 100644
--- a/nullability_verification/pointer_nullability_diagnosis.cc
+++ b/nullability_verification/pointer_nullability_diagnosis.cc
@@ -23,6 +23,7 @@
 using ast_matchers::MatchFinder;
 using dataflow::CFGMatchSwitchBuilder;
 using dataflow::Environment;
+using dataflow::TransferStateForDiagnostics;
 
 namespace {
 
@@ -41,19 +42,19 @@
          isNullableOrUntracked(E, Env);
 }
 
-llvm::Optional<CFGElement> diagnoseDereference(const UnaryOperator* UnaryOp,
-                                               const MatchFinder::MatchResult&,
-                                               const Environment& Env) {
-  if (isNullableOrUntracked(UnaryOp->getSubExpr(), Env)) {
+llvm::Optional<CFGElement> diagnoseDereference(
+    const UnaryOperator* UnaryOp, const MatchFinder::MatchResult&,
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
+  if (isNullableOrUntracked(UnaryOp->getSubExpr(), State.Env)) {
     return llvm::Optional<CFGElement>(CFGStmt(UnaryOp));
   }
   return llvm::None;
 }
 
-llvm::Optional<CFGElement> diagnoseArrow(const MemberExpr* MemberExpr,
-                                         const MatchFinder::MatchResult& Result,
-                                         const Environment& Env) {
-  if (isNullableOrUntracked(MemberExpr->getBase(), Env)) {
+llvm::Optional<CFGElement> diagnoseArrow(
+    const MemberExpr* MemberExpr, const MatchFinder::MatchResult& Result,
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
+  if (isNullableOrUntracked(MemberExpr->getBase(), State.Env)) {
     return llvm::Optional<CFGElement>(CFGStmt(MemberExpr));
   }
   return llvm::None;
@@ -80,7 +81,7 @@
 // interpreted as a function type (e.g. a pointer to a function pointer).
 llvm::Optional<CFGElement> diagnoseCallExpr(
     const CallExpr* CE, const MatchFinder::MatchResult& Result,
-    const Environment& Env) {
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
   auto* Callee = CE->getCalleeDecl();
   if (!Callee) return llvm::None;
 
@@ -95,56 +96,60 @@
     Args = Args.drop_front();
   }
 
-  return isIncompatibleArgumentList(ParamTypes, Args, Env, *Result.Context)
+  return isIncompatibleArgumentList(ParamTypes, Args, State.Env,
+                                    *Result.Context)
              ? llvm::Optional<CFGElement>(CFGStmt(CE))
              : llvm::None;
 }
 
 llvm::Optional<CFGElement> diagnoseConstructExpr(
     const CXXConstructExpr* CE, const MatchFinder::MatchResult& Result,
-    const Environment& Env) {
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
   auto ConstructorParamTypes = CE->getConstructor()
                                    ->getType()
                                    ->getAs<FunctionProtoType>()
                                    ->getParamTypes();
   ArrayRef<const Expr*> ConstructorArgs(CE->getArgs(), CE->getNumArgs());
-  return isIncompatibleArgumentList(ConstructorParamTypes, ConstructorArgs, Env,
-                                    *Result.Context)
+  return isIncompatibleArgumentList(ConstructorParamTypes, ConstructorArgs,
+                                    State.Env, *Result.Context)
              ? llvm::Optional<CFGElement>(CFGStmt(CE))
              : llvm::None;
 }
 
 llvm::Optional<CFGElement> diagnoseReturn(
     const ReturnStmt* RS, const MatchFinder::MatchResult& Result,
-    const Environment& Env) {
-  auto ReturnType = cast<FunctionDecl>(Env.getDeclCtx())->getReturnType();
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
+  auto ReturnType = cast<FunctionDecl>(State.Env.getDeclCtx())->getReturnType();
   assert(ReturnType->isPointerType());
 
   auto* ReturnExpr = RS->getRetValue();
   assert(ReturnExpr->getType()->isPointerType());
 
-  return isIncompatibleAssignment(ReturnType, ReturnExpr, Env, *Result.Context)
+  return isIncompatibleAssignment(ReturnType, ReturnExpr, State.Env,
+                                  *Result.Context)
              ? llvm::Optional<CFGElement>(CFGStmt(RS))
              : llvm::None;
 }
 
 llvm::Optional<CFGElement> diagnoseMemberInitializer(
     const CXXCtorInitializer* CI, const MatchFinder::MatchResult& Result,
-    const Environment& Env) {
+    const TransferStateForDiagnostics<PointerNullabilityLattice>& State) {
   assert(CI->isAnyMemberInitializer());
   auto MemberType = CI->getAnyMember()->getType();
   if (!MemberType->isAnyPointerType()) {
     return llvm::None;
   }
   auto MemberInitExpr = CI->getInit();
-  return isIncompatibleAssignment(MemberType, MemberInitExpr, Env,
+  return isIncompatibleAssignment(MemberType, MemberInitExpr, State.Env,
                                   *Result.Context)
              ? llvm::Optional<CFGElement>(CFGInitializer(CI))
              : llvm::None;
 }
 
 auto buildDiagnoser() {
-  return CFGMatchSwitchBuilder<const Environment, llvm::Optional<CFGElement>>()
+  return CFGMatchSwitchBuilder<const dataflow::TransferStateForDiagnostics<
+                                   PointerNullabilityLattice>,
+                               llvm::Optional<CFGElement>>()
       // (*)
       .CaseOfCFGStmt<UnaryOperator>(isPointerDereference(), diagnoseDereference)
       // (->)
diff --git a/nullability_verification/pointer_nullability_diagnosis.h b/nullability_verification/pointer_nullability_diagnosis.h
index 5547e77..a4785f3 100644
--- a/nullability_verification/pointer_nullability_diagnosis.h
+++ b/nullability_verification/pointer_nullability_diagnosis.h
@@ -5,6 +5,7 @@
 #ifndef CRUBIT_NULLABILITY_VERIFICATION_POINTER_NULLABILITY_DIAGNOSIS_H_
 #define CRUBIT_NULLABILITY_VERIFICATION_POINTER_NULLABILITY_DIAGNOSIS_H_
 
+#include "nullability_verification/pointer_nullability_lattice.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Analysis/FlowSensitive/CFGMatchSwitch.h"
@@ -30,14 +31,17 @@
   ///
   /// TODO(b/233582219): Extend diagnosis to return more information, e.g. the
   /// type of violation.
-  llvm::Optional<CFGElement> diagnose(const CFGElement* Elt, ASTContext& Ctx,
-                                      const dataflow::Environment& Env) {
-    return Diagnoser(*Elt, Ctx, Env);
+  llvm::Optional<CFGElement> diagnose(
+      const CFGElement* Elt, ASTContext& Ctx,
+      const dataflow::TransferStateForDiagnostics<PointerNullabilityLattice>&
+          State) {
+    return Diagnoser(*Elt, Ctx, State);
   }
 
  private:
-  dataflow::CFGMatchSwitch<const dataflow::Environment,
-                           llvm::Optional<CFGElement>>
+  dataflow::CFGMatchSwitch<
+      const dataflow::TransferStateForDiagnostics<PointerNullabilityLattice>,
+      llvm::Optional<CFGElement>>
       Diagnoser;
 };
 
diff --git a/nullability_verification/pointer_nullability_verification_test.cc b/nullability_verification/pointer_nullability_verification_test.cc
index 9de3e88..722532b 100644
--- a/nullability_verification/pointer_nullability_verification_test.cc
+++ b/nullability_verification/pointer_nullability_verification_test.cc
@@ -21,6 +21,7 @@
 namespace {
 
 using dataflow::Environment;
+using dataflow::TransferStateForDiagnostics;
 using dataflow::TypeErasedDataflowAnalysisState;
 using dataflow::test::AnalysisInputs;
 using dataflow::test::AnalysisOutputs;
@@ -38,16 +39,15 @@
               [](ASTContext &ASTCtx, Environment &) {
                 return PointerNullabilityAnalysis(ASTCtx);
               })
-              .withPostVisitCFG(
-                  [&Diagnostics, &Diagnoser](
-                      ASTContext &Ctx, const CFGElement &Elt,
-                      const TypeErasedDataflowAnalysisState &State) {
-                    auto EltDiagnostics =
-                        Diagnoser.diagnose(&Elt, Ctx, State.Env);
-                    if (EltDiagnostics.has_value()) {
-                      Diagnostics.push_back(EltDiagnostics.value());
-                    }
-                  })
+              .withPostVisitCFG([&Diagnostics, &Diagnoser](
+                                    ASTContext &Ctx, const CFGElement &Elt,
+                                    const TransferStateForDiagnostics<
+                                        PointerNullabilityLattice> &State) {
+                auto EltDiagnostics = Diagnoser.diagnose(&Elt, Ctx, State);
+                if (EltDiagnostics.has_value()) {
+                  Diagnostics.push_back(EltDiagnostics.value());
+                }
+              })
               .withASTBuildArgs({"-fsyntax-only", "-std=c++17",
                                  "-Wno-unused-value", "-Wno-nonnull"}),
           [&Diagnostics](