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](