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
@@ -219,7 +219,9 @@
     std::function<Diagnosis<typename AnalysisT::Lattice, llvm::DenseSet<Diag>>(
         ASTContext &)>
         MakeDiagnosis,
-    std::function<void(llvm::DenseSet<Diag>, ASTContext &)> Expectations,
+    std::function<void(llvm::DenseMap<const clang::Stmt *, std::string> &,
+                       llvm::DenseSet<Diag> &, ASTContext &)>
+        Expectations,
     ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
   return checkDataflowHelper(
@@ -229,7 +231,8 @@
         auto Diags = diagnoseCFG(
             AnalysisResults.CFCtx, AnalysisResults.BlockStates,
             AnalysisResults.Env, AnalysisResults.Analysis, Diagnose);
-        Expectations(Diags, AnalysisResults.Context);
+        Expectations(AnalysisResults.Annotations, Diags,
+                     AnalysisResults.Context);
       },
       Args, VirtualMappedFiles);
 }
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,7 +11,6 @@
 #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"
@@ -28,6 +27,8 @@
 using namespace dataflow;
 using namespace test;
 
+using ::testing::ContainerEq;
+
 // FIXME: Move header definitions in separate file(s).
 static constexpr char CSDtdDefHeader[] = R"(
 #ifndef CSTDDEF_H
@@ -1179,12 +1180,6 @@
 } // namespace base
 )";
 
-/// Converts `L` to string.
-static std::string ConvertToString(const llvm::DenseSet<SourceLocation> &L,
-                                   const ASTContext &Ctx) {
-  return L.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) {
@@ -1205,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);
 
@@ -1260,10 +1251,30 @@
                 return Diagnosis.diagnose(Stmt, Env);
               };
             },
-            [&MatchesLatticeChecks](llvm::DenseSet<SourceLocation> Locs,
-                                    ASTContext &Ctx) {
-              std::string StringifiedDiags = ConvertToString(Locs, Ctx);
-              EXPECT_THAT(StringifiedDiags, MatchesLatticeChecks);
+            [](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. 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);
@@ -1282,65 +1293,55 @@
     });
 
 TEST_P(UncheckedOptionalAccessTest, EmptyFunctionBody) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     void target() {
       (void)0;
-      /*[[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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "unsafe: input.cc:5:8");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, UnwrapUsingOperatorArrowNoCheck) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1349,13 +1350,11 @@
     };
 
     void target($ns::$optional<Foo> opt) {
-      opt->foo();
-      /*[[check]]*/
+      opt->foo(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:9:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1364,107 +1363,91 @@
     };
 
     void target($ns::$optional<Foo> opt) {
-      std::move(opt)->foo();
-      /*[[check]]*/
+      std::move(opt)->foo(); // [[unsafe]]
     }
-  )",
-      "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]]*/
       }
     }
-  )",
-                         "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]]*/
       }
     }
-  )",
-                         "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]]*/
     }
-  )",
-      "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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "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]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1472,12 +1455,10 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1487,12 +1468,10 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place, 3, false);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1502,46 +1481,38 @@
     void target() {
       $ns::$optional<Foo> opt($ns::in_place, {3});
       opt.value();
-      /*[[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]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target() {
       $ns::$optional<int> opt = $ns::$optional<int>(21);
       opt.value();
-      /*[[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]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1551,12 +1522,10 @@
     void target() {
       $ns::$optional<MyString> opt("foo");
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1568,12 +1537,10 @@
     void target() {
       $ns::$optional<Bar> opt(Make<Foo>());
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1583,14 +1550,12 @@
     void target() {
       $ns::$optional<Foo> opt(3);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                        "safe");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ConvertibleOptionalConstructor) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1602,13 +1567,11 @@
 
     void target() {
       $ns::$optional<Bar> opt(Make<$ns::$optional<Foo>>());
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:12:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1620,13 +1583,11 @@
 
     void target() {
       $ns::$optional<Bar> opt(Make<$ns::$optional<Foo>>());
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:12:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -1639,13 +1600,11 @@
     void target() {
       $ns::$optional<Foo> opt1 = $ns::nullopt;
       $ns::$optional<Bar> opt2(opt1);
-      opt2.value();
-      /*[[check]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:13:7");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1658,12 +1617,10 @@
       $ns::$optional<Foo> opt1(Make<Foo>());
       $ns::$optional<Bar> opt2(opt1);
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1676,25 +1633,21 @@
       $ns::$optional<Foo> opt1(Make<Foo>());
       $ns::$optional<Bar> opt2(opt1);
       opt2.value();
-      /*[[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]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1704,12 +1657,10 @@
     void target() {
       $ns::$optional<Foo> opt = $ns::make_optional<Foo>(21, 22);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
@@ -1720,99 +1671,83 @@
       char a = 'a';
       $ns::$optional<Foo> opt = $ns::make_optional<Foo>({a});
       opt.value();
-      /*[[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]]*/
     }
-  )",
-                         "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",
-      "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",
-      "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",
-      "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",
-      "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"
 
@@ -1820,42 +1755,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",
-      "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]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target($ns::$optional<int> *opt) {
       opt->emplace(0);
       opt->value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
   // FIXME: Add tests that call `emplace` in conditional branches:
-  //  ExpectLatticeChecksFor(
+  //  ExpectDiagnosticsFor(
   //      R"(
   //    #include "unchecked_optional_access_test.h"
   //
@@ -1865,46 +1793,39 @@
   //      }
   //      if (b) {
   //        opt.value();
-  //        /*[[check-1]]*/
   //      } else {
-  //        opt.value();
-  //        /*[[check-2]]*/
+  //        opt.value(); // [[unsafe]]
   //      }
   //    }
-  //  )",
-  //      "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]]
     }
-  )",
-      "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]]
       }
     }
-  )",
-      "unsafe: input.cc:7:9");
+  )");
 
   // FIXME: Add tests that call `reset` in conditional branches:
-  //  ExpectLatticeChecksFor(
+  //  ExpectDiagnosticsFor(
   //      R"(
   //    #include "unchecked_optional_access_test.h"
   //
@@ -1914,19 +1835,16 @@
   //        opt.reset();
   //      }
   //      if (b) {
-  //        opt.value();
-  //        /*[[check-1]]*/
+  //        opt.value(); // [[unsafe]]
   //      } else {
   //        opt.value();
-  //        /*[[check-2]]*/
   //      }
   //    }
-  //  )",
-  //      "safe");
+  //  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ValueAssignment) {
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1935,12 +1853,10 @@
       $ns::$optional<Foo> opt;
       opt = Foo();
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
@@ -1949,12 +1865,10 @@
       $ns::$optional<Foo> opt;
       (opt = Foo()).value();
       (void)0;
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1965,12 +1879,10 @@
       $ns::$optional<MyString> opt;
       opt = "foo";
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
@@ -1980,14 +1892,12 @@
     void target() {
       $ns::$optional<MyString> opt;
       (opt = "foo").value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalConversionAssignment) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2002,12 +1912,10 @@
       $ns::$optional<Bar> opt2;
       opt2 = opt1;
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-      "safe");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2022,14 +1930,12 @@
       $ns::$optional<Bar> opt2;
       if (opt2.has_value()) {
         opt2 = opt1;
-        opt2.value();
-        /*[[check]]*/
+        opt2.value(); // [[unsafe]]
       }
     }
-  )",
-      "unsafe: input.cc:15:9");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2044,41 +1950,35 @@
       $ns::$optional<Bar> opt2;
       (opt2 = opt1).value();
       (void)0;
-      /*[[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]]
     }
-  )",
-      "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]]
     }
-  )",
-      "unsafe: input.cc:6:7");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalSwap) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2089,15 +1989,12 @@
       opt1.swap(opt2);
 
       opt1.value();
-      /*[[check-1]]*/
 
-      opt2.value();
-      /*[[check-2]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:13:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2108,17 +2005,14 @@
       opt2.swap(opt1);
 
       opt1.value();
-      /*[[check-3]]*/
 
-      opt2.value();
-      /*[[check-4]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:13:7");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, StdSwap) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2129,15 +2023,12 @@
       std::swap(opt1, opt2);
 
       opt1.value();
-      /*[[check-1]]*/
 
-      opt2.value();
-      /*[[check-2]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:13:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2148,19 +2039,16 @@
       std::swap(opt2, opt1);
 
       opt1.value();
-      /*[[check-3]]*/
 
-      opt2.value();
-      /*[[check-4]]*/
+      opt2.value(); // [[unsafe]]
     }
-  )",
-      "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"
 
@@ -2177,16 +2065,13 @@
     void target() {
       smart_ptr<Foo> foo;
       *foo->opt;
-      /*[[check-1]]*/
       *(*foo).opt;
-      /*[[check-2]]*/
     }
-  )",
-      "safe");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, CallReturningOptional) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2195,12 +2080,10 @@
     void target() {
       $ns::$optional<int> opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-1]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:9:7");
-  ExpectLatticeChecksFor(
+  )");
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2209,13 +2092,11 @@
     void target() {
       $ns::$optional<int> opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-2]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:9:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2225,13 +2106,11 @@
     void target() {
       IntOpt opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-3]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:10:7");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2241,16 +2120,14 @@
     void target() {
       IntOpt opt = 0;
       opt = MakeOpt();
-      opt.value();
-      /*[[check-4]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "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"
 
@@ -2258,18 +2135,16 @@
     using MyOptional = $ns::$optional<T>;
 
     void target(MyOptional<int> opt) {
-      opt.value();
-      /*[[check]]*/
+      opt.value(); // [[unsafe]]
     }
-  )",
-      "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"
 
@@ -2278,14 +2153,12 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && *foo) {
         foo->value();
-        /*[[access]]*/
       }
     }
-  )",
-      "safe");
+  )");
 
   // Mutation is supported for nested values.
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2294,18 +2167,16 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && *foo) {
         foo->reset();
-        foo->value();
-        /*[[reset]]*/
+        foo->value(); // [[unsafe]]
       }
     }
-  )",
-      "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"
 
@@ -2316,18 +2187,16 @@
     void target($ns::$optional<Foo> foo) {
       if (foo && foo->opt) {
         foo->opt.value();
-        /*[[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"
 
@@ -2346,15 +2215,13 @@
       }
       // Now we merge the two values. UncheckedOptionalAccessModel::merge() will
       // throw away the "value" property.
-      foo->value();
-      /*[[merge]]*/
+      foo->value(); // [[unsafe]]
     }
-  )",
-      "unsafe: input.cc:19:7");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, AssignThroughLvalueReferencePtr) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
@@ -2366,56 +2233,48 @@
     void target() {
       smart_ptr<$ns::$optional<float>> x;
       *x = $ns::nullopt;
-      (*x).value();
-      /*[[check]]*/
+      (*x).value(); // [[unsafe]]
     }
-  )",
-      "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",
-                         "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",
-                         "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",
-      "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) {
@@ -2423,41 +2282,35 @@
       if (opt.has_value()) b = true;
       if (b) {
         opt.value();
-        /*[[check-4]]*/
       }
     }
-  )code",
-                         "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]]*/
         }
       }
     }
-  )",
-                         "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]]*/
         }
       }
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
@@ -2470,30 +2323,26 @@
       }
       if (opt2.has_value()) {
         opt1.value();
-        /*[[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]]*/
   //       }
   //     }
   //   }
-  // )",
-  //                     "safe");
+  // )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, JoinDistinctValues) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2506,16 +2355,13 @@
       }
       if (opt.has_value()) {
         opt.value();
-        /*[[check-1]]*/
       } else {
-        opt.value();
-        /*[[check-2]]*/
+        opt.value(); // [[unsafe]]
       }
     }
-  )code",
-      "unsafe: input.cc:15:9");
+  )code");
 
-  ExpectLatticeChecksFor(R"code(
+  ExpectDiagnosticsFor(R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
@@ -2528,12 +2374,10 @@
         if (!opt.has_value()) return;
       }
       opt.value();
-      /*[[check-3]]*/
     }
-  )code",
-                         "safe");
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2545,13 +2389,11 @@
       } else {
         opt = Make<$ns::$optional<int>>();
       }
-      opt.value();
-      /*[[check-4]]*/
+      opt.value(); // [[unsafe]]
     }
-  )code",
-      "unsafe: input.cc:12:7");
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2563,12 +2405,10 @@
         opt = 2;
       }
       opt.value();
-      /*[[check-5]]*/
     }
-  )code",
-      "safe");
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
@@ -2579,75 +2419,65 @@
       } else {
         opt = Make<$ns::$optional<int>>();
       }
-      opt.value();
-      /*[[check-6]]*/
+      opt.value(); // [[unsafe]]
     }
-  )code",
-      "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]]*/
       }
     }
-  )",
-                         "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;
       }
     }
-  )",
-                         "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>>();
       }
     }
-  )",
-      "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;
       }
     }
-  )",
-      "unsafe: input.cc:7:9");
+  )");
 }
 
 // FIXME: Add support for: