1 files changed, 9 insertions, 5 deletions
diff --git a/clang/lib/Analysis/LifetimeSafety/Dataflow.h b/clang/lib/Analysis/LifetimeSafety/Dataflow.h
index 2f7bcb6..de821bb 100644
--- a/clang/lib/Analysis/LifetimeSafety/Dataflow.h
+++ b/clang/lib/Analysis/LifetimeSafety/Dataflow.h
@@ -67,10 +67,10 @@ private:
   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 @@ public:
     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 @@ public:
   }
 
 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 @@ private:
     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);
       }
     }