[clang][dataflow] Add test for crash repro and clean up const accessor handling

Add test for llvm#125589

The crash is actually incidentally fixed by
llvm#128437 since it added a branch
for the reference case and would no longer fall through when the return
type is a reference to a pointer.

Clean up a bit as well:
- make the fallback for early returns more consistent (check if
  returning optional and call transfer function for that case)
- check RecordLoc == nullptr in one place

Add some init for the reference to pointer/bool cases.

Author: jvoung

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

@@ -551,15 +551,17 @@ void transferCallReturningOptional(const CallExpr *E,
   setHasValue(*Loc, State.Env.makeAtomicBoolValue(), State.Env);
 }
 
-void handleConstMemberCall(const CallExpr *CE,
+bool handleConstMemberCall(const CallExpr *CE,
                            dataflow::RecordStorageLocation *RecordLoc,
                            const MatchFinder::MatchResult &Result,
                            LatticeTransferState &State) {
+  if (RecordLoc == nullptr) return false;
+
   // If the const method returns an optional or reference to an optional.
-  if (RecordLoc != nullptr && isSupportedOptionalType(CE->getType())) {
+  if (isSupportedOptionalType(CE->getType())) {
     const FunctionDecl *DirectCallee = CE->getDirectCallee();
     if (DirectCallee == nullptr)
-      return;
+      return false;
     StorageLocation &Loc =
         State.Lattice.getOrCreateConstMethodReturnStorageLocation(
             *RecordLoc, DirectCallee, State.Env, [&](StorageLocation &Loc) {
@@ -577,57 +579,65 @@ void handleConstMemberCall(const CallExpr *CE,
       auto &ResultLoc = State.Env.getResultObjectLocation(*CE);
       copyRecord(cast<RecordStorageLocation>(Loc), ResultLoc, State.Env);
     }
-    return;
+    return true;
   }
 
+  bool IsBooleanOrPointer = CE->getType()->isBooleanType() ||
+                            CE->getType()->isPointerType();
+
   // Cache if the const method returns a reference
-  if (RecordLoc != nullptr && CE->isGLValue()) {
+  if (CE->isGLValue()) {
     const FunctionDecl *DirectCallee = CE->getDirectCallee();
     if (DirectCallee == nullptr)
-      return;
+      return false;
 
     StorageLocation &Loc =
         State.Lattice.getOrCreateConstMethodReturnStorageLocation(
             *RecordLoc, DirectCallee, State.Env, [&](StorageLocation &Loc) {
-              // no-op
+              if (IsBooleanOrPointer) {
+                State.Env.setValue(Loc, *State.Env.createValue(CE->getType()));
+              }
             });
 
     State.Env.setStorageLocation(*CE, Loc);
-    return;
-  }
-
-  // Cache if the const method returns a boolean or pointer type.
-  // We may decide to cache other return types in the future.
-  if (RecordLoc != nullptr &&
-      (CE->getType()->isBooleanType() || CE->getType()->isPointerType())) {
+    return true;
+  } else if (IsBooleanOrPointer) {
+    // Cache if the const method returns a boolean or pointer type.
     Value *Val = State.Lattice.getOrCreateConstMethodReturnValue(*RecordLoc, CE,
                                                                  State.Env);
     if (Val == nullptr)
-      return;
+      return false;
     State.Env.setValue(*CE, *Val);
-    return;
+    return true;
   }
 
-  // Perform default handling if the call returns an optional
-  // but wasn't handled above (if RecordLoc is nullptr).
-  if (isSupportedOptionalType(CE->getType())) {
-    transferCallReturningOptional(CE, Result, State);
-  }
+  return false;
 }
 
-void transferValue_ConstMemberCall(const CXXMemberCallExpr *MCE,
-                                   const MatchFinder::MatchResult &Result,
-                                   LatticeTransferState &State) {
-  handleConstMemberCall(
-      MCE, dataflow::getImplicitObjectLocation(*MCE, State.Env), Result, State);
+void transferConstMemberCall(const CXXMemberCallExpr *MCE,
+                             const MatchFinder::MatchResult &Result,
+                             LatticeTransferState &State) {
+  if (!handleConstMemberCall(
+          MCE, dataflow::getImplicitObjectLocation(*MCE, State.Env), Result,
+          State)) {
+    // Perform default handling if the call returns an optional,
+    // but wasn't handled.
+    if (isSupportedOptionalType(MCE->getType()))
+      transferCallReturningOptional(MCE, Result, State);
+  }
 }
 
-void transferValue_ConstMemberOperatorCall(
-    const CXXOperatorCallExpr *OCE, const MatchFinder::MatchResult &Result,
-    LatticeTransferState &State) {
+void transferConstMemberOperatorCall(const CXXOperatorCallExpr *OCE,
+                                     const MatchFinder::MatchResult &Result,
+                                     LatticeTransferState &State) {
   auto *RecordLoc = cast_or_null<dataflow::RecordStorageLocation>(
       State.Env.getStorageLocation(*OCE->getArg(0)));
-  handleConstMemberCall(OCE, RecordLoc, Result, State);
+  if (!handleConstMemberCall(OCE, RecordLoc, Result, State)) {
+    // Perform default handling if the call returns an optional,
+    // but wasn't handled.
+    if (isSupportedOptionalType(OCE->getType()))
+      transferCallReturningOptional(OCE, Result, State);
+  }
 }
 
 void handleNonConstMemberCall(const CallExpr *CE,

clang/unittests/Analysis/FlowSensitive/UncheckedOptionalAccessModelTest.cpp

@@ -3771,6 +3771,30 @@ TEST_P(UncheckedOptionalAccessTest, ConstPointerAccessorWithModInBetween) {
                        /*IgnoreSmartPointerDereference=*/false);
 }
 
+TEST_P(UncheckedOptionalAccessTest, ConstPointerRefAccessor) {
+  ExpectDiagnosticsFor(R"cc(
+    #include "unchecked_optional_access_test.h"
+
+    struct A {
+      $ns::$optional<int> x;
+    };
+
+    struct PtrWrapper {
+      A*& getPtrRef() const;
+      void reset(A*);
+    };
+
+    void target(PtrWrapper p) {
+      if (p.getPtrRef()->x) {
+        *p.getPtrRef()->x;
+        p.reset(nullptr);
+        *p.getPtrRef()->x;  // [[unsafe]]
+      }
+    }
+  )cc",
+                       /*IgnoreSmartPointerDereference=*/false);
+}
+
 TEST_P(UncheckedOptionalAccessTest, SmartPointerAccessorMixed) {
   ExpectDiagnosticsFor(R"cc(
      #include "unchecked_optional_access_test.h"