[clang][dataflow] Use annotations for optional diagnostic tests

[samestep]

Summary

Followup to D127898. This patch makes the optional model tests much nicer by removing the second parameter (list of expected diagnostics) entirely, and instead just checking the set of annotation line numbers against the set of diagnostics (so you only write annotations on exactly the set of lines where you expect a diagnostic).

Currently this just reuses the existing buildStatementToAnnotationMapping which already gets called in checkDataflowHelper, so all annotations are required to be ranges and attached to specific statements. @gribozavr2 originally suggested just letting the annotations exist à la carte and using their line numbers directly, which I think would be a better idea since it would allow us to deal with statements spanning multiple lines. That would take more implementation effort though, so I'm leaving it to a followup.

[samestep]

• Merge branch 'diagnose-api' into diagnose-test-annotations

[gribozavr2]

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

1258

1264–1267

But maybe there's an easy way to do that -- change checkDataflowDiagnosis pass down the llvm::Annotations objects to this callback? You could even make that change in the earlier patch in the stack.

[ymandel]

This looks quite nice. I really like how you solved the problem of diagnostics representation/checking.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

 template <typename AnalysisT, typename Diag>
 llvm::Error checkDataflowDiagnosis(
     llvm::StringRef Code,
     ast_matchers::internal::Matcher<FunctionDecl> FuncMatcher,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     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(
       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(Diags, AnalysisResults.Context);
+        Expectations(AnalysisResults.Annotations, Diags,
+                     AnalysisResults.Context);
       },
       Args, VirtualMappedFiles);
 }
 
 /// Returns the `ValueDecl` for the given identifier.
 ///
 /// Requirements:
 ///
 ///  `Name` must be unique in `ASTCtx`.
 const ValueDecl *findValueDecl(ASTContext &ASTCtx, llvm::StringRef Name);
 
 } // namespace test
 } // namespace dataflow
 } // namespace clang
 
 #endif // LLVM_CLANG_ANALYSIS_FLOW_SENSITIVE_TESTING_SUPPORT_H_

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

 //===- UncheckedOptionalAccessModelTest.cpp -------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 // FIXME: Move this to clang/unittests/Analysis/FlowSensitive/Models.
 
 #include "clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h"
 #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"
 #include "gtest/gtest.h"
 #include <string>
 #include <utility>
 #include <vector>
 
 using namespace clang;
 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
 #define CSTDDEF_H
 
 namespace std {
 
 typedef decltype(sizeof(char)) size_t;

 
 template <typename T, typename U, typename... Args>
 constexpr Optional<T> make_optional(std::initializer_list<U> il,
                                     Args&&... args);
 
 } // 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) {
   size_t Pos = 0;
   while (true) {
     Pos = S.find(Pattern, Pos);
     if (Pos == std::string::npos)
       break;
     S.replace(Pos, Pattern.size(), Replacement);
   }
 }
 
 struct OptionalTypeIdentifier {
   std::string NamespaceName;
   std::string TypeName;
 };
 
 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);
 
     std::vector<std::pair<std::string, std::string>> Headers;
     Headers.emplace_back("cstddef.h", CSDtdDefHeader);
     Headers.emplace_back("std_initializer_list.h", StdInitializerListHeader);
     Headers.emplace_back("std_string.h", StdStringHeader);
     Headers.emplace_back("std_type_traits.h", StdTypeTraitsHeader);
llvm::Error Error =
               UncheckedOptionalAccessDiagnosis Diagnosis(Ctx, Options);
               return [Diagnosis = std::move(Diagnosis)](
                          const Stmt *Stmt,
                          const UncheckedOptionalAccessModel::Lattice &,
                          const Environment &Env) mutable {
                 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);
     if (Error)
       FAIL() << llvm::toString(std::move(Error));
   }
 };
 
 INSTANTIATE_TEST_SUITE_P(
     UncheckedOptionalUseTestInst, UncheckedOptionalAccessTest,
     ::testing::Values(OptionalTypeIdentifier{"std", "optional"},
                       OptionalTypeIdentifier{"absl", "optional"},
                       OptionalTypeIdentifier{"base", "Optional"}),
     [](const ::testing::TestParamInfo<OptionalTypeIdentifier> &Info) {
       return Info.param.NamespaceName;
     });
 
 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"
 
     struct Foo {
       void foo();
     };
 
     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"
 
     struct Foo {
       void foo();
     };
 
     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 {};
 
     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 {
       explicit Foo(int, bool);
     };
 
     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 {
       explicit Foo(std::initializer_list<int>);
     };
 
     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 {
       MyString(const char*);
     };
 
     void target() {
       $ns::$optional<MyString> opt("foo");
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $ns::$optional<Bar> opt(Make<Foo>());
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {
       explicit Foo(int);
     };
 
     void target() {
       $ns::$optional<Foo> opt(3);
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                        "safe");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, ConvertibleOptionalConstructor) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     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"
 
     struct Foo {};
 
     struct Bar {
       explicit Bar(const Foo&);
     };
 
     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"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     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 {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $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 {};
 
     struct Bar {
       explicit Bar(const Foo&);
     };
 
     void target() {
       $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 {
       Foo(int, int);
     };
 
     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 {
       constexpr Foo(std::initializer_list<char>);
     };
 
     void target() {
       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"
 
     void target($ns::$optional<int> p) {
       $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"
   //
   //    void target($ns::$optional<int> opt, bool b) {
   //      if (b) {
   //        opt.emplace(0);
   //      }
   //      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"
   //
   //    void target(bool b) {
   //      $ns::$optional<int> opt = $ns::make_optional(0);
   //      if (b) {
   //        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 {};
 
     void target() {
       $ns::$optional<Foo> opt;
       opt = Foo();
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
 
     void target() {
       $ns::$optional<Foo> opt;
       (opt = Foo()).value();
       (void)0;
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
       MyString(const char*);
     };
 
     void target() {
       $ns::$optional<MyString> opt;
       opt = "foo";
       opt.value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 
-  ExpectLatticeChecksFor(R"(
+  ExpectDiagnosticsFor(R"(
     #include "unchecked_optional_access_test.h"
 
     struct MyString {
       MyString(const char*);
     };
 
     void target() {
       $ns::$optional<MyString> opt;
       (opt = "foo").value();
-      /*[[check]]*/
     }
-  )",
-                         "safe");
+  )");
 }
 
 TEST_P(UncheckedOptionalAccessTest, OptionalConversionAssignment) {
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $ns::$optional<Foo> opt1 = Foo();
       $ns::$optional<Bar> opt2;
       opt2 = opt1;
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-      "safe");
+  )");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"(
     #include "unchecked_optional_access_test.h"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $ns::$optional<Foo> opt1;
       $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"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $ns::$optional<Foo> opt1 = Foo();
       $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"
 
     void target() {
       $ns::$optional<int> opt1 = $ns::nullopt;
       $ns::$optional<int> opt2 = 3;
 
       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"
 
     void target() {
       $ns::$optional<int> opt1 = $ns::nullopt;
       $ns::$optional<int> opt2 = 3;
 
       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"
 
     void target() {
       $ns::$optional<int> opt1 = $ns::nullopt;
       $ns::$optional<int> opt2 = 3;
 
       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"
 
     void target() {
       $ns::$optional<int> opt1 = $ns::nullopt;
       $ns::$optional<int> opt2 = 3;
 
       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"
 
     template <typename T>
     struct smart_ptr {
       T& operator*() &;
       T* operator->();
     };
 
     struct Foo {
       $ns::$optional<int> opt;
     };
 
     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"
 
     $ns::$optional<int> MakeOpt();
 
     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"
 
     const $ns::$optional<int>& MakeOpt();
 
     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"
 
     using IntOpt = $ns::$optional<int>;
     IntOpt MakeOpt();
 
     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"
 
     using IntOpt = $ns::$optional<int>;
     const IntOpt& MakeOpt();
 
     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"
 
     template <typename T>
     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"
 
     using Foo = $ns::$optional<std::string>;
 
     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"
 
     using Foo = $ns::$optional<std::string>;
 
     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"
 
     struct Foo {
       $ns::$optional<std::string> opt;
     };
 
     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"
 
     using Foo = $ns::$optional<std::string>;
 
     void target($ns::$optional<Foo> foo, bool b) {
       if (!foo.has_value()) return;
       if (b) {
         if (!foo->has_value()) return;
         // We have created `foo.value()`.
         foo->value();
       } else {
         if (!foo->has_value()) return;
         // We have created `foo.value()` again, in a different environment.
         foo->value();
       }
       // 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"
 
     template <typename T>
     struct smart_ptr {
       typename std::add_lvalue_reference<T>::type operator*() &;
     };
 
     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) {
       if (b) return;
       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) {
       $ns::$optional<int> opt1 = $ns::nullopt;
       $ns::$optional<int> opt2;
       if (b) {
         opt2 = $ns::nullopt;
       } else {
         opt2 = $ns::nullopt;
       }
       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"
 
     void target(bool b) {
       $ns::$optional<int> opt;
       if (b) {
         opt = Make<$ns::$optional<int>>();
       } else {
         opt = Make<$ns::$optional<int>>();
       }
       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) {
       $ns::$optional<int> opt;
       if (b) {
         opt = Make<$ns::$optional<int>>();
         if (!opt.has_value()) return;
       } else {
         opt = Make<$ns::$optional<int>>();
         if (!opt.has_value()) return;
       }
       opt.value();
-      /*[[check-3]]*/
     }
-  )code",
-                         "safe");
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
       $ns::$optional<int> opt;
       if (b) {
         opt = Make<$ns::$optional<int>>();
         if (!opt.has_value()) return;
       } 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"
 
     void target(bool b) {
       $ns::$optional<int> opt;
       if (b) {
         opt = 1;
       } else {
         opt = 2;
       }
       opt.value();
-      /*[[check-5]]*/
     }
-  )code",
-      "safe");
+  )code");
 
-  ExpectLatticeChecksFor(
+  ExpectDiagnosticsFor(
       R"code(
     #include "unchecked_optional_access_test.h"
 
     void target(bool b) {
       $ns::$optional<int> opt;
       if (b) {
         opt = 1;
       } 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:
 // - constructors (copy, move)
 // - assignment operators (default, copy, move)
 // - invalidation (passing optional by non-const reference/pointer)