[LifetimeSafety] Optimize fact storage with IDs and vector-based lookup

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

@@ -16,13 +16,17 @@
 
 #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 "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include <cstdint>
 
 namespace clang::lifetimes::internal {
+
+using FactID = utils::ID<struct FactTag>;
+
 /// An abstract base class for a single, atomic lifetime-relevant event.
 class Fact {
 
@@ -48,6 +52,7 @@ class Fact {
 
 private:
   Kind K;
+  FactID ID;
 
 protected:
   Fact(Kind K) : K(K) {}
@@ -56,6 +61,9 @@ class Fact {
   virtual ~Fact() = default;
   Kind getKind() const { return K; }
 
+  void setID(FactID ID) { this->ID = ID; }
+  FactID getID() const { return ID; }
+
   template <typename T> const T *getAs() const {
     if (T::classof(this))
       return static_cast<const T *>(this);
@@ -183,22 +191,26 @@ class TestPointFact : public Fact {
 
 class FactManager {
 public:
+  void init(const CFG &Cfg) {
+    assert(BlockToFacts.empty() && "FactManager already initialized");
+    BlockToFacts.resize(Cfg.getNumBlockIDs());
+  }
+
   llvm::ArrayRef<const Fact *> getFacts(const CFGBlock *B) const {
-    auto It = BlockToFactsMap.find(B);
-    if (It != BlockToFactsMap.end())
-      return It->second;
-    return {};
+    return BlockToFacts[B->getBlockID()];
   }
 
   void addBlockFacts(const CFGBlock *B, llvm::ArrayRef<Fact *> NewFacts) {
     if (!NewFacts.empty())
-      BlockToFactsMap[B].assign(NewFacts.begin(), NewFacts.end());
+      BlockToFacts[B->getBlockID()].assign(NewFacts.begin(), NewFacts.end());
   }
 
   template <typename FactType, typename... Args>
   FactType *createFact(Args &&...args) {
     void *Mem = FactAllocator.Allocate<FactType>();
-    return new (Mem) FactType(std::forward<Args>(args)...);
+    FactType *Res = new (Mem) FactType(std::forward<Args>(args)...);
+    Res->setID(NextFactID++);
+    return Res;
   }
 
   void dump(const CFG &Cfg, AnalysisDeclContext &AC) const;
@@ -214,16 +226,19 @@ class FactManager {
   /// \note This is intended for testing only.
   llvm::StringMap<ProgramPoint> getTestPoints() const;
 
+  unsigned getNumFacts() const { return NextFactID.Value; }
+
   LoanManager &getLoanMgr() { return LoanMgr; }
   const LoanManager &getLoanMgr() const { return LoanMgr; }
   OriginManager &getOriginMgr() { return OriginMgr; }
   const OriginManager &getOriginMgr() const { return OriginMgr; }
 
 private:
+  FactID NextFactID{0};
   LoanManager LoanMgr;
   OriginManager OriginMgr;
-  llvm::DenseMap<const clang::CFGBlock *, llvm::SmallVector<const Fact *>>
-      BlockToFactsMap;
+  /// Facts for each CFG block, indexed by block ID.
+  llvm::SmallVector<llvm::SmallVector<const Fact *>> BlockToFacts;
   llvm::BumpPtrAllocator FactAllocator;
 };
 } // namespace clang::lifetimes::internal

clang/lib/Analysis/LifetimeSafety/Dataflow.h

@@ -67,10 +67,10 @@ class DataflowAnalysis {
   llvm::DenseMap<const CFGBlock *, Lattice> InStates;
   /// The dataflow state after a basic block is processed.
   llvm::DenseMap<const CFGBlock *, Lattice> OutStates;
-  /// The dataflow state at a Program Point.
+  /// Dataflow state at each program point, indexed by Fact ID.
   /// In a forward analysis, this is the state after the Fact at that point has
   /// been applied, while in a backward analysis, it is the state before.
-  llvm::DenseMap<ProgramPoint, Lattice> PerPointStates;
+  llvm::SmallVector<Lattice> PointToState;
 
   static constexpr bool isForward() { return Dir == Direction::Forward; }
 
@@ -86,6 +86,8 @@ class DataflowAnalysis {
     Derived &D = static_cast<Derived &>(*this);
     llvm::TimeTraceScope Time(D.getAnalysisName());
 
+    PointToState.resize(FactMgr.getNumFacts());
+
     using Worklist =
         std::conditional_t<Dir == Direction::Forward, ForwardDataflowWorklist,
                            BackwardDataflowWorklist>;
@@ -116,7 +118,9 @@ class DataflowAnalysis {
   }
 
 protected:
-  Lattice getState(ProgramPoint P) const { return PerPointStates.lookup(P); }
+  Lattice getState(ProgramPoint P) const {
+    return PointToState[P->getID().Value];
+  }
 
   std::optional<Lattice> getInState(const CFGBlock *B) const {
     auto It = InStates.find(B);
@@ -144,12 +148,12 @@ class DataflowAnalysis {
     if constexpr (isForward()) {
       for (const Fact *F : Facts) {
         State = transferFact(State, F);
-        PerPointStates[F] = State;
+        PointToState[F->getID().Value] = State;
       }
     } else {
       for (const Fact *F : llvm::reverse(Facts)) {
         // In backward analysis, capture the state before applying the fact.
-        PerPointStates[F] = State;
+        PointToState[F->getID().Value] = State;
         State = transferFact(State, F);
       }
     }

clang/lib/Analysis/LifetimeSafety/Facts.cpp

@@ -64,8 +64,8 @@ void TestPointFact::dump(llvm::raw_ostream &OS, const LoanManager &,
 
 llvm::StringMap<ProgramPoint> FactManager::getTestPoints() const {
   llvm::StringMap<ProgramPoint> AnnotationToPointMap;
-  for (const CFGBlock *Block : BlockToFactsMap.keys()) {
-    for (const Fact *F : getFacts(Block)) {
+  for (const auto &BlockFacts : BlockToFacts) {
+    for (const Fact *F : BlockFacts) {
       if (const auto *TPF = F->getAs<TestPointFact>()) {
         StringRef PointName = TPF->getAnnotation();
         assert(AnnotationToPointMap.find(PointName) ==
@@ -88,12 +88,9 @@ void FactManager::dump(const CFG &Cfg, AnalysisDeclContext &AC) const {
   // Print blocks in the order as they appear in code for a stable ordering.
   for (const CFGBlock *B : *AC.getAnalysis<PostOrderCFGView>()) {
     llvm::dbgs() << "  Block B" << B->getBlockID() << ":\n";
-    auto It = BlockToFactsMap.find(B);
-    if (It != BlockToFactsMap.end()) {
-      for (const Fact *F : It->second) {
-        llvm::dbgs() << "    ";
-        F->dump(llvm::dbgs(), LoanMgr, OriginMgr);
-      }
+    for (const Fact *F : getFacts(B)) {
+      llvm::dbgs() << "    ";
+      F->dump(llvm::dbgs(), LoanMgr, OriginMgr);
     }
     llvm::dbgs() << "  End of Block\n";
   }

clang/lib/Analysis/LifetimeSafety/LifetimeSafety.cpp

@@ -41,6 +41,7 @@ void LifetimeSafetyAnalysis::run() {
   const CFG &Cfg = *AC.getCFG();
   DEBUG_WITH_TYPE("PrintCFG", Cfg.dump(AC.getASTContext().getLangOpts(),
                                        /*ShowColors=*/true));
+  FactMgr.init(Cfg);
 
   FactsGenerator FactGen(FactMgr, AC);
   FactGen.run();