[LifetimeSafety] Support bidirectional dataflow analysis (#148967)

Generalize the dataflow analysis to support both forward and backward analyses.

Some program analyses would be expressed as backward dataflow problems (like liveness analysis). This change enables the framework to support both forward analyses (like the loan propagation analysis) and backward analyses with the same infrastructure.

Author: usx95

clang/lib/Analysis/LifetimeSafety.cpp

@@ -499,13 +499,16 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
 // ========================================================================= //
 //                         Generic Dataflow Analysis
 // ========================================================================= //
-/// A generic, policy-based driver for forward dataflow analyses. It combines
+
+enum class Direction { Forward, Backward };
+
+/// A generic, policy-based driver for dataflow analyses. It combines
 /// the dataflow runner and the transferer logic into a single class hierarchy.
 ///
 /// The derived class is expected to provide:
 /// - A `Lattice` type.
 /// - `StringRef getAnalysisName() const`
-/// - `Lattice getInitialState();` The initial state at the function entry.
+/// - `Lattice getInitialState();` The initial state of the analysis.
 /// - `Lattice join(Lattice, Lattice);` Merges states from multiple CFG paths.
 /// - `Lattice transfer(Lattice, const FactType&);` Defines how a single
 ///   lifetime-relevant `Fact` transforms the lattice state. Only overloads
@@ -514,18 +517,23 @@ class FactGenerator : public ConstStmtVisitor<FactGenerator> {
 /// \tparam Derived The CRTP derived class that implements the specific
 /// analysis.
 /// \tparam LatticeType The dataflow lattice used by the analysis.
+/// \tparam Dir The direction of the analysis (Forward or Backward).
 /// TODO: Maybe use the dataflow framework! The framework might need changes
 /// to support the current comparison done at block-entry.
-template <typename Derived, typename LatticeType> class DataflowAnalysis {
+template <typename Derived, typename LatticeType, Direction Dir>
+class DataflowAnalysis {
 public:
   using Lattice = LatticeType;
+  using Base = DataflowAnalysis<Derived, LatticeType, Dir>;
 
 private:
   const CFG &Cfg;
   AnalysisDeclContext &AC;
 
-  llvm::DenseMap<const CFGBlock *, Lattice> BlockEntryStates;
-  llvm::DenseMap<const CFGBlock *, Lattice> BlockExitStates;
+  llvm::DenseMap<const CFGBlock *, Lattice> InStates;
+  llvm::DenseMap<const CFGBlock *, Lattice> OutStates;
+
+  static constexpr bool isForward() { return Dir == Direction::Forward; }
 
 protected:
   FactManager &AllFacts;
@@ -539,75 +547,76 @@ template <typename Derived, typename LatticeType> class DataflowAnalysis {
     Derived &D = static_cast<Derived &>(*this);
     llvm::TimeTraceScope Time(D.getAnalysisName());
 
-    ForwardDataflowWorklist Worklist(Cfg, AC);
-    const CFGBlock *Entry = &Cfg.getEntry();
-    BlockEntryStates[Entry] = D.getInitialState();
-    Worklist.enqueueBlock(Entry);
-    llvm::SmallBitVector Visited;
-    Visited.resize(Cfg.getNumBlockIDs() + 1);
-
-    while (const CFGBlock *B = Worklist.dequeue()) {
-      Lattice EntryState = getEntryState(B);
-      Lattice ExitState = transferBlock(B, EntryState);
-      BlockExitStates[B] = ExitState;
-      Visited.set(B->getBlockID());
+    using Worklist =
+        std::conditional_t<Dir == Direction::Forward, ForwardDataflowWorklist,
+                           BackwardDataflowWorklist>;
+    Worklist W(Cfg, AC);
+
+    const CFGBlock *Start = isForward() ? &Cfg.getEntry() : &Cfg.getExit();
+    InStates[Start] = D.getInitialState();
+    W.enqueueBlock(Start);
 
-      for (const CFGBlock *Successor : B->succs()) {
-        Lattice OldSuccEntryState = getEntryState(Successor);
-        Lattice NewSuccEntryState = D.join(OldSuccEntryState, ExitState);
+    llvm::SmallBitVector Visited(Cfg.getNumBlockIDs() + 1);
 
-        // Enqueue the successor if its entry state has changed or if we have
+    while (const CFGBlock *B = W.dequeue()) {
+      Lattice StateIn = getInState(B);
+      Lattice StateOut = transferBlock(B, StateIn);
+      OutStates[B] = StateOut;
+      Visited.set(B->getBlockID());
+      for (const CFGBlock *AdjacentB : isForward() ? B->succs() : B->preds()) {
+        Lattice OldInState = getInState(AdjacentB);
+        Lattice NewInState = D.join(OldInState, StateOut);
+        // Enqueue the adjacent block if its in-state has changed or if we have
         // never visited it.
-        if (!Visited.test(Successor->getBlockID()) ||
-            NewSuccEntryState != OldSuccEntryState) {
-          BlockEntryStates[Successor] = NewSuccEntryState;
-          Worklist.enqueueBlock(Successor);
+        if (!Visited.test(AdjacentB->getBlockID()) ||
+            NewInState != OldInState) {
+          InStates[AdjacentB] = NewInState;
+          W.enqueueBlock(AdjacentB);
         }
       }
     }
   }
 
-  Lattice getEntryState(const CFGBlock *B) const {
-    return BlockEntryStates.lookup(B);
-  }
+  Lattice getInState(const CFGBlock *B) const { return InStates.lookup(B); }
 
-  Lattice getExitState(const CFGBlock *B) const {
-    return BlockExitStates.lookup(B);
-  }
+  Lattice getOutState(const CFGBlock *B) const { return OutStates.lookup(B); }
 
   void dump() const {
     const Derived *D = static_cast<const Derived *>(this);
     llvm::dbgs() << "==========================================\n";
     llvm::dbgs() << D->getAnalysisName() << " results:\n";
     llvm::dbgs() << "==========================================\n";
-    const CFGBlock &B = Cfg.getExit();
-    getExitState(&B).dump(llvm::dbgs());
+    const CFGBlock &B = isForward() ? Cfg.getExit() : Cfg.getEntry();
+    getOutState(&B).dump(llvm::dbgs());
   }
 
-private:
-  /// Computes the exit state of a block by applying all its facts sequentially
-  /// to a given entry state.
+  /// Computes the state at one end of a block by applying all its facts
+  /// sequentially to a given state from the other end.
   /// TODO: We might need to store intermediate states per-fact in the block for
   /// later analysis.
-  Lattice transferBlock(const CFGBlock *Block, Lattice EntryState) {
-    Lattice BlockState = EntryState;
-    for (const Fact *F : AllFacts.getFacts(Block)) {
-      BlockState = transferFact(BlockState, F);
-    }
-    return BlockState;
+  Lattice transferBlock(const CFGBlock *Block, Lattice State) {
+    auto Facts = AllFacts.getFacts(Block);
+    if constexpr (isForward())
+      for (const Fact *F : Facts)
+        State = transferFact(State, F);
+    else
+      for (const Fact *F : llvm::reverse(Facts))
+        State = transferFact(State, F);
+    return State;
   }
 
   Lattice transferFact(Lattice In, const Fact *F) {
-    Derived *d = static_cast<Derived *>(this);
+    assert(F);
+    Derived *D = static_cast<Derived *>(this);
     switch (F->getKind()) {
     case Fact::Kind::Issue:
-      return d->transfer(In, *F->getAs<IssueFact>());
+      return D->transfer(In, *F->getAs<IssueFact>());
     case Fact::Kind::Expire:
-      return d->transfer(In, *F->getAs<ExpireFact>());
+      return D->transfer(In, *F->getAs<ExpireFact>());
     case Fact::Kind::AssignOrigin:
-      return d->transfer(In, *F->getAs<AssignOriginFact>());
+      return D->transfer(In, *F->getAs<AssignOriginFact>());
     case Fact::Kind::ReturnOfOrigin:
-      return d->transfer(In, *F->getAs<ReturnOfOriginFact>());
+      return D->transfer(In, *F->getAs<ReturnOfOriginFact>());
     }
     llvm_unreachable("Unknown fact kind");
   }
@@ -715,7 +724,8 @@ struct LoanPropagationLattice {
 
 /// The analysis that tracks which loans belong to which origins.
 class LoanPropagationAnalysis
-    : public DataflowAnalysis<LoanPropagationAnalysis, LoanPropagationLattice> {
+    : public DataflowAnalysis<LoanPropagationAnalysis, LoanPropagationLattice,
+                              Direction::Forward> {
 
   LifetimeFactory &Factory;
 
@@ -724,7 +734,7 @@ class LoanPropagationAnalysis
                           LifetimeFactory &Factory)
       : DataflowAnalysis(C, AC, F), Factory(Factory) {}
 
-  using DataflowAnalysis<LoanPropagationAnalysis, Lattice>::transfer;
+  using Base::transfer;
 
   StringRef getAnalysisName() const { return "LoanPropagation"; }