[clang][dataflow] Use smart pointer caching in unchecked optional accessor

clang/include/clang/Analysis/FlowSensitive/CachedConstAccessorsLattice.h

@@ -13,7 +13,9 @@
 #ifndef LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_CACHED_CONST_ACCESSORS_LATTICE_H
 #define LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_CACHED_CONST_ACCESSORS_LATTICE_H
 
+#include "clang/AST/Decl.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/Type.h"
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #include "clang/Analysis/FlowSensitive/DataflowLattice.h"
 #include "clang/Analysis/FlowSensitive/StorageLocation.h"
@@ -71,10 +73,27 @@ template <typename Base> class CachedConstAccessorsLattice : public Base {
   /// Requirements:
   ///
   ///  - `CE` should return a location (GLValue or a record type).
+  ///
+  /// DEPRECATED: switch users to the below overload which takes Callee and Type
+  /// directly.
   StorageLocation *getOrCreateConstMethodReturnStorageLocation(
       const RecordStorageLocation &RecordLoc, const CallExpr *CE,
       Environment &Env, llvm::function_ref<void(StorageLocation &)> Initialize);
 
+  /// Creates or returns a previously created `StorageLocation` associated with
+  /// a const method call `obj.getFoo()` where `RecordLoc` is the
+  /// `RecordStorageLocation` of `obj`, `Callee` is the decl for `getFoo`.
+  ///
+  /// The callback `Initialize` runs on the storage location if newly created.
+  ///
+  /// Requirements:
+  ///
+  ///  - `Callee` should return a location (return type is a reference type or a
+  ///     record type).
+  StorageLocation &getOrCreateConstMethodReturnStorageLocation(
+      const RecordStorageLocation &RecordLoc, const FunctionDecl *Callee,
+      Environment &Env, llvm::function_ref<void(StorageLocation &)> Initialize);
+
   void clearConstMethodReturnValues(const RecordStorageLocation &RecordLoc) {
     ConstMethodReturnValues.erase(&RecordLoc);
   }
@@ -212,6 +231,27 @@ CachedConstAccessorsLattice<Base>::getOrCreateConstMethodReturnStorageLocation(
   return &Loc;
 }
 
+template <typename Base>
+StorageLocation &
+CachedConstAccessorsLattice<Base>::getOrCreateConstMethodReturnStorageLocation(
+    const RecordStorageLocation &RecordLoc, const FunctionDecl *Callee,
+    Environment &Env, llvm::function_ref<void(StorageLocation &)> Initialize) {
+  assert(Callee != nullptr);
+  QualType Type = Callee->getReturnType();
+  assert(!Type.isNull());
+  assert(Type->isReferenceType() || Type->isRecordType());
+  auto &ObjMap = ConstMethodReturnStorageLocations[&RecordLoc];
+  auto it = ObjMap.find(Callee);
+  if (it != ObjMap.end())
+    return *it->second;
+
+  StorageLocation &Loc = Env.createStorageLocation(Type.getNonReferenceType());
+  Initialize(Loc);
+
+  ObjMap.insert({Callee, &Loc});
+  return Loc;
+}
+
 } // namespace dataflow
 } // namespace clang
 

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

@@ -37,14 +37,13 @@ struct UncheckedOptionalAccessModelOptions {
   /// can't identify when their results are used safely (across calls),
   /// resulting in false positives in all such cases. Note: this option does not
   /// cover access through `operator[]`.
-  /// FIXME: we currently cache and equate the result of const accessors
-  /// returning pointers, so cover the case of operator-> followed by
-  /// operator->, which covers the common case of smart pointers. We also cover
-  /// some limited cases of returning references (if return type is an optional
-  /// type), so cover some cases of operator* followed by operator*. We don't
-  /// cover mixing operator-> and operator*. Once we are confident in this const
-  /// accessor caching, we shouldn't need the IgnoreSmartPointerDereference
-  /// option anymore.
+  ///
+  /// FIXME: we now cache and equate the result of const accessors
+  /// that look like unique_ptr, have both `->` (returning a pointer type) and
+  /// `*` (returning a reference type). This includes mixing `->` and
+  /// `*` in a sequence of calls as long as the object is not modified. Once we
+  /// are confident in this const accessor caching, we shouldn't need the
+  /// IgnoreSmartPointerDereference option anymore.
   bool IgnoreSmartPointerDereference = false;
 };
 

clang/include/clang/Analysis/FlowSensitive/SmartPointerAccessorCaching.h

@@ -27,8 +27,13 @@
 #include <cassert>
 
 #include "clang/AST/Decl.h"
+#include "clang/AST/Expr.h"
 #include "clang/AST/Stmt.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Analysis/FlowSensitive/MatchSwitch.h"
+#include "clang/Analysis/FlowSensitive/StorageLocation.h"
+#include "clang/Analysis/FlowSensitive/Value.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 
 namespace clang::dataflow {
 
@@ -58,6 +63,107 @@ ast_matchers::StatementMatcher isSmartPointerLikeOperatorArrow();
 ast_matchers::StatementMatcher isSmartPointerLikeValueMethodCall();
 ast_matchers::StatementMatcher isSmartPointerLikeGetMethodCall();
 
+// Common transfer functions.
+
+/// Returns the "canonical" callee for smart pointer operators (`*` and `->`)
+/// as a key for caching.
+///
+/// We choose `*` as the canonical one, since it needs a
+/// StorageLocation anyway.
+///
+/// Note: there may be multiple `operator*` (one const, one non-const).
+/// We pick the const one, which the above provided matchers require to exist.
+const FunctionDecl *
+getCanonicalSmartPointerLikeOperatorCallee(const CallExpr *CE);
+
+/// A transfer function for `operator*` (and `value`) calls that can be
+/// cached. Runs the `InitializeLoc` callback to initialize any new
+/// StorageLocations.
+///
+/// Requirements:
+///
+/// - LatticeT should use the `CachedConstAccessorsLattice` mixin.
+template <typename LatticeT>
+void transferSmartPointerLikeCachedDeref(
+    const CallExpr *DerefExpr, RecordStorageLocation *SmartPointerLoc,
+    TransferState<LatticeT> &State,
+    llvm::function_ref<void(StorageLocation &)> InitializeLoc);
+
+/// A transfer function for `operator->` (and `get`) calls that can be cached.
+/// Runs the `InitializeLoc` callback to initialize any new StorageLocations.
+///
+/// Requirements:
+///
+/// - LatticeT should use the `CachedConstAccessorsLattice` mixin.
+template <typename LatticeT>
+void transferSmartPointerLikeCachedGet(
+    const CallExpr *GetExpr, RecordStorageLocation *SmartPointerLoc,
+    TransferState<LatticeT> &State,
+    llvm::function_ref<void(StorageLocation &)> InitializeLoc);
+
+template <typename LatticeT>
+void transferSmartPointerLikeCachedDeref(
+    const CallExpr *DerefExpr, RecordStorageLocation *SmartPointerLoc,
+    TransferState<LatticeT> &State,
+    llvm::function_ref<void(StorageLocation &)> InitializeLoc) {
+  if (State.Env.getStorageLocation(*DerefExpr) != nullptr)
+    return;
+  if (SmartPointerLoc == nullptr)
+    return;
+
+  const FunctionDecl *Callee = DerefExpr->getDirectCallee();
+  if (Callee == nullptr)
+    return;
+  const FunctionDecl *CanonicalCallee =
+      getCanonicalSmartPointerLikeOperatorCallee(DerefExpr);
+  // This shouldn't happen, as we should at least find `Callee` itself.
+  assert(CanonicalCallee != nullptr);
+  if (CanonicalCallee != Callee) {
+    // When using the provided matchers, we should always get a reference to
+    // the same type.
+    assert(CanonicalCallee->getReturnType()->isReferenceType() &&
+           Callee->getReturnType()->isReferenceType());
+    assert(CanonicalCallee->getReturnType()
+               .getNonReferenceType()
+               ->getCanonicalTypeUnqualified() ==
+           Callee->getReturnType()
+               .getNonReferenceType()
+               ->getCanonicalTypeUnqualified());
+  }
+
+  StorageLocation &LocForValue =
+      State.Lattice.getOrCreateConstMethodReturnStorageLocation(
+          *SmartPointerLoc, CanonicalCallee, State.Env, InitializeLoc);
+  State.Env.setStorageLocation(*DerefExpr, LocForValue);
+}
+
+template <typename LatticeT>
+void transferSmartPointerLikeCachedGet(
+    const CallExpr *GetExpr, RecordStorageLocation *SmartPointerLoc,
+    TransferState<LatticeT> &State,
+    llvm::function_ref<void(StorageLocation &)> InitializeLoc) {
+  if (SmartPointerLoc == nullptr)
+    return;
+
+  const FunctionDecl *CanonicalCallee =
+      getCanonicalSmartPointerLikeOperatorCallee(GetExpr);
+
+  if (CanonicalCallee != nullptr) {
+    auto &LocForValue =
+        State.Lattice.getOrCreateConstMethodReturnStorageLocation(
+            *SmartPointerLoc, CanonicalCallee, State.Env, InitializeLoc);
+    State.Env.setValue(*GetExpr,
+                       State.Env.template create<PointerValue>(LocForValue));
+  } else {
+    // Otherwise, just cache the pointer value as if it was a const accessor.
+    Value *Val = State.Lattice.getOrCreateConstMethodReturnValue(
+        *SmartPointerLoc, GetExpr, State.Env);
+    if (Val == nullptr)
+      return;
+    State.Env.setValue(*GetExpr, *Val);
+  }
+}
+
 } // namespace clang::dataflow
 
 #endif // LLVM_CLANG_ANALYSIS_FLOWSENSITIVE_SMARTPOINTERACCESSORCACHING_H

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

@@ -25,8 +25,10 @@
 #include "clang/Analysis/FlowSensitive/DataflowEnvironment.h"
 #include "clang/Analysis/FlowSensitive/Formula.h"
 #include "clang/Analysis/FlowSensitive/RecordOps.h"
+#include "clang/Analysis/FlowSensitive/SmartPointerAccessorCaching.h"
 #include "clang/Analysis/FlowSensitive/StorageLocation.h"
 #include "clang/Analysis/FlowSensitive/Value.h"
+#include "clang/Basic/OperatorKinds.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
@@ -555,24 +557,25 @@ void handleConstMemberCall(const CallExpr *CE,
                            LatticeTransferState &State) {
   // If the const method returns an optional or reference to an optional.
   if (RecordLoc != nullptr && isSupportedOptionalType(CE->getType())) {
-    StorageLocation *Loc =
+    const FunctionDecl *DirectCallee = CE->getDirectCallee();
+    if (DirectCallee == nullptr)
+      return;
+    StorageLocation &Loc =
         State.Lattice.getOrCreateConstMethodReturnStorageLocation(
-            *RecordLoc, CE, State.Env, [&](StorageLocation &Loc) {
+            *RecordLoc, DirectCallee, State.Env, [&](StorageLocation &Loc) {
               setHasValue(cast<RecordStorageLocation>(Loc),
                           State.Env.makeAtomicBoolValue(), State.Env);
             });
-    if (Loc == nullptr)
-      return;
     if (CE->isGLValue()) {
       // If the call to the const method returns a reference to an optional,
       // link the call expression to the cached StorageLocation.
-      State.Env.setStorageLocation(*CE, *Loc);
+      State.Env.setStorageLocation(*CE, Loc);
     } else {
       // If the call to the const method returns an optional by value, we
       // need to use CopyRecord to link the optional to the result object
       // of the call expression.
       auto &ResultLoc = State.Env.getResultObjectLocation(*CE);
-      copyRecord(*cast<RecordStorageLocation>(Loc), ResultLoc, State.Env);
+      copyRecord(cast<RecordStorageLocation>(Loc), ResultLoc, State.Env);
     }
     return;
   }
@@ -1031,6 +1034,48 @@ auto buildTransferMatchSwitch() {
             transferOptionalAndValueCmp(Cmp, Cmp->getArg(1), State.Env);
           })
 
+      // Smart-pointer-like operator* and operator-> calls that may look like
+      // const accessors (below) but need special handling to allow mixing
+      // the accessor calls.
+      .CaseOfCFGStmt<CXXOperatorCallExpr>(
+          isSmartPointerLikeOperatorStar(),
+          [](const CXXOperatorCallExpr *E,
+             const MatchFinder::MatchResult &Result,
+             LatticeTransferState &State) {
+            transferSmartPointerLikeCachedDeref(
+                E,
+                dyn_cast_or_null<RecordStorageLocation>(
+                    getLocBehindPossiblePointer(*E->getArg(0), State.Env)),
+                State, [](StorageLocation &Loc) {});
+          })
+      .CaseOfCFGStmt<CXXOperatorCallExpr>(
+          isSmartPointerLikeOperatorArrow(),
+          [](const CXXOperatorCallExpr *E,
+             const MatchFinder::MatchResult &Result,
+             LatticeTransferState &State) {
+            transferSmartPointerLikeCachedGet(
+                E,
+                dyn_cast_or_null<RecordStorageLocation>(
+                    getLocBehindPossiblePointer(*E->getArg(0), State.Env)),
+                State, [](StorageLocation &Loc) {});
+          })
+      .CaseOfCFGStmt<CXXMemberCallExpr>(
+          isSmartPointerLikeValueMethodCall(),
+          [](const CXXMemberCallExpr *E, const MatchFinder::MatchResult &Result,
+             LatticeTransferState &State) {
+            transferSmartPointerLikeCachedDeref(
+                E, getImplicitObjectLocation(*E, State.Env), State,
+                [](StorageLocation &Loc) {});
+          })
+      .CaseOfCFGStmt<CXXMemberCallExpr>(
+          isSmartPointerLikeGetMethodCall(),
+          [](const CXXMemberCallExpr *E, const MatchFinder::MatchResult &Result,
+             LatticeTransferState &State) {
+            transferSmartPointerLikeCachedGet(
+                E, getImplicitObjectLocation(*E, State.Env), State,
+                [](StorageLocation &Loc) {});
+          })
+
       // const accessor calls
       .CaseOfCFGStmt<CXXMemberCallExpr>(isZeroParamConstMemberCall(),
                                         transferValue_ConstMemberCall)

clang/lib/Analysis/FlowSensitive/SmartPointerAccessorCaching.cpp

@@ -132,6 +132,7 @@ ast_matchers::StatementMatcher isSmartPointerLikeOperatorArrow() {
       callee(cxxMethodDecl(parameterCountIs(0), returns(pointerType()),
                            ofClass(smartPointerClassWithGetOrValue()))));
 }
+
 ast_matchers::StatementMatcher isSmartPointerLikeValueMethodCall() {
   return cxxMemberCallExpr(callee(
       cxxMethodDecl(parameterCountIs(0), returns(referenceType()),
@@ -144,4 +145,24 @@ ast_matchers::StatementMatcher isSmartPointerLikeGetMethodCall() {
                     ofClass(smartPointerClassWithGet()))));
 }
 
+const FunctionDecl *
+getCanonicalSmartPointerLikeOperatorCallee(const CallExpr *CE) {
+  const FunctionDecl *CanonicalCallee = nullptr;
+  const CXXMethodDecl *Callee =
+      cast_or_null<CXXMethodDecl>(CE->getDirectCallee());
+  if (Callee == nullptr)
+    return nullptr;
+  const CXXRecordDecl *RD = Callee->getParent();
+  if (RD == nullptr)
+    return nullptr;
+  for (const auto *MD : RD->methods()) {
+    if (MD->getOverloadedOperator() == OO_Star && MD->isConst() &&
+        MD->getNumParams() == 0 && MD->getReturnType()->isReferenceType()) {
+      CanonicalCallee = MD;
+      break;
+    }
+  }
+  return CanonicalCallee;
+}
+
 } // namespace clang::dataflow

clang/unittests/Analysis/FlowSensitive/CachedConstAccessorsLatticeTest.cpp

@@ -148,6 +148,35 @@ TEST_F(CachedConstAccessorsLatticeTest, SameLocBeforeClearOrDiffAfterClear) {
   EXPECT_NE(Loc3, Loc2);
 }
 
+TEST_F(CachedConstAccessorsLatticeTest,
+       SameLocBeforeClearOrDiffAfterClearWithCallee) {
+  CommonTestInputs Inputs;
+  auto *CE = Inputs.CallRef;
+  RecordStorageLocation Loc(Inputs.SType, RecordStorageLocation::FieldToLoc(),
+                            {});
+
+  LatticeT Lattice;
+  auto NopInit = [](StorageLocation &) {};
+  const FunctionDecl *Callee = CE->getDirectCallee();
+  ASSERT_NE(Callee, nullptr);
+  StorageLocation &Loc1 = Lattice.getOrCreateConstMethodReturnStorageLocation(
+      Loc, Callee, Env, NopInit);
+  auto NotCalled = [](StorageLocation &) {
+    ASSERT_TRUE(false) << "Not reached";
+  };
+  StorageLocation &Loc2 = Lattice.getOrCreateConstMethodReturnStorageLocation(
+      Loc, Callee, Env, NotCalled);
+
+  EXPECT_EQ(&Loc1, &Loc2);
+
+  Lattice.clearConstMethodReturnStorageLocations(Loc);
+  StorageLocation &Loc3 = Lattice.getOrCreateConstMethodReturnStorageLocation(
+      Loc, Callee, Env, NopInit);
+
+  EXPECT_NE(&Loc3, &Loc1);
+  EXPECT_NE(&Loc3, &Loc2);
+}
+
 TEST_F(CachedConstAccessorsLatticeTest,
        SameStructValBeforeClearOrDiffAfterClear) {
   TestAST AST(R"cpp(