diff --git a/clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h b/clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h
--- a/clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h
+++ b/clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h
@@ -48,6 +48,16 @@
   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>;
 
diff --git a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
--- a/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
+++ b/clang/unittests/Analysis/FlowSensitive/TestingSupport.h
@@ -30,6 +30,7 @@
 #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"
@@ -116,11 +117,26 @@
     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 =
+        [=](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) && {
@@ -148,8 +164,9 @@
   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.
@@ -226,11 +243,15 @@
                      const TypeErasedDataflowAnalysisState &)>
       PostVisitCFGClosure = nullptr;
   if (AI.PostVisitCFG) {
-    PostVisitCFGClosure =
-        [&AI, &Context](const CFGElement &Element,
-                        const TypeErasedDataflowAnalysisState &State) {
-          AI.PostVisitCFG(Context, Element, State);
-        };
+    PostVisitCFGClosure = [&AI, &Context](
+                              const CFGElement &Element,
+                              const TypeErasedDataflowAnalysisState &State) {
+      AI.PostVisitCFG(Context, Element,
+                      TransferStateForDiagnostics<typename AnalysisT::Lattice>(
+                          llvm::any_cast<const typename AnalysisT::Lattice &>(
+                              State.Lattice.Value),
+                          State.Env));
+    };
   }
 
   // Additional test setup.
@@ -326,28 +347,28 @@
   // 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) {
-    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}});
-    (void)_;
-    (void)InsertSuccess;
-    assert(InsertSuccess);
-  };
+  auto PostVisitCFG =
+      [&StmtToAnnotations, &AnnotationStates,
+       PrevPostVisitCFG = std::move(AI.PostVisitCFG)](
+          ASTContext &Ctx, const CFGElement &Elt,
+          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 [_, 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))
diff --git a/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp b/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
--- a/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
+++ b/clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp
@@ -1314,7 +1314,8 @@
                 [&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));