Identify DeclRefExpr as a use of an origin

clang/lib/Analysis/LifetimeSafety.cpp

@@ -8,6 +8,7 @@
 #include "clang/Analysis/Analyses/LifetimeSafety.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/AST/Type.h"
 #include "clang/Analysis/Analyses/PostOrderCFGView.h"
@@ -316,6 +317,7 @@ public:
 class UseFact : public Fact {
   OriginID UsedOrigin;
   const Expr *UseExpr;
+  bool IsWritten = false;
 
 public:
   static bool classof(const Fact *F) { return F->getKind() == Kind::Use; }
@@ -325,11 +327,13 @@ public:
 
   OriginID getUsedOrigin() const { return UsedOrigin; }
   const Expr *getUseExpr() const { return UseExpr; }
+  void markAsWritten() { IsWritten = true; }
+  bool isWritten() const { return IsWritten; }
 
   void dump(llvm::raw_ostream &OS, const OriginManager &OM) const override {
     OS << "Use (";
     OM.dump(getUsedOrigin(), OS);
-    OS << ")\n";
+    OS << " " << (isWritten() ? "Write" : "Read") << ")\n";
   }
 };
 
@@ -403,29 +407,20 @@ private:
   llvm::BumpPtrAllocator FactAllocator;
 };
 
-class FactGenerator : public ConstStmtVisitor<FactGenerator> {
-  using Base = ConstStmtVisitor<FactGenerator>;
+class FactGeneratorVisitor : public ConstStmtVisitor<FactGeneratorVisitor> {
+  using Base = ConstStmtVisitor<FactGeneratorVisitor>;
 
 public:
-  FactGenerator(FactManager &FactMgr, AnalysisDeclContext &AC)
-      : FactMgr(FactMgr), AC(AC) {}
+  FactGeneratorVisitor(FactManager &FactMgr) : FactMgr(FactMgr) {}
 
-  void run() {
-    llvm::TimeTraceScope TimeProfile("FactGenerator");
-    // Iterate through the CFG blocks in reverse post-order to ensure that
-    // initializations and destructions are processed in the correct sequence.
-    for (const CFGBlock *Block : *AC.getAnalysis<PostOrderCFGView>()) {
-      CurrentBlockFacts.clear();
-      for (unsigned I = 0; I < Block->size(); ++I) {
-        const CFGElement &Element = Block->Elements[I];
-        if (std::optional<CFGStmt> CS = Element.getAs<CFGStmt>())
-          Visit(CS->getStmt());
-        else if (std::optional<CFGAutomaticObjDtor> DtorOpt =
-                     Element.getAs<CFGAutomaticObjDtor>())
-          handleDestructor(*DtorOpt);
-      }
-      FactMgr.addBlockFacts(Block, CurrentBlockFacts);
-    }
+  void startBlock(const CFGBlock *Block) {
+    CurrentBlock = Block;
+    CurrentBlockFacts.clear();
+  }
+
+  void endBlock() {
+    FactMgr.addBlockFacts(CurrentBlock, CurrentBlockFacts);
+    startBlock(nullptr);
   }
 
   void VisitDeclStmt(const DeclStmt *DS) {
@@ -436,6 +431,8 @@ public:
             addAssignOriginFact(*VD, *InitExpr);
   }
 
+  void VisitDeclRefExpr(const DeclRefExpr *DRE) { handleUse(DRE); }
+
   void VisitCXXNullPtrLiteralExpr(const CXXNullPtrLiteralExpr *N) {
     /// TODO: Handle nullptr expr as a special 'null' loan. Uninitialized
     /// pointers can use the same type of loan.
@@ -445,7 +442,6 @@ public:
   void VisitImplicitCastExpr(const ImplicitCastExpr *ICE) {
     if (!hasOrigin(ICE->getType()))
       return;
-    Visit(ICE->getSubExpr());
     // An ImplicitCastExpr node itself gets an origin, which flows from the
     // origin of its sub-expression (after stripping its own parens/casts).
     // TODO: Consider if this is actually useful in practice. Alternatively, we
@@ -470,10 +466,6 @@ public:
           }
         }
       }
-    } else if (UO->getOpcode() == UO_Deref) {
-      // This is a pointer use, like '*p'.
-      OriginID OID = FactMgr.getOriginMgr().get(*UO->getSubExpr());
-      CurrentBlockFacts.push_back(FactMgr.createFact<UseFact>(OID, UO));
     }
   }
 
@@ -488,20 +480,13 @@ public:
   }
 
   void VisitBinaryOperator(const BinaryOperator *BO) {
-    if (BO->isAssignmentOp()) {
-      const Expr *LHSExpr = BO->getLHS();
-      const Expr *RHSExpr = BO->getRHS();
+    if (BO->isAssignmentOp())
+      handleAssignment(BO->getLHS(), BO->getRHS());
+  }
 
-      // We are interested in assignments like `ptr1 = ptr2` or `ptr = &var`
-      // LHS must be a pointer/reference type that can be an origin.
-      // RHS must also represent an origin (either another pointer/ref or an
-      // address-of).
-      if (const auto *DRE_LHS = dyn_cast<DeclRefExpr>(LHSExpr))
-        if (const auto *VD_LHS =
-                dyn_cast<ValueDecl>(DRE_LHS->getDecl()->getCanonicalDecl());
-            VD_LHS && hasOrigin(VD_LHS->getType()))
-          addAssignOriginFact(*VD_LHS, *RHSExpr);
-    }
+  void VisitCXXOperatorCallExpr(const CXXOperatorCallExpr *OCE) {
+    if (OCE->isAssignmentOp() && OCE->getNumArgs() == 2)
+      handleAssignment(OCE->getArg(0), OCE->getArg(1));
   }
 
   void VisitCXXFunctionalCastExpr(const CXXFunctionalCastExpr *FCE) {
@@ -513,18 +498,6 @@ public:
     Base::VisitCXXFunctionalCastExpr(FCE);
   }
 
-private:
-  // Check if a type has an origin.
-  bool hasOrigin(QualType QT) { return QT->isPointerOrReferenceType(); }
-
-  template <typename Destination, typename Source>
-  void addAssignOriginFact(const Destination &D, const Source &S) {
-    OriginID DestOID = FactMgr.getOriginMgr().getOrCreate(D);
-    OriginID SrcOID = FactMgr.getOriginMgr().get(S);
-    CurrentBlockFacts.push_back(
-        FactMgr.createFact<AssignOriginFact>(DestOID, SrcOID));
-  }
-
   void handleDestructor(const CFGAutomaticObjDtor &DtorOpt) {
     /// TODO: Also handle trivial destructors (e.g., for `int`
     /// variables) which will never have a CFGAutomaticObjDtor node.
@@ -547,6 +520,18 @@ private:
     }
   }
 
+private:
+  // Check if a type has an origin.
+  bool hasOrigin(QualType QT) { return QT->isPointerOrReferenceType(); }
+
+  template <typename Destination, typename Source>
+  void addAssignOriginFact(const Destination &D, const Source &S) {
+    OriginID DestOID = FactMgr.getOriginMgr().getOrCreate(D);
+    OriginID SrcOID = FactMgr.getOriginMgr().get(S);
+    CurrentBlockFacts.push_back(
+        FactMgr.createFact<AssignOriginFact>(DestOID, SrcOID));
+  }
+
   /// Checks if the expression is a `void("__lifetime_test_point_...")` cast.
   /// If so, creates a `TestPointFact` and returns true.
   bool VisitTestPoint(const CXXFunctionalCastExpr *FCE) {
@@ -568,9 +553,101 @@ private:
     return false;
   }
 
+  void handleAssignment(const Expr *LHSExpr, const Expr *RHSExpr) {
+    // Find the underlying variable declaration for the left-hand side.
+    if (const auto *DRE_LHS =
+            dyn_cast<DeclRefExpr>(LHSExpr->IgnoreParenImpCasts())) {
+      markUseAsWrite(DRE_LHS);
+      if (const auto *VD_LHS = dyn_cast<ValueDecl>(DRE_LHS->getDecl()))
+        if (hasOrigin(VD_LHS->getType()))
+          // We are interested in assignments like `ptr1 = ptr2` or `ptr = &var`
+          // LHS must be a pointer/reference type that can be an origin.
+          // RHS must also represent an origin (either another pointer/ref or an
+          // address-of).
+          addAssignOriginFact(*VD_LHS, *RHSExpr);
+    }
+  }
+
+  // A DeclRefExpr is a use of the referenced decl. It is checked for
+  // use-after-free unless it is being written to (e.g. on the left-hand side
+  // of an assignment).
+  void handleUse(const DeclRefExpr *DRE) {
+    const auto *VD = dyn_cast<ValueDecl>(DRE->getDecl());
+    if (VD && hasOrigin(VD->getType())) {
+      OriginID OID = FactMgr.getOriginMgr().get(*VD);
+      UseFact *UF = FactMgr.createFact<UseFact>(OID, DRE);
+      CurrentBlockFacts.push_back(UF);
+      assert(!UseFacts.contains(DRE));
+      UseFacts[DRE] = UF;
+    }
+  }
+
+  void markUseAsWrite(const DeclRefExpr *DRE) {
+    assert(UseFacts.contains(DRE));
+    UseFacts[DRE]->markAsWritten();
+  }
+
   FactManager &FactMgr;
-  AnalysisDeclContext &AC;
+  const CFGBlock *CurrentBlock = nullptr;
   llvm::SmallVector<Fact *> CurrentBlockFacts;
+  // To distinguish between reads and writes for use-after-free checks, this map
+  // stores the `UseFact` for each `DeclRefExpr`. We initially identify all
+  // `DeclRefExpr`s as "read" uses. When an assignment is processed, the use
+  // corresponding to the left-hand side is updated to be a "write", thereby
+  // exempting it from the check.
+  llvm::DenseMap<const DeclRefExpr *, UseFact *> UseFacts;
+};
+
+class FactGenerator : public RecursiveASTVisitor<FactGenerator> {
+public:
+  FactGenerator(FactManager &FactMgr, AnalysisDeclContext &AC)
+      : FG(FactMgr), AC(AC) {}
+
+  bool shouldTraversePostOrder() const { return true; }
+
+  void run() {
+    llvm::TimeTraceScope TimeProfile("FactGenerator");
+    // Iterate through the CFG blocks in reverse post-order to ensure that
+    // initializations and destructions are processed in the correct sequence.
+    for (const CFGBlock *Block : *AC.getAnalysis<PostOrderCFGView>()) {
+      FactGeneratorBlockRAII BlockGenerator(FG, Block);
+      for (const CFGElement &Element : *Block) {
+        if (std::optional<CFGStmt> CS = Element.getAs<CFGStmt>())
+          TraverseStmt(const_cast<Stmt *>(CS->getStmt()));
+        else if (std::optional<CFGAutomaticObjDtor> DtorOpt =
+                     Element.getAs<CFGAutomaticObjDtor>())
+          FG.handleDestructor(*DtorOpt);
+      }
+    }
+  }
+
+  bool TraverseStmt(Stmt *S) {
+    // Avoid re-visiting nodes to not create duplicate facts.
+    if (!S || !VisitedStmts.insert(S).second)
+      return true;
+    return RecursiveASTVisitor::TraverseStmt(S);
+  }
+
+  bool VisitStmt(Stmt *S) {
+    FG.Visit(S);
+    return true; // Continue traversing to children.
+  }
+
+private:
+  struct FactGeneratorBlockRAII {
+    FactGeneratorBlockRAII(FactGeneratorVisitor &FG, const CFGBlock *Block)
+        : FG(FG) {
+      FG.startBlock(Block);
+    }
+    ~FactGeneratorBlockRAII() { FG.endBlock(); }
+
+  private:
+    FactGeneratorVisitor &FG;
+  };
+
+  FactGeneratorVisitor FG;
+  AnalysisDeclContext &AC;
+  llvm::DenseSet<const Stmt *> VisitedStmts;
 };
 
 // ========================================================================= //
@@ -1033,7 +1110,8 @@ public:
   /// graph. It determines if the loans held by the used origin have expired
   /// at the point of use.
   void checkUse(const UseFact *UF) {
-
+    if (UF->isWritten())
+      return;
     OriginID O = UF->getUsedOrigin();
 
     // Get the set of loans that the origin might hold at this program point.
@@ -1116,8 +1194,8 @@ void LifetimeSafetyAnalysis::run() {
   DEBUG_WITH_TYPE("PrintCFG", Cfg.dump(AC.getASTContext().getLangOpts(),
                                        /*ShowColors=*/true));
 
-  FactGenerator FactGen(*FactMgr, AC);
-  FactGen.run();
+  FactGenerator FG(*FactMgr, AC);
+  FG.run();
   DEBUG_WITH_TYPE("LifetimeFacts", FactMgr->dump(Cfg, AC));
 
   /// TODO(opt): Consider optimizing individual blocks before running the

clang/test/Sema/warn-lifetime-safety-dataflow.cpp

@@ -293,3 +293,29 @@ void pointer_indirection() {
   int *q = *pp;
 // CHECK:   AssignOrigin (Dest: {{[0-9]+}} (Decl: q), Src: {{[0-9]+}} (Expr: ImplicitCastExpr))
 }
+
+// CHECK-LABEL: Function: test_use_facts
+void usePointer(MyObj*);
+void test_use_facts() {
+  // CHECK: Block B{{[0-9]+}}:
+  MyObj x;
+  MyObj *p;
+  p = &x;
+  // CHECK:   Use ([[O_P:[0-9]+]] (Decl: p) Write)
+  (void)*p;
+  // CHECK:   Use ([[O_P]] (Decl: p) Read)
+  usePointer(p);
+  // CHECK:   Use ([[O_P]] (Decl: p) Read)
+  p->id = 1;
+  // CHECK:   Use ([[O_P]] (Decl: p) Read)
+
+
+  MyObj* q;
+  q = p;
+  // CHECK:   Use ([[O_P]] (Decl: p) Read)
+  // CHECK:   Use ([[O_Q:[0-9]+]] (Decl: q) Write)
+  usePointer(q);
+  // CHECK:   Use ([[O_Q]] (Decl: q) Read)
+  q->id = 2;
+  // CHECK:   Use ([[O_Q]] (Decl: q) Read)
+}

clang/unittests/Analysis/LifetimeSafetyTest.cpp

@@ -11,6 +11,7 @@
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Testing/TestAST.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <optional>
@@ -20,6 +21,7 @@ namespace clang::lifetimes::internal {
 namespace {
 
 using namespace ast_matchers;
+using ::testing::SizeIs;
 using ::testing::UnorderedElementsAreArray;
 
 // A helper class to run the full lifetime analysis on a piece of code
@@ -96,21 +98,18 @@ public:
     return OID;
   }
 
-  std::optional<LoanID> getLoanForVar(llvm::StringRef VarName) {
+  std::vector<LoanID> getLoansForVar(llvm::StringRef VarName) {
     auto *VD = findDecl<VarDecl>(VarName);
-    if (!VD)
-      return std::nullopt;
+    if (!VD) {
+      ADD_FAILURE() << "No VarDecl found for '" << VarName << "'";
+      return {};
+    }
     std::vector<LoanID> LID = Analysis.getLoanIDForVar(VD);
     if (LID.empty()) {
       ADD_FAILURE() << "Loan for '" << VarName << "' not found.";
-      return std::nullopt;
+      return {};
     }
-    // TODO: Support retrieving more than one loans to a var.
-    if (LID.size() > 1) {
-      ADD_FAILURE() << "More than 1 loans found for '" << VarName;
-      return std::nullopt;
-    }
-    return LID[0];
+    return LID;
   }
 
   std::optional<LoanSet> getLoansAtPoint(OriginID OID,
@@ -121,13 +120,12 @@ public:
     return Analysis.getLoansAtPoint(OID, PP);
   }
 
-  std::optional<llvm::DenseSet<LoanID>>
+  std::optional<std::vector<LoanID>>
   getExpiredLoansAtPoint(llvm::StringRef Annotation) {
     ProgramPoint PP = Runner.getProgramPoint(Annotation);
     if (!PP)
       return std::nullopt;
-    auto Expired = Analysis.getExpiredLoansAtPoint(PP);
-    return llvm::DenseSet<LoanID>{Expired.begin(), Expired.end()};
+    return Analysis.getExpiredLoansAtPoint(PP);
   }
 
 private:
@@ -197,12 +195,13 @@ MATCHER_P2(HasLoansToImpl, LoanVars, Annotation, "") {
 
   std::vector<LoanID> ExpectedLoans;
   for (const auto &LoanVar : LoanVars) {
-    std::optional<LoanID> ExpectedLIDOpt = Info.Helper.getLoanForVar(LoanVar);
-    if (!ExpectedLIDOpt) {
+    std::vector<LoanID> ExpectedLIDs = Info.Helper.getLoansForVar(LoanVar);
+    if (ExpectedLIDs.empty()) {
       *result_listener << "could not find loan for var '" << LoanVar << "'";
       return false;
     }
-    ExpectedLoans.push_back(*ExpectedLIDOpt);
+    ExpectedLoans.insert(ExpectedLoans.end(), ExpectedLIDs.begin(),
+                         ExpectedLIDs.end());
   }
 
   return ExplainMatchResult(UnorderedElementsAreArray(ExpectedLoans),
@@ -221,17 +220,17 @@ MATCHER_P(AreExpiredAt, Annotation, "") {
                      << Annotation << "'";
     return false;
   }
-  std::vector<LoanID> ActualExpiredLoans(ActualExpiredSetOpt->begin(),
-                                         ActualExpiredSetOpt->end());
+  std::vector<LoanID> ActualExpiredLoans = *ActualExpiredSetOpt;
   std::vector<LoanID> ExpectedExpiredLoans;
   for (const auto &VarName : Info.LoanVars) {
-    auto LoanIDOpt = Helper.getLoanForVar(VarName);
-    if (!LoanIDOpt) {
+    auto LoanIDs = Helper.getLoansForVar(VarName);
+    if (LoanIDs.empty()) {
       *result_listener << "could not find a loan for variable '" << VarName
                        << "'";
       return false;
     }
-    ExpectedExpiredLoans.push_back(*LoanIDOpt);
+    ExpectedExpiredLoans.insert(ExpectedExpiredLoans.end(), LoanIDs.begin(),
+                                LoanIDs.end());
   }
   return ExplainMatchResult(UnorderedElementsAreArray(ExpectedExpiredLoans),
                             ActualExpiredLoans, result_listener);
@@ -730,5 +729,17 @@ TEST_F(LifetimeAnalysisTest, ReassignedPointerThenOriginalExpires) {
   EXPECT_THAT(LoansTo({"s1", "s2"}), AreExpiredAt("p_after_s1_expires"));
 }
 
+TEST_F(LifetimeAnalysisTest, NoDuplicateLoansForImplicitCastToConst) {
+  SetupTest(R"(
+    void target() {
+      MyObj a;
+      const MyObj* p = &a;
+      const MyObj* q = &a;
+      POINT(at_end);
+    }
+  )");
+  EXPECT_THAT(Helper->getLoansForVar("a"), SizeIs(2));
+}
+
 } // anonymous namespace
 } // namespace clang::lifetimes::internal