[LifetimeSafety] Optimize loan propagation by separating persistent and block-local origins

clang/include/clang/Analysis/Analyses/LifetimeSafety/Facts.h

@@ -144,19 +144,18 @@ class ReturnOfOriginFact : public Fact {
 
 class UseFact : public Fact {
   const Expr *UseExpr;
+  OriginID OID;
   // True if this use is a write operation (e.g., left-hand side of assignment).
   // Write operations are exempted from use-after-free checks.
   bool IsWritten = false;
 
 public:
   static bool classof(const Fact *F) { return F->getKind() == Kind::Use; }
 
-  UseFact(const Expr *UseExpr) : Fact(Kind::Use), UseExpr(UseExpr) {}
+  UseFact(const Expr *UseExpr, OriginManager &OM)
+      : Fact(Kind::Use), UseExpr(UseExpr), OID(OM.get(*UseExpr)) {}
 
-  OriginID getUsedOrigin(const OriginManager &OM) const {
-    // TODO: Remove const cast and make OriginManager::get as const.
-    return const_cast<OriginManager &>(OM).get(*UseExpr);
-  }
+  OriginID getUsedOrigin() const { return OID; }
   const Expr *getUseExpr() const { return UseExpr; }
   void markAsWritten() { IsWritten = true; }
   bool isWritten() const { return IsWritten; }

clang/include/clang/Analysis/Analyses/LifetimeSafety/Origins.h

@@ -74,6 +74,8 @@ class OriginManager {
 
   OriginID getOrCreate(const ValueDecl &D);
 
+  unsigned getNumOrigins() const { return NextOriginID.Value; }
+
   void dump(OriginID OID, llvm::raw_ostream &OS) const;
 
 private:

clang/lib/Analysis/LifetimeSafety/Facts.cpp

@@ -53,7 +53,7 @@ void ReturnOfOriginFact::dump(llvm::raw_ostream &OS, const LoanManager &,
 void UseFact::dump(llvm::raw_ostream &OS, const LoanManager &,
                    const OriginManager &OM) const {
   OS << "Use (";
-  OM.dump(getUsedOrigin(OM), OS);
+  OM.dump(getUsedOrigin(), OS);
   OS << ", " << (isWritten() ? "Write" : "Read") << ")\n";
 }
 

clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp

@@ -333,7 +333,7 @@ void FactsGenerator::handleAssignment(const Expr *LHSExpr,
 // (e.g. on the left-hand side of an assignment).
 void FactsGenerator::handleUse(const DeclRefExpr *DRE) {
   if (isPointerType(DRE->getType())) {
-    UseFact *UF = FactMgr.createFact<UseFact>(DRE);
+    UseFact *UF = FactMgr.createFact<UseFact>(DRE, FactMgr.getOriginMgr());
     CurrentBlockFacts.push_back(UF);
     assert(!UseFacts.contains(DRE));
     UseFacts[DRE] = UF;

clang/lib/Analysis/LifetimeSafety/LiveOrigins.cpp

@@ -111,7 +111,7 @@ class AnalysisImpl
   /// dominates this program point. A write operation kills the liveness of
   /// the origin since it overwrites the value.
   Lattice transfer(Lattice In, const UseFact &UF) {
-    OriginID OID = UF.getUsedOrigin(FactMgr.getOriginMgr());
+    OriginID OID = UF.getUsedOrigin();
     // Write kills liveness.
     if (UF.isWritten())
       return Lattice(Factory.remove(In.LiveOrigins, OID));

clang/lib/Analysis/LifetimeSafety/LoanPropagation.cpp

@@ -5,36 +5,114 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-#include "clang/Analysis/Analyses/LifetimeSafety/LoanPropagation.h"
-#include "Dataflow.h"
+#include <cassert>
 #include <memory>
 
+#include "Dataflow.h"
+#include "clang/Analysis/Analyses/LifetimeSafety/Facts.h"
+#include "clang/Analysis/Analyses/LifetimeSafety/LoanPropagation.h"
+#include "clang/Analysis/Analyses/LifetimeSafety/Loans.h"
+#include "clang/Analysis/Analyses/LifetimeSafety/Origins.h"
+#include "clang/Analysis/Analyses/LifetimeSafety/Utils.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
+#include "clang/Analysis/CFG.h"
+#include "clang/Basic/LLVM.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/TimeProfiler.h"
+#include "llvm/Support/raw_ostream.h"
+
 namespace clang::lifetimes::internal {
+
+// Prepass to find persistent origins. An origin is persistent if it is
+// referenced in more than one basic block.
+static llvm::BitVector computePersistentOrigins(const FactManager &FactMgr,
+                                                const CFG &C) {
+  llvm::TimeTraceScope("ComputePersistentOrigins");
+  unsigned NumOrigins = FactMgr.getOriginMgr().getNumOrigins();
+  llvm::BitVector PersistentOrigins(NumOrigins);
+
+  llvm::SmallVector<const CFGBlock *> OriginToFirstSeenBlock(NumOrigins,
+                                                             nullptr);
+  for (const CFGBlock *B : C) {
+    for (const Fact *F : FactMgr.getFacts(B)) {
+      auto CheckOrigin = [&](OriginID OID) {
+        if (PersistentOrigins.test(OID.Value))
+          return;
+        auto &FirstSeenBlock = OriginToFirstSeenBlock[OID.Value];
+        if (FirstSeenBlock == nullptr)
+          FirstSeenBlock = B;
+        if (FirstSeenBlock != B) {
+          // We saw this origin in more than one block.
+          PersistentOrigins.set(OID.Value);
+        }
+      };
+
+      switch (F->getKind()) {
+      case Fact::Kind::Issue:
+        CheckOrigin(F->getAs<IssueFact>()->getOriginID());
+        break;
+      case Fact::Kind::OriginFlow: {
+        const auto *OF = F->getAs<OriginFlowFact>();
+        CheckOrigin(OF->getDestOriginID());
+        CheckOrigin(OF->getSrcOriginID());
+        break;
+      }
+      case Fact::Kind::ReturnOfOrigin:
+        CheckOrigin(F->getAs<ReturnOfOriginFact>()->getReturnedOriginID());
+        break;
+      case Fact::Kind::Use:
+        CheckOrigin(F->getAs<UseFact>()->getUsedOrigin());
+        break;
+      case Fact::Kind::Expire:
+      case Fact::Kind::TestPoint:
+        break;
+      }
+    }
+  }
+  return PersistentOrigins;
+}
+
 namespace {
+
 /// Represents the dataflow lattice for loan propagation.
 ///
 /// This lattice tracks which loans each origin may hold at a given program
 /// point.The lattice has a finite height: An origin's loan set is bounded by
 /// the total number of loans in the function.
-/// TODO(opt): To reduce the lattice size, propagate origins of declarations,
-/// not expressions, because expressions are not visible across blocks.
 struct Lattice {
   /// The map from an origin to the set of loans it contains.
-  OriginLoanMap Origins = OriginLoanMap(nullptr);
-
-  explicit Lattice(const OriginLoanMap &S) : Origins(S) {}
+  /// Origins that appear in multiple blocks. Participates in join operations.
+  OriginLoanMap PersistentOrigins = OriginLoanMap(nullptr);
+  /// Origins confined to a single block. Discarded at block boundaries.
+  OriginLoanMap BlockLocalOrigins = OriginLoanMap(nullptr);
+
+  explicit Lattice(const OriginLoanMap &Persistent,
+                   const OriginLoanMap &BlockLocal)
+      : PersistentOrigins(Persistent), BlockLocalOrigins(BlockLocal) {}
   Lattice() = default;
 
   bool operator==(const Lattice &Other) const {
-    return Origins == Other.Origins;
+    return PersistentOrigins == Other.PersistentOrigins &&
+           BlockLocalOrigins == Other.BlockLocalOrigins;
   }
   bool operator!=(const Lattice &Other) const { return !(*this == Other); }
 
   void dump(llvm::raw_ostream &OS) const {
     OS << "LoanPropagationLattice State:\n";
-    if (Origins.isEmpty())
+    OS << " Persistent Origins:\n";
+    if (PersistentOrigins.isEmpty())
       OS << "  <empty>\n";
-    for (const auto &Entry : Origins) {
+    for (const auto &Entry : PersistentOrigins) {
+      if (Entry.second.isEmpty())
+        OS << "  Origin " << Entry.first << " contains no loans\n";
+      for (const LoanID &LID : Entry.second)
+        OS << "  Origin " << Entry.first << " contains Loan " << LID << "\n";
+    }
+    OS << " Block-Local Origins:\n";
+    if (BlockLocalOrigins.isEmpty())
+      OS << "  <empty>\n";
+    for (const auto &Entry : BlockLocalOrigins) {
       if (Entry.second.isEmpty())
         OS << "  Origin " << Entry.first << " contains no loans\n";
       for (const LoanID &LID : Entry.second)
@@ -50,7 +128,8 @@ class AnalysisImpl
                OriginLoanMap::Factory &OriginLoanMapFactory,
                LoanSet::Factory &LoanSetFactory)
       : DataflowAnalysis(C, AC, F), OriginLoanMapFactory(OriginLoanMapFactory),
-        LoanSetFactory(LoanSetFactory) {}
+        LoanSetFactory(LoanSetFactory),
+        PersistentOrigins(computePersistentOrigins(F, C)) {}
 
   using Base::transfer;
 
@@ -59,10 +138,10 @@ class AnalysisImpl
   Lattice getInitialState() { return Lattice{}; }
 
   /// Merges two lattices by taking the union of loans for each origin.
-  // TODO(opt): Keep the state small by removing origins which become dead.
+  /// Only persistent origins are joined; block-local origins are discarded.
   Lattice join(Lattice A, Lattice B) {
     OriginLoanMap JoinedOrigins = utils::join(
-        A.Origins, B.Origins, OriginLoanMapFactory,
+        A.PersistentOrigins, B.PersistentOrigins, OriginLoanMapFactory,
         [&](const LoanSet *S1, const LoanSet *S2) {
           assert((S1 || S2) && "unexpectedly merging 2 empty sets");
           if (!S1)
@@ -74,16 +153,15 @@ class AnalysisImpl
         // Asymmetric join is a performance win. For origins present only on one
         // branch, the loan set can be carried over as-is.
         utils::JoinKind::Asymmetric);
-    return Lattice(JoinedOrigins);
+    return Lattice(JoinedOrigins, OriginLoanMapFactory.getEmptyMap());
   }
 
   /// A new loan is issued to the origin. Old loans are erased.
   Lattice transfer(Lattice In, const IssueFact &F) {
     OriginID OID = F.getOriginID();
     LoanID LID = F.getLoanID();
-    return Lattice(OriginLoanMapFactory.add(
-        In.Origins, OID,
-        LoanSetFactory.add(LoanSetFactory.getEmptySet(), LID)));
+    LoanSet NewLoans = LoanSetFactory.add(LoanSetFactory.getEmptySet(), LID);
+    return setLoans(In, OID, NewLoans);
   }
 
   /// A flow from source to destination. If `KillDest` is true, this replaces
@@ -98,22 +176,41 @@ class AnalysisImpl
     LoanSet SrcLoans = getLoans(In, SrcOID);
     LoanSet MergedLoans = utils::join(DestLoans, SrcLoans, LoanSetFactory);
 
-    return Lattice(OriginLoanMapFactory.add(In.Origins, DestOID, MergedLoans));
+    return setLoans(In, DestOID, MergedLoans);
   }
 
   LoanSet getLoans(OriginID OID, ProgramPoint P) const {
     return getLoans(getState(P), OID);
   }
 
 private:
+  /// Returns true if the origin is persistent (referenced in multiple blocks).
+  bool isPersistent(OriginID OID) const {
+    return PersistentOrigins.test(OID.Value);
+  }
+
+  Lattice setLoans(Lattice L, OriginID OID, LoanSet Loans) {
+    if (isPersistent(OID))
+      return Lattice(OriginLoanMapFactory.add(L.PersistentOrigins, OID, Loans),
+                     L.BlockLocalOrigins);
+    return Lattice(L.PersistentOrigins,
+                   OriginLoanMapFactory.add(L.BlockLocalOrigins, OID, Loans));
+  }
+
   LoanSet getLoans(Lattice L, OriginID OID) const {
-    if (auto *Loans = L.Origins.lookup(OID))
+    const OriginLoanMap *Map =
+        isPersistent(OID) ? &L.PersistentOrigins : &L.BlockLocalOrigins;
+    if (auto *Loans = Map->lookup(OID))
       return *Loans;
     return LoanSetFactory.getEmptySet();
   }
 
   OriginLoanMap::Factory &OriginLoanMapFactory;
   LoanSet::Factory &LoanSetFactory;
+  /// Boolean vector indexed by origin ID. If true, the origin appears in
+  /// multiple basic blocks and must participate in join operations. If false,
+  /// the origin is block-local and can be discarded at block boundaries.
+  llvm::BitVector PersistentOrigins;
 };
 } // namespace
 

clang/unittests/Analysis/LifetimeSafetyTest.cpp

@@ -543,7 +543,6 @@ TEST_F(LifetimeAnalysisTest, PointersInACycle) {
   EXPECT_THAT(Origin("p1"), HasLoansTo({"v1", "v2", "v3"}, "after_loop"));
   EXPECT_THAT(Origin("p2"), HasLoansTo({"v1", "v2", "v3"}, "after_loop"));
   EXPECT_THAT(Origin("p3"), HasLoansTo({"v1", "v2", "v3"}, "after_loop"));
-  EXPECT_THAT(Origin("temp"), HasLoansTo({"v1", "v2", "v3"}, "after_loop"));
 }
 
 TEST_F(LifetimeAnalysisTest, PointersAndExpirationInACycle) {