1 files changed, 9 insertions, 16 deletions
diff --git a/clang/unittests/Analysis/LifetimeSafetyTest.cpp b/clang/unittests/Analysis/LifetimeSafetyTest.cpp
index 0c05184..169b2d2 100644
--- a/clang/unittests/Analysis/LifetimeSafetyTest.cpp
+++ b/clang/unittests/Analysis/LifetimeSafetyTest.cpp
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "clang/Analysis/Analyses/LifetimeSafety/LifetimeSafety.h"
+#include "clang/Analysis/Analyses/LifetimeSafety.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Testing/TestAST.h"
@@ -63,7 +63,7 @@ public:
     Analysis = std::make_unique<LifetimeSafetyAnalysis>(*AnalysisCtx, nullptr);
     Analysis->run();
 
-    AnnotationToPointMap = Analysis->getFactManager().getTestPoints();
+    AnnotationToPointMap = Analysis->getTestPoints();
   }
 
   LifetimeSafetyAnalysis &getAnalysis() { return *Analysis; }
@@ -98,11 +98,10 @@ public:
     auto *VD = findDecl<ValueDecl>(VarName);
     if (!VD)
       return std::nullopt;
-    // This assumes the OriginManager's `get` can find an existing origin.
-    // We might need a `find` method on OriginManager to avoid `getOrCreate`
-    // logic in a const-query context if that becomes an issue.
-    return const_cast<OriginManager &>(Analysis.getFactManager().getOriginMgr())
-        .get(*VD);
+    auto OID = Analysis.getOriginIDForDecl(VD);
+    if (!OID)
+      ADD_FAILURE() << "Origin for '" << VarName << "' not found.";
+    return OID;
   }
 
   std::vector<LoanID> getLoansForVar(llvm::StringRef VarName) {
@@ -111,10 +110,7 @@ public:
       ADD_FAILURE() << "Failed to find VarDecl for '" << VarName << "'";
       return {};
     }
-    std::vector<LoanID> LID;
-    for (const Loan &L : Analysis.getFactManager().getLoanMgr().getLoans())
-      if (L.Path.D == VD)
-        LID.push_back(L.ID);
+    std::vector<LoanID> LID = Analysis.getLoanIDForVar(VD);
     if (LID.empty()) {
       ADD_FAILURE() << "Loan for '" << VarName << "' not found.";
       return {};
@@ -127,7 +123,7 @@ public:
     ProgramPoint PP = Runner.getProgramPoint(Annotation);
     if (!PP)
       return std::nullopt;
-    return Analysis.getLoanPropagation().getLoans(OID, PP);
+    return Analysis.getLoansAtPoint(OID, PP);
   }
 
   std::optional<std::vector<std::pair<OriginID, LivenessKind>>>
@@ -135,10 +131,7 @@ public:
     ProgramPoint PP = Runner.getProgramPoint(Annotation);
     if (!PP)
       return std::nullopt;
-    std::vector<std::pair<OriginID, LivenessKind>> Result;
-    for (auto &[OID, Info] : Analysis.getLiveOrigins().getLiveOriginsAt(PP))
-      Result.push_back({OID, Info.Kind});
-    return Result;
+    return Analysis.getLiveOriginsAtPoint(PP);
   }
 
 private: