[clang][dataflow] Add API to separate analysis from diagnosis

[samestep]

Summary

This patch adds an optional PostVisitStmt parameter to the runTypeErasedDataflowAnalysis function, which does one more pass over all statements in the CFG after a fixpoint is reached. It then defines a diagnose method for the optional model in a new UncheckedOptionalAccessDiagnosis class, but only integrates that into the tests and not the actual optional check for clang-tidy. That will be done in a followup patch.

The primary motivation is to separate the implementation of the unchecked optional access check into two parts, to allow for further refactoring of just the model part later, while leaving the checking part alone. Currently there is duplication between the transferUnwrapCall and diagnoseUnwrapCall functions, but that will be dealt with in the followup.

Because diagnostics are now all gathered into one collection rather than being populated at each program point like when computing a fixpoint, this patch removes the usage of Pair and UnorderedElementsAre from the optional model tests, and instead modifies all their expectations to simply check the stringified set of diagnostics against a single string, either "safe" or some concatenation of "unsafe: input.cc:y:x". This is not ideal as it loses any connection to the /*[[check]]*/ annotations in the source strings, but it does still retain the source locations from the diagnostic strings themselves.

[ymandel]

clang/include/clang/Analysis/FlowSensitive/DataflowAnalysis.h

35 ↗

(On Diff #437299)

I'm not sure that this is valid style. It's essentially a unit type, which I quite like, but I've never seen this done before.

64 ↗

(On Diff #437299)

Please describe this parameter.

103 ↗

(On Diff #437299)

The virtual function seems out of place here, given that everything else is statically resolved. Either drop virtual or move to the dyn. typed super class.

That said, I'm not clear on why this isn't a free function. I imagine a model of creating the analysis, running it and that results in blocks annotated with Environments and custome lattices. Then, you run any function you want over that using some visitor template helper. To account for state, we either make it an explicit parameter or take a std::function.

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

52 ↗

(On Diff #437299)

this implies that we want to support some form of accumulation. I'd argue that any accumulation should already occur in the lattice (or the environment) and this should be a straight read of the state and some output. So, the TransferState struct should be enough.

[tschuett]

C++ seems to be tied to templates. I wonder if super classes resp. inheritance would make your life easier. You can explicitly define the API and documentation. You could provide default implementations.

[xazax.hun]

clang-tools-extra/clang-tidy/bugprone/UncheckedOptionalAccessCheck.cpp

66 ↗

(On Diff #437299)

While doing yet another iteration after we reached the fixed point is a valid approach, many analyses have a monotonic property when it comes to emitting diagnostics. In this case, when the analysis discovered that an optional access is unsafe in one of the iterations I would not expect it to become safe in subsequent iterations. In most of the cases this makes collecting diagnostics eagerly a viable option and saves us from doing one more traversal of the CFG. On the other hand, we need to be careful to not to emit duplicate diagnostics.

I wonder whether, from a user point of view, not having to do one more iteration is a more ergonomic interface.

[samestep]

clang-tools-extra/clang-tidy/bugprone/UncheckedOptionalAccessCheck.cpp

66 ↗

(On Diff #437299)

@xazax.hun Agreed, I don't think for this particular analysis there's a risk of diagnostics being prematurely added and then becoming invalid as the analysis progresses. However, the broader goal of this patch is to separate the information-gathering part of the analysis to the diagnostic-flagging part of the analysis, with the hope that we can maybe derive the former automatically while still allowing the latter to be written manually. It seems to me that it would be infeasible to automatically derive both.

[samestep]

Only add the new API, don't change the check

[sgatev]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

26 ↗

(On Diff #438891)

Let's add some documentation for Diagnosis and diagnoseCFG.

27 ↗

(On Diff #438891)

#include <functional>

32 ↗

(On Diff #438891)

#include <vector> and #include "llvm/ADT/Optional.h"

49 ↗

(On Diff #438891)

#include <utility>

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

Move this to Diagnosis.h?

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

Move this to a new UncheckedOptionalAccessDiagnosis.(h,cpp)?

[samestep]

Comment Diagnosis.h and add missing headers

[samestep]

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

I'd prefer to keep it in MatchSwitch.h alongside TransferState, unless we plan to, e.g., move that type to DataflowAnalysis.h.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

OK! I'll do that next.

[samestep]

Try to fix the broken patch

[samestep]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

Actually, @sgatev would that go under FlowSensitive/Models/ or just under FlowSensitive/?

[sgatev]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

I suggest keeping it under FlowSensitive/Models/ for now. We can change that at a later point if there's a better alternative.

[samestep]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

Where should I put the UncheckedOptionalAccessModelOptions type and ignorableOptional function that need to be shared between both the model and the diagnosis?

[samestep]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

Same question for the following other helper functions:

• hasOptionalType
• isOptionalMemberCallWithName
• isOptionalOperatorCallWithName

Given how much code is shared between the two, I'm really not sure whether it's the best idea to move this into a separate file in this patch...

[xazax.hun]

clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp

546

I think this should be moved closer to its use.

684–685

This logic is duplicated between the check and the diagnosis. I think it would be nice to factor this out to avoid these two parts getting out of sync.

[samestep]

• Merge branch 'main' into diagnose-api
• Incorporate Gábor's suggestions

[sgatev]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

57 ↗

(On Diff #439372)

Better to remove this and rely on NRVO? https://en.cppreference.com/w/cpp/language/copy_elision

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

Fair enough, but generally I see TransferState and DiagnoseState as very specific to the dataflow analysis framework whereas MatchSwitchBuilder seems to be a bit more generic so I'd consider separating them.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

I suggest either declaring those in the header and using them from UncheckedOptionalAccessDiagnosis.cpp or moving them to a separate file and including them both in the model and in the diagnosis. In general I think it's fair if the diagnosis depends on the model, but not the other way around. It's perfectly fine if we do this refactoring in a follow up.

[samestep]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

57 ↗

(On Diff #439372)

Ah OK, I'm not a C++ expert so I wasn't sure when to write std::move and when not to. Will do, thanks!

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

Sure, if people think it makes sense to put those two types elsewhere then we can do that in a followup.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

76

Yep, makes sense; I just didn't want to make declarations for all those in a header file in this patch.

[samestep]

• Remove unnecessary std::move

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

38 ↗

(On Diff #439372)

This function seems pretty general and not necessarily tied to diagnostics.

WDYT about using "visitor" in the name?

For example, "visitDataflowAnalysisResults".

The "Diagnosis" would also become "DataflowAnalysisResultVisitor".

42 ↗

(On Diff #439372)

I'm leaning towards std::vector<Diag> instead of a set, to avoid requiring that Diag is hashable. Users who need hash-based deduplication can easily do it themselves (but why would they have duplication? we only visit each Stmt once)

57 ↗

(On Diff #439372)

https://stackoverflow.com/questions/14856344/when-should-stdmove-be-used-on-a-function-return-value

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

For now, since there is only a single user, I'd actually prefer to move it into UncheckedOptionalAccessDiagnosis itself. It is a trivial component, basically a std::pair, it seems that exposing it as a public API might not be that helpful.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

78

"Diagnosis" sounds like the result. Should this be a "Diagnoser"?

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

78–79

It is better to name a function with a verb. WDYT about "VerifyResults"?

If you agree, please also update the functions below.

147

189

Can we use diagnoseCFG to implement the loop above?

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

1841

The original was unsafe, is it an intentional change?

[xazax.hun]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

38–42 ↗

(On Diff #439411)

I feel a bit uneasy about this API. The need of helpers in the tests is a signal that this API might not be ergonomic as is. One needs to pass a lot of stuff in, and it is really error prone. For example, the user might get confused what Environment to pass in. Or one might pass the wrong Analysis if there are multiple similar ones in scope. I wonder if there is a way to reduce the ceremony needed to get diagnostics out. How about adding an optional argument to runDataflowAnalysis function to collect the results? I wonder if that would be a bit cleaner. You do not need to change the return type of that function, the collector object could take a reference to the llvm::DenseSet<Diag> that it populates.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

85

Do we expect this to modify the ASTContext? Could this be a const ref instead?

[xazax.hun]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

38–42 ↗

(On Diff #439411)

Moreover, if we pass the diagnostic function to runDataflowAnalysis, the fact whether we collect diagnostic during the fixed-point iteration or after that in a separate pass becomes an implementation detail. We can easily change back and forth without affecting the callers. The current API mandates an implementation approach and makes it harder to change in the future.

[samestep]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

38–42 ↗

(On Diff #439411)

I think this makes a lot of sense. I agree that the current set of parameters is unwieldy. I hadn't thought about your point of hiding the separate pass as an implementation detail, but that also makes sense to me. @gribozavr2 @ymandel @sgatev thoughts?

38 ↗

(On Diff #439372)

Sure, that sounds fine to me. What do you think about @xazax.hun's suggested alternate API? If we go with that, would we rename it to say "visitor" regardless?

42 ↗

(On Diff #439372)

This makes a lot of sense, I'll make that change. Somehow along the way I forgot that the whole reason we even needed a set was because we were using it as a lattice when computing a fixpoint, but you're right that since we're pulling this out of the fixpoint iteration, we a vector would do fine.

57 ↗

(On Diff #439372)

Thanks!

clang/include/clang/Analysis/FlowSensitive/MatchSwitch.h

49 ↗

(On Diff #438891)

That's also fine; I can make that change here.

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

78

Sure, I can change it to say "Diagnoser" instead (unless we want to just replace it with "Visitor").

85

I'd guess it shouldn't; I was mostly just copying the ASTContext & field from the DataflowAnalysis class, which is not const. I can try changing it to const here, though.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

78–79

Sounds good to me, I can do that.

189

Possibly! I'll look into it.

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

1841

Oh! Good catch, I didn't notice that. Will fix.

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/Diagnosis.h

38–42 ↗

(On Diff #439411)

I think that's a good idea, it will make a simpler API for typical usage scenarios. I still think there might be space for an API like diagnoseCFG (e.g., if you want to run multiple visitations of the analysis results), but we don't have a use case for it so far.

[samestep]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

85

@xazax.hun Coming back to this today, I remembered why it has to be ASTContext & and not const ASTContext &: because MatchSwitch requires it to not be const.

[xazax.hun]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

85

Thanks! Do you mean MatchSwitch cannot be refactored to use const ASTContext&? I quickly skimmed through the code and it looks like almost every instance of ASTContext is non-const. So I think it is ok to leave as is for this PR. I am also not sure if making it const at more places would bring any significant value.

[samestep]

• Merge branch 'main' into diagnose-api
• WIP

[samestep]

• Merge branch 'main' into diagnose-api

[samestep]

• Bring back Diagnosis.h

[samestep]

• Implement Gábor's suggestion
• Merge branch 'main' into diagnose-api

[ymandel]

clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h

128

Please update comment to mention this new param (and that its optional)

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

65–73

nit: since this includes both inputs and results, maybe AnalysisData?

114

nit: i prefer explicit comparisons, so MakePostVisitStmt != nullptr.

161

172–191

any chance that this could be folded into MakePostVisitStmt?

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

1249

[xazax.hun]

Thanks, LGTM!

[gribozavr2]

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

78

Not done?

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

144–145

155

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

2335–2336

[sgatev]

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

401

Let's harmonize the arguments to transferBlock and runTypeErasedDataflowAnalysis. I think in both cases we only need access to the internal Stmt so there's no need to pass CFGStmt to transferBlock. This way we can pass PostVisitStmt directly here without having to wrap it in a lambda.

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

78–82

Why is this a function that returns a function? Couldn't it be simply a void function, like VerifyResults?

[samestep]

Fixed all the minor nits and responded to some comments.

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

401

Makes sense to me; should I do that in a followup?

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

78–82

It needs to take in the ASTContext, and it needs to be able to close over state; see the one we pass from UncheckedOptionalAccessModelTest.cpp.

172–191

Looking into this now.

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

2335–2336

I don't think this is correct, is it?

[samestep]

• Merge branch 'main' into diagnose-api
• Fix minor nits

[samestep]

• Tidy AnalysisData struct and PostVisitStmt param

clang/include/clang/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.h

 #define CLANG_ANALYSIS_FLOWSENSITIVE_MODELS_UNCHECKEDOPTIONALACCESSMODEL_H
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Stmt.h"
 #include "clang/Analysis/FlowSensitive/DataflowAnalysis.h"
 #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 {
 
 // FIXME: Explore using an allowlist-approach, where constructs supported by the
 // analysis are always enabled and additional constructs are enabled through the
 // `Options`.
 struct UncheckedOptionalAccessModelOptions {
public:
   bool merge(QualType Type, const Value &Val1, const Environment &Env1,
              const Value &Val2, const Environment &Env2, Value &MergedVal,
              Environment &MergedEnv) override;
 
 private:
   MatchSwitch<TransferState<SourceLocationsLattice>> TransferMatchSwitch;
 };
 
+class UncheckedOptionalAccessDiagnoser {
+public:
+  UncheckedOptionalAccessDiagnoser(
+      UncheckedOptionalAccessModelOptions Options = {});
+
+  std::vector<SourceLocation> diagnose(ASTContext &Context, const Stmt *Stmt,
+                                       const Environment &Env);
+
+private:
+  MatchSwitch<const Environment, std::vector<SourceLocation>>
+      DiagnoseMatchSwitch;
+};
+
 } // namespace dataflow
 } // namespace clang
 
 #endif // CLANG_ANALYSIS_FLOWSENSITIVE_MODELS_UNCHECKEDOPTIONALACCESSMODEL_H

clang/include/clang/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h

TypeErasedDataflowAnalysisState transferBlock(
     std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates,
     const CFGBlock &Block, const Environment &InitEnv,
     TypeErasedDataflowAnalysis &Analysis,
     std::function<void(const CFGStmt &,
                        const TypeErasedDataflowAnalysisState &)>
         HandleTransferredStmt = nullptr);
 
 /// Performs dataflow analysis and returns a mapping from basic block IDs to
-/// dataflow analysis states that model the respective basic blocks. Indices
-/// of the returned vector correspond to basic block IDs. Returns an error if
-/// the dataflow analysis cannot be performed successfully.
+/// dataflow analysis states that model the respective basic blocks. Indices of
+/// the returned vector correspond to basic block IDs. Returns an error if the
+/// dataflow analysis cannot be performed successfully. Otherwise, calls
+/// `PostVisitStmt` on each statement with the final analysis results at that
+/// program point.
 llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
-runTypeErasedDataflowAnalysis(const ControlFlowContext &CFCtx,
-                              TypeErasedDataflowAnalysis &Analysis,
-                              const Environment &InitEnv);
+runTypeErasedDataflowAnalysis(
+    const ControlFlowContext &CFCtx, TypeErasedDataflowAnalysis &Analysis,
+    const Environment &InitEnv,
+    std::function<void(const Stmt *, const TypeErasedDataflowAnalysisState &)>
+        PostVisitStmt = nullptr);
 
 } // namespace dataflow
 } // namespace clang
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_TYPEERASEDDATAFLOWANALYSIS_H

clang/lib/Analysis/FlowSensitive/Models/UncheckedOptionalAccessModel.cpp

 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/Stmt.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #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 {
 namespace {
 
 using namespace ::clang::ast_matchers;
 using LatticeTransferState = TransferState<SourceLocationsLattice>;
 
void transferStdSwapCall(const CallExpr *E, const MatchFinder::MatchResult &,
 
   auto *OptionalLoc2 =
       State.Env.getStorageLocation(*E->getArg(1), SkipPast::Reference);
   assert(OptionalLoc2 != nullptr);
 
   transferSwap(*OptionalLoc1, *OptionalLoc2, State);
 }
 
 llvm::Optional<StatementMatcher>
 ignorableOptional(const UncheckedOptionalAccessModelOptions &Options) {
   if (Options.IgnoreSmartPointerDereference)
     return memberExpr(hasObjectExpression(ignoringParenImpCasts(
         cxxOperatorCallExpr(anyOf(hasOverloadedOperatorName("->"),
                                   hasOverloadedOperatorName("*")),
                             unless(hasArgument(0, expr(hasOptionalType())))))));
   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
   // lot of duplicated work (e.g. string comparisons), consider providing APIs
   // that avoid it through memoization.
   auto IgnorableOptional = ignorableOptional(Options);
   return MatchSwitchBuilder<LatticeTransferState>()
       // Attach a symbolic "has_value" state to optional values that we see for
return MatchSwitchBuilder<LatticeTransferState>()
       // optional::operator=
       .CaseOf<CXXOperatorCallExpr>(isOptionalValueOrConversionAssignment(),
                                    transferValueOrConversionAssignment)
       .CaseOf<CXXOperatorCallExpr>(isOptionalNulloptAssignment(),
                                    transferNulloptAssignment)
 
       // 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))),
+      .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"),
                                  transferOptionalHasValueCall)
 
       // optional::operator bool
       .CaseOf<CXXMemberCallExpr>(isOptionalMemberCallWithName("operator bool"),
                                  transferOptionalHasValueCall)
return MatchSwitchBuilder<LatticeTransferState>()
 
       // returns optional
       .CaseOf<CallExpr>(isCallReturningOptional(),
                         transferCallReturningOptional)
 
       .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
 UncheckedOptionalAccessModel::optionalClassDecl() {
   return optionalClass();
 }
 
 UncheckedOptionalAccessModel::UncheckedOptionalAccessModel(
bool UncheckedOptionalAccessModel::merge(QualType Type, const Value &Val1,
   if (isNonEmptyOptional(Val1, Env1) && isNonEmptyOptional(Val2, Env2))
     MergedEnv.addToFlowCondition(HasValueVal);
   else if (isEmptyOptional(Val1, Env1) && isEmptyOptional(Val2, Env2))
     MergedEnv.addToFlowCondition(MergedEnv.makeNot(HasValueVal));
   setHasValue(MergedVal, HasValueVal);
   return true;
 }
 
+UncheckedOptionalAccessDiagnoser::UncheckedOptionalAccessDiagnoser(
+    UncheckedOptionalAccessModelOptions Options)
+    : DiagnoseMatchSwitch(buildDiagnoseMatchSwitch(Options)) {}
+
+std::vector<SourceLocation> UncheckedOptionalAccessDiagnoser::diagnose(
+    ASTContext &Context, const Stmt *Stmt, const Environment &Env) {
+  return DiagnoseMatchSwitch(*Stmt, Context, Env);
+}
+
 } // namespace dataflow
 } // namespace clang

clang/lib/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.cpp

default:
       // FIXME: Evaluate other kinds of `CFGElement`.
       break;
     }
   }
   return State;
 }
 
 llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
-runTypeErasedDataflowAnalysis(const ControlFlowContext &CFCtx,
-                              TypeErasedDataflowAnalysis &Analysis,
-                              const Environment &InitEnv) {
+runTypeErasedDataflowAnalysis(
+    const ControlFlowContext &CFCtx, TypeErasedDataflowAnalysis &Analysis,
+    const Environment &InitEnv,
+    std::function<void(const Stmt *, const TypeErasedDataflowAnalysisState &)>
+        PostVisitStmt) {
   PostOrderCFGView POV(&CFCtx.getCFG());
   ForwardDataflowWorklist Worklist(CFCtx.getCFG(), &POV);
 
   std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> BlockStates;
   BlockStates.resize(CFCtx.getCFG().size(), llvm::None);
 
   // The entry basic block doesn't contain statements so it can be skipped.
   const CFGBlock &Entry = CFCtx.getCFG().getEntry();
while (const CFGBlock *Block = Worklist.dequeue()) {
     if (Block->hasNoReturnElement())
       continue;
 
     Worklist.enqueueSuccessors(Block);
   }
   // FIXME: Consider evaluating unreachable basic blocks (those that have a
   // state set to `llvm::None` at this point) to also analyze dead code.
 
+  if (PostVisitStmt) {
+    for (const CFGBlock *Block : CFCtx.getCFG()) {
+      // Skip blocks that were not evaluated.
+      if (!BlockStates[Block->getBlockID()])
+        continue;
+      transferBlock(
+          CFCtx, BlockStates, *Block, InitEnv, Analysis,
+          [&PostVisitStmt](const clang::CFGStmt &Stmt,
+                           const TypeErasedDataflowAnalysisState &State) {
+            PostVisitStmt(Stmt.getStmt(), State);
+          });
+    }
+  }
+
   return BlockStates;
 }
 
 } // namespace dataflow
 } // namespace clang

clang/unittests/Analysis/FlowSensitive/TestingSupport.h

 namespace test {
 
 // Returns assertions based on annotations that are present after statements in
 // `AnnotatedCode`.
 llvm::Expected<llvm::DenseMap<const Stmt *, std::string>>
 buildStatementToAnnotationMapping(const FunctionDecl *Func,
                                   llvm::Annotations AnnotatedCode);
 
-// Runs dataflow on the body of the function that matches `func_matcher` in code
-// snippet `code`. Requires: `Analysis` contains a type `Lattice`.
+struct AnalysisData {
+  ASTContext &ASTCtx;
+  const ControlFlowContext &CFCtx;
+  const Environment &Env;
+  TypeErasedDataflowAnalysis &Analysis;
+  llvm::DenseMap<const clang::Stmt *, std::string> &Annotations;
+  std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>> &BlockStates;
+};
+
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code,
-    ast_matchers::internal::Matcher<FunctionDecl> FuncMatcher,
+    ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
     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,
+    std::function<void(ASTContext &, const Stmt *,
+                       const TypeErasedDataflowAnalysisState &)>
+        PostVisitStmt,
+    std::function<void(AnalysisData)> VerifyResults, ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
-  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
-
   llvm::Annotations AnnotatedCode(Code);
   auto Unit = tooling::buildASTFromCodeWithArgs(
       AnnotatedCode.code(), Args, "input.cc", "clang-dataflow-test",
       std::make_shared<PCHContainerOperations>(),
       tooling::getClangStripDependencyFileAdjuster(), VirtualMappedFiles);
   auto &Context = Unit->getASTContext();
 
   if (Context.getDiagnostics().getClient()->getNumErrors() != 0) {
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Source file has syntax or type errors, "
                                       "they were printed to the test log");
   }
 
   const FunctionDecl *F = ast_matchers::selectFirst<FunctionDecl>(
       "target",
-      ast_matchers::match(
-          ast_matchers::functionDecl(ast_matchers::isDefinition(), FuncMatcher)
+      ast_matchers::match(ast_matchers::functionDecl(
+                              ast_matchers::isDefinition(), TargetFuncMatcher)
                               .bind("target"),
                           Context));
   if (F == nullptr)
     return llvm::make_error<llvm::StringError>(
         llvm::errc::invalid_argument, "Could not find target function.");
 
   auto CFCtx = ControlFlowContext::build(F, F->getBody(), &F->getASTContext());
   if (!CFCtx)
     return CFCtx.takeError();
 
   DataflowAnalysisContext DACtx(std::make_unique<WatchedLiteralsSolver>());
   Environment Env(DACtx, *F);
   auto Analysis = MakeAnalysis(Context, Env);
 
+  std::function<void(const Stmt *, const TypeErasedDataflowAnalysisState &)>
+      PostVisitStmtClosure = nullptr;
+  if (PostVisitStmt != nullptr) {
+    PostVisitStmtClosure = [&PostVisitStmt, &Context](
+                               const Stmt *Stmt,
+                               const TypeErasedDataflowAnalysisState &State) {
+      PostVisitStmt(Context, Stmt, State);
+    };
+  }
+
   llvm::Expected<llvm::DenseMap<const clang::Stmt *, std::string>>
       StmtToAnnotations = buildStatementToAnnotationMapping(F, AnnotatedCode);
   if (!StmtToAnnotations)
     return StmtToAnnotations.takeError();
   auto &Annotations = *StmtToAnnotations;
 
   llvm::Expected<std::vector<llvm::Optional<TypeErasedDataflowAnalysisState>>>
-      MaybeBlockStates = runTypeErasedDataflowAnalysis(*CFCtx, Analysis, Env);
+      MaybeBlockStates = runTypeErasedDataflowAnalysis(*CFCtx, Analysis, Env,
+                                                       PostVisitStmtClosure);
   if (!MaybeBlockStates)
     return MaybeBlockStates.takeError();
   auto &BlockStates = *MaybeBlockStates;
 
-  if (BlockStates.empty()) {
-    Expectations({}, Context);
-    return llvm::Error::success();
-  }
+  AnalysisData AnalysisData{Context,  *CFCtx,      Env,
+                            Analysis, Annotations, BlockStates};
+  VerifyResults(AnalysisData);
 
+// Runs dataflow on the body of the function that matches `TargetFuncMatcher` in
+// code snippet `Code`. Requires: `AnalysisT` contains a type `Lattice`.
+template <typename AnalysisT>
+llvm::Error checkDataflow(
+    llvm::StringRef Code,
+    ast_matchers::internal::Matcher<FunctionDecl> TargetFuncMatcher,
+    std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
+    std::function<void(
+        llvm::ArrayRef<std::pair<
+            std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
+        ASTContext &)>
+        VerifyResults,
+    ArrayRef<std::string> Args,
+    const tooling::FileContentMappings &VirtualMappedFiles = {}) {
+  using StateT = DataflowAnalysisState<typename AnalysisT::Lattice>;
+
+  return checkDataflow(
+      Code, std::move(TargetFuncMatcher), std::move(MakeAnalysis),
+      /*PostVisitStmt=*/nullptr,
+      [&VerifyResults](AnalysisData AnalysisData) {
+        if (AnalysisData.BlockStates.empty()) {
+          VerifyResults({}, AnalysisData.ASTCtx);
+          return;
+        }
+
+        auto &Annotations = AnalysisData.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 : CFCtx->getCFG()) {
+        for (const CFGBlock *Block : AnalysisData.CFCtx.getCFG()) {
           // Skip blocks that were not evaluated.
-    if (!BlockStates[Block->getBlockID()])
+          if (!AnalysisData.BlockStates[Block->getBlockID()])
             continue;
 
           transferBlock(
-        *CFCtx, BlockStates, *Block, Env, Analysis,
-        [&Results, &Annotations](const clang::CFGStmt &Stmt,
-                                 const TypeErasedDataflowAnalysisState &State) {
+              AnalysisData.CFCtx, AnalysisData.BlockStates, *Block,
+              AnalysisData.Env, AnalysisData.Analysis,
+              [&Results,
+               &Annotations](const clang::CFGStmt &Stmt,
                 auto It = Annotations.find(Stmt.getStmt());
                 if (It == Annotations.end())
                   return;
-          auto *Lattice =
-              llvm::any_cast<typename AnalysisT::Lattice>(&State.Lattice.Value);
+                auto *Lattice = llvm::any_cast<typename AnalysisT::Lattice>(
+                    &State.Lattice.Value);
                 Results.emplace_back(It->second, StateT{*Lattice, State.Env});
               });
         }
-  Expectations(Results, Context);
-  return llvm::Error::success();
+        VerifyResults(Results, AnalysisData.ASTCtx);
+      },
+      Args, VirtualMappedFiles);
 }
 
 // Runs dataflow on the body of the function named `target_fun` in code snippet
 // `code`.
 template <typename AnalysisT>
 llvm::Error checkDataflow(
     llvm::StringRef Code, llvm::StringRef TargetFun,
     std::function<AnalysisT(ASTContext &, Environment &)> MakeAnalysis,
     std::function<void(
         llvm::ArrayRef<std::pair<
             std::string, DataflowAnalysisState<typename AnalysisT::Lattice>>>,
         ASTContext &)>
-        Expectations,
+        VerifyResults,
     ArrayRef<std::string> Args,
     const tooling::FileContentMappings &VirtualMappedFiles = {}) {
   return checkDataflow(Code, ast_matchers::hasName(TargetFun),
-                       std::move(MakeAnalysis), std::move(Expectations), Args,
+                       std::move(MakeAnalysis), std::move(VerifyResults), 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/Analysis/FlowSensitive/TypeErasedDataflowAnalysis.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.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::Pair;
-using ::testing::UnorderedElementsAre;
+using ::testing::ContainerEq;
 
 // FIXME: Move header definitions in separate file(s).
 static constexpr char CSDtdDefHeader[] = R"(
 #ifndef CSTDDEF_H
 #define CSTDDEF_H
 
 namespace std {
 

 
 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 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) {
   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);
Headers.emplace_back("unchecked_optional_access_test.h", R"(
       #include "std_string.h"
       #include "std_utility.h"
 
       template <typename T>
       T Make();
     )");
     const tooling::FileContentMappings FileContents(Headers.begin(),
                                                     Headers.end());
+    UncheckedOptionalAccessModelOptions Options{
+        /*IgnoreSmartPointerDereference=*/true};
+    std::vector<SourceLocation> Diagnostics;
     llvm::Error Error = checkDataflow<UncheckedOptionalAccessModel>(
         SourceCode, FuncMatcher,
-        [](ASTContext &Ctx, Environment &) {
-          return UncheckedOptionalAccessModel(
-              Ctx, UncheckedOptionalAccessModelOptions{
-                       /*IgnoreSmartPointerDereference=*/true});
+        [Options](ASTContext &Ctx, Environment &) {
+          return UncheckedOptionalAccessModel(Ctx, Options);
         },
-        [&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));
+        [&Diagnostics, Diagnoser = UncheckedOptionalAccessDiagnoser(Options)](
+            ASTContext &Ctx, const Stmt *Stmt,
+            const TypeErasedDataflowAnalysisState &State) mutable {
+          auto StmtDiagnostics = Diagnoser.diagnose(Ctx, Stmt, State.Env);
+          llvm::move(StmtDiagnostics, std::back_inserter(Diagnostics));
+        },
+        [&Diagnostics](AnalysisData AnalysisData) {
+          auto &SrcMgr = AnalysisData.ASTCtx.getSourceManager();
+
+          llvm::DenseSet<unsigned> AnnotationLines;
+          for (const auto &Pair : AnalysisData.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 : Diagnostics) {
+            DiagnosticLines.insert(SrcMgr.getPresumedLineNumber(Loc));
           }
-          EXPECT_THAT(CheckToStringifiedLatticeMap, MatchesLatticeChecks);
+
+          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]]*/
     }
-  )",
-                         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"
 
     struct Foo {
       void foo();
     };
 
     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"
 
     struct Foo {
       void foo();
     };
 
     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 {};
 
     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 {
       explicit Foo(int, bool);
     };
 
     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 {
       explicit Foo(std::initializer_list<int>);
     };
 
     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 {
       MyString(const char*);
     };
 
     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 {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     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 {
       explicit Foo(int);
     };
 
     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"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     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"
 
     struct Foo {};
 
     struct Bar {
       explicit Bar(const Foo&);
     };
 
     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"
 
     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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check", "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]]*/
     }
-  )",
-                         UnorderedElementsAre(Pair("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]]*/
     }
-  )",
-                         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 {
       Foo(int, int);
     };
 
     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 {
       constexpr Foo(std::initializer_list<char>);
     };
 
     void target() {
       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"
 
     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",
-      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"
   //
   //    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]]
   //      }
   //    }
-  //  )",
-  //      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"
   //
   //    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]]*/
   //      }
   //    }
-  //  )",
-  //      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 {};
 
     void target() {
       $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 {};
 
     void target() {
       $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 {
       MyString(const char*);
     };
 
     void target() {
       $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 {
       MyString(const char*);
     };
 
     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"
 
     struct Foo {};
 
     struct Bar {
       Bar(const Foo&);
     };
 
     void target() {
       $ns::$optional<Foo> opt1 = Foo();
       $ns::$optional<Bar> opt2;
       opt2 = opt1;
       opt2.value();
-      /*[[check]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("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]]
       }
     }
-  )",
-      UnorderedElementsAre(Pair("check", "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]]*/
     }
-  )",
-      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"
 
     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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"),
-                           Pair("check-2", "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]]
     }
-  )",
-      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"
 
     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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"),
-                           Pair("check-2", "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]]
     }
-  )",
-      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"
 
     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]]*/
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "safe"), Pair("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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-1", "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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-2", "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]]
     }
-  )",
-      UnorderedElementsAre(Pair("check-3", "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]]
     }
-  )",
-      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"
 
     template <typename T>
     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"
 
     using Foo = $ns::$optional<std::string>;
 
     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"
 
     using Foo = $ns::$optional<std::string>;
 
     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"
 
     struct Foo {
       $ns::$optional<std::string> opt;
     };
 
     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"
 
     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]]
     }
-  )",
-      UnorderedElementsAre(Pair("merge", "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]]
     }
-  )",
-      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) {
       if (b) return;
       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) {
       $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]]*/
       }
     }
-  )",
-                         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"
 
     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",
-      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) {
       $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",
-                         UnorderedElementsAre(Pair("check-3", "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",
-      UnorderedElementsAre(Pair("check-4", "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",
-      UnorderedElementsAre(Pair("check-5", "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",
-      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:
 // - constructors (copy, move)
 // - assignment operators (default, copy, move)
 // - invalidation (passing optional by non-const reference/pointer)