diff --git a/clang/docs/tools/clang-formatted-files.txt b/clang/docs/tools/clang-formatted-files.txt
--- a/clang/docs/tools/clang-formatted-files.txt
+++ b/clang/docs/tools/clang-formatted-files.txt
@@ -129,6 +129,7 @@
 clang/include/clang/Analysis/FlowSensitive/DataflowEnvironment.h
 clang/include/clang/Analysis/FlowSensitive/DataflowLattice.h
 clang/include/clang/Analysis/FlowSensitive/DataflowWorklist.h
+clang/include/clang/Analysis/FlowSensitive/Diagnosis.h
 clang/include/clang/Analysis/FlowSensitive/MapLattice.h
 clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h
 clang/include/clang/Analysis/FlowSensitive/Solver.h
diff --git a/clang/include/clang/Analysis/FlowSensitive/Diagnosis.h b/clang/include/clang/Analysis/FlowSensitive/Diagnosis.h
new file mode 100644
--- /dev/null
+++ b/clang/include/clang/Analysis/FlowSensitive/Diagnosis.h
@@ -0,0 +1,68 @@
+//===- Diagnosis.h ----------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines a helper function for gathering diagnostics using
+//  the results of a dataflow analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DIAGNOSIS_H
+#define LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DIAGNOSIS_H
+
+#include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
+#include "clang/Analysis/FlowSensitive/DataflowAnalysis.h"
+#include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
+#include "clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h"
+#include "llvm/ADT/Optional.h"
+#include <functional>
+#include <utility>
+#include <vector>
+
+namespace clang {
+namespace dataflow {
+
+/// Collects diagnostics from all blocks in a CFG, given some dataflow analysis
+/// results and a `Diagnose` function which can be run on individual statements.
+///
+/// The `CFCtx`, `Analysyis`, and `InitEnv` arguments should be the same
+/// arguments passed to `runDataflowAnalysis`, and the `BlockStates` argument
+/// should be its returned value (extracted from the `llvm::Expected`).
+template <typename AnalysisT, typename DiagnosticT>
+std::vector<DiagnosticT> diagnoseCFG(
+    const ControlFlowContext &CFCtx, AnalysisT &Analysis,
+    const Environment &InitEnv,
+    std::vector<
+        llvm::Optional<DataflowAnalysisState<typename AnalysisT::Lattice>>>
+        &BlockStates,
+    std::function<std::vector<DiagnosticT>(
+        const Stmt *, const typename AnalysisT::Lattice &, const Environment &)>
+        Diagnose) {
+  std::vector<DiagnosticT> Diagnostics;
+  for (const CFGBlock *Block : CFCtx.getCFG()) {
+    // Skip blocks that were not evaluated.
+    if (!BlockStates[Block->getBlockID()])
+      continue;
+    transferBlock(CFCtx, BlockStates, *Block, InitEnv, Analysis,
+                  [&Diagnose,
+                   &Diagnostics](const clang::CFGStmt &Stmt,
+                                 const TypeErasedDataflowAnalysisState &State) {
+                    auto *L = llvm::any_cast<typename AnalysisT::Lattice>(
+                        &State.Lattice.Value);
+                    auto BlockDiagnostics =
+                        Diagnose(Stmt.getStmt(), *L, State.Env);
+                    std::move(BlockDiagnostics.begin(), BlockDiagnostics.end(),
+                              std::back_inserter(Diagnostics));
+                  });
+  }
+  return Diagnostics;
+}
+
+} // namespace dataflow
+} // namespace clang
+
+#endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_DIAGNOSIS_H
diff --git a/clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h b/clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h
--- a/clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h
+++ b/clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h
@@ -20,6 +20,8 @@
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #include "clang/Analysis/FlowSensitive/MatchSwitch.h"
 #include "clang/Analysis/FlowSensitive/SourceLocationsLattice.h"
+#include "clang/Basic/SourceLocation.h"
+#include <vector>
 
 namespace clang {
 namespace dataflow {
@@ -71,6 +73,20 @@
   MatchSwitch<TransferState<SourceLocationsLattice>> TransferMatchSwitch;
 };
 
+class UncheckedOptionalAccessDiagnosis {
+public:
+  UncheckedOptionalAccessDiagnosis(
+      ASTContext &AstContext, UncheckedOptionalAccessModelOptions Options = {});
+
+  std::vector<SourceLocation> diagnose(const Stmt *Stmt,
+                                       const Environment &Env);
+
+private:
+  ASTContext &Context;
+  MatchSwitch<const Environment, std::vector<SourceLocation>>
+      DiagnoseMatchSwitch;
+};
+
 } // namespace dataflow
 } // namespace clang
 
diff --git a/clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp b/clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp
--- a/clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp
+++ b/clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp
@@ -22,11 +22,13 @@
 #include "clang/Analysis/FlowSensitive/MatchSwitch.h"
 #include "clang/Analysis/FlowSensitive/SourceLocationsLattice.h"
 #include "clang/Analysis/FlowSensitive/Value.h"
+#include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include <cassert>
 #include <memory>
 #include <utility>
+#include <vector>
 
 namespace clang {
 namespace dataflow {
@@ -551,6 +553,17 @@
   return llvm::None;
 }
 
+StatementMatcher
+valueCall(llvm::Optional<StatementMatcher> &IgnorableOptional) {
+  return isOptionalMemberCallWithName("value", IgnorableOptional);
+}
+
+StatementMatcher
+valueOperatorCall(llvm::Optional<StatementMatcher> &IgnorableOptional) {
+  return expr(anyOf(isOptionalOperatorCallWithName("*", IgnorableOptional),
+                    isOptionalOperatorCallWithName("->", IgnorableOptional)));
+}
+
 auto buildTransferMatchSwitch(
     const UncheckedOptionalAccessModelOptions &Options) {
   // FIXME: Evaluate the efficiency of matchers. If using matchers results in a
@@ -592,20 +605,18 @@
 
       // optional::value
       .CaseOf<CXXMemberCallExpr>(
-          isOptionalMemberCallWithName("value", IgnorableOptional),
+          valueCall(IgnorableOptional),
           [](const CXXMemberCallExpr *E, const MatchFinder::MatchResult &,
              LatticeTransferState &State) {
             transferUnwrapCall(E, E->getImplicitObjectArgument(), State);
           })
 
       // optional::operator*, optional::operator->
-      .CaseOf<CallExpr>(
-          expr(anyOf(isOptionalOperatorCallWithName("*", IgnorableOptional),
-                     isOptionalOperatorCallWithName("->", IgnorableOptional))),
-          [](const CallExpr *E, const MatchFinder::MatchResult &,
-             LatticeTransferState &State) {
-            transferUnwrapCall(E, E->getArg(0), State);
-          })
+      .CaseOf<CallExpr>(valueOperatorCall(IgnorableOptional),
+                        [](const CallExpr *E, const MatchFinder::MatchResult &,
+                           LatticeTransferState &State) {
+                          transferUnwrapCall(E, E->getArg(0), State);
+                        })
 
       // optional::has_value
       .CaseOf<CXXMemberCallExpr>(isOptionalMemberCallWithName("has_value"),
@@ -653,6 +664,49 @@
       .Build();
 }
 
+std::vector<SourceLocation> diagnoseUnwrapCall(const Expr *UnwrapExpr,
+                                               const Expr *ObjectExpr,
+                                               const Environment &Env) {
+  if (auto *OptionalVal =
+          Env.getValue(*ObjectExpr, SkipPast::ReferenceThenPointer)) {
+    auto *Prop = OptionalVal->getProperty("has_value");
+    if (auto *HasValueVal = cast_or_null<BoolValue>(Prop)) {
+      if (Env.flowConditionImplies(*HasValueVal))
+        return {};
+    }
+  }
+
+  // Record that this unwrap is *not* provably safe.
+  // FIXME: include either the name of the optional (if applicable) or a source
+  // range of the access for easier interpretation of the result.
+  return {ObjectExpr->getBeginLoc()};
+}
+
+auto buildDiagnoseMatchSwitch(
+    const UncheckedOptionalAccessModelOptions &Options) {
+  // FIXME: Evaluate the efficiency of matchers. If using matchers results in a
+  // lot of duplicated work (e.g. string comparisons), consider providing APIs
+  // that avoid it through memoization.
+  auto IgnorableOptional = ignorableOptional(Options);
+  return MatchSwitchBuilder<const Environment, std::vector<SourceLocation>>()
+      // optional::value
+      .CaseOf<CXXMemberCallExpr>(
+          valueCall(IgnorableOptional),
+          [](const CXXMemberCallExpr *E, const MatchFinder::MatchResult &,
+             const Environment &Env) {
+            return diagnoseUnwrapCall(E, E->getImplicitObjectArgument(), Env);
+          })
+
+      // optional::operator*, optional::operator->
+      .CaseOf<CallExpr>(
+          valueOperatorCall(IgnorableOptional),
+          [](const CallExpr *E, const MatchFinder::MatchResult &,
+             const Environment &Env) {
+            return diagnoseUnwrapCall(E, E->getArg(0), Env);
+          })
+      .Build();
+}
+
 } // namespace
 
 ast_matchers::DeclarationMatcher
@@ -699,5 +753,16 @@
   return true;
 }
 
+UncheckedOptionalAccessDiagnosis::UncheckedOptionalAccessDiagnosis(
+    ASTContext &AstContext, UncheckedOptionalAccessModelOptions Options)
+    : Context(AstContext),
+      DiagnoseMatchSwitch(buildDiagnoseMatchSwitch(Options)) {}
+
+std::vector<SourceLocation>
+UncheckedOptionalAccessDiagnosis::diagnose(const Stmt *Stmt,
+                                           const Environment &Env) {
+  return DiagnoseMatchSwitch(*Stmt, Context, Env);
+}
+
 } // namespace dataflow
 } // namespace clang
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
@@ -23,6 +23,8 @@
 #include "clang/Analysis/FlowSensitive/ControlFlowContext.h"
 #include "clang/Analysis/FlowSensitive/DataflowAnalysis.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
+#include "clang/Analysis/FlowSensitive/Diagnosis.h"
+#include "clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h"
 #include "clang/Analysis/FlowSensitive/WatchedLiteralsSolver.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Serialization/PCHContainerOperations.h"
@@ -62,21 +64,33 @@
 buildStatementToAnnotationMapping(const FunctionDecl *Func,
                                   llvm::Annotations AnnotatedCode);
 
+struct AnalysisResults {
+  AnalysisResults(
+      ASTContext &Context, const ControlFlowContext &CFCtx,
+      const Environment &Env, TypeErasedDataflowAnalysis &Analysis,
+      llvm::DenseMap<const clang::Stmt *, std::string> &Annotations,
+      std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates)
+      : Context(Context), CFCtx(CFCtx), Env(Env), Analysis(Analysis),
+        Annotations(Annotations), BlockStates(BlockStates) {}
+
+  ASTContext &Context;
+  const ControlFlowContext &CFCtx;
+  const Environment &Env;
+  TypeErasedDataflowAnalysis &Analysis;
+  llvm::DenseMap<const clang::Stmt *, std::string> &Annotations;
+  std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates;
+};
+
 // Runs dataflow on the body of the function that matches `func_matcher` in code
 // snippet `code`. Requires: `Analysis` contains a type `Lattice`.
 template <typename AnalysisT>
-llvm::Error checkDataflow(
+llvm::Error checkDataflowHelper(
     llvm::StringRef Code,
     ast_matchers::internal::Matcher<FunctionDecl> FuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
-    std::function<void(
-        llvm::ArrayRef<std::pair<
-            std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
-        ASTContext &)>
-        Expectations,
+    std::function<void(AnalysisResults)> Expectations,
     ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
 
   llvm::Annotations AnnotatedCode(Code);
   auto Unit = tooling::buildASTFromCodeWithArgs(
@@ -121,35 +135,63 @@
     return MaybeBlockStates.takeError();
   auto &BlockStates = *MaybeBlockStates;
 
-  if (BlockStates.empty()) {
-    Expectations({}, Context);
-    return llvm::Error::success();
-  }
-
-  // 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 : CFCtx->getCFG()) {
-    // Skip blocks that were not evaluated.
-    if (!BlockStates[Block->getBlockID()])
-      continue;
-
-    transferBlock(
-        *CFCtx, BlockStates, *Block, Env, Analysis,
-        [&Results, &Annotations](const clang::CFGStmt &Stmt,
-                                 const TypeErasedDataflowAnalysisState &State) {
-          auto It = Annotations.find(Stmt.getStmt());
-          if (It == Annotations.end())
-            return;
-          auto *Lattice =
-              llvm::any_cast<typename AnalysisT::Lattice>(&State.Lattice.Value);
-          Results.emplace_back(It->second, StateT{*Lattice, State.Env});
-        });
-  }
-  Expectations(Results, Context);
+  AnalysisResults AnalysisResults(Context, *CFCtx, Env, Analysis, Annotations,
+                                  BlockStates);
+  Expectations(AnalysisResults);
   return llvm::Error::success();
 }
 
+template <typename AnalysisT>
+llvm::Error checkDataflow(
+    llvm::StringRef Code,
+    ast_matchers::internal::Matcher<FunctionDecl> FuncMatcher,
+    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
+    std::function<void(
+        llvm::ArrayRef<std::pair<
+            std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
+        ASTContext &)>
+        Expectations,
+    ArrayRef<std::string> Args,
+    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
+  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
+
+  return checkDataflowHelper(
+      Code, std::move(FuncMatcher), std::move(MakeAnalysis),
+      [&Expectations](AnalysisResults AnalysisResults) {
+        if (AnalysisResults.BlockStates.empty()) {
+          Expectations({}, AnalysisResults.Context);
+          return;
+        }
+
+        auto &Annotations = AnalysisResults.Annotations;
+
+        // 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 : AnalysisResults.CFCtx.getCFG()) {
+          // Skip blocks that were not evaluated.
+          if (!AnalysisResults.BlockStates[Block->getBlockID()])
+            continue;
+
+          transferBlock(
+              AnalysisResults.CFCtx, AnalysisResults.BlockStates, *Block,
+              AnalysisResults.Env, AnalysisResults.Analysis,
+              [&Results,
+               &Annotations](const clang::CFGStmt &Stmt,
+                             const TypeErasedDataflowAnalysisState &State) {
+                auto It = Annotations.find(Stmt.getStmt());
+                if (It == Annotations.end())
+                  return;
+                auto *Lattice = llvm::any_cast<typename AnalysisT::Lattice>(
+                    &State.Lattice.Value);
+                Results.emplace_back(It->second, StateT{*Lattice, State.Env});
+              });
+        }
+        Expectations(Results, AnalysisResults.Context);
+      },
+      Args, VirtualMappedFiles);
+}
+
 // Runs dataflow on the body of the function named `target_fun` in code snippet
 // `code`.
 template <typename AnalysisT>
@@ -168,6 +210,33 @@
                        VirtualMappedFiles);
 }
 
+template <typename AnalysisT, typename DiagnosticT>
+llvm::Error checkDataflowDiagnosis(
+    llvm::StringRef Code,
+    ast_matchers::internal::Matcher<FunctionDecl> FuncMatcher,
+    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
+    std::function<std::function<std::vector<DiagnosticT>(
+        const Stmt *, const typename AnalysisT::Lattice &,
+        const Environment &)>(ASTContext &)>
+        MakeDiagnosis,
+    std::function<void(llvm::DenseMap<const clang::Stmt *, std::string> &,
+                       std::vector<DiagnosticT>, ASTContext &)>
+        Expectations,
+    ArrayRef<std::string> Args,
+    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
+  return checkDataflowHelper(
+      Code, std::move(FuncMatcher), std::move(MakeAnalysis),
+      [&MakeDiagnosis, &Expectations](AnalysisResults AnalysisResults) {
+        auto Diagnose = MakeDiagnosis(AnalysisResults.Context);
+        auto Diags = diagnoseCFG(
+            AnalysisResults.CFCtx, AnalysisResults.BlockStates,
+            AnalysisResults.Env, AnalysisResults.Analysis, Diagnose);
+        Expectations(AnalysisResults.Annotations, Diags,
+                     AnalysisResults.Context);
+      },
+      Args, VirtualMappedFiles);
+}
+
 /// Returns the `ValueDecl` for the given identifier.
 ///
 /// Requirements:
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
@@ -11,9 +11,10 @@
 #include "TestingSupport.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
-#include "clang/Analysis/FlowSensitive/SourceLocationsLattice.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
@@ -26,8 +27,7 @@
 using namespace dataflow;
 using namespace test;
 
-using ::testing::Pair;
-using ::testing::UnorderedElementsAre;
+using ::testing::ContainerEq;
 
 // FIXME: Move header definitions in separate file(s).
 static constexpr char CSDtdDefHeader[] = R"(
@@ -1180,13 +1180,6 @@
 } // namespace base
 )";
 
-/// Converts `L` to string.
-static std::string ConvertToString(const SourceLocationsLattice &L,
-                                   const ASTContext &Ctx) {
-  return L.getSourceLocations().empty() ? "safe"
-                                        : "unsafe: " + DebugString(L, Ctx);
-}
-
 /// Replaces all occurrences of `Pattern` in `S` with `Replacement`.
 static void ReplaceAllOccurrences(std::string &S, const std::string &Pattern,
                                   const std::string &Replacement) {
@@ -1207,18 +1200,14 @@
 class UncheckedOptionalAccessTest
     : public ::testing::TestWithParam<OptionalTypeIdentifier> {
 protected:
-  template <typename LatticeChecksMatcher>
-  void ExpectLatticeChecksFor(std::string SourceCode,
-                              LatticeChecksMatcher MatchesLatticeChecks) {
-    ExpectLatticeChecksFor(SourceCode, ast_matchers::hasName("target"),
-                           MatchesLatticeChecks);
+  void ExpectDiagnosticsFor(std::string SourceCode) {
+    ExpectDiagnosticsFor(SourceCode, ast_matchers::hasName("target"));
   }
 
 private:
-  template <typename FuncDeclMatcher, typename LatticeChecksMatcher>
-  void ExpectLatticeChecksFor(std::string SourceCode,
-                              FuncDeclMatcher FuncMatcher,
-                              LatticeChecksMatcher MatchesLatticeChecks) {
+  template <typename FuncDeclMatcher>
+  void ExpectDiagnosticsFor(std::string SourceCode,
+                            FuncDeclMatcher FuncMatcher) {
     ReplaceAllOccurrences(SourceCode, "$ns", GetParam().NamespaceName);
     ReplaceAllOccurrences(SourceCode, "$optional", GetParam().TypeName);
 
@@ -1245,29 +1234,51 @@
     )");
     const tooling::FileContentMappings FileContents(Headers.begin(),
                                                     Headers.end());
-    llvm::Error Error = checkDataflow<UncheckedOptionalAccessModel>(
-        SourceCode, FuncMatcher,
-        [](ASTContext &Ctx, Environment &) {
-          return UncheckedOptionalAccessModel(
-              Ctx, UncheckedOptionalAccessModelOptions{
-                       /*IgnoreSmartPointerDereference=*/true});
-        },
-        [&MatchesLatticeChecks](
-            llvm::ArrayRef<std::pair<
-                std::string, DataflowAnalysisState<SourceLocationsLattice>>>
-                CheckToLatticeMap,
-            ASTContext &Ctx) {
-          // FIXME: Consider using a matcher instead of translating
-          // `CheckToLatticeMap` to `CheckToStringifiedLatticeMap`.
-          std::vector<std::pair<std::string, std::string>>
-              CheckToStringifiedLatticeMap;
-          for (const auto &E : CheckToLatticeMap) {
-            CheckToStringifiedLatticeMap.emplace_back(
-                E.first, ConvertToString(E.second.Lattice, Ctx));
-          }
-          EXPECT_THAT(CheckToStringifiedLatticeMap, MatchesLatticeChecks);
-        },
-        {"-fsyntax-only", "-std=c++17", "-Wno-undefined-inline"}, FileContents);
+    UncheckedOptionalAccessModelOptions Options{
+        /*IgnoreSmartPointerDereference=*/true};
+    llvm::Error Error =
+        checkDataflowDiagnosis<UncheckedOptionalAccessModel, SourceLocation>(
+            SourceCode, FuncMatcher,
+            [Options](ASTContext &Ctx, Environment &) {
+              return UncheckedOptionalAccessModel(Ctx, Options);
+            },
+            [Options](ASTContext &Ctx) {
+              UncheckedOptionalAccessDiagnosis Diagnosis(Ctx, Options);
+              return [Diagnosis = std::move(Diagnosis)](
+                         const Stmt *Stmt,
+                         const UncheckedOptionalAccessModel::Lattice &,
+                         const Environment &Env) mutable {
+                return Diagnosis.diagnose(Stmt, Env);
+              };
+            },
+            [](llvm::DenseMap<const Stmt *, std::string> &Annotations,
+               llvm::DenseSet<SourceLocation> &Locs, ASTContext &Ctx) {
+              auto &SrcMgr = Ctx.getSourceManager();
+
+              llvm::DenseSet<unsigned> AnnotationLines;
+              for (const auto &Pair : Annotations) {
+                auto *Stmt = Pair.getFirst();
+                AnnotationLines.insert(
+                    SrcMgr.getPresumedLineNumber(Stmt->getBeginLoc()));
+                // We add both the begin and end locations, so that if the
+                // statement spans multiple lines then the test will fail.
+                //
+                // FIXME: Going forward, we should change this to instead just
+                // get the single line number from the annotation itself, rather
+                // than looking at the statement it's attached to.
+                AnnotationLines.insert(
+                    SrcMgr.getPresumedLineNumber(Stmt->getEndLoc()));
+              }
+
+              llvm::DenseSet<unsigned> DiagnosticLines;
+              for (SourceLocation &Loc : Locs) {
+                DiagnosticLines.insert(SrcMgr.getPresumedLineNumber(Loc));
+              }
+
+              EXPECT_THAT(DiagnosticLines, ContainerEq(AnnotationLines));
+            },
+            {"-fsyntax-only", "-std=c++17", "-Wno-undefined-inline"},
+            FileContents);
     if (Error)
       FAIL() << llvm::toString(std::move(Error));
   }
@@ -1283,65 +1294,55 @@
     });
 
 TEST_P(UncheckedOptionalAccessTest, EmptyFunctionBody) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     void target() {
       (void)0;
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UnwrapUsingValueNoCheck) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
-      std::move(opt).value();
-      /*[[check]]*/
+      std::move(opt).value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UnwrapUsingOperatorStarNoCheck) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
-      *opt;
-      /*[[check]]*/
+      *opt; // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:8")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
-      *std::move(opt);
-      /*[[check]]*/
+      *std::move(opt); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:8")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UnwrapUsingOperatorArrowNoCheck) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1350,13 +1351,11 @@
     };
 
     void target($ns::$optional<Foo> opt) {
-      opt->foo();
-      /*[[check]]*/
+      opt->foo(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:9:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1365,107 +1364,91 @@
     };
 
     void target($ns::$optional<Foo> opt) {
-      std::move(opt)->foo();
-      /*[[check]]*/
+      std::move(opt)->foo(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:9:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, HasValueCheck) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
       if (opt.has_value()) {
         opt.value();
-        /*[[check]]*/
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OperatorBoolCheck) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
       if (opt) {
         opt.value();
-        /*[[check]]*/
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UnwrapFunctionCallResultNoCheck) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
-      Make<$ns::$optional<int>>().value();
+      Make<$ns::$optional<int>>().value(); // [[unsafe]]
       (void)0;
-      /*[[check]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
-      std::move(opt).value();
-      /*[[check]]*/
+      std::move(opt).value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:5:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, DefaultConstructor) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt;
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:6:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, NulloptConstructor) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt($ns::nullopt);
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:6:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, InPlaceConstructor) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt($ns::in_place, 3);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1473,12 +1456,10 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1488,12 +1469,10 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place, 3, false);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1503,46 +1482,38 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place, {3});
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ValueConstructor) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt(21);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = $ns::$optional<int>(21);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
-  ExpectLatticeChecksFor(R"(
+  )");
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<$ns::$optional<int>> opt(Make<$ns::$optional<int>>());
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1552,12 +1523,10 @@
     void target() {
       $ns::$optional<MyString> opt("foo");
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1569,12 +1538,10 @@
     void target() {
       $ns::$optional<Bar> opt(Make<Foo>());
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1584,14 +1551,12 @@
     void target() {
       $ns::$optional<Foo> opt(3);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ConvertibleOptionalConstructor) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1603,13 +1568,11 @@
 
     void target() {
       $ns::$optional<Bar> opt(Make<$ns::$optional<Foo>>());
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:12:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1621,13 +1584,11 @@
 
     void target() {
       $ns::$optional<Bar> opt(Make<$ns::$optional<Foo>>());
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:12:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1640,13 +1601,11 @@
     void target() {
       $ns::$optional<Foo> opt1 = $ns::nullopt;
       $ns::$optional<Bar> opt2(opt1);
-      opt2.value();
-      /*[[check]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:13:7")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1659,12 +1618,10 @@
       $ns::$optional<Foo> opt1(Make<Foo>());
       $ns::$optional<Bar> opt2(opt1);
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1677,25 +1634,21 @@
       $ns::$optional<Foo> opt1(Make<Foo>());
       $ns::$optional<Bar> opt2(opt1);
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, MakeOptional) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = $ns::make_optional(0);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1705,12 +1658,10 @@
     void target() {
       $ns::$optional<Foo> opt = $ns::make_optional<Foo>(21, 22);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1721,104 +1672,83 @@
       char a = 'a';
       $ns::$optional<Foo> opt = $ns::make_optional<Foo>({a});
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ValueOr) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt;
       opt.value_or(0);
       (void)0;
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ValueOrComparison) {
   // Pointers.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int*> opt) {
       if (opt.value_or(nullptr) != nullptr) {
         opt.value();
-        /*[[check-ptrs-1]]*/
       } else {
-        opt.value();
-        /*[[check-ptrs-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-ptrs-1", "safe"),
-                           Pair("check-ptrs-2", "unsafe: input.cc:9:9")));
+  )code");
 
   // Integers.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> opt) {
       if (opt.value_or(0) != 0) {
         opt.value();
-        /*[[check-ints-1]]*/
       } else {
-        opt.value();
-        /*[[check-ints-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-ints-1", "safe"),
-                           Pair("check-ints-2", "unsafe: input.cc:9:9")));
+  )code");
 
   // Strings.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<std::string> opt) {
       if (!opt.value_or("").empty()) {
         opt.value();
-        /*[[check-strings-1]]*/
       } else {
-        opt.value();
-        /*[[check-strings-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-strings-1", "safe"),
-                           Pair("check-strings-2", "unsafe: input.cc:9:9")));
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<std::string> opt) {
       if (opt.value_or("") != "") {
         opt.value();
-        /*[[check-strings-neq-1]]*/
       } else {
-        opt.value();
-        /*[[check-strings-neq-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(
-          Pair("check-strings-neq-1", "safe"),
-          Pair("check-strings-neq-2", "unsafe: input.cc:9:9")));
+  )code");
 
   // Pointer-to-optional.
   //
   // FIXME: make `opt` a parameter directly, once we ensure that all `optional`
   // values have a `has_value` property.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -1826,43 +1756,35 @@
       $ns::$optional<int> *opt = &p;
       if (opt->value_or(0) != 0) {
         opt->value();
-        /*[[check-pto-1]]*/
       } else {
-        opt->value();
-        /*[[check-pto-2]]*/
+        opt->value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-pto-1", "safe"),
-                           Pair("check-pto-2", "unsafe: input.cc:10:9")));
+  )code");
 }
 
 TEST_P(UncheckedOptionalAccessTest, Emplace) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt;
       opt.emplace(0);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> *opt) {
       opt->emplace(0);
       opt->value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
   // FIXME: Add tests that call `emplace` in conditional branches:
-  //  ExpectLatticeChecksFor(
+  //  ExpectDiagnosticsFor(
   //      R"(
   //    #include "unchecked_optional_access_test.h"
   //
@@ -1872,47 +1794,39 @@
   //      }
   //      if (b) {
   //        opt.value();
-  //        /*[[check-1]]*/
   //      } else {
-  //        opt.value();
-  //        /*[[check-2]]*/
+  //        opt.value(); // [[unsafe]]
   //      }
   //    }
-  //  )",
-  //      UnorderedElementsAre(Pair("check-1", "safe"),
-  //                           Pair("check-2", "unsafe: input.cc:12:9")));
+  //  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, Reset) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = $ns::make_optional(0);
       opt.reset();
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:7:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> &opt) {
       if (opt.has_value()) {
         opt.reset();
-        opt.value();
-        /*[[check]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:7:9")));
+  )");
 
   // FIXME: Add tests that call `reset` in conditional branches:
-  //  ExpectLatticeChecksFor(
+  //  ExpectDiagnosticsFor(
   //      R"(
   //    #include "unchecked_optional_access_test.h"
   //
@@ -1922,20 +1836,16 @@
   //        opt.reset();
   //      }
   //      if (b) {
-  //        opt.value();
-  //        /*[[check-1]]*/
+  //        opt.value(); // [[unsafe]]
   //      } else {
   //        opt.value();
-  //        /*[[check-2]]*/
   //      }
   //    }
-  //  )",
-  //      UnorderedElementsAre(Pair("check-1", "unsafe: input.cc:10:9"),
-  //                           Pair("check-2", "safe")));
+  //  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ValueAssignment) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1944,12 +1854,10 @@
       $ns::$optional<Foo> opt;
       opt = Foo();
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1958,12 +1866,10 @@
       $ns::$optional<Foo> opt;
       (opt = Foo()).value();
       (void)0;
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1974,12 +1880,10 @@
       $ns::$optional<MyString> opt;
       opt = "foo";
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1989,14 +1893,12 @@
     void target() {
       $ns::$optional<MyString> opt;
       (opt = "foo").value();
-      /*[[check]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalConversionAssignment) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2011,12 +1913,10 @@
       $ns::$optional<Bar> opt2;
       opt2 = opt1;
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2031,14 +1931,12 @@
       $ns::$optional<Bar> opt2;
       if (opt2.has_value()) {
         opt2 = opt1;
-        opt2.value();
-        /*[[check]]*/
+        opt2.value(); // [[unsafe]]
       }
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:15:9")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2053,41 +1951,35 @@
       $ns::$optional<Bar> opt2;
       (opt2 = opt1).value();
       (void)0;
-      /*[[check]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("check", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, NulloptAssignment) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
       opt = $ns::nullopt;
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:7:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
-      (opt = $ns::nullopt).value();
-      /*[[check]]*/
+      (opt = $ns::nullopt).value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:6:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalSwap) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2098,16 +1990,12 @@
       opt1.swap(opt2);
 
       opt1.value();
-      /*[[check-1]]*/
 
-      opt2.value();
-      /*[[check-2]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"),
-                           Pair("check-2", "unsafe: input.cc:13:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2118,18 +2006,14 @@
       opt2.swap(opt1);
 
       opt1.value();
-      /*[[check-3]]*/
 
-      opt2.value();
-      /*[[check-4]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-3", "safe"),
-                           Pair("check-4", "unsafe: input.cc:13:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, StdSwap) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2140,16 +2024,12 @@
       std::swap(opt1, opt2);
 
       opt1.value();
-      /*[[check-1]]*/
 
-      opt2.value();
-      /*[[check-2]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"),
-                           Pair("check-2", "unsafe: input.cc:13:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2160,20 +2040,16 @@
       std::swap(opt2, opt1);
 
       opt1.value();
-      /*[[check-3]]*/
 
-      opt2.value();
-      /*[[check-4]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-3", "safe"),
-                           Pair("check-4", "unsafe: input.cc:13:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UniquePtrToStructWithOptionalField) {
   // We suppress diagnostics for values reachable from smart pointers (other
   // than `optional` itself).
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2190,16 +2066,13 @@
     void target() {
       smart_ptr<Foo> foo;
       *foo->opt;
-      /*[[check-1]]*/
       *(*foo).opt;
-      /*[[check-2]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"), Pair("check-2", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, CallReturningOptional) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2208,12 +2081,10 @@
     void target() {
       $ns::$optional<int> opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-1]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "unsafe: input.cc:9:7")));
-  ExpectLatticeChecksFor(
+  )");
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2222,13 +2093,11 @@
     void target() {
       $ns::$optional<int> opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-2]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-2", "unsafe: input.cc:9:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2238,13 +2107,11 @@
     void target() {
       IntOpt opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-3]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-3", "unsafe: input.cc:10:7")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2254,16 +2121,14 @@
     void target() {
       IntOpt opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-4]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-4", "unsafe: input.cc:10:7")));
+  )");
 }
 
 // Verifies that the model sees through aliases.
 TEST_P(UncheckedOptionalAccessTest, WithAlias) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2271,18 +2136,16 @@
     using MyOptional = $ns::$optional<T>;
 
     void target(MyOptional<int> opt) {
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:8:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalValueOptional) {
   // Basic test that nested values are populated.  We nest an optional because
   // its easy to use in a test, but the type of the nested value shouldn't
   // matter.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2291,14 +2154,12 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && *foo) {
         foo->value();
-        /*[[access]]*/
       }
     }
-  )",
-      UnorderedElementsAre(Pair("access", "safe")));
+  )");
 
   // Mutation is supported for nested values.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2307,18 +2168,16 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && *foo) {
         foo->reset();
-        foo->value();
-        /*[[reset]]*/
+        foo->value(); // [[unsafe]]
       }
     }
-  )",
-      UnorderedElementsAre(Pair("reset", "unsafe: input.cc:9:9")));
+  )");
 }
 
 // Tests that structs can be nested. We use an optional field because its easy
 // to use in a test, but the type of the field shouldn't matter.
 TEST_P(UncheckedOptionalAccessTest, OptionalValueStruct) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2329,18 +2188,16 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && foo->opt) {
         foo->opt.value();
-        /*[[access]]*/
       }
     }
-  )",
-      UnorderedElementsAre(Pair("access", "safe")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalValueInitialization) {
   // FIXME: Fix when to initialize `value`. All unwrapping should be safe in
   // this example, but `value` initialization is done multiple times during the
   // fixpoint iterations and joining the environment won't correctly merge them.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2359,15 +2216,13 @@
       }
       // Now we merge the two values. UncheckedOptionalAccessModel::merge() will
       // throw away the "value" property.
-      foo->value();
-      /*[[merge]]*/
+      foo->value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("merge", "unsafe: input.cc:19:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, AssignThroughLvalueReferencePtr) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2379,56 +2234,48 @@
     void target() {
       smart_ptr<$ns::$optional<float>> x;
       *x = $ns::nullopt;
-      (*x).value();
-      /*[[check]]*/
+      (*x).value(); // [[unsafe]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "unsafe: input.cc:12:7")));
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, CorrelatedBranches) {
-  ExpectLatticeChecksFor(R"code(
+  ExpectDiagnosticsFor(R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
       if (b || opt.has_value()) {
         if (!b) {
           opt.value();
-          /*[[check-1]]*/
         }
       }
     }
-  )code",
-                         UnorderedElementsAre(Pair("check-1", "safe")));
+  )code");
 
-  ExpectLatticeChecksFor(R"code(
+  ExpectDiagnosticsFor(R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
       if (b && !opt.has_value()) return;
       if (b) {
         opt.value();
-        /*[[check-2]]*/
       }
     }
-  )code",
-                         UnorderedElementsAre(Pair("check-2", "safe")));
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
       if (opt.has_value()) b = true;
       if (b) {
-        opt.value();
-        /*[[check-3]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-3", "unsafe: input.cc:7:9")));
+  )code");
 
-  ExpectLatticeChecksFor(R"code(
+  ExpectDiagnosticsFor(R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
@@ -2436,41 +2283,35 @@
       if (opt.has_value()) b = true;
       if (b) {
         opt.value();
-        /*[[check-4]]*/
       }
     }
-  )code",
-                         UnorderedElementsAre(Pair("check-4", "safe")));
+  )code");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
       if (opt.has_value() == b) {
         if (b) {
           opt.value();
-          /*[[check-5]]*/
         }
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check-5", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b, $ns::$optional<int> opt) {
       if (opt.has_value() != b) {
         if (!b) {
           opt.value();
-          /*[[check-6]]*/
         }
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check-6", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
@@ -2483,30 +2324,26 @@
       }
       if (opt2.has_value()) {
         opt1.value();
-        /*[[check]]*/
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check", "safe")));
+  )");
 
   // FIXME: Add support for operator==.
-  // ExpectLatticeChecksFor(R"(
+  // ExpectDiagnosticsFor(R"(
   //   #include "unchecked_optional_access_test.h"
   //
   //   void target($ns::$optional<int> opt1, $ns::$optional<int> opt2) {
   //     if (opt1 == opt2) {
   //       if (opt1.has_value()) {
   //         opt2.value();
-  //         /*[[check-7]]*/
   //       }
   //     }
   //   }
-  // )",
-  //                     UnorderedElementsAre(Pair("check-7", "safe")));
+  // )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, JoinDistinctValues) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2519,17 +2356,13 @@
       }
       if (opt.has_value()) {
         opt.value();
-        /*[[check-1]]*/
       } else {
-        opt.value();
-        /*[[check-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      UnorderedElementsAre(Pair("check-1", "safe"),
-                           Pair("check-2", "unsafe: input.cc:15:9")));
+  )code");
 
-  ExpectLatticeChecksFor(R"code(
+  ExpectDiagnosticsFor(R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
@@ -2542,12 +2375,10 @@
         if (!opt.has_value()) return;
       }
       opt.value();
-      /*[[check-3]]*/
     }
-  )code",
-                         UnorderedElementsAre(Pair("check-3", "safe")));
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2559,13 +2390,11 @@
       } else {
         opt = Make<$ns::$optional<int>>();
       }
-      opt.value();
-      /*[[check-4]]*/
+      opt.value(); // [[unsafe]]
     }
-  )code",
-      UnorderedElementsAre(Pair("check-4", "unsafe: input.cc:12:7")));
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2577,12 +2406,10 @@
         opt = 2;
       }
       opt.value();
-      /*[[check-5]]*/
     }
-  )code",
-      UnorderedElementsAre(Pair("check-5", "safe")));
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2593,75 +2420,65 @@
       } else {
         opt = Make<$ns::$optional<int>>();
       }
-      opt.value();
-      /*[[check-6]]*/
+      opt.value(); // [[unsafe]]
     }
-  )code",
-      UnorderedElementsAre(Pair("check-6", "unsafe: input.cc:11:7")));
+  )code");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ReassignValueInLoop) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
       while (Make<bool>()) {
         opt.value();
-        /*[[check-1]]*/
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check-1", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
       while (Make<bool>()) {
         opt.value();
-        /*[[check-2]]*/
 
         opt = Make<$ns::$optional<int>>();
         if (!opt.has_value()) return;
       }
     }
-  )",
-                         UnorderedElementsAre(Pair("check-2", "safe")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
       while (Make<bool>()) {
-        opt.value();
-        /*[[check-3]]*/
+        opt.value(); // [[unsafe]]
 
         opt = Make<$ns::$optional<int>>();
       }
     }
-  )",
-      UnorderedElementsAre(Pair("check-3", "unsafe: input.cc:7:9")));
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = 3;
       while (Make<bool>()) {
-        opt.value();
-        /*[[check-4]]*/
+        opt.value(); // [[unsafe]]
 
         opt = Make<$ns::$optional<int>>();
         if (!opt.has_value()) continue;
       }
     }
-  )",
-      UnorderedElementsAre(Pair("check-4", "unsafe: input.cc:7:9")));
+  )");
 }
 
 // FIXME: Add support for: