[clang] [libc] [llvm] [Analysis] Add interprocedural heap provenance analysis and baremetal fortification fallback (PR #205909)

Schrodinger ZHU Yifan via cfe-commits cfe-commits at lists.llvm.org
Mon Jun 29 11:56:09 PDT 2026


https://github.com/SchrodingerZhu updated https://github.com/llvm/llvm-project/pull/205909

>From 00edecd208f2b7ed0e884f34661927d374e32863 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 12 Jun 2026 09:16:48 -0700
Subject: [PATCH 01/18] fix 203411. This PR addresses the problem that EAGAIN
 may be treated as timeout in mutex and rwlock. Two changes are applied: 1.
 timeout site always explicitly check for timeout now to make the logic more
 robust; 2. the futex wait now discards the error of EAGAIN/EWOULDBLOCK and
 returns 0; We don't distinguish waking up from signal and waking up from
 mismatch for the following 3 reasons: - We have userspace guard to avoid
 futex syscall if we already know value would match, it seems awkward to make
 that check returns error, as we may wake up and loop back to the check, where
 signal is consumed but we still return error....; - futex syscall can
 spuriously wake up anyway, there is no way to tell whether the signal is
 "indeed" consumed; - other platforms like darwin does not distinguish these
 states either. Assisted-by: Gemini powered automation tools
 (human-in-the-loop).

TAG=agy
CONV=5e48f6bc-22e7-4614-a572-5ed16ff5f9a9
---
 libc/src/__support/threads/linux/futex_utils.h               | 5 +++++
 libc/src/__support/threads/raw_mutex.h                       | 4 ++--
 libc/src/__support/threads/raw_rwlock.h                      | 5 +++--
 .../integration/src/__support/threads/futex_requeue_test.cpp | 2 +-
 4 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/libc/src/__support/threads/linux/futex_utils.h b/libc/src/__support/threads/linux/futex_utils.h
index ff6b5d526a3c1..aa3499dd8a7b7 100644
--- a/libc/src/__support/threads/linux/futex_utils.h
+++ b/libc/src/__support/threads/linux/futex_utils.h
@@ -66,6 +66,11 @@ class Futex : public cpp::Atomic<FutexWordType> {
       if (ret == -EINTR)
         continue;
 
+      // EAGAIN and EWOULDBLOCK will be treated as a normal finish, as the value
+      // has changed.
+      if (ret == -EAGAIN || ret == -EWOULDBLOCK)
+        return 0;
+
       if (ret < 0)
         return cpp::unexpected(-ret);
       return ret;
diff --git a/libc/src/__support/threads/raw_mutex.h b/libc/src/__support/threads/raw_mutex.h
index e216e0e95b77c..262899db24009 100644
--- a/libc/src/__support/threads/raw_mutex.h
+++ b/libc/src/__support/threads/raw_mutex.h
@@ -83,8 +83,8 @@ class RawMutex {
           futex.exchange(IN_CONTENTION, cpp::MemoryOrder::ACQUIRE) == UNLOCKED)
         return true;
       // Contention persists. Park the thread and wait for further notification.
-      if (!futex.wait(IN_CONTENTION, timeout, is_pshared).has_value() &&
-          timeout.has_value())
+      if (ErrorOr<int> res = futex.wait(IN_CONTENTION, timeout, is_pshared);
+          !res.has_value() && res.error() == ETIMEDOUT)
         return false;
 
       // Continue to spin after waking up.
diff --git a/libc/src/__support/threads/raw_rwlock.h b/libc/src/__support/threads/raw_rwlock.h
index 987cfa4f82152..c0140b213a088 100644
--- a/libc/src/__support/threads/raw_rwlock.h
+++ b/libc/src/__support/threads/raw_rwlock.h
@@ -400,8 +400,9 @@ class RawRwLock {
       // reached.
       bool timeout_flag = false;
       if (!old.can_acquire<role>(get_preference())) {
-        auto wait_result = queue.wait<role>(serial_number, timeout, is_pshared);
-        timeout_flag = (!wait_result.has_value() && timeout.has_value());
+        ErrorOr<int> wait_result = queue.wait<role>(serial_number, timeout, is_pshared);
+        timeout_flag =
+            (!wait_result.has_value() && wait_result.error() == ETIMEDOUT);
       }
 
       // Phase 7: unregister ourselves as a pending reader/writer.
diff --git a/libc/test/integration/src/__support/threads/futex_requeue_test.cpp b/libc/test/integration/src/__support/threads/futex_requeue_test.cpp
index b6ce58d2d1da8..ab13d36a79922 100644
--- a/libc/test/integration/src/__support/threads/futex_requeue_test.cpp
+++ b/libc/test/integration/src/__support/threads/futex_requeue_test.cpp
@@ -29,7 +29,7 @@ struct SharedState {
 void wait_until_zero(LIBC_NAMESPACE::Futex &futex) {
   while (futex.load() != 0) {
     auto wait_result = futex.wait(1, LIBC_NAMESPACE::cpp::nullopt, false);
-    ASSERT_TRUE(wait_result.has_value() || wait_result.error() == EAGAIN);
+    ASSERT_TRUE(wait_result.has_value());
   }
 }
 

>From 4a93f639f6cf00156a946b32e527922d424f6d03 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 12 Jun 2026 09:20:04 -0700
Subject: [PATCH 02/18] format

---
 libc/src/__support/threads/raw_rwlock.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/libc/src/__support/threads/raw_rwlock.h b/libc/src/__support/threads/raw_rwlock.h
index c0140b213a088..b08d8e4c0768a 100644
--- a/libc/src/__support/threads/raw_rwlock.h
+++ b/libc/src/__support/threads/raw_rwlock.h
@@ -400,7 +400,8 @@ class RawRwLock {
       // reached.
       bool timeout_flag = false;
       if (!old.can_acquire<role>(get_preference())) {
-        ErrorOr<int> wait_result = queue.wait<role>(serial_number, timeout, is_pshared);
+        ErrorOr<int> wait_result =
+            queue.wait<role>(serial_number, timeout, is_pshared);
         timeout_flag =
             (!wait_result.has_value() && wait_result.error() == ETIMEDOUT);
       }

>From 8b78b46cb12d66e7f87b439f7936ab9eb2d987a6 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 09:43:01 -0700
Subject: [PATCH 03/18] [Analysis][HeapProvenance] add a pass to analyze
 pointer's heap provenance

TAG=agy

CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 .../llvm/Analysis/HeapProvenanceAnalysis.h    | 114 ++++
 llvm/lib/Analysis/CMakeLists.txt              |   1 +
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 583 ++++++++++++++++++
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 llvm/lib/Passes/PassRegistry.def              |   3 +
 5 files changed, 702 insertions(+)
 create mode 100644 llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
 create mode 100644 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp

diff --git a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
new file mode 100644
index 0000000000000..f31f2312ec69c
--- /dev/null
+++ b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
@@ -0,0 +1,114 @@
+#ifndef LLVM_ANALYSIS_HEAPPROVENANCEANALYSIS_H
+#define LLVM_ANALYSIS_HEAPPROVENANCEANALYSIS_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/Value.h"
+#include <string>
+#include <vector>
+
+namespace llvm {
+
+class HeapProvenanceAnalysisResult {
+public:
+  struct ProvenanceInfo {
+    enum Kind {
+      Uninit = 0,
+      HeapChunkPtr,
+      RecoverableHeapChunkPtr,
+      Unknown
+    } State = Uninit;
+
+    enum Direction {
+      None = 0,
+      Forward = 1,
+      Backward = 2,
+      Both = 3
+    } Dir = None;
+
+    int64_t ConstOffset = 0;
+    std::vector<std::string> SymOffsets;
+    std::string CustomExpr;
+
+    bool isUninit() const { return State == Uninit; }
+    bool isValid() const {
+      return State == HeapChunkPtr || State == RecoverableHeapChunkPtr;
+    }
+
+    std::string getDirectionStr() const {
+      if (Dir == Forward)
+        return " [forward: from alloc]";
+      if (Dir == Backward)
+        return " [backward: into dealloc]";
+      if (Dir == Both)
+        return " [both: alloc & dealloc]";
+      return "";
+    }
+
+    std::string getExpr() const {
+      if (State == Uninit)
+        return "Uninit";
+      if (State == Unknown)
+        return "Unknown";
+      if (!CustomExpr.empty())
+        return CustomExpr;
+      if (State == HeapChunkPtr && ConstOffset == 0 && SymOffsets.empty())
+        return "head";
+
+      std::string Res = "head";
+      if (ConstOffset > 0)
+        Res += " + " + std::to_string(ConstOffset);
+      else if (ConstOffset < 0)
+        Res += " - " + std::to_string(-ConstOffset);
+
+      for (const auto &S : SymOffsets) {
+        if (!S.empty() && S[0] == '-')
+          Res += " " + S;
+        else
+          Res += " + " + S;
+      }
+      return Res;
+    }
+  };
+
+private:
+  DenseMap<const Value *, ProvenanceInfo> ValueMap;
+
+public:
+  void setInfo(const Value *V, const ProvenanceInfo &Info) {
+    ValueMap[V] = Info;
+  }
+  const ProvenanceInfo &getInfo(const Value *V) const {
+    static ProvenanceInfo Empty;
+    auto It = ValueMap.find(V);
+    return It != ValueMap.end() ? It->second : Empty;
+  }
+  DenseMap<const Value *, ProvenanceInfo> &getMap() { return ValueMap; }
+  const DenseMap<const Value *, ProvenanceInfo> &getMap() const {
+    return ValueMap;
+  }
+};
+
+class HeapProvenanceAnalysis
+    : public AnalysisInfoMixin<HeapProvenanceAnalysis> {
+  friend AnalysisInfoMixin<HeapProvenanceAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  using Result = HeapProvenanceAnalysisResult;
+  Result run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+class HeapProvenancePrinterPass
+    : public PassInfoMixin<HeapProvenancePrinterPass> {
+  raw_ostream &OS;
+
+public:
+  explicit HeapProvenancePrinterPass(raw_ostream &OS) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_ANALYSIS_HEAPPROVENANCEANALYSIS_H
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index f3586c66cb056..6cc8b44e84e84 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -146,6 +146,7 @@ add_llvm_component_library(LLVMAnalysis
   ScalarEvolutionDivision.cpp
   ScalarEvolutionNormalization.cpp
   StaticDataProfileInfo.cpp
+  HeapProvenanceAnalysis.cpp
   StackLifetime.cpp
   StackSafetyAnalysis.cpp
   StructuralHash.cpp
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
new file mode 100644
index 0000000000000..5f6c2575fed33
--- /dev/null
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -0,0 +1,583 @@
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/GetElementPtrTypeIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Operator.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace llvm;
+
+AnalysisKey HeapProvenanceAnalysis::Key;
+
+static bool isAllocFunc(StringRef Name) {
+  return Name == "malloc" || Name == "calloc" || Name == "realloc" ||
+         Name == "aligned_alloc" || Name == "strdup" || Name == "strndup" ||
+         Name == "g_malloc" || Name == "g_malloc0" || Name == "g_realloc" ||
+         Name == "g_strdup" || Name.starts_with("_Zn");
+}
+
+static bool isFreeFunc(StringRef Name) {
+  return Name == "free" || Name == "realloc" || Name == "cfree" ||
+         Name == "g_free" || Name.starts_with("_Zd");
+}
+
+static std::string getValueOperandName(const Value *V) {
+  std::string Str;
+  raw_string_ostream OS(Str);
+  V->printAsOperand(OS, false);
+  return Str;
+}
+
+static void addSymOffset(std::vector<std::string> &Syms, const std::string &S) {
+  StringRef SR(S);
+  std::string NegS = (SR.size() > 3 && SR.starts_with("-(") && SR.ends_with(")"))
+                         ? S.substr(2, S.size() - 3)
+                         : "-(" + S + ")";
+  for (auto It = Syms.begin(); It != Syms.end(); ++It) {
+    if (*It == NegS) {
+      Syms.erase(It);
+      return;
+    }
+  }
+  Syms.push_back(S);
+}
+
+static void extractGEPOffsets(const GEPOperator *GEP, const DataLayout &DL,
+                              int64_t &ConstOff,
+                              std::vector<std::string> &SymOffs) {
+  APInt APIntOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+  if (GEP->accumulateConstantOffset(DL, APIntOffset)) {
+    ConstOff += APIntOffset.getSExtValue();
+    return;
+  }
+  for (gep_type_iterator GTI = gep_type_begin(GEP), GTE = gep_type_end(GEP);
+       GTI != GTE; ++GTI) {
+    Value *Op = GTI.getOperand();
+    if (StructType *STy = GTI.getStructTypeOrNull()) {
+      ConstantInt *OpC = cast<ConstantInt>(Op);
+      const StructLayout *SL = DL.getStructLayout(STy);
+      ConstOff += SL->getElementOffset(OpC->getZExtValue());
+    } else {
+      TypeSize Stride = GTI.getSequentialElementStride(DL);
+      int64_t StrideVal = Stride.getKnownMinValue();
+      if (ConstantInt *OpC = dyn_cast<ConstantInt>(Op)) {
+        ConstOff += OpC->getSExtValue() * StrideVal;
+      } else {
+        std::string OpName = getValueOperandName(Op);
+        if (StrideVal == 1)
+          addSymOffset(SymOffs, OpName);
+        else
+          addSymOffset(SymOffs, OpName + " * " + std::to_string(StrideVal));
+      }
+    }
+  }
+}
+
+static HeapProvenanceAnalysisResult::ProvenanceInfo
+addOffset(const HeapProvenanceAnalysisResult::ProvenanceInfo &BaseInfo,
+          int64_t C, const std::vector<std::string> &Syms) {
+  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
+  if (!BaseInfo.isValid())
+    return BaseInfo;
+  ProvenanceInfo Res = BaseInfo;
+  Res.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+  if (!Res.CustomExpr.empty()) {
+    std::string OffStr;
+    if (C > 0)
+      OffStr += " + " + std::to_string(C);
+    else if (C < 0)
+      OffStr += " - " + std::to_string(-C);
+    for (const auto &S : Syms)
+      OffStr += " + " + S;
+    Res.CustomExpr = "(" + Res.CustomExpr + ")" + OffStr;
+    return Res;
+  }
+  Res.ConstOffset += C;
+  for (const auto &S : Syms)
+    addSymOffset(Res.SymOffsets, S);
+  if (Res.ConstOffset == 0 && Res.SymOffsets.empty())
+    Res.State = ProvenanceInfo::HeapChunkPtr;
+  return Res;
+}
+
+static HeapProvenanceAnalysisResult::ProvenanceInfo
+subOffset(const HeapProvenanceAnalysisResult::ProvenanceInfo &IInfo, int64_t C,
+          const std::vector<std::string> &Syms) {
+  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
+  if (!IInfo.isValid())
+    return IInfo;
+  ProvenanceInfo Res = IInfo;
+  Res.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+  if (!Res.CustomExpr.empty()) {
+    std::string OffStr;
+    if (C > 0)
+      OffStr += " - " + std::to_string(C);
+    else if (C < 0)
+      OffStr += " + " + std::to_string(-C);
+    for (const auto &S : Syms)
+      OffStr += " - (" + S + ")";
+    Res.CustomExpr = "(" + Res.CustomExpr + ")" + OffStr;
+    return Res;
+  }
+  Res.ConstOffset -= C;
+  for (const auto &S : Syms)
+    addSymOffset(Res.SymOffsets, "-(" + S + ")");
+  if (Res.ConstOffset == 0 && Res.SymOffsets.empty())
+    Res.State = ProvenanceInfo::HeapChunkPtr;
+  return Res;
+}
+
+static bool mergeInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
+                      const HeapProvenanceAnalysisResult::ProvenanceInfo &Src) {
+  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
+  if (Src.State == ProvenanceInfo::Uninit)
+    return false;
+
+  auto MergedDir = static_cast<ProvenanceInfo::Direction>(Dest.Dir | Src.Dir);
+
+  if (Dest.State == ProvenanceInfo::Uninit) {
+    Dest = Src;
+    Dest.Dir = MergedDir;
+    return true;
+  }
+  if (Dest.State == ProvenanceInfo::Unknown) {
+    if (Dest.Dir != MergedDir) {
+      Dest.Dir = MergedDir;
+      return true;
+    }
+    return false;
+  }
+  if (Src.State == ProvenanceInfo::Unknown) {
+    if (Dest.Dir != MergedDir) {
+      Dest.Dir = MergedDir;
+      return true;
+    }
+    return false;
+  }
+
+  bool Changed = false;
+  if (Dest.Dir != MergedDir) {
+    Dest.Dir = MergedDir;
+    Changed = true;
+  }
+
+  if (Dest.getExpr() == Src.getExpr()) {
+    if (Src.State == ProvenanceInfo::RecoverableHeapChunkPtr &&
+        Dest.State != ProvenanceInfo::RecoverableHeapChunkPtr) {
+      Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+      Changed = true;
+    }
+    return Changed;
+  }
+
+  if (Dest.CustomExpr == "head + dynamic_offset")
+    return Changed;
+
+  if (!Dest.CustomExpr.empty()) {
+    Dest.CustomExpr = "head + dynamic_offset";
+    Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+    return true;
+  }
+
+  Dest.CustomExpr = "PHI(" + Dest.getExpr() + ", " + Src.getExpr() + ")";
+  Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+  return true;
+}
+
+static bool updateInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
+                       const HeapProvenanceAnalysisResult::ProvenanceInfo &NewVal) {
+  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
+  if (Dest.State == NewVal.State && Dest.Dir == NewVal.Dir &&
+      Dest.ConstOffset == NewVal.ConstOffset &&
+      Dest.SymOffsets == NewVal.SymOffsets &&
+      Dest.CustomExpr == NewVal.CustomExpr)
+    return false;
+
+  if (!Dest.CustomExpr.empty() && !NewVal.CustomExpr.empty() &&
+      Dest.CustomExpr != NewVal.CustomExpr) {
+    std::string Widened = "head + dynamic_offset";
+    if (Dest.CustomExpr != Widened) {
+      Dest = NewVal;
+      Dest.CustomExpr = Widened;
+      return true;
+    }
+    return false;
+  }
+
+  Dest = NewVal;
+  return true;
+}
+
+static HeapProvenanceAnalysisResult::ProvenanceInfo
+keepForwardOnly(const HeapProvenanceAnalysisResult::ProvenanceInfo &Info) {
+  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
+  auto Res = Info;
+  Res.Dir = static_cast<ProvenanceInfo::Direction>(Info.Dir &
+                                                     ProvenanceInfo::Forward);
+  if (Res.Dir == ProvenanceInfo::None)
+    Res.State = ProvenanceInfo::Uninit;
+  return Res;
+}
+
+HeapProvenanceAnalysis::Result
+HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  Result Res;
+  const DataLayout &DL = M.getDataLayout();
+
+  for (Function &F : M) {
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        if (auto *CB = dyn_cast<CallBase>(&I)) {
+          Function *Callee = CB->getCalledFunction();
+          if (!Callee) {
+            Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
+            Callee = dyn_cast<Function>(CalledVal);
+          }
+          if (Callee) {
+            StringRef Name = Callee->getName();
+            if (isAllocFunc(Name)) {
+              Result::ProvenanceInfo Info;
+              Info.State = Result::ProvenanceInfo::HeapChunkPtr;
+              Info.Dir = Result::ProvenanceInfo::Forward;
+              mergeInfo(Res.getMap()[&I], Info);
+            }
+            if (isFreeFunc(Name)) {
+              if (CB->arg_size() > 0) {
+                Value *ArgVal = CB->getArgOperand(0);
+                Result::ProvenanceInfo Info;
+                Info.State = Result::ProvenanceInfo::HeapChunkPtr;
+                Info.Dir = Result::ProvenanceInfo::Backward;
+                mergeInfo(Res.getMap()[ArgVal], Info);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  bool Changed = true;
+  int MaxIters = 50;
+  while (Changed && MaxIters-- > 0) {
+    Changed = false;
+
+    // Forward propagation
+    for (Function &F : M) {
+      for (BasicBlock &BB : F) {
+        for (Instruction &I : BB) {
+          if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
+            Value *Base = GEP->getPointerOperand();
+            auto BaseInfo = keepForwardOnly(Res.getInfo(Base));
+            if (BaseInfo.isValid()) {
+              auto NewInfo = BaseInfo;
+              NewInfo.State = Result::ProvenanceInfo::RecoverableHeapChunkPtr;
+              int64_t ConstOff = 0;
+              std::vector<std::string> SymOffs;
+              extractGEPOffsets(GEP, DL, ConstOff, SymOffs);
+              NewInfo = addOffset(NewInfo, ConstOff, SymOffs);
+              if (mergeInfo(Res.getMap()[&I], NewInfo))
+                Changed = true;
+            }
+          } else if (auto *BC = dyn_cast<BitCastOperator>(&I)) {
+            Value *Op = BC->getOperand(0);
+            auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+            if (OpInfo.isValid()) {
+              if (mergeInfo(Res.getMap()[&I], OpInfo))
+                Changed = true;
+            }
+          } else if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(&I)) {
+            Value *Op = ASC->getOperand(0);
+            auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+            if (OpInfo.isValid()) {
+              if (mergeInfo(Res.getMap()[&I], OpInfo))
+                Changed = true;
+            }
+          } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
+            Value *Op = ITP->getOperand(0);
+            if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
+              if (BO->getOpcode() == Instruction::Add ||
+                  BO->getOpcode() == Instruction::Sub) {
+                Value *LHS = BO->getOperand(0);
+                Value *RHS = BO->getOperand(1);
+                auto *PTI = dyn_cast<PtrToIntOperator>(LHS);
+                if (!PTI) {
+                  PTI = dyn_cast<PtrToIntOperator>(RHS);
+                  if (PTI && BO->getOpcode() == Instruction::Add)
+                    std::swap(LHS, RHS);
+                  else
+                    PTI = nullptr;
+                }
+                if (PTI) {
+                  Value *BasePtr = PTI->getOperand(0);
+                  auto BaseInfo = keepForwardOnly(Res.getInfo(BasePtr));
+                  if (BaseInfo.isValid()) {
+                    int64_t ConstOff = 0;
+                    std::vector<std::string> SymOffs;
+                    if (auto *CInt = dyn_cast<ConstantInt>(RHS)) {
+                      ConstOff = CInt->getSExtValue();
+                    } else {
+                      SymOffs.push_back(getValueOperandName(RHS));
+                    }
+                    if (BO->getOpcode() == Instruction::Sub) {
+                      ConstOff = -ConstOff;
+                      for (auto &S : SymOffs)
+                        S = "-(" + S + ")";
+                    }
+                    auto NewInfo = addOffset(BaseInfo, ConstOff, SymOffs);
+                    if (mergeInfo(Res.getMap()[&I], NewInfo))
+                      Changed = true;
+                  }
+                }
+              }
+            } else if (auto *PTI = dyn_cast<PtrToIntOperator>(Op)) {
+              auto BaseInfo = keepForwardOnly(Res.getInfo(PTI->getOperand(0)));
+              if (BaseInfo.isValid()) {
+                if (mergeInfo(Res.getMap()[&I], BaseInfo))
+                  Changed = true;
+              }
+            }
+          } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
+            Result::ProvenanceInfo Temp;
+            for (Value *InV : PHI->incoming_values()) {
+              if (isa<ConstantPointerNull>(InV) || isa<UndefValue>(InV))
+                continue;
+              auto InInfo = keepForwardOnly(Res.getInfo(InV));
+              if (InInfo.isValid())
+                mergeInfo(Temp, InInfo);
+            }
+            if (Temp.isValid() && updateInfo(Res.getMap()[&I], Temp))
+              Changed = true;
+          } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
+            Result::ProvenanceInfo Temp;
+            for (Value *Op : {Sel->getTrueValue(), Sel->getFalseValue()}) {
+              if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op))
+                continue;
+              auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+              if (OpInfo.isValid())
+                mergeInfo(Temp, OpInfo);
+            }
+            if (Temp.isValid() && updateInfo(Res.getMap()[&I], Temp))
+              Changed = true;
+          } else if (auto *CB = dyn_cast<CallBase>(&I)) {
+            Function *Callee = CB->getCalledFunction();
+            if (!Callee) {
+              Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
+              Callee = dyn_cast<Function>(CalledVal);
+            }
+            if (Callee && !Callee->isDeclaration()) {
+              unsigned ArgIdx = 0;
+              for (Argument &Arg : Callee->args()) {
+                if (ArgIdx < CB->arg_size()) {
+                  Value *ArgVal = CB->getArgOperand(ArgIdx);
+                  auto ArgInfo = keepForwardOnly(Res.getInfo(ArgVal));
+                  if (ArgInfo.isValid()) {
+                    if (mergeInfo(Res.getMap()[&Arg], ArgInfo))
+                      Changed = true;
+                  }
+                }
+                ArgIdx++;
+              }
+              if (I.getType()->isPointerTy()) {
+                for (BasicBlock &CalleeBB : *Callee) {
+                  if (auto *RI =
+                          dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
+                    Value *RetVal = RI->getReturnValue();
+                    if (RetVal) {
+                      auto RetInfo = keepForwardOnly(Res.getInfo(RetVal));
+                      if (RetInfo.isValid()) {
+                        if (mergeInfo(Res.getMap()[&I], RetInfo))
+                          Changed = true;
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+
+    // Backward propagation
+    for (Function &F : M) {
+      for (BasicBlock &BB : F) {
+        for (Instruction &I : BB) {
+          auto Info = Res.getInfo(&I);
+          if (!Info.isValid() ||
+              (Info.Dir & Result::ProvenanceInfo::Backward) == 0)
+            continue;
+
+          if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
+            Value *Base = GEP->getPointerOperand();
+            int64_t ConstOff = 0;
+            std::vector<std::string> SymOffs;
+            extractGEPOffsets(GEP, DL, ConstOff, SymOffs);
+            auto BaseNew = subOffset(Info, ConstOff, SymOffs);
+            BaseNew.Dir = Result::ProvenanceInfo::Backward;
+            if (mergeInfo(Res.getMap()[Base], BaseNew))
+              Changed = true;
+          } else if (auto *BC = dyn_cast<BitCastOperator>(&I)) {
+            Value *Op = BC->getOperand(0);
+            auto OpNew = Info;
+            OpNew.Dir = Result::ProvenanceInfo::Backward;
+            if (mergeInfo(Res.getMap()[Op], OpNew))
+              Changed = true;
+          } else if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(&I)) {
+            Value *Op = ASC->getOperand(0);
+            auto OpNew = Info;
+            OpNew.Dir = Result::ProvenanceInfo::Backward;
+            if (mergeInfo(Res.getMap()[Op], OpNew))
+              Changed = true;
+          } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
+            Value *Op = ITP->getOperand(0);
+            if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
+              if (BO->getOpcode() == Instruction::Add ||
+                  BO->getOpcode() == Instruction::Sub) {
+                Value *LHS = BO->getOperand(0);
+                Value *RHS = BO->getOperand(1);
+                auto *PTI = dyn_cast<PtrToIntOperator>(LHS);
+                if (!PTI) {
+                  PTI = dyn_cast<PtrToIntOperator>(RHS);
+                  if (PTI && BO->getOpcode() == Instruction::Add)
+                    std::swap(LHS, RHS);
+                  else
+                    PTI = nullptr;
+                }
+                if (PTI) {
+                  Value *BasePtr = PTI->getOperand(0);
+                  int64_t ConstOff = 0;
+                  std::vector<std::string> SymOffs;
+                  if (auto *CInt = dyn_cast<ConstantInt>(RHS)) {
+                    ConstOff = CInt->getSExtValue();
+                  } else {
+                    SymOffs.push_back(getValueOperandName(RHS));
+                  }
+                  if (BO->getOpcode() == Instruction::Sub) {
+                    ConstOff = -ConstOff;
+                    for (auto &S : SymOffs)
+                      S = "-(" + S + ")";
+                  }
+                  auto BaseNew = subOffset(Info, ConstOff, SymOffs);
+                  BaseNew.Dir = Result::ProvenanceInfo::Backward;
+                  if (mergeInfo(Res.getMap()[BasePtr], BaseNew))
+                    Changed = true;
+                }
+              }
+            } else if (auto *PTI = dyn_cast<PtrToIntOperator>(Op)) {
+              Value *BasePtr = PTI->getOperand(0);
+              auto BaseNew = Info;
+              BaseNew.Dir = Result::ProvenanceInfo::Backward;
+              if (mergeInfo(Res.getMap()[BasePtr], BaseNew))
+                Changed = true;
+            }
+          } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
+            for (Value *InV : PHI->incoming_values()) {
+              if (isa<ConstantPointerNull>(InV) || isa<UndefValue>(InV))
+                continue;
+              auto InNew = Info;
+              InNew.Dir = Result::ProvenanceInfo::Backward;
+              if (mergeInfo(Res.getMap()[InV], InNew))
+                Changed = true;
+            }
+          } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
+            for (Value *Op : {Sel->getTrueValue(), Sel->getFalseValue()}) {
+              if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op))
+                continue;
+              auto OpNew = Info;
+              OpNew.Dir = Result::ProvenanceInfo::Backward;
+              if (mergeInfo(Res.getMap()[Op], OpNew))
+                Changed = true;
+            }
+          } else if (auto *CB = dyn_cast<CallBase>(&I)) {
+            Function *Callee = CB->getCalledFunction();
+            if (!Callee) {
+              Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
+              Callee = dyn_cast<Function>(CalledVal);
+            }
+            if (Callee && !Callee->isDeclaration()) {
+              if (I.getType()->isPointerTy()) {
+                for (BasicBlock &CalleeBB : *Callee) {
+                  if (auto *RI =
+                          dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
+                    Value *RetVal = RI->getReturnValue();
+                    if (RetVal) {
+                      auto RetNew = Info;
+                      RetNew.Dir = Result::ProvenanceInfo::Backward;
+                      if (mergeInfo(Res.getMap()[RetVal], RetNew))
+                        Changed = true;
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+
+      for (Argument &Arg : F.args()) {
+        auto ArgInfo = Res.getInfo(&Arg);
+        if (!ArgInfo.isValid() ||
+            (ArgInfo.Dir & Result::ProvenanceInfo::Backward) == 0)
+          continue;
+        for (User *U : F.users()) {
+          if (auto *CB = dyn_cast<CallBase>(U)) {
+            if (CB->getCalledOperand()->stripPointerCasts() == &F ||
+                CB->getCalledFunction() == &F) {
+              unsigned ArgIdx = Arg.getArgNo();
+              if (ArgIdx < CB->arg_size()) {
+                Value *CallArg = CB->getArgOperand(ArgIdx);
+                auto CallArgNew = ArgInfo;
+                CallArgNew.Dir = Result::ProvenanceInfo::Backward;
+                if (mergeInfo(Res.getMap()[CallArg], CallArgNew))
+                  Changed = true;
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  return Res;
+}
+
+PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
+                                                 ModuleAnalysisManager &MAM) {
+  auto &Result = MAM.getResult<HeapProvenanceAnalysis>(M);
+  for (Function &F : M) {
+    if (F.isDeclaration())
+      continue;
+    OS << "Printing analysis 'Heap Provenance Analysis' for function '"
+       << F.getName() << "':\n";
+    for (Argument &Arg : F.args()) {
+      auto Info = Result.getInfo(&Arg);
+      if (Info.isValid()) {
+        OS << "  argument ";
+        Arg.printAsOperand(OS, false);
+        OS << ": "
+           << (Info.State ==
+                       HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr
+                   ? "HeapChunkPtr"
+                   : "RecoverableHeapChunkPtr")
+           << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
+      }
+    }
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        auto Info = Result.getInfo(&I);
+        if (Info.isValid()) {
+          OS << "  ";
+          I.printAsOperand(OS, false);
+          OS << ": "
+             << (Info.State ==
+                         HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr
+                     ? "HeapChunkPtr"
+                     : "RecoverableHeapChunkPtr")
+             << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
+        }
+      }
+    }
+  }
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 0f288a2124b7b..125eb5f97cda7 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -72,6 +72,7 @@
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
 #include "llvm/Analysis/ScalarEvolutionDivision.h"
 #include "llvm/Analysis/ScopedNoAliasAA.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/StackLifetime.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/StructuralHash.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 8f70c7cefa408..2962394549056 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -37,6 +37,7 @@ MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
 MODULE_ANALYSIS("runtime-libcall-info", RuntimeLibraryAnalysis())
+MODULE_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
@@ -144,6 +145,8 @@ MODULE_PASS("print-lcg-dot", LazyCallGraphDOTPrinterPass(errs()))
 MODULE_PASS("print-must-be-executed-contexts",
             MustBeExecutedContextPrinterPass(errs()))
 MODULE_PASS("print-profile-summary", ProfileSummaryPrinterPass(errs()))
+MODULE_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
+MODULE_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 MODULE_PASS("print-stack-safety", StackSafetyGlobalPrinterPass(errs()))
 MODULE_PASS("print<dxil-metadata>", DXILMetadataAnalysisPrinterPass(errs()))
 MODULE_PASS("print<dxil-resources>", DXILResourcePrinterPass(errs()))

>From 1ae91a4dd5d375b00bd64032950b9a5455b61ea2 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 11:29:48 -0700
Subject: [PATCH 04/18] [Analysis] add a fallback fortification method for
 baremetal heap

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 .../llvm/Analysis/HeapProvenanceAnalysis.h    |  1 +
 llvm/include/llvm/Analysis/MemoryBuiltins.h   |  1 +
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 36 +++----
 llvm/lib/Analysis/MemoryBuiltins.cpp          | 96 +++++++++++++++++++
 4 files changed, 113 insertions(+), 21 deletions(-)

diff --git a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
index f31f2312ec69c..033c1cde9b3f8 100644
--- a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
+++ b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
@@ -97,6 +97,7 @@ class HeapProvenanceAnalysis
 
 public:
   using Result = HeapProvenanceAnalysisResult;
+  static Result analyzeModule(Module &M);
   Result run(Module &M, ModuleAnalysisManager &MAM);
 };
 
diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 10a31973be7fa..5d37b9c0d0a87 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -357,6 +357,7 @@ class ObjectSizeOffsetEvaluator
   SmallPtrSet<Instruction *, 8> InsertedInstructions;
 
   SizeOffsetValue compute_(Value *V);
+  bool computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result);
 
 public:
   LLVM_ABI ObjectSizeOffsetEvaluator(const DataLayout &DL,
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 5f6c2575fed33..8711fe8957c93 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -187,7 +187,6 @@ static bool mergeInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
 
 static bool updateInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
                        const HeapProvenanceAnalysisResult::ProvenanceInfo &NewVal) {
-  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
   if (Dest.State == NewVal.State && Dest.Dir == NewVal.Dir &&
       Dest.ConstOffset == NewVal.ConstOffset &&
       Dest.SymOffsets == NewVal.SymOffsets &&
@@ -209,17 +208,6 @@ static bool updateInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
   return true;
 }
 
-static HeapProvenanceAnalysisResult::ProvenanceInfo
-keepForwardOnly(const HeapProvenanceAnalysisResult::ProvenanceInfo &Info) {
-  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
-  auto Res = Info;
-  Res.Dir = static_cast<ProvenanceInfo::Direction>(Info.Dir &
-                                                     ProvenanceInfo::Forward);
-  if (Res.Dir == ProvenanceInfo::None)
-    Res.State = ProvenanceInfo::Uninit;
-  return Res;
-}
-
 HeapProvenanceAnalysis::Result
 HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
   Result Res;
@@ -268,7 +256,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
         for (Instruction &I : BB) {
           if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
             Value *Base = GEP->getPointerOperand();
-            auto BaseInfo = keepForwardOnly(Res.getInfo(Base));
+            auto BaseInfo = Res.getInfo(Base);
             if (BaseInfo.isValid()) {
               auto NewInfo = BaseInfo;
               NewInfo.State = Result::ProvenanceInfo::RecoverableHeapChunkPtr;
@@ -281,14 +269,14 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
             }
           } else if (auto *BC = dyn_cast<BitCastOperator>(&I)) {
             Value *Op = BC->getOperand(0);
-            auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+            auto OpInfo = Res.getInfo(Op);
             if (OpInfo.isValid()) {
               if (mergeInfo(Res.getMap()[&I], OpInfo))
                 Changed = true;
             }
           } else if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(&I)) {
             Value *Op = ASC->getOperand(0);
-            auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+            auto OpInfo = Res.getInfo(Op);
             if (OpInfo.isValid()) {
               if (mergeInfo(Res.getMap()[&I], OpInfo))
                 Changed = true;
@@ -310,7 +298,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
                 }
                 if (PTI) {
                   Value *BasePtr = PTI->getOperand(0);
-                  auto BaseInfo = keepForwardOnly(Res.getInfo(BasePtr));
+                  auto BaseInfo = Res.getInfo(BasePtr);
                   if (BaseInfo.isValid()) {
                     int64_t ConstOff = 0;
                     std::vector<std::string> SymOffs;
@@ -331,7 +319,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
                 }
               }
             } else if (auto *PTI = dyn_cast<PtrToIntOperator>(Op)) {
-              auto BaseInfo = keepForwardOnly(Res.getInfo(PTI->getOperand(0)));
+              auto BaseInfo = Res.getInfo(PTI->getOperand(0));
               if (BaseInfo.isValid()) {
                 if (mergeInfo(Res.getMap()[&I], BaseInfo))
                   Changed = true;
@@ -342,7 +330,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
             for (Value *InV : PHI->incoming_values()) {
               if (isa<ConstantPointerNull>(InV) || isa<UndefValue>(InV))
                 continue;
-              auto InInfo = keepForwardOnly(Res.getInfo(InV));
+              auto InInfo = Res.getInfo(InV);
               if (InInfo.isValid())
                 mergeInfo(Temp, InInfo);
             }
@@ -353,7 +341,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
             for (Value *Op : {Sel->getTrueValue(), Sel->getFalseValue()}) {
               if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op))
                 continue;
-              auto OpInfo = keepForwardOnly(Res.getInfo(Op));
+              auto OpInfo = Res.getInfo(Op);
               if (OpInfo.isValid())
                 mergeInfo(Temp, OpInfo);
             }
@@ -370,7 +358,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
               for (Argument &Arg : Callee->args()) {
                 if (ArgIdx < CB->arg_size()) {
                   Value *ArgVal = CB->getArgOperand(ArgIdx);
-                  auto ArgInfo = keepForwardOnly(Res.getInfo(ArgVal));
+                  auto ArgInfo = Res.getInfo(ArgVal);
                   if (ArgInfo.isValid()) {
                     if (mergeInfo(Res.getMap()[&Arg], ArgInfo))
                       Changed = true;
@@ -384,7 +372,7 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
                           dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
                     Value *RetVal = RI->getReturnValue();
                     if (RetVal) {
-                      auto RetInfo = keepForwardOnly(Res.getInfo(RetVal));
+                      auto RetInfo = Res.getInfo(RetVal);
                       if (RetInfo.isValid()) {
                         if (mergeInfo(Res.getMap()[&I], RetInfo))
                           Changed = true;
@@ -542,6 +530,12 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
   return Res;
 }
 
+HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeModule(Module &M) {
+  HeapProvenanceAnalysis HPA;
+  ModuleAnalysisManager DummyMAM;
+  return HPA.run(M, DummyMAM);
+}
+
 PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
                                                  ModuleAnalysisManager &MAM) {
   auto &Result = MAM.getResult<HeapProvenanceAnalysis>(M);
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index a587f72fa9a18..c03071ed661d1 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/TargetFolder.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/Utils/Local.h"
@@ -1239,6 +1240,98 @@ ObjectSizeOffsetEvaluator::ObjectSizeOffsetEvaluator(
   // be different for later objects.
 }
 
+bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result) {
+  Module *M = nullptr;
+  if (auto *I = dyn_cast<Instruction>(V))
+    M = I->getModule();
+  else if (auto *A = dyn_cast<Argument>(V))
+    M = A->getParent()->getParent();
+  if (!M)
+    return false;
+
+  auto HPAResult = HeapProvenanceAnalysis::analyzeModule(*M);
+  auto Info = HPAResult.getInfo(V);
+  if (!Info.isValid())
+    return false;
+
+  if (auto *I = dyn_cast<Instruction>(V)) {
+    if (auto *PHI = dyn_cast<PHINode>(I))
+      Builder.SetInsertPoint(PHI->getParent(), PHI->getParent()->getFirstInsertionPt());
+    else if (I->isTerminator())
+      Builder.SetInsertPoint(I);
+    else if (I->getNextNode())
+      Builder.SetInsertPoint(I->getNextNode());
+    else
+      Builder.SetInsertPoint(I->getParent());
+  } else if (auto *A = dyn_cast<Argument>(V)) {
+    Builder.SetInsertPoint(&A->getParent()->getEntryBlock(),
+                           A->getParent()->getEntryBlock().getFirstInsertionPt());
+  }
+
+  Value *HeadVal = nullptr;
+  Value *OffsetVal = nullptr;
+
+  if (Info.State == HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr) {
+    HeadVal = V;
+    OffsetVal = Zero;
+  } else {
+    Value *Curr = V->stripPointerCasts();
+    while (Curr) {
+      auto CInfo = HPAResult.getInfo(Curr);
+      if (CInfo.State == HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr) {
+        HeadVal = Curr;
+        break;
+      }
+      if (auto *GEP = dyn_cast<GEPOperator>(Curr)) {
+        Curr = GEP->getPointerOperand()->stripPointerCasts();
+      } else if (auto *ITP = dyn_cast<IntToPtrInst>(Curr)) {
+        if (auto *BO = dyn_cast<BinaryOperator>(ITP->getOperand(0))) {
+          if (auto *PTI = dyn_cast<PtrToIntOperator>(BO->getOperand(0)))
+            Curr = PTI->getOperand(0)->stripPointerCasts();
+          else if (auto *PTI = dyn_cast<PtrToIntOperator>(BO->getOperand(1)))
+            Curr = PTI->getOperand(0)->stripPointerCasts();
+          else
+            break;
+        } else {
+          break;
+        }
+      } else {
+        break;
+      }
+    }
+
+    if (!HeadVal) {
+      if (Info.SymOffsets.empty()) {
+        HeadVal = Builder.CreateGEP(Builder.getInt8Ty(), V, ConstantInt::getSigned(IntTy, -Info.ConstOffset), "head.meta");
+        OffsetVal = ConstantInt::get(IntTy, Info.ConstOffset);
+      } else {
+        return false;
+      }
+    } else {
+      Value *VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
+      Value *HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
+      OffsetVal = Builder.CreateSub(VInt, HeadInt, "meta.offset");
+    }
+  }
+
+  FunctionCallee GetSizeFC = M->getOrInsertFunction(
+      "__heap_provanence_sanitizer_get_size", IntTy, Builder.getPtrTy());
+  if (Function *Fn = dyn_cast<Function>(GetSizeFC.getCallee()->stripPointerCasts())) {
+    Fn->setDoesNotThrow();
+    Fn->setOnlyReadsMemory();
+    Fn->setWillReturn();
+    Fn->setDoesNotFreeMemory();
+    Fn->addFnAttr(Attribute::NoCallback);
+    Fn->addFnAttr(Attribute::NoSync);
+  }
+  Value *ChunkPayloadSize = Builder.CreateCall(GetSizeFC, {HeadVal}, "chunk.size");
+  if (OffsetVal->getType() != IntTy)
+    OffsetVal = Builder.CreateZExtOrTrunc(OffsetVal, IntTy);
+
+  Result = SizeOffsetValue(ChunkPayloadSize, OffsetVal);
+  return true;
+}
+
 SizeOffsetValue ObjectSizeOffsetEvaluator::compute(Value *V) {
   // XXX - Are vectors of pointers possible here?
   IntTy = cast<IntegerType>(DL.getIndexType(V->getType()));
@@ -1246,6 +1339,9 @@ SizeOffsetValue ObjectSizeOffsetEvaluator::compute(Value *V) {
 
   SizeOffsetValue Result = compute_(V);
 
+  if (!Result.bothKnown())
+    computeFallbackHeapMetadata(V, Result);
+
   if (!Result.bothKnown()) {
     // Erase everything that was computed in this iteration from the cache, so
     // that no dangling references are left behind. We could be a bit smarter if

>From d595e4036d9cb7434659f4ad2b794ce21918ccd3 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:00:25 -0700
Subject: [PATCH 05/18] [Analysis] Refactor HeapProvenanceAnalysis to use LLVM
 SparseSolver

Redesign HeapProvenanceAnalysis ground-up into separate Forward and Backward analysis passes driven by LLVM's canonical sparse conditional dataflow driver (SparseSolver / AbstractLatticeFunction).

Update the lattice to directly track heap allocation heads (HeapChunkHead -> HeapChunkInterim -> Unknown) with support for PHI nodes and Select instructions. Update computeFallbackHeapMetadata to enforce local function scope checking to prevent cross-function SSA violations.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 .../llvm/Analysis/HeapProvenanceAnalysis.h    | 209 +++--
 llvm/include/llvm/Analysis/MemoryBuiltins.h   |  22 +-
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 792 ++++++++----------
 llvm/lib/Analysis/MemoryBuiltins.cpp          |  82 +-
 llvm/lib/Passes/PassRegistry.def              |   2 +
 .../Instrumentation/BoundsChecking.cpp        |  10 +-
 6 files changed, 538 insertions(+), 579 deletions(-)

diff --git a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
index 033c1cde9b3f8..7144888334bb7 100644
--- a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
+++ b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
@@ -6,95 +6,171 @@
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
 #include <string>
-#include <vector>
 
 namespace llvm {
 
-class HeapProvenanceAnalysisResult {
-public:
-  struct ProvenanceInfo {
-    enum Kind {
-      Uninit = 0,
-      HeapChunkPtr,
-      RecoverableHeapChunkPtr,
-      Unknown
-    } State = Uninit;
-
-    enum Direction {
-      None = 0,
-      Forward = 1,
-      Backward = 2,
-      Both = 3
-    } Dir = None;
-
-    int64_t ConstOffset = 0;
-    std::vector<std::string> SymOffsets;
-    std::string CustomExpr;
-
-    bool isUninit() const { return State == Uninit; }
-    bool isValid() const {
-      return State == HeapChunkPtr || State == RecoverableHeapChunkPtr;
+struct HeapProvenanceLattice {
+  enum class StateKind {
+    Uninit = 0,
+    HeapChunkHead,
+    HeapChunkInterim,
+    Unknown
+  } State = StateKind::Uninit;
+
+  enum Direction {
+    None = 0,
+    Forward = 1,
+    Backward = 2,
+    Both = 3
+  } Dir = None;
+
+  struct Payload {
+    enum class Kind { None, Ref, Select, Phi } K = Kind::None;
+    const Value *Val = nullptr;
+
+    bool operator==(const Payload &RHS) const {
+      return K == RHS.K && Val == RHS.Val;
     }
+    bool operator!=(const Payload &RHS) const { return !(*this == RHS); }
+  } HeadPayload;
 
-    std::string getDirectionStr() const {
-      if (Dir == Forward)
-        return " [forward: from alloc]";
-      if (Dir == Backward)
-        return " [backward: into dealloc]";
-      if (Dir == Both)
-        return " [both: alloc & dealloc]";
-      return "";
-    }
+  bool isUninit() const { return State == StateKind::Uninit; }
+  bool isValid() const {
+    return State == StateKind::HeapChunkHead || State == StateKind::HeapChunkInterim;
+  }
+  bool operator==(const HeapProvenanceLattice &RHS) const {
+    return State == RHS.State && Dir == RHS.Dir && HeadPayload == RHS.HeadPayload;
+  }
+  bool operator!=(const HeapProvenanceLattice &RHS) const {
+    return !(*this == RHS);
+  }
 
-    std::string getExpr() const {
-      if (State == Uninit)
-        return "Uninit";
-      if (State == Unknown)
-        return "Unknown";
-      if (!CustomExpr.empty())
-        return CustomExpr;
-      if (State == HeapChunkPtr && ConstOffset == 0 && SymOffsets.empty())
-        return "head";
-
-      std::string Res = "head";
-      if (ConstOffset > 0)
-        Res += " + " + std::to_string(ConstOffset);
-      else if (ConstOffset < 0)
-        Res += " - " + std::to_string(-ConstOffset);
-
-      for (const auto &S : SymOffsets) {
-        if (!S.empty() && S[0] == '-')
-          Res += " " + S;
-        else
-          Res += " + " + S;
-      }
-      return Res;
+  const Value *getHead() const {
+    if (isValid())
+      return HeadPayload.Val;
+    return nullptr;
+  }
+
+  std::string getDirectionStr() const {
+    if (Dir == Forward)
+      return " [forward: from alloc]";
+    if (Dir == Backward)
+      return " [backward: into dealloc]";
+    if (Dir == Both)
+      return " [both: alloc & dealloc]";
+    return "";
+  }
+
+  std::string getExpr() const {
+    if (State == StateKind::Uninit)
+      return "Uninit";
+    if (State == StateKind::Unknown)
+      return "Unknown";
+    if (HeadPayload.K == Payload::Kind::Ref && HeadPayload.Val) {
+      std::string Name = HeadPayload.Val->getName().str();
+      return Name.empty() ? "Ref" : ("Ref(" + Name + ")");
     }
-  };
+    if (HeadPayload.K == Payload::Kind::Select)
+      return "Select";
+    if (HeadPayload.K == Payload::Kind::Phi)
+      return "Phi";
+    return "Head";
+  }
+};
 
-private:
-  DenseMap<const Value *, ProvenanceInfo> ValueMap;
+class ForwardHeapProvenanceAnalysis;
+class BackwardHeapProvenanceAnalysis;
+class HeapProvenanceAnalysis;
 
+class ForwardHeapProvenanceAnalysisResult {
+  DenseMap<const Value *, HeapProvenanceLattice> ValueMap;
 public:
-  void setInfo(const Value *V, const ProvenanceInfo &Info) {
+  void setInfo(const Value *V, const HeapProvenanceLattice &Info) {
     ValueMap[V] = Info;
   }
-  const ProvenanceInfo &getInfo(const Value *V) const {
-    static ProvenanceInfo Empty;
+  const HeapProvenanceLattice &getInfo(const Value *V) const {
+    static HeapProvenanceLattice Empty;
     auto It = ValueMap.find(V);
     return It != ValueMap.end() ? It->second : Empty;
   }
-  DenseMap<const Value *, ProvenanceInfo> &getMap() { return ValueMap; }
-  const DenseMap<const Value *, ProvenanceInfo> &getMap() const {
+  DenseMap<const Value *, HeapProvenanceLattice> &getMap() { return ValueMap; }
+  const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
     return ValueMap;
   }
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
+};
+
+class ForwardHeapProvenanceAnalysis
+    : public AnalysisInfoMixin<ForwardHeapProvenanceAnalysis> {
+  friend AnalysisInfoMixin<ForwardHeapProvenanceAnalysis>;
+  static AnalysisKey Key;
+public:
+  using Result = ForwardHeapProvenanceAnalysisResult;
+  Result run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+class BackwardHeapProvenanceAnalysisResult {
+  DenseMap<const Value *, HeapProvenanceLattice> ValueMap;
+public:
+  void setInfo(const Value *V, const HeapProvenanceLattice &Info) {
+    ValueMap[V] = Info;
+  }
+  const HeapProvenanceLattice &getInfo(const Value *V) const {
+    static HeapProvenanceLattice Empty;
+    auto It = ValueMap.find(V);
+    return It != ValueMap.end() ? It->second : Empty;
+  }
+  DenseMap<const Value *, HeapProvenanceLattice> &getMap() { return ValueMap; }
+  const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
+    return ValueMap;
+  }
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
+};
+
+class BackwardHeapProvenanceAnalysis
+    : public AnalysisInfoMixin<BackwardHeapProvenanceAnalysis> {
+  friend AnalysisInfoMixin<BackwardHeapProvenanceAnalysis>;
+  static AnalysisKey Key;
+public:
+  using Result = BackwardHeapProvenanceAnalysisResult;
+  Result run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+// Combined wrapper for backward compatibility with ObjectSizeOffsetEvaluator
+class HeapProvenanceAnalysisResult {
+public:
+  using ProvenanceInfo = HeapProvenanceLattice;
+
+private:
+  ForwardHeapProvenanceAnalysisResult ForwardRes;
+  BackwardHeapProvenanceAnalysisResult BackwardRes;
+
+public:
+  HeapProvenanceAnalysisResult() = default;
+  HeapProvenanceAnalysisResult(ForwardHeapProvenanceAnalysisResult F,
+                               BackwardHeapProvenanceAnalysisResult B)
+      : ForwardRes(std::move(F)), BackwardRes(std::move(B)) {}
+
+  const ProvenanceInfo &getInfo(const Value *V) const {
+    const auto &F = ForwardRes.getInfo(V);
+    if (F.isValid())
+      return F;
+    return BackwardRes.getInfo(V);
+  }
+
+  const ForwardHeapProvenanceAnalysisResult &getForwardResult() const { return ForwardRes; }
+  const BackwardHeapProvenanceAnalysisResult &getBackwardResult() const { return BackwardRes; }
+
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
 };
 
 class HeapProvenanceAnalysis
     : public AnalysisInfoMixin<HeapProvenanceAnalysis> {
   friend AnalysisInfoMixin<HeapProvenanceAnalysis>;
   static AnalysisKey Key;
-
 public:
   using Result = HeapProvenanceAnalysisResult;
   static Result analyzeModule(Module &M);
@@ -104,7 +180,6 @@ class HeapProvenanceAnalysis
 class HeapProvenancePrinterPass
     : public PassInfoMixin<HeapProvenancePrinterPass> {
   raw_ostream &OS;
-
 public:
   explicit HeapProvenancePrinterPass(raw_ostream &OS) : OS(OS) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 5d37b9c0d0a87..088701b579006 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -185,18 +185,21 @@ LLVM_ABI std::optional<TypeSize> getBaseObjectSize(const Value *Ptr,
                                                    const TargetLibraryInfo *TLI,
                                                    ObjectSizeOpts Opts = {});
 
+class HeapProvenanceAnalysisResult;
+
 /// Try to turn a call to \@llvm.objectsize into an integer value of the given
 /// Type. Returns null on failure. If MustSucceed is true, this function will
 /// not return null, and may return conservative values governed by the second
 /// argument of the call to objectsize.
-LLVM_ABI Value *lowerObjectSizeCall(IntrinsicInst *ObjectSize,
-                                    const DataLayout &DL,
-                                    const TargetLibraryInfo *TLI,
-                                    bool MustSucceed);
+LLVM_ABI Value *lowerObjectSizeCall(
+    IntrinsicInst *ObjectSize, const DataLayout &DL,
+    const TargetLibraryInfo *TLI, bool MustSucceed,
+    const HeapProvenanceAnalysisResult *HPA = nullptr);
 LLVM_ABI Value *lowerObjectSizeCall(
     IntrinsicInst *ObjectSize, const DataLayout &DL,
     const TargetLibraryInfo *TLI, AAResults *AA, bool MustSucceed,
-    SmallVectorImpl<Instruction *> *InsertedInstructions = nullptr);
+    SmallVectorImpl<Instruction *> *InsertedInstructions = nullptr,
+    const HeapProvenanceAnalysisResult *HPA = nullptr);
 
 /// SizeOffsetType - A base template class for the object size visitors. Used
 /// here as a self-documenting way to handle the values rather than using a
@@ -355,15 +358,16 @@ class ObjectSizeOffsetEvaluator
   PtrSetTy SeenVals;
   ObjectSizeOpts EvalOpts;
   SmallPtrSet<Instruction *, 8> InsertedInstructions;
+  const HeapProvenanceAnalysisResult *HPA;
 
   SizeOffsetValue compute_(Value *V);
   bool computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result);
 
 public:
-  LLVM_ABI ObjectSizeOffsetEvaluator(const DataLayout &DL,
-                                     const TargetLibraryInfo *TLI,
-                                     LLVMContext &Context,
-                                     ObjectSizeOpts EvalOpts = {});
+  LLVM_ABI ObjectSizeOffsetEvaluator(
+      const DataLayout &DL, const TargetLibraryInfo *TLI, LLVMContext &Context,
+      ObjectSizeOpts EvalOpts = {},
+      const HeapProvenanceAnalysisResult *HPA = nullptr);
 
   static SizeOffsetValue unknown() { return SizeOffsetValue(); }
 
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 8711fe8957c93..ad4af26c8bdac 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -1,15 +1,45 @@
 #include "llvm/Analysis/HeapProvenanceAnalysis.h"
+#include "llvm/Analysis/SparsePropagation.h"
 #include "llvm/IR/Constants.h"
-#include "llvm/IR/DataLayout.h"
-#include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Support/raw_ostream.h"
 
+namespace llvm {
+template <> struct LatticeKeyInfo<const Value *> {
+  static inline Value *getValueFromLatticeKey(const Value *Key) {
+    return const_cast<Value *>(Key);
+  }
+  static inline const Value *getLatticeKeyFromValue(Value *V) {
+    return V;
+  }
+};
+} // namespace llvm
+
 using namespace llvm;
 
+AnalysisKey ForwardHeapProvenanceAnalysis::Key;
+AnalysisKey BackwardHeapProvenanceAnalysis::Key;
 AnalysisKey HeapProvenanceAnalysis::Key;
 
+bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                                     ModuleAnalysisManager::Invalidator &) {
+  auto PAC = PA.getChecker<ForwardHeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                                      ModuleAnalysisManager::Invalidator &) {
+  auto PAC = PA.getChecker<BackwardHeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                              ModuleAnalysisManager::Invalidator &) {
+  auto PAC = PA.getChecker<HeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
 static bool isAllocFunc(StringRef Name) {
   return Name == "malloc" || Name == "calloc" || Name == "realloc" ||
          Name == "aligned_alloc" || Name == "strdup" || Name == "strndup" ||
@@ -22,138 +52,33 @@ static bool isFreeFunc(StringRef Name) {
          Name == "g_free" || Name.starts_with("_Zd");
 }
 
-static std::string getValueOperandName(const Value *V) {
-  std::string Str;
-  raw_string_ostream OS(Str);
-  V->printAsOperand(OS, false);
-  return Str;
-}
-
-static void addSymOffset(std::vector<std::string> &Syms, const std::string &S) {
-  StringRef SR(S);
-  std::string NegS = (SR.size() > 3 && SR.starts_with("-(") && SR.ends_with(")"))
-                         ? S.substr(2, S.size() - 3)
-                         : "-(" + S + ")";
-  for (auto It = Syms.begin(); It != Syms.end(); ++It) {
-    if (*It == NegS) {
-      Syms.erase(It);
-      return;
-    }
-  }
-  Syms.push_back(S);
-}
-
-static void extractGEPOffsets(const GEPOperator *GEP, const DataLayout &DL,
-                              int64_t &ConstOff,
-                              std::vector<std::string> &SymOffs) {
-  APInt APIntOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
-  if (GEP->accumulateConstantOffset(DL, APIntOffset)) {
-    ConstOff += APIntOffset.getSExtValue();
-    return;
-  }
-  for (gep_type_iterator GTI = gep_type_begin(GEP), GTE = gep_type_end(GEP);
-       GTI != GTE; ++GTI) {
-    Value *Op = GTI.getOperand();
-    if (StructType *STy = GTI.getStructTypeOrNull()) {
-      ConstantInt *OpC = cast<ConstantInt>(Op);
-      const StructLayout *SL = DL.getStructLayout(STy);
-      ConstOff += SL->getElementOffset(OpC->getZExtValue());
-    } else {
-      TypeSize Stride = GTI.getSequentialElementStride(DL);
-      int64_t StrideVal = Stride.getKnownMinValue();
-      if (ConstantInt *OpC = dyn_cast<ConstantInt>(Op)) {
-        ConstOff += OpC->getSExtValue() * StrideVal;
-      } else {
-        std::string OpName = getValueOperandName(Op);
-        if (StrideVal == 1)
-          addSymOffset(SymOffs, OpName);
-        else
-          addSymOffset(SymOffs, OpName + " * " + std::to_string(StrideVal));
-      }
-    }
-  }
-}
-
-static HeapProvenanceAnalysisResult::ProvenanceInfo
-addOffset(const HeapProvenanceAnalysisResult::ProvenanceInfo &BaseInfo,
-          int64_t C, const std::vector<std::string> &Syms) {
-  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
-  if (!BaseInfo.isValid())
-    return BaseInfo;
-  ProvenanceInfo Res = BaseInfo;
-  Res.State = ProvenanceInfo::RecoverableHeapChunkPtr;
-  if (!Res.CustomExpr.empty()) {
-    std::string OffStr;
-    if (C > 0)
-      OffStr += " + " + std::to_string(C);
-    else if (C < 0)
-      OffStr += " - " + std::to_string(-C);
-    for (const auto &S : Syms)
-      OffStr += " + " + S;
-    Res.CustomExpr = "(" + Res.CustomExpr + ")" + OffStr;
-    return Res;
-  }
-  Res.ConstOffset += C;
-  for (const auto &S : Syms)
-    addSymOffset(Res.SymOffsets, S);
-  if (Res.ConstOffset == 0 && Res.SymOffsets.empty())
-    Res.State = ProvenanceInfo::HeapChunkPtr;
-  return Res;
-}
-
-static HeapProvenanceAnalysisResult::ProvenanceInfo
-subOffset(const HeapProvenanceAnalysisResult::ProvenanceInfo &IInfo, int64_t C,
-          const std::vector<std::string> &Syms) {
-  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
-  if (!IInfo.isValid())
-    return IInfo;
-  ProvenanceInfo Res = IInfo;
-  Res.State = ProvenanceInfo::RecoverableHeapChunkPtr;
-  if (!Res.CustomExpr.empty()) {
-    std::string OffStr;
-    if (C > 0)
-      OffStr += " - " + std::to_string(C);
-    else if (C < 0)
-      OffStr += " + " + std::to_string(-C);
-    for (const auto &S : Syms)
-      OffStr += " - (" + S + ")";
-    Res.CustomExpr = "(" + Res.CustomExpr + ")" + OffStr;
-    return Res;
-  }
-  Res.ConstOffset -= C;
-  for (const auto &S : Syms)
-    addSymOffset(Res.SymOffsets, "-(" + S + ")");
-  if (Res.ConstOffset == 0 && Res.SymOffsets.empty())
-    Res.State = ProvenanceInfo::HeapChunkPtr;
-  return Res;
-}
-
-static bool mergeInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
-                      const HeapProvenanceAnalysisResult::ProvenanceInfo &Src) {
-  using ProvenanceInfo = HeapProvenanceAnalysisResult::ProvenanceInfo;
-  if (Src.State == ProvenanceInfo::Uninit)
+static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
+                         const HeapProvenanceLattice &Src) {
+  using Lattice = HeapProvenanceLattice;
+  if (Src.State == Lattice::StateKind::Uninit)
     return false;
 
-  auto MergedDir = static_cast<ProvenanceInfo::Direction>(Dest.Dir | Src.Dir);
+  auto MergedDir = static_cast<Lattice::Direction>(Dest.Dir | Src.Dir);
 
-  if (Dest.State == ProvenanceInfo::Uninit) {
+  if (Dest.State == Lattice::StateKind::Uninit) {
     Dest = Src;
     Dest.Dir = MergedDir;
     return true;
   }
-  if (Dest.State == ProvenanceInfo::Unknown) {
+
+  if (Dest.State == Lattice::StateKind::Unknown) {
     if (Dest.Dir != MergedDir) {
       Dest.Dir = MergedDir;
       return true;
     }
     return false;
   }
-  if (Src.State == ProvenanceInfo::Unknown) {
-    if (Dest.Dir != MergedDir) {
-      Dest.Dir = MergedDir;
-      return true;
-    }
-    return false;
+
+  if (Src.State == Lattice::StateKind::Unknown) {
+    Dest.State = Lattice::StateKind::Unknown;
+    Dest.Dir = MergedDir;
+    Dest.HeadPayload = {Lattice::Payload::Kind::None, nullptr};
+    return true;
   }
 
   bool Changed = false;
@@ -162,81 +87,126 @@ static bool mergeInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
     Changed = true;
   }
 
-  if (Dest.getExpr() == Src.getExpr()) {
-    if (Src.State == ProvenanceInfo::RecoverableHeapChunkPtr &&
-        Dest.State != ProvenanceInfo::RecoverableHeapChunkPtr) {
-      Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+  if (Dest.State == Src.State && Dest.HeadPayload == Src.HeadPayload)
+    return Changed;
+
+  if (isa_and_nonnull<PHINode>(Target)) {
+    Lattice::Payload NewPayload{Lattice::Payload::Kind::Phi, Target};
+    if (Dest.State != Lattice::StateKind::HeapChunkInterim ||
+        Dest.HeadPayload != NewPayload) {
+      Dest.State = Lattice::StateKind::HeapChunkInterim;
+      Dest.HeadPayload = NewPayload;
       Changed = true;
     }
     return Changed;
   }
 
-  if (Dest.CustomExpr == "head + dynamic_offset")
+  if (isa_and_nonnull<SelectInst>(Target)) {
+    Lattice::Payload NewPayload{Lattice::Payload::Kind::Select, Target};
+    if (Dest.State != Lattice::StateKind::HeapChunkInterim ||
+        Dest.HeadPayload != NewPayload) {
+      Dest.State = Lattice::StateKind::HeapChunkInterim;
+      Dest.HeadPayload = NewPayload;
+      Changed = true;
+    }
     return Changed;
+  }
 
-  if (!Dest.CustomExpr.empty()) {
-    Dest.CustomExpr = "head + dynamic_offset";
-    Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
+  if (Dest.HeadPayload.Val != Src.HeadPayload.Val) {
+    Dest.State = Lattice::StateKind::Unknown;
+    Dest.HeadPayload = {Lattice::Payload::Kind::None, nullptr};
     return true;
   }
 
-  Dest.CustomExpr = "PHI(" + Dest.getExpr() + ", " + Src.getExpr() + ")";
-  Dest.State = ProvenanceInfo::RecoverableHeapChunkPtr;
-  return true;
+  return Changed;
 }
 
-static bool updateInfo(HeapProvenanceAnalysisResult::ProvenanceInfo &Dest,
-                       const HeapProvenanceAnalysisResult::ProvenanceInfo &NewVal) {
-  if (Dest.State == NewVal.State && Dest.Dir == NewVal.Dir &&
-      Dest.ConstOffset == NewVal.ConstOffset &&
-      Dest.SymOffsets == NewVal.SymOffsets &&
-      Dest.CustomExpr == NewVal.CustomExpr)
-    return false;
+class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapProvenanceLattice> {
+  DenseMap<const Value *, HeapProvenanceLattice> Seeds;
+public:
+  ForwardLatticeFunc()
+      : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
 
-  if (!Dest.CustomExpr.empty() && !NewVal.CustomExpr.empty() &&
-      Dest.CustomExpr != NewVal.CustomExpr) {
-    std::string Widened = "head + dynamic_offset";
-    if (Dest.CustomExpr != Widened) {
-      Dest = NewVal;
-      Dest.CustomExpr = Widened;
-      return true;
-    }
-    return false;
+  void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
+    mergeLattice(V, Seeds[V], Info);
+  }
+  bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
+  HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
+    mergeLattice(nullptr, X, Y);
+    return X;
   }
+  HeapProvenanceLattice ComputeLatticeVal(const Value *Key) override {
+    auto It = Seeds.find(Key);
+    if (It != Seeds.end()) return It->second;
+    return getUndefVal();
+  }
+  void ComputeInstructionState(Instruction &I,
+                               SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
+                               SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
+    auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
+      if (!Target || isa<ConstantPointerNull>(Target) || isa<UndefValue>(Target)) return;
+      auto &Dest = ChangedValues[Target];
+      if (Dest.isUninit()) {
+        auto Existing = SS.getExistingValueState(Target);
+        if (Existing.isValid()) Dest = Existing;
+      }
+      mergeLattice(Target, Dest, NewI);
+    };
 
-  Dest = NewVal;
-  return true;
-}
+    auto SeedIt = Seeds.find(&I);
+    if (SeedIt != Seeds.end()) MergeInto(&I, SeedIt->second);
 
-HeapProvenanceAnalysis::Result
-HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
-  Result Res;
-  const DataLayout &DL = M.getDataLayout();
+    for (const Use &U : I.operands()) {
+      const Value *Op = U.get();
+      auto OpSeedIt = Seeds.find(Op);
+      if (OpSeedIt != Seeds.end()) MergeInto(Op, OpSeedIt->second);
 
-  for (Function &F : M) {
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        if (auto *CB = dyn_cast<CallBase>(&I)) {
-          Function *Callee = CB->getCalledFunction();
-          if (!Callee) {
-            Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
-            Callee = dyn_cast<Function>(CalledVal);
-          }
-          if (Callee) {
-            StringRef Name = Callee->getName();
-            if (isAllocFunc(Name)) {
-              Result::ProvenanceInfo Info;
-              Info.State = Result::ProvenanceInfo::HeapChunkPtr;
-              Info.Dir = Result::ProvenanceInfo::Forward;
-              mergeInfo(Res.getMap()[&I], Info);
+      auto OpInfo = SS.getExistingValueState(Op);
+      if (!OpInfo.isValid() || !(OpInfo.Dir & HeapProvenanceLattice::Forward)) continue;
+
+      HeapProvenanceLattice NewI = OpInfo;
+      if (NewI.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
+        NewI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+        NewI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
+      }
+
+      if (isa<GEPOperator>(&I) || isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
+          isa<IntToPtrInst>(&I) || isa<PHINode>(&I) || isa<SelectInst>(&I)) {
+        MergeInto(&I, NewI);
+      } else if (auto *CB = dyn_cast<CallBase>(&I)) {
+        Function *Callee = CB->getCalledFunction();
+        if (!Callee) Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+        if (Callee && !Callee->isDeclaration()) {
+          for (unsigned idx = 0, e = CB->arg_size(); idx != e; ++idx) {
+            if (CB->getArgOperand(idx) == Op && idx < Callee->arg_size()) {
+              Argument *Arg = Callee->getArg(idx);
+              HeapProvenanceLattice ArgI = NewI;
+              ArgI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+              ArgI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Arg};
+              MergeInto(Arg, ArgI);
             }
-            if (isFreeFunc(Name)) {
-              if (CB->arg_size() > 0) {
-                Value *ArgVal = CB->getArgOperand(0);
-                Result::ProvenanceInfo Info;
-                Info.State = Result::ProvenanceInfo::HeapChunkPtr;
-                Info.Dir = Result::ProvenanceInfo::Backward;
-                mergeInfo(Res.getMap()[ArgVal], Info);
+          }
+        }
+      }
+    }
+
+    if (auto *RI = dyn_cast<ReturnInst>(&I)) {
+      if (Value *RetVal = RI->getReturnValue()) {
+        auto RetInfo = SS.getExistingValueState(RetVal);
+        if (RetInfo.isValid() && (RetInfo.Dir & HeapProvenanceLattice::Forward)) {
+          HeapProvenanceLattice NewI = RetInfo;
+          if (NewI.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
+            NewI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+            NewI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, RetVal};
+          }
+          const Function *F = RI->getFunction();
+          for (const User *U : F->users()) {
+            if (auto *CB = dyn_cast<CallBase>(U)) {
+              if (CB->getCalledFunction() == F || CB->getCalledOperand()->stripPointerCasts() == F) {
+                HeapProvenanceLattice CBI = NewI;
+                CBI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+                CBI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, CB};
+                MergeInto(CB, CBI);
               }
             }
           }
@@ -244,141 +214,107 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
       }
     }
   }
+};
 
-  bool Changed = true;
-  int MaxIters = 50;
-  while (Changed && MaxIters-- > 0) {
-    Changed = false;
-
-    // Forward propagation
-    for (Function &F : M) {
-      for (BasicBlock &BB : F) {
-        for (Instruction &I : BB) {
-          if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
-            Value *Base = GEP->getPointerOperand();
-            auto BaseInfo = Res.getInfo(Base);
-            if (BaseInfo.isValid()) {
-              auto NewInfo = BaseInfo;
-              NewInfo.State = Result::ProvenanceInfo::RecoverableHeapChunkPtr;
-              int64_t ConstOff = 0;
-              std::vector<std::string> SymOffs;
-              extractGEPOffsets(GEP, DL, ConstOff, SymOffs);
-              NewInfo = addOffset(NewInfo, ConstOff, SymOffs);
-              if (mergeInfo(Res.getMap()[&I], NewInfo))
-                Changed = true;
-            }
-          } else if (auto *BC = dyn_cast<BitCastOperator>(&I)) {
-            Value *Op = BC->getOperand(0);
-            auto OpInfo = Res.getInfo(Op);
-            if (OpInfo.isValid()) {
-              if (mergeInfo(Res.getMap()[&I], OpInfo))
-                Changed = true;
-            }
-          } else if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(&I)) {
-            Value *Op = ASC->getOperand(0);
-            auto OpInfo = Res.getInfo(Op);
-            if (OpInfo.isValid()) {
-              if (mergeInfo(Res.getMap()[&I], OpInfo))
-                Changed = true;
-            }
-          } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
-            Value *Op = ITP->getOperand(0);
-            if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
-              if (BO->getOpcode() == Instruction::Add ||
-                  BO->getOpcode() == Instruction::Sub) {
-                Value *LHS = BO->getOperand(0);
-                Value *RHS = BO->getOperand(1);
-                auto *PTI = dyn_cast<PtrToIntOperator>(LHS);
-                if (!PTI) {
-                  PTI = dyn_cast<PtrToIntOperator>(RHS);
-                  if (PTI && BO->getOpcode() == Instruction::Add)
-                    std::swap(LHS, RHS);
-                  else
-                    PTI = nullptr;
-                }
-                if (PTI) {
-                  Value *BasePtr = PTI->getOperand(0);
-                  auto BaseInfo = Res.getInfo(BasePtr);
-                  if (BaseInfo.isValid()) {
-                    int64_t ConstOff = 0;
-                    std::vector<std::string> SymOffs;
-                    if (auto *CInt = dyn_cast<ConstantInt>(RHS)) {
-                      ConstOff = CInt->getSExtValue();
-                    } else {
-                      SymOffs.push_back(getValueOperandName(RHS));
-                    }
-                    if (BO->getOpcode() == Instruction::Sub) {
-                      ConstOff = -ConstOff;
-                      for (auto &S : SymOffs)
-                        S = "-(" + S + ")";
-                    }
-                    auto NewInfo = addOffset(BaseInfo, ConstOff, SymOffs);
-                    if (mergeInfo(Res.getMap()[&I], NewInfo))
-                      Changed = true;
-                  }
-                }
-              }
-            } else if (auto *PTI = dyn_cast<PtrToIntOperator>(Op)) {
-              auto BaseInfo = Res.getInfo(PTI->getOperand(0));
-              if (BaseInfo.isValid()) {
-                if (mergeInfo(Res.getMap()[&I], BaseInfo))
-                  Changed = true;
+class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapProvenanceLattice> {
+  DenseMap<const Value *, HeapProvenanceLattice> Seeds;
+public:
+  BackwardLatticeFunc()
+      : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
+
+  void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
+    mergeLattice(V, Seeds[V], Info);
+  }
+  bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
+  HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
+    mergeLattice(nullptr, X, Y);
+    return X;
+  }
+  HeapProvenanceLattice ComputeLatticeVal(const Value *Key) override {
+    auto It = Seeds.find(Key);
+    if (It != Seeds.end()) return It->second;
+    return getUndefVal();
+  }
+  void ComputeInstructionState(Instruction &I,
+                               SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
+                               SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
+    auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
+      if (!Target || isa<ConstantPointerNull>(Target) || isa<UndefValue>(Target)) return;
+      auto &Dest = ChangedValues[Target];
+      if (Dest.isUninit()) {
+        auto Existing = SS.getExistingValueState(Target);
+        if (Existing.isValid()) Dest = Existing;
+      }
+      mergeLattice(Target, Dest, NewI);
+    };
+
+    auto SeedIt = Seeds.find(&I);
+    if (SeedIt != Seeds.end()) MergeInto(&I, SeedIt->second);
+
+    for (const Use &U : I.operands()) {
+      const Value *Op = U.get();
+      auto OpSeedIt = Seeds.find(Op);
+      if (OpSeedIt != Seeds.end()) MergeInto(Op, OpSeedIt->second);
+    }
+
+    auto Info = SS.getExistingValueState(&I);
+    if (Info.isValid() && (Info.Dir & HeapProvenanceLattice::Backward)) {
+      HeapProvenanceLattice BackInfo = Info;
+      BackInfo.Dir = HeapProvenanceLattice::Backward;
+      if (BackInfo.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
+        BackInfo.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+        BackInfo.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
+      }
+
+      if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
+        MergeInto(GEP->getPointerOperand(), BackInfo);
+      } else if (auto *BC = dyn_cast<BitCastInst>(&I)) {
+        MergeInto(BC->getOperand(0), BackInfo);
+      } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&I)) {
+        MergeInto(ASC->getOperand(0), BackInfo);
+      } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
+        if (auto *PTI = dyn_cast<PtrToIntOperator>(ITP->getOperand(0)))
+          MergeInto(PTI->getOperand(0), BackInfo);
+      } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
+        for (Value *InV : PHI->incoming_values())
+          MergeInto(InV, BackInfo);
+      } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
+        MergeInto(Sel->getTrueValue(), BackInfo);
+        MergeInto(Sel->getFalseValue(), BackInfo);
+      } else if (auto *CB = dyn_cast<CallBase>(&I)) {
+        Function *Callee = CB->getCalledFunction();
+        if (!Callee) Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+        if (Callee && !Callee->isDeclaration() && I.getType()->isPointerTy()) {
+          for (BasicBlock &CalleeBB : *Callee) {
+            if (auto *RI = dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
+              if (Value *RetVal = RI->getReturnValue()) {
+                HeapProvenanceLattice RetI = BackInfo;
+                RetI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+                RetI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, RetVal};
+                MergeInto(RetVal, RetI);
               }
             }
-          } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
-            Result::ProvenanceInfo Temp;
-            for (Value *InV : PHI->incoming_values()) {
-              if (isa<ConstantPointerNull>(InV) || isa<UndefValue>(InV))
-                continue;
-              auto InInfo = Res.getInfo(InV);
-              if (InInfo.isValid())
-                mergeInfo(Temp, InInfo);
-            }
-            if (Temp.isValid() && updateInfo(Res.getMap()[&I], Temp))
-              Changed = true;
-          } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
-            Result::ProvenanceInfo Temp;
-            for (Value *Op : {Sel->getTrueValue(), Sel->getFalseValue()}) {
-              if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op))
-                continue;
-              auto OpInfo = Res.getInfo(Op);
-              if (OpInfo.isValid())
-                mergeInfo(Temp, OpInfo);
-            }
-            if (Temp.isValid() && updateInfo(Res.getMap()[&I], Temp))
-              Changed = true;
-          } else if (auto *CB = dyn_cast<CallBase>(&I)) {
-            Function *Callee = CB->getCalledFunction();
-            if (!Callee) {
-              Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
-              Callee = dyn_cast<Function>(CalledVal);
-            }
-            if (Callee && !Callee->isDeclaration()) {
-              unsigned ArgIdx = 0;
-              for (Argument &Arg : Callee->args()) {
+          }
+        }
+      }
+    }
+
+    for (const Use &U : I.operands()) {
+      if (auto *Arg = dyn_cast<Argument>(U.get())) {
+        auto ArgInfo = SS.getExistingValueState(Arg);
+        if (ArgInfo.isValid() && (ArgInfo.Dir & HeapProvenanceLattice::Backward)) {
+          const Function *F = Arg->getParent();
+          for (const User *Usr : F->users()) {
+            if (auto *CB = dyn_cast<CallBase>(Usr)) {
+              if (CB->getCalledFunction() == F || CB->getCalledOperand()->stripPointerCasts() == F) {
+                unsigned ArgIdx = Arg->getArgNo();
                 if (ArgIdx < CB->arg_size()) {
-                  Value *ArgVal = CB->getArgOperand(ArgIdx);
-                  auto ArgInfo = Res.getInfo(ArgVal);
-                  if (ArgInfo.isValid()) {
-                    if (mergeInfo(Res.getMap()[&Arg], ArgInfo))
-                      Changed = true;
-                  }
-                }
-                ArgIdx++;
-              }
-              if (I.getType()->isPointerTy()) {
-                for (BasicBlock &CalleeBB : *Callee) {
-                  if (auto *RI =
-                          dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
-                    Value *RetVal = RI->getReturnValue();
-                    if (RetVal) {
-                      auto RetInfo = Res.getInfo(RetVal);
-                      if (RetInfo.isValid()) {
-                        if (mergeInfo(Res.getMap()[&I], RetInfo))
-                          Changed = true;
-                      }
-                    }
-                  }
+                  Value *CallerArg = CB->getArgOperand(ArgIdx);
+                  HeapProvenanceLattice CallerI = ArgInfo;
+                  CallerI.Dir = HeapProvenanceLattice::Backward;
+                  CallerI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+                  CallerI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, CallerArg};
+                  MergeInto(CallerArg, CallerI);
                 }
               }
             }
@@ -386,154 +322,114 @@ HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
         }
       }
     }
+  }
+};
 
-    // Backward propagation
-    for (Function &F : M) {
-      for (BasicBlock &BB : F) {
-        for (Instruction &I : BB) {
-          auto Info = Res.getInfo(&I);
-          if (!Info.isValid() ||
-              (Info.Dir & Result::ProvenanceInfo::Backward) == 0)
-            continue;
-
-          if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
-            Value *Base = GEP->getPointerOperand();
-            int64_t ConstOff = 0;
-            std::vector<std::string> SymOffs;
-            extractGEPOffsets(GEP, DL, ConstOff, SymOffs);
-            auto BaseNew = subOffset(Info, ConstOff, SymOffs);
-            BaseNew.Dir = Result::ProvenanceInfo::Backward;
-            if (mergeInfo(Res.getMap()[Base], BaseNew))
-              Changed = true;
-          } else if (auto *BC = dyn_cast<BitCastOperator>(&I)) {
-            Value *Op = BC->getOperand(0);
-            auto OpNew = Info;
-            OpNew.Dir = Result::ProvenanceInfo::Backward;
-            if (mergeInfo(Res.getMap()[Op], OpNew))
-              Changed = true;
-          } else if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(&I)) {
-            Value *Op = ASC->getOperand(0);
-            auto OpNew = Info;
-            OpNew.Dir = Result::ProvenanceInfo::Backward;
-            if (mergeInfo(Res.getMap()[Op], OpNew))
-              Changed = true;
-          } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
-            Value *Op = ITP->getOperand(0);
-            if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
-              if (BO->getOpcode() == Instruction::Add ||
-                  BO->getOpcode() == Instruction::Sub) {
-                Value *LHS = BO->getOperand(0);
-                Value *RHS = BO->getOperand(1);
-                auto *PTI = dyn_cast<PtrToIntOperator>(LHS);
-                if (!PTI) {
-                  PTI = dyn_cast<PtrToIntOperator>(RHS);
-                  if (PTI && BO->getOpcode() == Instruction::Add)
-                    std::swap(LHS, RHS);
-                  else
-                    PTI = nullptr;
-                }
-                if (PTI) {
-                  Value *BasePtr = PTI->getOperand(0);
-                  int64_t ConstOff = 0;
-                  std::vector<std::string> SymOffs;
-                  if (auto *CInt = dyn_cast<ConstantInt>(RHS)) {
-                    ConstOff = CInt->getSExtValue();
-                  } else {
-                    SymOffs.push_back(getValueOperandName(RHS));
-                  }
-                  if (BO->getOpcode() == Instruction::Sub) {
-                    ConstOff = -ConstOff;
-                    for (auto &S : SymOffs)
-                      S = "-(" + S + ")";
-                  }
-                  auto BaseNew = subOffset(Info, ConstOff, SymOffs);
-                  BaseNew.Dir = Result::ProvenanceInfo::Backward;
-                  if (mergeInfo(Res.getMap()[BasePtr], BaseNew))
-                    Changed = true;
-                }
-              }
-            } else if (auto *PTI = dyn_cast<PtrToIntOperator>(Op)) {
-              Value *BasePtr = PTI->getOperand(0);
-              auto BaseNew = Info;
-              BaseNew.Dir = Result::ProvenanceInfo::Backward;
-              if (mergeInfo(Res.getMap()[BasePtr], BaseNew))
-                Changed = true;
-            }
-          } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
-            for (Value *InV : PHI->incoming_values()) {
-              if (isa<ConstantPointerNull>(InV) || isa<UndefValue>(InV))
-                continue;
-              auto InNew = Info;
-              InNew.Dir = Result::ProvenanceInfo::Backward;
-              if (mergeInfo(Res.getMap()[InV], InNew))
-                Changed = true;
-            }
-          } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
-            for (Value *Op : {Sel->getTrueValue(), Sel->getFalseValue()}) {
-              if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op))
-                continue;
-              auto OpNew = Info;
-              OpNew.Dir = Result::ProvenanceInfo::Backward;
-              if (mergeInfo(Res.getMap()[Op], OpNew))
-                Changed = true;
-            }
-          } else if (auto *CB = dyn_cast<CallBase>(&I)) {
-            Function *Callee = CB->getCalledFunction();
-            if (!Callee) {
-              Value *CalledVal = CB->getCalledOperand()->stripPointerCasts();
-              Callee = dyn_cast<Function>(CalledVal);
-            }
-            if (Callee && !Callee->isDeclaration()) {
-              if (I.getType()->isPointerTy()) {
-                for (BasicBlock &CalleeBB : *Callee) {
-                  if (auto *RI =
-                          dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
-                    Value *RetVal = RI->getReturnValue();
-                    if (RetVal) {
-                      auto RetNew = Info;
-                      RetNew.Dir = Result::ProvenanceInfo::Backward;
-                      if (mergeInfo(Res.getMap()[RetVal], RetNew))
-                        Changed = true;
-                    }
-                  }
-                }
-              }
-            }
+ForwardHeapProvenanceAnalysis::Result
+ForwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  Result Res;
+  ForwardLatticeFunc Lattice;
+
+  for (Function &F : M) {
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        if (auto *CB = dyn_cast<CallBase>(&I)) {
+          Function *Callee = CB->getCalledFunction();
+          if (!Callee)
+            Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+          if (Callee && isAllocFunc(Callee->getName())) {
+            HeapProvenanceLattice Info;
+            Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+            Info.Dir = HeapProvenanceLattice::Forward;
+            Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
+            Lattice.addSeed(&I, Info);
           }
         }
       }
+    }
+  }
 
-      for (Argument &Arg : F.args()) {
-        auto ArgInfo = Res.getInfo(&Arg);
-        if (!ArgInfo.isValid() ||
-            (ArgInfo.Dir & Result::ProvenanceInfo::Backward) == 0)
-          continue;
-        for (User *U : F.users()) {
-          if (auto *CB = dyn_cast<CallBase>(U)) {
-            if (CB->getCalledOperand()->stripPointerCasts() == &F ||
-                CB->getCalledFunction() == &F) {
-              unsigned ArgIdx = Arg.getArgNo();
-              if (ArgIdx < CB->arg_size()) {
-                Value *CallArg = CB->getArgOperand(ArgIdx);
-                auto CallArgNew = ArgInfo;
-                CallArgNew.Dir = Result::ProvenanceInfo::Backward;
-                if (mergeInfo(Res.getMap()[CallArg], CallArgNew))
-                  Changed = true;
-              }
-            }
+  SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
+  for (Function &F : M)
+    if (!F.isDeclaration())
+      Solver.MarkBlockExecutable(&F.front());
+
+  Solver.Solve();
+
+  for (Function &F : M) {
+    for (Argument &Arg : F.args()) {
+      auto LV = Solver.getExistingValueState(&Arg);
+      if (LV.isValid()) Res.setInfo(&Arg, LV);
+    }
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        auto LV = Solver.getExistingValueState(&I);
+        if (LV.isValid()) Res.setInfo(&I, LV);
+      }
+    }
+  }
+  return Res;
+}
+
+BackwardHeapProvenanceAnalysis::Result
+BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  Result Res;
+  BackwardLatticeFunc Lattice;
+
+  for (Function &F : M) {
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        if (auto *CB = dyn_cast<CallBase>(&I)) {
+          Function *Callee = CB->getCalledFunction();
+          if (!Callee)
+            Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+          if (Callee && isFreeFunc(Callee->getName()) && CB->arg_size() > 0) {
+            Value *ArgVal = CB->getArgOperand(0);
+            HeapProvenanceLattice Info;
+            Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+            Info.Dir = HeapProvenanceLattice::Backward;
+            Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, ArgVal};
+            Lattice.addSeed(ArgVal, Info);
           }
         }
       }
     }
   }
 
+  SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
+  for (Function &F : M)
+    if (!F.isDeclaration())
+      Solver.MarkBlockExecutable(&F.front());
+
+  Solver.Solve();
+
+  for (Function &F : M) {
+    for (Argument &Arg : F.args()) {
+      auto LV = Solver.getExistingValueState(&Arg);
+      if (LV.isValid()) Res.setInfo(&Arg, LV);
+    }
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        auto LV = Solver.getExistingValueState(&I);
+        if (LV.isValid()) Res.setInfo(&I, LV);
+      }
+    }
+  }
   return Res;
 }
 
 HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeModule(Module &M) {
-  HeapProvenanceAnalysis HPA;
   ModuleAnalysisManager DummyMAM;
-  return HPA.run(M, DummyMAM);
+  ForwardHeapProvenanceAnalysis ForwardHPA;
+  BackwardHeapProvenanceAnalysis BackwardHPA;
+  return Result(ForwardHPA.run(M, DummyMAM), BackwardHPA.run(M, DummyMAM));
+}
+
+HeapProvenanceAnalysis::Result
+HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  auto ForwardRes = MAM.getResult<ForwardHeapProvenanceAnalysis>(M);
+  auto BackwardRes = MAM.getResult<BackwardHeapProvenanceAnalysis>(M);
+  return Result(std::move(ForwardRes), std::move(BackwardRes));
 }
 
 PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
@@ -550,10 +446,9 @@ PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
         OS << "  argument ";
         Arg.printAsOperand(OS, false);
         OS << ": "
-           << (Info.State ==
-                       HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr
-                   ? "HeapChunkPtr"
-                   : "RecoverableHeapChunkPtr")
+           << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
+                   ? "HeapChunkHead"
+                   : "HeapChunkInterim")
            << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
       }
     }
@@ -564,10 +459,9 @@ PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
           OS << "  ";
           I.printAsOperand(OS, false);
           OS << ": "
-             << (Info.State ==
-                         HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr
-                     ? "HeapChunkPtr"
-                     : "RecoverableHeapChunkPtr")
+             << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
+                     ? "HeapChunkHead"
+                     : "HeapChunkInterim")
              << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
         }
       }
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index c03071ed661d1..e72e4192f4d3a 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -635,15 +635,17 @@ std::optional<TypeSize> llvm::getBaseObjectSize(const Value *Ptr,
 Value *llvm::lowerObjectSizeCall(IntrinsicInst *ObjectSize,
                                  const DataLayout &DL,
                                  const TargetLibraryInfo *TLI,
-                                 bool MustSucceed) {
+                                 bool MustSucceed,
+                                 const HeapProvenanceAnalysisResult *HPA) {
   return lowerObjectSizeCall(ObjectSize, DL, TLI, /*AAResults=*/nullptr,
-                             MustSucceed);
+                             MustSucceed, /*InsertedInstructions=*/nullptr, HPA);
 }
 
 Value *llvm::lowerObjectSizeCall(
     IntrinsicInst *ObjectSize, const DataLayout &DL,
     const TargetLibraryInfo *TLI, AAResults *AA, bool MustSucceed,
-    SmallVectorImpl<Instruction *> *InsertedInstructions) {
+    SmallVectorImpl<Instruction *> *InsertedInstructions,
+    const HeapProvenanceAnalysisResult *HPA) {
   assert(ObjectSize->getIntrinsicID() == Intrinsic::objectsize &&
          "ObjectSize must be a call to llvm.objectsize!");
 
@@ -673,7 +675,7 @@ Value *llvm::lowerObjectSizeCall(
       return ConstantInt::get(ResultType, Size);
   } else {
     LLVMContext &Ctx = ObjectSize->getFunction()->getContext();
-    ObjectSizeOffsetEvaluator Eval(DL, TLI, Ctx, EvalOptions);
+    ObjectSizeOffsetEvaluator Eval(DL, TLI, Ctx, EvalOptions, HPA);
     SizeOffsetValue SizeOffsetPair = Eval.compute(ObjectSize->getArgOperand(0));
 
     if (SizeOffsetPair != ObjectSizeOffsetEvaluator::unknown()) {
@@ -1230,12 +1232,12 @@ SizeOffsetValue::SizeOffsetValue(const SizeOffsetWeakTrackingVH &SOT)
 
 ObjectSizeOffsetEvaluator::ObjectSizeOffsetEvaluator(
     const DataLayout &DL, const TargetLibraryInfo *TLI, LLVMContext &Context,
-    ObjectSizeOpts EvalOpts)
+    ObjectSizeOpts EvalOpts, const HeapProvenanceAnalysisResult *HPA)
     : DL(DL), TLI(TLI), Context(Context),
       Builder(Context, TargetFolder(DL),
               IRBuilderCallbackInserter(
                   [&](Instruction *I) { InsertedInstructions.insert(I); })),
-      EvalOpts(EvalOpts) {
+      EvalOpts(EvalOpts), HPA(HPA) {
   // IntTy and Zero must be set for each compute() since the address space may
   // be different for later objects.
 }
@@ -1246,11 +1248,10 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     M = I->getModule();
   else if (auto *A = dyn_cast<Argument>(V))
     M = A->getParent()->getParent();
-  if (!M)
+  if (!M || !HPA)
     return false;
 
-  auto HPAResult = HeapProvenanceAnalysis::analyzeModule(*M);
-  auto Info = HPAResult.getInfo(V);
+  auto Info = HPA->getInfo(V);
   if (!Info.isValid())
     return false;
 
@@ -1268,54 +1269,33 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
                            A->getParent()->getEntryBlock().getFirstInsertionPt());
   }
 
-  Value *HeadVal = nullptr;
-  Value *OffsetVal = nullptr;
+  Value *HeadVal = const_cast<Value *>(Info.getHead());
+  if (!HeadVal)
+    return false;
+
+  auto getFn = [](const Value *X) -> const Function * {
+    if (auto *I = dyn_cast<Instruction>(X))
+      return I->getFunction();
+    if (auto *A = dyn_cast<Argument>(X))
+      return A->getParent();
+    return nullptr;
+  };
+  const Function *VF = getFn(V);
+  const Function *HF = getFn(HeadVal);
+  if (VF && HF && VF != HF)
+    return false;
 
-  if (Info.State == HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr) {
-    HeadVal = V;
+  Value *OffsetVal = nullptr;
+  if (HeadVal == V) {
     OffsetVal = Zero;
   } else {
-    Value *Curr = V->stripPointerCasts();
-    while (Curr) {
-      auto CInfo = HPAResult.getInfo(Curr);
-      if (CInfo.State == HeapProvenanceAnalysisResult::ProvenanceInfo::HeapChunkPtr) {
-        HeadVal = Curr;
-        break;
-      }
-      if (auto *GEP = dyn_cast<GEPOperator>(Curr)) {
-        Curr = GEP->getPointerOperand()->stripPointerCasts();
-      } else if (auto *ITP = dyn_cast<IntToPtrInst>(Curr)) {
-        if (auto *BO = dyn_cast<BinaryOperator>(ITP->getOperand(0))) {
-          if (auto *PTI = dyn_cast<PtrToIntOperator>(BO->getOperand(0)))
-            Curr = PTI->getOperand(0)->stripPointerCasts();
-          else if (auto *PTI = dyn_cast<PtrToIntOperator>(BO->getOperand(1)))
-            Curr = PTI->getOperand(0)->stripPointerCasts();
-          else
-            break;
-        } else {
-          break;
-        }
-      } else {
-        break;
-      }
-    }
-
-    if (!HeadVal) {
-      if (Info.SymOffsets.empty()) {
-        HeadVal = Builder.CreateGEP(Builder.getInt8Ty(), V, ConstantInt::getSigned(IntTy, -Info.ConstOffset), "head.meta");
-        OffsetVal = ConstantInt::get(IntTy, Info.ConstOffset);
-      } else {
-        return false;
-      }
-    } else {
-      Value *VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
-      Value *HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
-      OffsetVal = Builder.CreateSub(VInt, HeadInt, "meta.offset");
-    }
+    Value *VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
+    Value *HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
+    OffsetVal = Builder.CreateSub(VInt, HeadInt, "meta.offset");
   }
 
   FunctionCallee GetSizeFC = M->getOrInsertFunction(
-      "__heap_provanence_sanitizer_get_size", IntTy, Builder.getPtrTy());
+      "malloc_usable_size", IntTy, Builder.getPtrTy());
   if (Function *Fn = dyn_cast<Function>(GetSizeFC.getCallee()->stripPointerCasts())) {
     Fn->setDoesNotThrow();
     Fn->setOnlyReadsMemory();
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 2962394549056..a0225f4f8094d 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -38,6 +38,8 @@ MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
 MODULE_ANALYSIS("runtime-libcall-info", RuntimeLibraryAnalysis())
 MODULE_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
+MODULE_ANALYSIS("forward-heap-provenance", ForwardHeapProvenanceAnalysis())
+MODULE_ANALYSIS("backward-heap-provenance", BackwardHeapProvenanceAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
diff --git a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
index fb7e58f4632ef..4c95d5ba79490 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetFolder.h"
@@ -189,7 +190,8 @@ getRuntimeCallName(const BoundsCheckingPass::Options::Runtime &Opts) {
 
 static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
                               ScalarEvolution &SE,
-                              const BoundsCheckingPass::Options &Opts) {
+                              const BoundsCheckingPass::Options &Opts,
+                              const HeapProvenanceAnalysisResult *HPA) {
   if (F.hasFnAttribute(Attribute::NoSanitizeBounds))
     return false;
 
@@ -197,7 +199,7 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
   ObjectSizeOpts EvalOpts;
   EvalOpts.RoundToAlign = true;
   EvalOpts.EvalMode = ObjectSizeOpts::Mode::ExactUnderlyingSizeAndOffset;
-  ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(), EvalOpts);
+  ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(), EvalOpts, HPA);
 
   // check HANDLE_MEMORY_INST in include/llvm/Instruction.def for memory
   // touching instructions
@@ -295,8 +297,10 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
 PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
+  auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+  auto *HPA = MAMProxy.getCachedResult<HeapProvenanceAnalysis>(*F.getParent());
 
-  if (!addBoundsChecking(F, TLI, SE, Opts))
+  if (!addBoundsChecking(F, TLI, SE, Opts, HPA))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();

>From cc15bc15849e772afb7de4372ea6414eb4dd14a8 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:02:52 -0700
Subject: [PATCH 06/18] [Analysis][HeapProvenance] Support integer pointer
 arithmetic in HPA

Extend Forward and Backward HeapProvenanceAnalysis transfer functions to preserve heap pointer provenance across integer manipulations (PtrToIntInst, IntToPtrInst, BinaryOperator).

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index ad4af26c8bdac..974ee64bb275e 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -170,8 +170,12 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
         NewI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
       }
 
-      if (isa<GEPOperator>(&I) || isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
-          isa<IntToPtrInst>(&I) || isa<PHINode>(&I) || isa<SelectInst>(&I)) {
+      if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
+        if (GEP->getPointerOperand() == Op)
+          MergeInto(&I, NewI);
+      } else if (isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
+                 isa<PtrToIntInst>(&I) || isa<IntToPtrInst>(&I) ||
+                 isa<BinaryOperator>(&I) || isa<PHINode>(&I) || isa<SelectInst>(&I)) {
         MergeInto(&I, NewI);
       } else if (auto *CB = dyn_cast<CallBase>(&I)) {
         Function *Callee = CB->getCalledFunction();
@@ -273,8 +277,12 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
       } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&I)) {
         MergeInto(ASC->getOperand(0), BackInfo);
       } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
-        if (auto *PTI = dyn_cast<PtrToIntOperator>(ITP->getOperand(0)))
-          MergeInto(PTI->getOperand(0), BackInfo);
+        MergeInto(ITP->getOperand(0), BackInfo);
+      } else if (auto *PTI = dyn_cast<PtrToIntInst>(&I)) {
+        MergeInto(PTI->getOperand(0), BackInfo);
+      } else if (auto *BO = dyn_cast<BinaryOperator>(&I)) {
+        for (Value *Op : BO->operands())
+          MergeInto(Op, BackInfo);
       } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
         for (Value *InV : PHI->incoming_values())
           MergeInto(InV, BackInfo);

>From aa2955092e0caa7215cba718aa77ffb1cf95ed22 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:06:08 -0700
Subject: [PATCH 07/18] [Analysis][HeapProvenance] Enforce direct state update
 for linear derivations

Ensure linear pointer derivations (GEP, Cast, PtrToInt, IntToPtr, BinaryOperator) directly take their operand's state rather than merging across dataflow solver iterations. This guarantees clean monotonic transitions when upstream PHI nodes resolve across branches.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 974ee64bb275e..48001d9a82098 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -162,7 +162,7 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
       if (OpSeedIt != Seeds.end()) MergeInto(Op, OpSeedIt->second);
 
       auto OpInfo = SS.getExistingValueState(Op);
-      if (!OpInfo.isValid() || !(OpInfo.Dir & HeapProvenanceLattice::Forward)) continue;
+      if (OpInfo.isUninit() || !(OpInfo.Dir & HeapProvenanceLattice::Forward)) continue;
 
       HeapProvenanceLattice NewI = OpInfo;
       if (NewI.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
@@ -172,10 +172,12 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
 
       if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
         if (GEP->getPointerOperand() == Op)
-          MergeInto(&I, NewI);
+          ChangedValues[&I] = NewI;
       } else if (isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
                  isa<PtrToIntInst>(&I) || isa<IntToPtrInst>(&I) ||
-                 isa<BinaryOperator>(&I) || isa<PHINode>(&I) || isa<SelectInst>(&I)) {
+                 isa<BinaryOperator>(&I)) {
+        ChangedValues[&I] = NewI;
+      } else if (isa<PHINode>(&I) || isa<SelectInst>(&I)) {
         MergeInto(&I, NewI);
       } else if (auto *CB = dyn_cast<CallBase>(&I)) {
         Function *Callee = CB->getCalledFunction();

>From b37da713029d816553b5ec02e0e10c89bc20d795 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:31:28 -0700
Subject: [PATCH 08/18] [Clang][Sanitizer] Split bounds checking pipeline to
 run HPA after mem2reg

When -fsanitize=bounds is enabled, split the module optimization pipeline by scheduling an early simplification callback that runs PromotePass (mem2reg) across functions first, followed by RequireAnalysisPass<HeapProvenanceAnalysis, llvm::Module>(). This ensures HPA computes heap provenance on clean SSA form and remains cached when BoundsCheckingPass executes late in the scalar optimization pipeline.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 clang/lib/CodeGen/BackendUtil.cpp            | 12 +++++++++++-
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp |  9 +++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index a46a25c4492f2..27d0505600387 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -68,6 +68,7 @@
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
 #include "llvm/Transforms/Instrumentation/GCOVProfiler.h"
@@ -90,6 +91,7 @@
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Utils/Debugify.h"
+#include "llvm/Transforms/Utils/Mem2Reg.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <limits>
 #include <memory>
@@ -1040,7 +1042,14 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
 
     // Register callbacks to schedule sanitizer passes at the appropriate part
     // of the pipeline.
-    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds))
+    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds)) {
+      PB.registerPipelineEarlySimplificationEPCallback(
+          [](ModulePassManager &MPM, OptimizationLevel Level, ThinOrFullLTOPhase) {
+            FunctionPassManager FPM;
+            FPM.addPass(PromotePass());
+            MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+            MPM.addPass(RequireAnalysisPass<HeapProvenanceAnalysis, llvm::Module>());
+          });
       PB.registerScalarOptimizerLateEPCallback([this](FunctionPassManager &FPM,
                                                       OptimizationLevel Level) {
         BoundsCheckingPass::Options Options;
@@ -1067,6 +1076,7 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
         }
         FPM.addPass(BoundsCheckingPass(Options));
       });
+    }
 
     if (!IsThinLTOPostLink) {
       // Most sanitizers only run during PreLink stage.
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 48001d9a82098..72d28c0307552 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -24,20 +24,17 @@ AnalysisKey HeapProvenanceAnalysis::Key;
 
 bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                                      ModuleAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<ForwardHeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+  return false;
 }
 
 bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                                       ModuleAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<BackwardHeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+  return false;
 }
 
 bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                               ModuleAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<HeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+  return false;
 }
 
 static bool isAllocFunc(StringRef Name) {

>From 63e4d3f714ebda53158904c9960ded2fb826b3a3 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:32:43 -0700
Subject: [PATCH 09/18] [Analysis][HeapProvenance] Restore standard stateless
 analysis invalidation logic

With the Clang pass pipeline split architecture properly executing HPA after mem2reg on functions, standard stateless pass invalidation check (!PAC.preservedWhenStateless()) suffices to preserve HPA cache across subsequent optimization steps.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 72d28c0307552..48001d9a82098 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -24,17 +24,20 @@ AnalysisKey HeapProvenanceAnalysis::Key;
 
 bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                                      ModuleAnalysisManager::Invalidator &) {
-  return false;
+  auto PAC = PA.getChecker<ForwardHeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
 }
 
 bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                                       ModuleAnalysisManager::Invalidator &) {
-  return false;
+  auto PAC = PA.getChecker<BackwardHeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
 }
 
 bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
                                               ModuleAnalysisManager::Invalidator &) {
-  return false;
+  auto PAC = PA.getChecker<HeapProvenanceAnalysis>();
+  return !PAC.preservedWhenStateless();
 }
 
 static bool isAllocFunc(StringRef Name) {

>From 5173a64f7ab270e6f90471709930a9b8f06d8aba Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 15:46:55 -0700
Subject: [PATCH 10/18] [Analysis][HeapProvenance] Fix backward provenance
 propagation and dominance insertion

1. In BackwardLatticeFunc, inspect ChangedValues and SparseSolver existing states when computing instruction states for operands, allowing backward provenance from sinks (e.g. free) to propagate upward to arguments and downward to derived pointer uses (GEPs).
2. In computeFallbackHeapMetadata, promote HeadInst to the top of the entry block when operands permit, guaranteeing SSA dominance over all instrumented accesses.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 21 +++++++++++++++++---
 llvm/lib/Analysis/MemoryBuiltins.cpp         | 14 +++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 48001d9a82098..31c76e7b24429 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -259,11 +259,26 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
 
     for (const Use &U : I.operands()) {
       const Value *Op = U.get();
-      auto OpSeedIt = Seeds.find(Op);
-      if (OpSeedIt != Seeds.end()) MergeInto(Op, OpSeedIt->second);
+      auto OpInfo = ChangedValues[Op];
+      if (OpInfo.isUninit())
+        OpInfo = SS.getExistingValueState(Op);
+      if (OpInfo.isUninit()) {
+        auto OpSeedIt = Seeds.find(Op);
+        if (OpSeedIt != Seeds.end()) OpInfo = OpSeedIt->second;
+      }
+      if (OpInfo.isValid() && (OpInfo.Dir & HeapProvenanceLattice::Backward)) {
+        HeapProvenanceLattice FwdFromBack = OpInfo;
+        if (FwdFromBack.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
+          FwdFromBack.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
+          FwdFromBack.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
+        }
+        MergeInto(&I, FwdFromBack);
+      }
     }
 
-    auto Info = SS.getExistingValueState(&I);
+    auto Info = ChangedValues[&I];
+    if (Info.isUninit())
+      Info = SS.getExistingValueState(&I);
     if (Info.isValid() && (Info.Dir & HeapProvenanceLattice::Backward)) {
       HeapProvenanceLattice BackInfo = Info;
       BackInfo.Dir = HeapProvenanceLattice::Backward;
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index e72e4192f4d3a..d9a49e0fb1ce9 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1285,6 +1285,20 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
   if (VF && HF && VF != HF)
     return false;
 
+  if (auto *HeadInst = dyn_cast<Instruction>(HeadVal)) {
+    if (VF && HF && VF == HF) {
+      bool CanMoveToEntry = true;
+      for (Value *Op : HeadInst->operands()) {
+        if (isa<Instruction>(Op)) {
+          CanMoveToEntry = false;
+          break;
+        }
+      }
+      if (CanMoveToEntry)
+        HeadInst->moveBefore(const_cast<Function *>(VF)->getEntryBlock().getFirstInsertionPt());
+    }
+  }
+
   Value *OffsetVal = nullptr;
   if (HeadVal == V) {
     OffsetVal = Zero;

>From 86fb14e1f125290f300a4d0d11cc2cc6fed5ed9c Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 25 Jun 2026 16:07:19 -0700
Subject: [PATCH 11/18] [Clang][Sanitizer][HeapProvenance] Group bounds check
 pipeline and separate HPA passes

1. In BackendUtil.cpp, group mem2reg, Forward/Backward HPA, and BoundsCheckingPass inside registerOptimizerLastEPCallback. This ensures all standard mutator passes complete before HPA runs, eliminating dangling pointers in HPA cache during bounds instrumentation.
2. Separate ForwardHeapProvenanceAnalysis and BackwardHeapProvenanceAnalysis into independent queries in ObjectSizeOffsetEvaluator and BoundsCheckingPass.
3. Restrict derivation instruction cloning in computeFallbackHeapMetadata to pure stateless linear derivations (GEP, Cast, BinaryOperator), avoiding CGSCC call graph corruption.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 clang/lib/CodeGen/BackendUtil.cpp             | 22 +++++++++-------
 llvm/include/llvm/Analysis/MemoryBuiltins.h   |  8 ++++--
 llvm/lib/Analysis/MemoryBuiltins.cpp          | 26 +++++++++++++------
 .../Instrumentation/BoundsChecking.cpp        | 10 ++++---
 4 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index 27d0505600387..710c9fee6141b 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -1043,15 +1043,15 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
     // Register callbacks to schedule sanitizer passes at the appropriate part
     // of the pipeline.
     if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds)) {
-      PB.registerPipelineEarlySimplificationEPCallback(
-          [](ModulePassManager &MPM, OptimizationLevel Level, ThinOrFullLTOPhase) {
-            FunctionPassManager FPM;
-            FPM.addPass(PromotePass());
-            MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
-            MPM.addPass(RequireAnalysisPass<HeapProvenanceAnalysis, llvm::Module>());
-          });
-      PB.registerScalarOptimizerLateEPCallback([this](FunctionPassManager &FPM,
-                                                      OptimizationLevel Level) {
+      PB.registerOptimizerLastEPCallback([this](ModulePassManager &MPM,
+                                                OptimizationLevel Level,
+                                                ThinOrFullLTOPhase Phase) {
+        FunctionPassManager FPM;
+        FPM.addPass(PromotePass());
+        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+        MPM.addPass(RequireAnalysisPass<ForwardHeapProvenanceAnalysis, llvm::Module>());
+        MPM.addPass(RequireAnalysisPass<BackwardHeapProvenanceAnalysis, llvm::Module>());
+
         BoundsCheckingPass::Options Options;
         if (CodeGenOpts.SanitizeSkipHotCutoffs[SanitizerKind::SO_LocalBounds] ||
             ClSanitizeGuardChecks) {
@@ -1074,7 +1074,9 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
               static_cast<bool>(CodeGenOpts.SanitizeHandlerPreserveAllRegs),
           };
         }
-        FPM.addPass(BoundsCheckingPass(Options));
+        FunctionPassManager BoundsFPM;
+        BoundsFPM.addPass(BoundsCheckingPass(Options));
+        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(BoundsFPM)));
       });
     }
 
diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 088701b579006..d500310bfa71b 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -49,6 +49,8 @@ class SelectInst;
 class Type;
 class UndefValue;
 class Value;
+class ForwardHeapProvenanceAnalysisResult;
+class BackwardHeapProvenanceAnalysisResult;
 
 /// Tests if a value is a call or invoke to a library function that
 /// allocates or reallocates memory (either malloc, calloc, realloc, or strdup
@@ -358,7 +360,8 @@ class ObjectSizeOffsetEvaluator
   PtrSetTy SeenVals;
   ObjectSizeOpts EvalOpts;
   SmallPtrSet<Instruction *, 8> InsertedInstructions;
-  const HeapProvenanceAnalysisResult *HPA;
+  const ForwardHeapProvenanceAnalysisResult *FwdHPA;
+  const BackwardHeapProvenanceAnalysisResult *BwdHPA;
 
   SizeOffsetValue compute_(Value *V);
   bool computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result);
@@ -367,7 +370,8 @@ class ObjectSizeOffsetEvaluator
   LLVM_ABI ObjectSizeOffsetEvaluator(
       const DataLayout &DL, const TargetLibraryInfo *TLI, LLVMContext &Context,
       ObjectSizeOpts EvalOpts = {},
-      const HeapProvenanceAnalysisResult *HPA = nullptr);
+      const ForwardHeapProvenanceAnalysisResult *FwdHPA = nullptr,
+      const BackwardHeapProvenanceAnalysisResult *BwdHPA = nullptr);
 
   static SizeOffsetValue unknown() { return SizeOffsetValue(); }
 
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index d9a49e0fb1ce9..da8ad567988f6 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -675,7 +675,9 @@ Value *llvm::lowerObjectSizeCall(
       return ConstantInt::get(ResultType, Size);
   } else {
     LLVMContext &Ctx = ObjectSize->getFunction()->getContext();
-    ObjectSizeOffsetEvaluator Eval(DL, TLI, Ctx, EvalOptions, HPA);
+    ObjectSizeOffsetEvaluator Eval(DL, TLI, Ctx, EvalOptions,
+                                   HPA ? &HPA->getForwardResult() : nullptr,
+                                   HPA ? &HPA->getBackwardResult() : nullptr);
     SizeOffsetValue SizeOffsetPair = Eval.compute(ObjectSize->getArgOperand(0));
 
     if (SizeOffsetPair != ObjectSizeOffsetEvaluator::unknown()) {
@@ -1232,12 +1234,13 @@ SizeOffsetValue::SizeOffsetValue(const SizeOffsetWeakTrackingVH &SOT)
 
 ObjectSizeOffsetEvaluator::ObjectSizeOffsetEvaluator(
     const DataLayout &DL, const TargetLibraryInfo *TLI, LLVMContext &Context,
-    ObjectSizeOpts EvalOpts, const HeapProvenanceAnalysisResult *HPA)
+    ObjectSizeOpts EvalOpts, const ForwardHeapProvenanceAnalysisResult *FwdHPA,
+    const BackwardHeapProvenanceAnalysisResult *BwdHPA)
     : DL(DL), TLI(TLI), Context(Context),
       Builder(Context, TargetFolder(DL),
               IRBuilderCallbackInserter(
                   [&](Instruction *I) { InsertedInstructions.insert(I); })),
-      EvalOpts(EvalOpts), HPA(HPA) {
+      EvalOpts(EvalOpts), FwdHPA(FwdHPA), BwdHPA(BwdHPA) {
   // IntTy and Zero must be set for each compute() since the address space may
   // be different for later objects.
 }
@@ -1248,10 +1251,12 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     M = I->getModule();
   else if (auto *A = dyn_cast<Argument>(V))
     M = A->getParent()->getParent();
-  if (!M || !HPA)
+  if (!M || (!FwdHPA && !BwdHPA))
     return false;
 
-  auto Info = HPA->getInfo(V);
+  HeapProvenanceLattice Info;
+  if (FwdHPA) Info = FwdHPA->getInfo(V);
+  if (!Info.isValid() && BwdHPA) Info = BwdHPA->getInfo(V);
   if (!Info.isValid())
     return false;
 
@@ -1286,7 +1291,9 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     return false;
 
   if (auto *HeadInst = dyn_cast<Instruction>(HeadVal)) {
-    if (VF && HF && VF == HF) {
+    if (VF && HF && VF == HF &&
+        (isa<GetElementPtrInst>(HeadInst) || isa<CastInst>(HeadInst) || isa<BinaryOperator>(HeadInst)) &&
+        HeadInst->getParent() != &VF->getEntryBlock()) {
       bool CanMoveToEntry = true;
       for (Value *Op : HeadInst->operands()) {
         if (isa<Instruction>(Op)) {
@@ -1294,8 +1301,11 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
           break;
         }
       }
-      if (CanMoveToEntry)
-        HeadInst->moveBefore(const_cast<Function *>(VF)->getEntryBlock().getFirstInsertionPt());
+      if (CanMoveToEntry) {
+        Instruction *Cloned = HeadInst->clone();
+        Cloned->insertBefore(const_cast<Function *>(VF)->getEntryBlock().getFirstInsertionPt());
+        HeadVal = Cloned;
+      }
     }
   }
 
diff --git a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
index 4c95d5ba79490..c6d55a42349ee 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -191,7 +191,8 @@ getRuntimeCallName(const BoundsCheckingPass::Options::Runtime &Opts) {
 static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
                               ScalarEvolution &SE,
                               const BoundsCheckingPass::Options &Opts,
-                              const HeapProvenanceAnalysisResult *HPA) {
+                              const ForwardHeapProvenanceAnalysisResult *FwdHPA,
+                              const BackwardHeapProvenanceAnalysisResult *BwdHPA) {
   if (F.hasFnAttribute(Attribute::NoSanitizeBounds))
     return false;
 
@@ -199,7 +200,7 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
   ObjectSizeOpts EvalOpts;
   EvalOpts.RoundToAlign = true;
   EvalOpts.EvalMode = ObjectSizeOpts::Mode::ExactUnderlyingSizeAndOffset;
-  ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(), EvalOpts, HPA);
+  ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(), EvalOpts, FwdHPA, BwdHPA);
 
   // check HANDLE_MEMORY_INST in include/llvm/Instruction.def for memory
   // touching instructions
@@ -298,9 +299,10 @@ PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
-  auto *HPA = MAMProxy.getCachedResult<HeapProvenanceAnalysis>(*F.getParent());
+  auto *FwdHPA = MAMProxy.getCachedResult<ForwardHeapProvenanceAnalysis>(*F.getParent());
+  auto *BwdHPA = MAMProxy.getCachedResult<BackwardHeapProvenanceAnalysis>(*F.getParent());
 
-  if (!addBoundsChecking(F, TLI, SE, Opts, HPA))
+  if (!addBoundsChecking(F, TLI, SE, Opts, FwdHPA, BwdHPA))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();

>From 7fac5f942ec16a2589435122e68b1a57cfdf0f33 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 26 Jun 2026 09:16:07 -0700
Subject: [PATCH 12/18] [Clang][Sanitizer][HeapProvenance] Check interim
 pointer bounds and verify use list dominance

1. Emit conditional size selection for interim pointers in fallback metadata calculation: if pointer is greater than or equal to chunk head, use malloc_usable_size callback; otherwise select 0.
2. Ensure instructions have a valid use list (preventing ConstantData assertion failures) before tracking or merging lattice states in HPA solver.
3. Automatically adjust insertion point after derivation instruction in computeFallbackHeapMetadata to maintain strict SSA dominance.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 16 ++++++++---
 llvm/lib/Analysis/MemoryBuiltins.cpp         | 30 ++++++++++++++++++--
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 31c76e7b24429..6c992d0914d3f 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -128,7 +128,11 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
       : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
 
   void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
-    mergeLattice(V, Seeds[V], Info);
+    if (V && V->hasUseList())
+      mergeLattice(V, Seeds[V], Info);
+  }
+  bool IsUntrackedValue(const Value *Key) override {
+    return !Key || !Key->hasUseList();
   }
   bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
   HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
@@ -144,7 +148,7 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
                                SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
                                SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
     auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
-      if (!Target || isa<ConstantPointerNull>(Target) || isa<UndefValue>(Target)) return;
+      if (!Target || !Target->hasUseList()) return;
       auto &Dest = ChangedValues[Target];
       if (Dest.isUninit()) {
         auto Existing = SS.getExistingValueState(Target);
@@ -229,7 +233,11 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
       : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
 
   void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
-    mergeLattice(V, Seeds[V], Info);
+    if (V && V->hasUseList())
+      mergeLattice(V, Seeds[V], Info);
+  }
+  bool IsUntrackedValue(const Value *Key) override {
+    return !Key || !Key->hasUseList();
   }
   bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
   HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
@@ -245,7 +253,7 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
                                SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
                                SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
     auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
-      if (!Target || isa<ConstantPointerNull>(Target) || isa<UndefValue>(Target)) return;
+      if (!Target || !Target->hasUseList()) return;
       auto &Dest = ChangedValues[Target];
       if (Dest.isUninit()) {
         auto Existing = SS.getExistingValueState(Target);
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index da8ad567988f6..68c0f36163b8d 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1309,12 +1309,32 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     }
   }
 
+  if (auto *FinalHeadInst = dyn_cast<Instruction>(HeadVal)) {
+    if (auto *VI = dyn_cast<Instruction>(V)) {
+      if (FinalHeadInst->getParent() == VI->getParent() &&
+          !FinalHeadInst->comesBefore(VI))
+        return false;
+    }
+    if (FinalHeadInst->getParent() == Builder.GetInsertBlock()) {
+      if (!FinalHeadInst->comesBefore(&*Builder.GetInsertPoint())) {
+        if (FinalHeadInst->getNextNode())
+          Builder.SetInsertPoint(FinalHeadInst->getNextNode());
+        else
+          Builder.SetInsertPoint(FinalHeadInst->getParent());
+      }
+    } else if (FinalHeadInst->getParent() != &VF->getEntryBlock()) {
+      return false;
+    }
+  }
+
   Value *OffsetVal = nullptr;
+  Value *VInt = nullptr;
+  Value *HeadInt = nullptr;
   if (HeadVal == V) {
     OffsetVal = Zero;
   } else {
-    Value *VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
-    Value *HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
+    VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
+    HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
     OffsetVal = Builder.CreateSub(VInt, HeadInt, "meta.offset");
   }
 
@@ -1328,7 +1348,11 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     Fn->addFnAttr(Attribute::NoCallback);
     Fn->addFnAttr(Attribute::NoSync);
   }
-  Value *ChunkPayloadSize = Builder.CreateCall(GetSizeFC, {HeadVal}, "chunk.size");
+  Value *ChunkPayloadSize = Builder.CreateCall(GetSizeFC, {HeadVal}, "meta.size");
+  if (HeadVal != V) {
+    Value *Cond = Builder.CreateICmpUGE(VInt, HeadInt, "ptr.ge.head");
+    ChunkPayloadSize = Builder.CreateSelect(Cond, ChunkPayloadSize, Zero, "chunk.size");
+  }
   if (OffsetVal->getType() != IntTy)
     OffsetVal = Builder.CreateZExtOrTrunc(OffsetVal, IntTy);
 

>From 67c1a1b06470cfff6c99a8b53e501dd1a54bc755 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 26 Jun 2026 09:36:51 -0700
Subject: [PATCH 13/18] [Clang][Sanitizer][HeapProvenance] Convert HPA to
 FunctionAnalysis and align pipeline with upstream main

1. Remove Interprocedural Analysis (IPA) loops across callers and callees in Forward and Backward HPA lattice transfer functions.
2. Convert ForwardHeapProvenanceAnalysis and BackwardHeapProvenanceAnalysis from ModuleAnalysis to FunctionAnalysis, reducing memory overhead and time complexity to lightweight intraprocedural limits.
3. Restrict backward FwdFromBack derivation rules strictly to GEP pointer operands, casts, PHIs, and Selects, preventing LoadInst from inheriting address provenance.
4. Revert custom BackendUtil pipeline splitting and register BoundsCheckingPass inside registerScalarOptimizerLateEPCallback exactly aligning with LLVM upstream main. Querying HPA via FunctionAnalysisManager computes function provenance on demand right before bounds check instrumentation.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 clang/lib/CodeGen/BackendUtil.cpp             |  20 +-
 .../llvm/Analysis/HeapProvenanceAnalysis.h    |  22 +-
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 266 ++++++------------
 llvm/lib/Passes/PassRegistry.def              |  10 +-
 .../Instrumentation/BoundsChecking.cpp        |   7 +-
 5 files changed, 116 insertions(+), 209 deletions(-)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index 710c9fee6141b..a46a25c4492f2 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -68,7 +68,6 @@
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
-#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
 #include "llvm/Transforms/Instrumentation/GCOVProfiler.h"
@@ -91,7 +90,6 @@
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Utils/Debugify.h"
-#include "llvm/Transforms/Utils/Mem2Reg.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <limits>
 #include <memory>
@@ -1042,16 +1040,9 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
 
     // Register callbacks to schedule sanitizer passes at the appropriate part
     // of the pipeline.
-    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds)) {
-      PB.registerOptimizerLastEPCallback([this](ModulePassManager &MPM,
-                                                OptimizationLevel Level,
-                                                ThinOrFullLTOPhase Phase) {
-        FunctionPassManager FPM;
-        FPM.addPass(PromotePass());
-        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
-        MPM.addPass(RequireAnalysisPass<ForwardHeapProvenanceAnalysis, llvm::Module>());
-        MPM.addPass(RequireAnalysisPass<BackwardHeapProvenanceAnalysis, llvm::Module>());
-
+    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds))
+      PB.registerScalarOptimizerLateEPCallback([this](FunctionPassManager &FPM,
+                                                      OptimizationLevel Level) {
         BoundsCheckingPass::Options Options;
         if (CodeGenOpts.SanitizeSkipHotCutoffs[SanitizerKind::SO_LocalBounds] ||
             ClSanitizeGuardChecks) {
@@ -1074,11 +1065,8 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
               static_cast<bool>(CodeGenOpts.SanitizeHandlerPreserveAllRegs),
           };
         }
-        FunctionPassManager BoundsFPM;
-        BoundsFPM.addPass(BoundsCheckingPass(Options));
-        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(BoundsFPM)));
+        FPM.addPass(BoundsCheckingPass(Options));
       });
-    }
 
     if (!IsThinLTOPostLink) {
       // Most sanitizers only run during PreLink stage.
diff --git a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
index 7144888334bb7..2dad1d62c7d8b 100644
--- a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
+++ b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
@@ -97,8 +97,8 @@ class ForwardHeapProvenanceAnalysisResult {
   const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
     return ValueMap;
   }
-  bool invalidate(Module &, const PreservedAnalyses &PA,
-                  ModuleAnalysisManager::Invalidator &);
+  bool invalidate(Function &, const PreservedAnalyses &PA,
+                  FunctionAnalysisManager::Invalidator &);
 };
 
 class ForwardHeapProvenanceAnalysis
@@ -107,7 +107,7 @@ class ForwardHeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = ForwardHeapProvenanceAnalysisResult;
-  Result run(Module &M, ModuleAnalysisManager &MAM);
+  Result run(Function &F, FunctionAnalysisManager &FAM);
 };
 
 class BackwardHeapProvenanceAnalysisResult {
@@ -125,8 +125,8 @@ class BackwardHeapProvenanceAnalysisResult {
   const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
     return ValueMap;
   }
-  bool invalidate(Module &, const PreservedAnalyses &PA,
-                  ModuleAnalysisManager::Invalidator &);
+  bool invalidate(Function &, const PreservedAnalyses &PA,
+                  FunctionAnalysisManager::Invalidator &);
 };
 
 class BackwardHeapProvenanceAnalysis
@@ -135,7 +135,7 @@ class BackwardHeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = BackwardHeapProvenanceAnalysisResult;
-  Result run(Module &M, ModuleAnalysisManager &MAM);
+  Result run(Function &F, FunctionAnalysisManager &FAM);
 };
 
 // Combined wrapper for backward compatibility with ObjectSizeOffsetEvaluator
@@ -163,8 +163,8 @@ class HeapProvenanceAnalysisResult {
   const ForwardHeapProvenanceAnalysisResult &getForwardResult() const { return ForwardRes; }
   const BackwardHeapProvenanceAnalysisResult &getBackwardResult() const { return BackwardRes; }
 
-  bool invalidate(Module &, const PreservedAnalyses &PA,
-                  ModuleAnalysisManager::Invalidator &);
+  bool invalidate(Function &, const PreservedAnalyses &PA,
+                  FunctionAnalysisManager::Invalidator &);
 };
 
 class HeapProvenanceAnalysis
@@ -173,8 +173,8 @@ class HeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = HeapProvenanceAnalysisResult;
-  static Result analyzeModule(Module &M);
-  Result run(Module &M, ModuleAnalysisManager &MAM);
+  static Result analyzeFunction(Function &F);
+  Result run(Function &F, FunctionAnalysisManager &FAM);
 };
 
 class HeapProvenancePrinterPass
@@ -182,7 +182,7 @@ class HeapProvenancePrinterPass
   raw_ostream &OS;
 public:
   explicit HeapProvenancePrinterPass(raw_ostream &OS) : OS(OS) {}
-  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
 };
 
 } // namespace llvm
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index 6c992d0914d3f..b72478860e8e0 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -22,20 +22,20 @@ AnalysisKey ForwardHeapProvenanceAnalysis::Key;
 AnalysisKey BackwardHeapProvenanceAnalysis::Key;
 AnalysisKey HeapProvenanceAnalysis::Key;
 
-bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                                     ModuleAnalysisManager::Invalidator &) {
+bool ForwardHeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
+                                                     FunctionAnalysisManager::Invalidator &) {
   auto PAC = PA.getChecker<ForwardHeapProvenanceAnalysis>();
   return !PAC.preservedWhenStateless();
 }
 
-bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                                      ModuleAnalysisManager::Invalidator &) {
+bool BackwardHeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
+                                                      FunctionAnalysisManager::Invalidator &) {
   auto PAC = PA.getChecker<BackwardHeapProvenanceAnalysis>();
   return !PAC.preservedWhenStateless();
 }
 
-bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                              ModuleAnalysisManager::Invalidator &) {
+bool HeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
+                                              FunctionAnalysisManager::Invalidator &) {
   auto PAC = PA.getChecker<HeapProvenanceAnalysis>();
   return !PAC.preservedWhenStateless();
 }
@@ -183,44 +183,6 @@ class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPro
         ChangedValues[&I] = NewI;
       } else if (isa<PHINode>(&I) || isa<SelectInst>(&I)) {
         MergeInto(&I, NewI);
-      } else if (auto *CB = dyn_cast<CallBase>(&I)) {
-        Function *Callee = CB->getCalledFunction();
-        if (!Callee) Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-        if (Callee && !Callee->isDeclaration()) {
-          for (unsigned idx = 0, e = CB->arg_size(); idx != e; ++idx) {
-            if (CB->getArgOperand(idx) == Op && idx < Callee->arg_size()) {
-              Argument *Arg = Callee->getArg(idx);
-              HeapProvenanceLattice ArgI = NewI;
-              ArgI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-              ArgI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Arg};
-              MergeInto(Arg, ArgI);
-            }
-          }
-        }
-      }
-    }
-
-    if (auto *RI = dyn_cast<ReturnInst>(&I)) {
-      if (Value *RetVal = RI->getReturnValue()) {
-        auto RetInfo = SS.getExistingValueState(RetVal);
-        if (RetInfo.isValid() && (RetInfo.Dir & HeapProvenanceLattice::Forward)) {
-          HeapProvenanceLattice NewI = RetInfo;
-          if (NewI.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
-            NewI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-            NewI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, RetVal};
-          }
-          const Function *F = RI->getFunction();
-          for (const User *U : F->users()) {
-            if (auto *CB = dyn_cast<CallBase>(U)) {
-              if (CB->getCalledFunction() == F || CB->getCalledOperand()->stripPointerCasts() == F) {
-                HeapProvenanceLattice CBI = NewI;
-                CBI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-                CBI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, CB};
-                MergeInto(CB, CBI);
-              }
-            }
-          }
-        }
       }
     }
   }
@@ -280,7 +242,16 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
           FwdFromBack.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
           FwdFromBack.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
         }
-        MergeInto(&I, FwdFromBack);
+        if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
+          if (GEP->getPointerOperand() == Op)
+            ChangedValues[&I] = FwdFromBack;
+        } else if (isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
+                   isa<PtrToIntInst>(&I) || isa<IntToPtrInst>(&I) ||
+                   isa<BinaryOperator>(&I)) {
+          ChangedValues[&I] = FwdFromBack;
+        } else if (isa<PHINode>(&I) || isa<SelectInst>(&I)) {
+          MergeInto(&I, FwdFromBack);
+        }
       }
     }
 
@@ -314,170 +285,133 @@ class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapPr
       } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
         MergeInto(Sel->getTrueValue(), BackInfo);
         MergeInto(Sel->getFalseValue(), BackInfo);
-      } else if (auto *CB = dyn_cast<CallBase>(&I)) {
-        Function *Callee = CB->getCalledFunction();
-        if (!Callee) Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-        if (Callee && !Callee->isDeclaration() && I.getType()->isPointerTy()) {
-          for (BasicBlock &CalleeBB : *Callee) {
-            if (auto *RI = dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
-              if (Value *RetVal = RI->getReturnValue()) {
-                HeapProvenanceLattice RetI = BackInfo;
-                RetI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-                RetI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, RetVal};
-                MergeInto(RetVal, RetI);
-              }
-            }
-          }
-        }
-      }
-    }
-
-    for (const Use &U : I.operands()) {
-      if (auto *Arg = dyn_cast<Argument>(U.get())) {
-        auto ArgInfo = SS.getExistingValueState(Arg);
-        if (ArgInfo.isValid() && (ArgInfo.Dir & HeapProvenanceLattice::Backward)) {
-          const Function *F = Arg->getParent();
-          for (const User *Usr : F->users()) {
-            if (auto *CB = dyn_cast<CallBase>(Usr)) {
-              if (CB->getCalledFunction() == F || CB->getCalledOperand()->stripPointerCasts() == F) {
-                unsigned ArgIdx = Arg->getArgNo();
-                if (ArgIdx < CB->arg_size()) {
-                  Value *CallerArg = CB->getArgOperand(ArgIdx);
-                  HeapProvenanceLattice CallerI = ArgInfo;
-                  CallerI.Dir = HeapProvenanceLattice::Backward;
-                  CallerI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-                  CallerI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, CallerArg};
-                  MergeInto(CallerArg, CallerI);
-                }
-              }
-            }
-          }
-        }
       }
     }
   }
 };
 
 ForwardHeapProvenanceAnalysis::Result
-ForwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+ForwardHeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
   Result Res;
   ForwardLatticeFunc Lattice;
 
-  for (Function &F : M) {
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        if (auto *CB = dyn_cast<CallBase>(&I)) {
-          Function *Callee = CB->getCalledFunction();
-          if (!Callee)
-            Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-          if (Callee && isAllocFunc(Callee->getName())) {
-            HeapProvenanceLattice Info;
-            Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
-            Info.Dir = HeapProvenanceLattice::Forward;
-            Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
-            Lattice.addSeed(&I, Info);
-          }
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        Function *Callee = CB->getCalledFunction();
+        if (!Callee)
+          Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+        if (Callee && isAllocFunc(Callee->getName())) {
+          HeapProvenanceLattice Info;
+          Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+          Info.Dir = HeapProvenanceLattice::Forward;
+          Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
+          Lattice.addSeed(&I, Info);
         }
       }
     }
   }
 
   SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
-  for (Function &F : M)
-    if (!F.isDeclaration())
-      Solver.MarkBlockExecutable(&F.front());
+  if (!F.isDeclaration())
+    Solver.MarkBlockExecutable(&F.front());
 
   Solver.Solve();
 
-  for (Function &F : M) {
-    for (Argument &Arg : F.args()) {
-      auto LV = Solver.getExistingValueState(&Arg);
-      if (LV.isValid()) Res.setInfo(&Arg, LV);
-    }
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        auto LV = Solver.getExistingValueState(&I);
-        if (LV.isValid()) Res.setInfo(&I, LV);
-      }
+  for (Argument &Arg : F.args()) {
+    auto LV = Solver.getExistingValueState(&Arg);
+    if (LV.isValid()) Res.setInfo(&Arg, LV);
+  }
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      auto LV = Solver.getExistingValueState(&I);
+      if (LV.isValid()) Res.setInfo(&I, LV);
     }
   }
   return Res;
 }
 
 BackwardHeapProvenanceAnalysis::Result
-BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+BackwardHeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
   Result Res;
   BackwardLatticeFunc Lattice;
 
-  for (Function &F : M) {
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        if (auto *CB = dyn_cast<CallBase>(&I)) {
-          Function *Callee = CB->getCalledFunction();
-          if (!Callee)
-            Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-          if (Callee && isFreeFunc(Callee->getName()) && CB->arg_size() > 0) {
-            Value *ArgVal = CB->getArgOperand(0);
-            HeapProvenanceLattice Info;
-            Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
-            Info.Dir = HeapProvenanceLattice::Backward;
-            Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, ArgVal};
-            Lattice.addSeed(ArgVal, Info);
-          }
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        Function *Callee = CB->getCalledFunction();
+        if (!Callee)
+          Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+        if (Callee && isFreeFunc(Callee->getName()) && CB->arg_size() > 0) {
+          Value *ArgVal = CB->getArgOperand(0);
+          HeapProvenanceLattice Info;
+          Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+          Info.Dir = HeapProvenanceLattice::Backward;
+          Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, ArgVal};
+          Lattice.addSeed(ArgVal, Info);
         }
       }
     }
   }
 
   SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
-  for (Function &F : M)
-    if (!F.isDeclaration())
-      Solver.MarkBlockExecutable(&F.front());
+  if (!F.isDeclaration())
+    Solver.MarkBlockExecutable(&F.front());
 
   Solver.Solve();
 
-  for (Function &F : M) {
-    for (Argument &Arg : F.args()) {
-      auto LV = Solver.getExistingValueState(&Arg);
-      if (LV.isValid()) Res.setInfo(&Arg, LV);
-    }
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        auto LV = Solver.getExistingValueState(&I);
-        if (LV.isValid()) Res.setInfo(&I, LV);
-      }
+  for (Argument &Arg : F.args()) {
+    auto LV = Solver.getExistingValueState(&Arg);
+    if (LV.isValid()) Res.setInfo(&Arg, LV);
+  }
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      auto LV = Solver.getExistingValueState(&I);
+      if (LV.isValid()) Res.setInfo(&I, LV);
     }
   }
   return Res;
 }
 
-HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeModule(Module &M) {
-  ModuleAnalysisManager DummyMAM;
+HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeFunction(Function &F) {
+  FunctionAnalysisManager DummyFAM;
   ForwardHeapProvenanceAnalysis ForwardHPA;
   BackwardHeapProvenanceAnalysis BackwardHPA;
-  return Result(ForwardHPA.run(M, DummyMAM), BackwardHPA.run(M, DummyMAM));
+  return Result(ForwardHPA.run(F, DummyFAM), BackwardHPA.run(F, DummyFAM));
 }
 
 HeapProvenanceAnalysis::Result
-HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
-  auto ForwardRes = MAM.getResult<ForwardHeapProvenanceAnalysis>(M);
-  auto BackwardRes = MAM.getResult<BackwardHeapProvenanceAnalysis>(M);
+HeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
+  auto ForwardRes = FAM.getResult<ForwardHeapProvenanceAnalysis>(F);
+  auto BackwardRes = FAM.getResult<BackwardHeapProvenanceAnalysis>(F);
   return Result(std::move(ForwardRes), std::move(BackwardRes));
 }
 
-PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
-                                                 ModuleAnalysisManager &MAM) {
-  auto &Result = MAM.getResult<HeapProvenanceAnalysis>(M);
-  for (Function &F : M) {
-    if (F.isDeclaration())
-      continue;
-    OS << "Printing analysis 'Heap Provenance Analysis' for function '"
-       << F.getName() << "':\n";
-    for (Argument &Arg : F.args()) {
-      auto Info = Result.getInfo(&Arg);
+PreservedAnalyses HeapProvenancePrinterPass::run(Function &F,
+                                                 FunctionAnalysisManager &FAM) {
+  auto &Result = FAM.getResult<HeapProvenanceAnalysis>(F);
+  if (F.isDeclaration())
+    return PreservedAnalyses::all();
+  OS << "Printing analysis 'Heap Provenance Analysis' for function '"
+     << F.getName() << "':\n";
+  for (Argument &Arg : F.args()) {
+    auto Info = Result.getInfo(&Arg);
+    if (Info.isValid()) {
+      OS << "  argument ";
+      Arg.printAsOperand(OS, false);
+      OS << ": "
+         << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
+                 ? "HeapChunkHead"
+                 : "HeapChunkInterim")
+         << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
+    }
+  }
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      auto Info = Result.getInfo(&I);
       if (Info.isValid()) {
-        OS << "  argument ";
-        Arg.printAsOperand(OS, false);
+        OS << "  ";
+        I.printAsOperand(OS, false);
         OS << ": "
            << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
                    ? "HeapChunkHead"
@@ -485,20 +419,6 @@ PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
            << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
       }
     }
-    for (BasicBlock &BB : F) {
-      for (Instruction &I : BB) {
-        auto Info = Result.getInfo(&I);
-        if (Info.isValid()) {
-          OS << "  ";
-          I.printAsOperand(OS, false);
-          OS << ": "
-             << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
-                     ? "HeapChunkHead"
-                     : "HeapChunkInterim")
-             << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
-        }
-      }
-    }
   }
   return PreservedAnalyses::all();
 }
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index a0225f4f8094d..08d10cbaa1257 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -37,9 +37,6 @@ MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
 MODULE_ANALYSIS("runtime-libcall-info", RuntimeLibraryAnalysis())
-MODULE_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
-MODULE_ANALYSIS("forward-heap-provenance", ForwardHeapProvenanceAnalysis())
-MODULE_ANALYSIS("backward-heap-provenance", BackwardHeapProvenanceAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
@@ -147,8 +144,6 @@ MODULE_PASS("print-lcg-dot", LazyCallGraphDOTPrinterPass(errs()))
 MODULE_PASS("print-must-be-executed-contexts",
             MustBeExecutedContextPrinterPass(errs()))
 MODULE_PASS("print-profile-summary", ProfileSummaryPrinterPass(errs()))
-MODULE_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
-MODULE_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 MODULE_PASS("print-stack-safety", StackSafetyGlobalPrinterPass(errs()))
 MODULE_PASS("print<dxil-metadata>", DXILMetadataAnalysisPrinterPass(errs()))
 MODULE_PASS("print<dxil-resources>", DXILResourcePrinterPass(errs()))
@@ -375,6 +370,9 @@ FUNCTION_ANALYSIS("ephemerals", EphemeralValuesAnalysis())
 FUNCTION_ANALYSIS("func-properties", FunctionPropertiesAnalysis())
 FUNCTION_ANALYSIS("machine-function-info", MachineFunctionAnalysis(*TM))
 FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis())
+FUNCTION_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
+FUNCTION_ANALYSIS("forward-heap-provenance", ForwardHeapProvenanceAnalysis())
+FUNCTION_ANALYSIS("backward-heap-provenance", BackwardHeapProvenanceAnalysis())
 FUNCTION_ANALYSIS("last-run-tracking", LastRunTrackingAnalysis())
 FUNCTION_ANALYSIS("lazy-value-info", LazyValueAnalysis())
 FUNCTION_ANALYSIS("loops", LoopAnalysis())
@@ -527,6 +525,8 @@ FUNCTION_PASS("print<demanded-bits>", DemandedBitsPrinterPass(errs()))
 FUNCTION_PASS("print<domfrontier>", DominanceFrontierPrinterPass(errs()))
 FUNCTION_PASS("print<domtree>", DominatorTreePrinterPass(errs()))
 FUNCTION_PASS("print<func-properties>", FunctionPropertiesPrinterPass(errs()))
+FUNCTION_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
+FUNCTION_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 FUNCTION_PASS("print<inline-cost>", InlineCostAnnotationPrinterPass(errs()))
 FUNCTION_PASS("print<lazy-value-info>", LazyValueInfoPrinterPass(errs()))
 FUNCTION_PASS("print<loops>", LoopPrinterPass(errs()))
diff --git a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
index c6d55a42349ee..4925fdcc99803 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -298,11 +298,10 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
 PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
-  auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
-  auto *FwdHPA = MAMProxy.getCachedResult<ForwardHeapProvenanceAnalysis>(*F.getParent());
-  auto *BwdHPA = MAMProxy.getCachedResult<BackwardHeapProvenanceAnalysis>(*F.getParent());
+  auto &FwdHPA = AM.getResult<ForwardHeapProvenanceAnalysis>(F);
+  auto &BwdHPA = AM.getResult<BackwardHeapProvenanceAnalysis>(F);
 
-  if (!addBoundsChecking(F, TLI, SE, Opts, FwdHPA, BwdHPA))
+  if (!addBoundsChecking(F, TLI, SE, Opts, &FwdHPA, &BwdHPA))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();

>From bd08bdf3c92f559a9f401f08c6d9247ce0c83c8d Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 26 Jun 2026 14:16:26 -0700
Subject: [PATCH 14/18] [Clang][Sanitizer][HeapProvenance] Rename interim to
 interior, analyze libcalls via TLI, and multi-split pipeline

1. Rename HeapChunkInterim to HeapChunkInterior across all HPA lattice states, propagation engines, and printer passes.
2. Replace hardcoded function string matching with TargetLibraryInfo libcall analysis (isAllocationFn and getFreedOperand) to identify heap allocation and deallocation sites.
3. Multi-split pipeline in BackendUtil.cpp so HeapProvenanceAnalysis runs after PromotePass (mem2reg) and feeds immutable cached results into BoundsCheckingPass, with proper analysis invalidation.
4. Maintain strict allocation head rooting across interprocedural call boundaries and intermediate derivations so computeFallbackHeapMetadata verifies VF == HF, eliminating false-positive traps on global variables and helper functions.
5. Add safe non-heap fallback sizing selecting all-ones integer when malloc_usable_size returns 0 at runtime.
6. Document SSA dominance analysis and hoisting criteria in HeapProvenanceAnalysis.h without claiming false CFG post-dominance for backward deallocation flow.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 clang/lib/CodeGen/BackendUtil.cpp             |  31 +-
 .../llvm/Analysis/HeapProvenanceAnalysis.h    |  87 +++-
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 490 ++++++++----------
 llvm/lib/Analysis/MemoryBuiltins.cpp          |  70 ++-
 llvm/lib/Passes/PassRegistry.def              |  10 +-
 .../Instrumentation/BoundsChecking.cpp        |  11 +-
 6 files changed, 357 insertions(+), 342 deletions(-)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index a46a25c4492f2..f5ccd9ad22b56 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -66,6 +66,7 @@
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
@@ -90,6 +91,7 @@
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Utils/Debugify.h"
+#include "llvm/Transforms/Utils/Mem2Reg.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <limits>
 #include <memory>
@@ -1040,9 +1042,14 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
 
     // Register callbacks to schedule sanitizer passes at the appropriate part
     // of the pipeline.
-    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds))
-      PB.registerScalarOptimizerLateEPCallback([this](FunctionPassManager &FPM,
-                                                      OptimizationLevel Level) {
+    if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds)) {
+      auto AddBoundsCheckingPipeline = [this](ModulePassManager &MPM) {
+        FunctionPassManager FPM;
+        FPM.addPass(PromotePass());
+        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+        MPM.addPass(RequireAnalysisPass<ForwardHeapProvenanceAnalysis, llvm::Module>());
+        MPM.addPass(RequireAnalysisPass<BackwardHeapProvenanceAnalysis, llvm::Module>());
+
         BoundsCheckingPass::Options Options;
         if (CodeGenOpts.SanitizeSkipHotCutoffs[SanitizerKind::SO_LocalBounds] ||
             ClSanitizeGuardChecks) {
@@ -1065,8 +1072,22 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
               static_cast<bool>(CodeGenOpts.SanitizeHandlerPreserveAllRegs),
           };
         }
-        FPM.addPass(BoundsCheckingPass(Options));
-      });
+        FunctionPassManager BoundsFPM;
+        BoundsFPM.addPass(BoundsCheckingPass(Options));
+        MPM.addPass(createModuleToFunctionPassAdaptor(std::move(BoundsFPM)));
+      };
+
+      PB.registerPipelineStartEPCallback(
+          [AddBoundsCheckingPipeline](ModulePassManager &MPM, OptimizationLevel Level) {
+            if (Level == OptimizationLevel::O0)
+              AddBoundsCheckingPipeline(MPM);
+          });
+      PB.registerOptimizerLastEPCallback(
+          [AddBoundsCheckingPipeline](ModulePassManager &MPM, OptimizationLevel Level,
+                                      ThinOrFullLTOPhase Phase) {
+            AddBoundsCheckingPipeline(MPM);
+          });
+    }
 
     if (!IsThinLTOPostLink) {
       // Most sanitizers only run during PreLink stage.
diff --git a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
index 2dad1d62c7d8b..d937933ec296e 100644
--- a/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
+++ b/llvm/include/llvm/Analysis/HeapProvenanceAnalysis.h
@@ -3,17 +3,58 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
 #include <string>
 
 namespace llvm {
 
+/// ============================================================================
+/// Heap Provenance Analysis (HPA) Theory & Dominance Properties
+/// ============================================================================
+///
+/// 1. Forward vs. Backward Propagation & Dominance Relations:
+///    Let A -> B denote that value A donates its heap provenance to value B.
+///
+///    - Forward Propagation (A --hpa_fwd--> B):
+///      For unary SSA derivations (e.g., B = GEP(A), BitCast(A), AddrSpaceCast(A)),
+///      B is computed directly from operand A along def-use chains, so A strictly
+///      dominates B (A dom B).
+///      Note on Control Flow Merges (PHINode / SelectInst): When B selects between
+///      multiple incoming heads A1, A2, ..., no single antecedent Ai dominates B across
+///      all CFG paths. Therefore, HPA transitions the HeadPayload at PHI/Select nodes
+///      to B itself ({Kind::Phi/Select, B}). Because B trivially dominates all its
+///      own uses, runtime instrumentation safely queries malloc_usable_size(B).
+///
+///    - Backward Propagation (A --hpa_bwd--> B):
+///      A is a deallocation site (e.g., call to free operand B), and B is the
+///      pointer definition being deallocated. In SSA form, because deallocation A
+///      consumes operand B, definition B strictly dominates A (B dom A).
+///      Note: A does NOT post-dominate B in the CFG, because B may return, escape,
+///      or pass to other calls without being freed locally on all paths. Backward
+///      HPA infers a flow-insensitive property: if B is freed anywhere, it must be
+///      a valid heap pointer across its entire live range dominated by B.
+///
+/// 2. GEP & Pointer Arithmetic Variations:
+///    For B = GEP(A, offset), forward propagation maintains A dom B. When bounds
+///    checking instruments a memory access at instruction I using derived pointer
+///    B, HPA traces B back to root allocation head A. Because A dom B and B dom I
+///    (or B == I operand), the root head A strictly dominates the access I.
+///
+/// 3. Hoisting & SSA Insertion Point Criteria:
+///    When BoundsCheckingPass instruments an access on pointer V with root head
+///    HeadVal, the runtime size query (malloc_usable_size(HeadVal)) is positioned
+///    safely in SSA form using InsertAfter on the latest definition between V
+///    and HeadVal. This hoists the metadata computation above the access while
+///    guaranteeing strict SSA dominance across PHIs and Select merges without
+///    breaking basic block terminators.
+/// ============================================================================
 struct HeapProvenanceLattice {
   enum class StateKind {
     Uninit = 0,
     HeapChunkHead,
-    HeapChunkInterim,
+    HeapChunkInterior,
     Unknown
   } State = StateKind::Uninit;
 
@@ -36,7 +77,7 @@ struct HeapProvenanceLattice {
 
   bool isUninit() const { return State == StateKind::Uninit; }
   bool isValid() const {
-    return State == StateKind::HeapChunkHead || State == StateKind::HeapChunkInterim;
+    return State == StateKind::HeapChunkHead || State == StateKind::HeapChunkInterior;
   }
   bool operator==(const HeapProvenanceLattice &RHS) const {
     return State == RHS.State && Dir == RHS.Dir && HeadPayload == RHS.HeadPayload;
@@ -91,14 +132,22 @@ class ForwardHeapProvenanceAnalysisResult {
   const HeapProvenanceLattice &getInfo(const Value *V) const {
     static HeapProvenanceLattice Empty;
     auto It = ValueMap.find(V);
-    return It != ValueMap.end() ? It->second : Empty;
+    if (It != ValueMap.end())
+      return It->second;
+    if (auto *GEP = dyn_cast<GEPOperator>(V))
+      return getInfo(GEP->getPointerOperand());
+    if (auto *BC = dyn_cast<BitCastOperator>(V))
+      return getInfo(BC->getOperand(0));
+    if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(V))
+      return getInfo(ASC->getOperand(0));
+    return Empty;
   }
   DenseMap<const Value *, HeapProvenanceLattice> &getMap() { return ValueMap; }
   const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
     return ValueMap;
   }
-  bool invalidate(Function &, const PreservedAnalyses &PA,
-                  FunctionAnalysisManager::Invalidator &);
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
 };
 
 class ForwardHeapProvenanceAnalysis
@@ -107,7 +156,7 @@ class ForwardHeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = ForwardHeapProvenanceAnalysisResult;
-  Result run(Function &F, FunctionAnalysisManager &FAM);
+  Result run(Module &M, ModuleAnalysisManager &MAM);
 };
 
 class BackwardHeapProvenanceAnalysisResult {
@@ -119,14 +168,22 @@ class BackwardHeapProvenanceAnalysisResult {
   const HeapProvenanceLattice &getInfo(const Value *V) const {
     static HeapProvenanceLattice Empty;
     auto It = ValueMap.find(V);
-    return It != ValueMap.end() ? It->second : Empty;
+    if (It != ValueMap.end())
+      return It->second;
+    if (auto *GEP = dyn_cast<GEPOperator>(V))
+      return getInfo(GEP->getPointerOperand());
+    if (auto *BC = dyn_cast<BitCastOperator>(V))
+      return getInfo(BC->getOperand(0));
+    if (auto *ASC = dyn_cast<AddrSpaceCastOperator>(V))
+      return getInfo(ASC->getOperand(0));
+    return Empty;
   }
   DenseMap<const Value *, HeapProvenanceLattice> &getMap() { return ValueMap; }
   const DenseMap<const Value *, HeapProvenanceLattice> &getMap() const {
     return ValueMap;
   }
-  bool invalidate(Function &, const PreservedAnalyses &PA,
-                  FunctionAnalysisManager::Invalidator &);
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
 };
 
 class BackwardHeapProvenanceAnalysis
@@ -135,7 +192,7 @@ class BackwardHeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = BackwardHeapProvenanceAnalysisResult;
-  Result run(Function &F, FunctionAnalysisManager &FAM);
+  Result run(Module &M, ModuleAnalysisManager &MAM);
 };
 
 // Combined wrapper for backward compatibility with ObjectSizeOffsetEvaluator
@@ -163,8 +220,8 @@ class HeapProvenanceAnalysisResult {
   const ForwardHeapProvenanceAnalysisResult &getForwardResult() const { return ForwardRes; }
   const BackwardHeapProvenanceAnalysisResult &getBackwardResult() const { return BackwardRes; }
 
-  bool invalidate(Function &, const PreservedAnalyses &PA,
-                  FunctionAnalysisManager::Invalidator &);
+  bool invalidate(Module &, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &);
 };
 
 class HeapProvenanceAnalysis
@@ -173,8 +230,8 @@ class HeapProvenanceAnalysis
   static AnalysisKey Key;
 public:
   using Result = HeapProvenanceAnalysisResult;
-  static Result analyzeFunction(Function &F);
-  Result run(Function &F, FunctionAnalysisManager &FAM);
+  static Result analyzeModule(Module &M);
+  Result run(Module &M, ModuleAnalysisManager &MAM);
 };
 
 class HeapProvenancePrinterPass
@@ -182,7 +239,7 @@ class HeapProvenancePrinterPass
   raw_ostream &OS;
 public:
   explicit HeapProvenancePrinterPass(raw_ostream &OS) : OS(OS) {}
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 };
 
 } // namespace llvm
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index b72478860e8e0..b27bf30caf3ec 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -1,55 +1,40 @@
 #include "llvm/Analysis/HeapProvenanceAnalysis.h"
-#include "llvm/Analysis/SparsePropagation.h"
+#include "llvm/Analysis/MemoryBuiltins.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Support/raw_ostream.h"
 
-namespace llvm {
-template <> struct LatticeKeyInfo<const Value *> {
-  static inline Value *getValueFromLatticeKey(const Value *Key) {
-    return const_cast<Value *>(Key);
-  }
-  static inline const Value *getLatticeKeyFromValue(Value *V) {
-    return V;
-  }
-};
-} // namespace llvm
-
 using namespace llvm;
 
 AnalysisKey ForwardHeapProvenanceAnalysis::Key;
 AnalysisKey BackwardHeapProvenanceAnalysis::Key;
 AnalysisKey HeapProvenanceAnalysis::Key;
 
-bool ForwardHeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
-                                                     FunctionAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<ForwardHeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                                     ModuleAnalysisManager::Invalidator &) {
+  return false;
 }
 
-bool BackwardHeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
-                                                      FunctionAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<BackwardHeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                                      ModuleAnalysisManager::Invalidator &) {
+  return false;
 }
 
-bool HeapProvenanceAnalysisResult::invalidate(Function &, const PreservedAnalyses &PA,
-                                              FunctionAnalysisManager::Invalidator &) {
-  auto PAC = PA.getChecker<HeapProvenanceAnalysis>();
-  return !PAC.preservedWhenStateless();
+bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
+                                              ModuleAnalysisManager::Invalidator &) {
+  return false;
 }
 
-static bool isAllocFunc(StringRef Name) {
-  return Name == "malloc" || Name == "calloc" || Name == "realloc" ||
-         Name == "aligned_alloc" || Name == "strdup" || Name == "strndup" ||
-         Name == "g_malloc" || Name == "g_malloc0" || Name == "g_realloc" ||
-         Name == "g_strdup" || Name.starts_with("_Zn");
+static bool isAllocLibCall(const Value *V, const TargetLibraryInfo *TLI) {
+  return isAllocationFn(V, TLI);
 }
 
-static bool isFreeFunc(StringRef Name) {
-  return Name == "free" || Name == "realloc" || Name == "cfree" ||
-         Name == "g_free" || Name.starts_with("_Zd");
+static Value *getFreeLibCallOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
+  return getFreedOperand(CB, TLI);
 }
 
 static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
@@ -92,9 +77,9 @@ static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
 
   if (isa_and_nonnull<PHINode>(Target)) {
     Lattice::Payload NewPayload{Lattice::Payload::Kind::Phi, Target};
-    if (Dest.State != Lattice::StateKind::HeapChunkInterim ||
+    if (Dest.State != Lattice::StateKind::HeapChunkInterior ||
         Dest.HeadPayload != NewPayload) {
-      Dest.State = Lattice::StateKind::HeapChunkInterim;
+      Dest.State = Lattice::StateKind::HeapChunkInterior;
       Dest.HeadPayload = NewPayload;
       Changed = true;
     }
@@ -103,9 +88,9 @@ static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
 
   if (isa_and_nonnull<SelectInst>(Target)) {
     Lattice::Payload NewPayload{Lattice::Payload::Kind::Select, Target};
-    if (Dest.State != Lattice::StateKind::HeapChunkInterim ||
+    if (Dest.State != Lattice::StateKind::HeapChunkInterior ||
         Dest.HeadPayload != NewPayload) {
-      Dest.State = Lattice::StateKind::HeapChunkInterim;
+      Dest.State = Lattice::StateKind::HeapChunkInterior;
       Dest.HeadPayload = NewPayload;
       Changed = true;
     }
@@ -121,304 +106,237 @@ static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
   return Changed;
 }
 
-class ForwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapProvenanceLattice> {
-  DenseMap<const Value *, HeapProvenanceLattice> Seeds;
-public:
-  ForwardLatticeFunc()
-      : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
-
-  void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
-    if (V && V->hasUseList())
-      mergeLattice(V, Seeds[V], Info);
-  }
-  bool IsUntrackedValue(const Value *Key) override {
-    return !Key || !Key->hasUseList();
-  }
-  bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
-  HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
-    mergeLattice(nullptr, X, Y);
-    return X;
-  }
-  HeapProvenanceLattice ComputeLatticeVal(const Value *Key) override {
-    auto It = Seeds.find(Key);
-    if (It != Seeds.end()) return It->second;
-    return getUndefVal();
-  }
-  void ComputeInstructionState(Instruction &I,
-                               SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
-                               SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
-    auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
-      if (!Target || !Target->hasUseList()) return;
-      auto &Dest = ChangedValues[Target];
-      if (Dest.isUninit()) {
-        auto Existing = SS.getExistingValueState(Target);
-        if (Existing.isValid()) Dest = Existing;
-      }
-      mergeLattice(Target, Dest, NewI);
-    };
-
-    auto SeedIt = Seeds.find(&I);
-    if (SeedIt != Seeds.end()) MergeInto(&I, SeedIt->second);
-
-    for (const Use &U : I.operands()) {
-      const Value *Op = U.get();
-      auto OpSeedIt = Seeds.find(Op);
-      if (OpSeedIt != Seeds.end()) MergeInto(Op, OpSeedIt->second);
-
-      auto OpInfo = SS.getExistingValueState(Op);
-      if (OpInfo.isUninit() || !(OpInfo.Dir & HeapProvenanceLattice::Forward)) continue;
-
-      HeapProvenanceLattice NewI = OpInfo;
-      if (NewI.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
-        NewI.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-        NewI.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
-      }
+static bool mergeIntoMap(DenseMap<const Value *, HeapProvenanceLattice> &Map,
+                         const Value *Target,
+                         const HeapProvenanceLattice &Src) {
+  if (!Target || !Target->hasUseList())
+    return false;
+  return mergeLattice(Target, Map[Target], Src);
+}
 
-      if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
-        if (GEP->getPointerOperand() == Op)
-          ChangedValues[&I] = NewI;
-      } else if (isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
-                 isa<PtrToIntInst>(&I) || isa<IntToPtrInst>(&I) ||
-                 isa<BinaryOperator>(&I)) {
-        ChangedValues[&I] = NewI;
-      } else if (isa<PHINode>(&I) || isa<SelectInst>(&I)) {
-        MergeInto(&I, NewI);
+ForwardHeapProvenanceAnalysis::Result
+ForwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  Result Res;
+  SmallVector<const Value *, 64> Worklist;
+  FunctionAnalysisManager &FAM =
+      MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+
+  for (Function &F : M) {
+    if (F.isDeclaration())
+      continue;
+    auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
+    for (Instruction &I : instructions(F)) {
+      if (isAllocLibCall(&I, &TLI)) {
+        HeapProvenanceLattice Info;
+        Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+        Info.Dir = HeapProvenanceLattice::Forward;
+        Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
+        Res.setInfo(&I, Info);
+        Worklist.push_back(&I);
       }
     }
   }
-};
 
-class BackwardLatticeFunc : public AbstractLatticeFunction<const Value *, HeapProvenanceLattice> {
-  DenseMap<const Value *, HeapProvenanceLattice> Seeds;
-public:
-  BackwardLatticeFunc()
-      : AbstractLatticeFunction(HeapProvenanceLattice(), HeapProvenanceLattice(), HeapProvenanceLattice()) {}
+  if (Worklist.empty())
+    return Res;
 
-  void addSeed(const Value *V, const HeapProvenanceLattice &Info) {
-    if (V && V->hasUseList())
-      mergeLattice(V, Seeds[V], Info);
-  }
-  bool IsUntrackedValue(const Value *Key) override {
-    return !Key || !Key->hasUseList();
-  }
-  bool IsSpecialCasedPHI(PHINode *PN) override { return true; }
-  HeapProvenanceLattice MergeValues(HeapProvenanceLattice X, HeapProvenanceLattice Y) override {
-    mergeLattice(nullptr, X, Y);
-    return X;
-  }
-  HeapProvenanceLattice ComputeLatticeVal(const Value *Key) override {
-    auto It = Seeds.find(Key);
-    if (It != Seeds.end()) return It->second;
-    return getUndefVal();
-  }
-  void ComputeInstructionState(Instruction &I,
-                               SmallDenseMap<const Value *, HeapProvenanceLattice, 16> &ChangedValues,
-                               SparseSolver<const Value *, HeapProvenanceLattice> &SS) override {
-    auto MergeInto = [&](const Value *Target, const HeapProvenanceLattice &NewI) {
-      if (!Target || !Target->hasUseList()) return;
-      auto &Dest = ChangedValues[Target];
-      if (Dest.isUninit()) {
-        auto Existing = SS.getExistingValueState(Target);
-        if (Existing.isValid()) Dest = Existing;
-      }
-      mergeLattice(Target, Dest, NewI);
-    };
-
-    auto SeedIt = Seeds.find(&I);
-    if (SeedIt != Seeds.end()) MergeInto(&I, SeedIt->second);
-
-    for (const Use &U : I.operands()) {
-      const Value *Op = U.get();
-      auto OpInfo = ChangedValues[Op];
-      if (OpInfo.isUninit())
-        OpInfo = SS.getExistingValueState(Op);
-      if (OpInfo.isUninit()) {
-        auto OpSeedIt = Seeds.find(Op);
-        if (OpSeedIt != Seeds.end()) OpInfo = OpSeedIt->second;
-      }
-      if (OpInfo.isValid() && (OpInfo.Dir & HeapProvenanceLattice::Backward)) {
-        HeapProvenanceLattice FwdFromBack = OpInfo;
-        if (FwdFromBack.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
-          FwdFromBack.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-          FwdFromBack.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, Op};
-        }
-        if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
-          if (GEP->getPointerOperand() == Op)
-            ChangedValues[&I] = FwdFromBack;
-        } else if (isa<BitCastInst>(&I) || isa<AddrSpaceCastInst>(&I) ||
-                   isa<PtrToIntInst>(&I) || isa<IntToPtrInst>(&I) ||
-                   isa<BinaryOperator>(&I)) {
-          ChangedValues[&I] = FwdFromBack;
-        } else if (isa<PHINode>(&I) || isa<SelectInst>(&I)) {
-          MergeInto(&I, FwdFromBack);
-        }
-      }
-    }
+  while (!Worklist.empty()) {
+    const Value *V = Worklist.pop_back_val();
+    HeapProvenanceLattice Info = Res.getInfo(V);
+    if (!Info.isValid() || !(Info.Dir & HeapProvenanceLattice::Forward))
+      continue;
 
-    auto Info = ChangedValues[&I];
-    if (Info.isUninit())
-      Info = SS.getExistingValueState(&I);
-    if (Info.isValid() && (Info.Dir & HeapProvenanceLattice::Backward)) {
-      HeapProvenanceLattice BackInfo = Info;
-      BackInfo.Dir = HeapProvenanceLattice::Backward;
-      if (BackInfo.State == HeapProvenanceLattice::StateKind::HeapChunkHead) {
-        BackInfo.State = HeapProvenanceLattice::StateKind::HeapChunkInterim;
-        BackInfo.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
-      }
+    HeapProvenanceLattice Derived = Info;
+    if (Derived.State == HeapProvenanceLattice::StateKind::HeapChunkHead)
+      Derived.State = HeapProvenanceLattice::StateKind::HeapChunkInterior;
 
-      if (auto *GEP = dyn_cast<GEPOperator>(&I)) {
-        MergeInto(GEP->getPointerOperand(), BackInfo);
-      } else if (auto *BC = dyn_cast<BitCastInst>(&I)) {
-        MergeInto(BC->getOperand(0), BackInfo);
-      } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&I)) {
-        MergeInto(ASC->getOperand(0), BackInfo);
-      } else if (auto *ITP = dyn_cast<IntToPtrInst>(&I)) {
-        MergeInto(ITP->getOperand(0), BackInfo);
-      } else if (auto *PTI = dyn_cast<PtrToIntInst>(&I)) {
-        MergeInto(PTI->getOperand(0), BackInfo);
-      } else if (auto *BO = dyn_cast<BinaryOperator>(&I)) {
-        for (Value *Op : BO->operands())
-          MergeInto(Op, BackInfo);
-      } else if (auto *PHI = dyn_cast<PHINode>(&I)) {
-        for (Value *InV : PHI->incoming_values())
-          MergeInto(InV, BackInfo);
-      } else if (auto *Sel = dyn_cast<SelectInst>(&I)) {
-        MergeInto(Sel->getTrueValue(), BackInfo);
-        MergeInto(Sel->getFalseValue(), BackInfo);
-      }
-    }
-  }
-};
-
-ForwardHeapProvenanceAnalysis::Result
-ForwardHeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
-  Result Res;
-  ForwardLatticeFunc Lattice;
-
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      if (auto *CB = dyn_cast<CallBase>(&I)) {
-        Function *Callee = CB->getCalledFunction();
-        if (!Callee)
-          Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-        if (Callee && isAllocFunc(Callee->getName())) {
-          HeapProvenanceLattice Info;
-          Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
-          Info.Dir = HeapProvenanceLattice::Forward;
-          Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, &I};
-          Lattice.addSeed(&I, Info);
+    for (const User *U : V->users()) {
+      if (auto *GEP = dyn_cast<GEPOperator>(U)) {
+        if (GEP->getPointerOperand() == V) {
+          if (mergeIntoMap(Res.getMap(), U, Derived))
+            Worklist.push_back(U);
+        }
+      } else if (isa<BitCastInst>(U) || isa<AddrSpaceCastInst>(U) ||
+                 isa<PtrToIntInst>(U) || isa<IntToPtrInst>(U) ||
+                 isa<BinaryOperator>(U)) {
+        if (mergeIntoMap(Res.getMap(), U, Derived))
+          Worklist.push_back(U);
+      } else if (isa<PHINode>(U) || isa<SelectInst>(U)) {
+        if (mergeIntoMap(Res.getMap(), U, Derived))
+          Worklist.push_back(U);
+      } else if (auto *CB = dyn_cast<CallBase>(U)) {
+        for (unsigned i = 0, e = CB->arg_size(); i != e; ++i) {
+          if (CB->getArgOperand(i) == V) {
+            Function *Callee = CB->getCalledFunction();
+            if (!Callee)
+              Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+            if (Callee && !Callee->isDeclaration() && i < Callee->arg_size()) {
+              Argument *Arg = Callee->getArg(i);
+              if (mergeIntoMap(Res.getMap(), Arg, Derived))
+                Worklist.push_back(Arg);
+            }
+          }
+        }
+      } else if (auto *RI = dyn_cast<ReturnInst>(U)) {
+        Function *Fn = const_cast<Function *>(RI->getParent()->getParent());
+        for (const User *FnU : Fn->users()) {
+          if (auto *Call = dyn_cast<CallBase>(FnU)) {
+            if (mergeIntoMap(Res.getMap(), Call, Derived))
+              Worklist.push_back(Call);
+          }
         }
       }
     }
   }
 
-  SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
-  if (!F.isDeclaration())
-    Solver.MarkBlockExecutable(&F.front());
-
-  Solver.Solve();
-
-  for (Argument &Arg : F.args()) {
-    auto LV = Solver.getExistingValueState(&Arg);
-    if (LV.isValid()) Res.setInfo(&Arg, LV);
-  }
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      auto LV = Solver.getExistingValueState(&I);
-      if (LV.isValid()) Res.setInfo(&I, LV);
-    }
-  }
   return Res;
 }
 
 BackwardHeapProvenanceAnalysis::Result
-BackwardHeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
+BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
   Result Res;
-  BackwardLatticeFunc Lattice;
-
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
+  SmallVector<const Value *, 64> Worklist;
+  FunctionAnalysisManager &FAM =
+      MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+
+  for (Function &F : M) {
+    if (F.isDeclaration())
+      continue;
+    auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
+    for (Instruction &I : instructions(F)) {
       if (auto *CB = dyn_cast<CallBase>(&I)) {
-        Function *Callee = CB->getCalledFunction();
-        if (!Callee)
-          Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-        if (Callee && isFreeFunc(Callee->getName()) && CB->arg_size() > 0) {
-          Value *ArgVal = CB->getArgOperand(0);
+        if (Value *ArgVal = getFreeLibCallOperand(CB, &TLI)) {
           HeapProvenanceLattice Info;
           Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
           Info.Dir = HeapProvenanceLattice::Backward;
           Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, ArgVal};
-          Lattice.addSeed(ArgVal, Info);
+          if (mergeIntoMap(Res.getMap(), ArgVal, Info))
+            Worklist.push_back(ArgVal);
         }
       }
     }
   }
 
-  SparseSolver<const Value *, HeapProvenanceLattice> Solver(&Lattice);
-  if (!F.isDeclaration())
-    Solver.MarkBlockExecutable(&F.front());
-
-  Solver.Solve();
-
-  for (Argument &Arg : F.args()) {
-    auto LV = Solver.getExistingValueState(&Arg);
-    if (LV.isValid()) Res.setInfo(&Arg, LV);
-  }
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      auto LV = Solver.getExistingValueState(&I);
-      if (LV.isValid()) Res.setInfo(&I, LV);
+  if (Worklist.empty())
+    return Res;
+
+  while (!Worklist.empty()) {
+    const Value *V = Worklist.pop_back_val();
+    HeapProvenanceLattice Info = Res.getInfo(V);
+    if (!Info.isValid() || !(Info.Dir & HeapProvenanceLattice::Backward))
+      continue;
+
+    HeapProvenanceLattice BackInfo = Info;
+    BackInfo.Dir = HeapProvenanceLattice::Backward;
+    if (BackInfo.State == HeapProvenanceLattice::StateKind::HeapChunkHead)
+      BackInfo.State = HeapProvenanceLattice::StateKind::HeapChunkInterior;
+
+    if (auto *GEP = dyn_cast<GEPOperator>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), GEP->getPointerOperand(), BackInfo))
+        Worklist.push_back(GEP->getPointerOperand());
+    } else if (auto *BC = dyn_cast<BitCastInst>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), BC->getOperand(0), BackInfo))
+        Worklist.push_back(BC->getOperand(0));
+    } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), ASC->getOperand(0), BackInfo))
+        Worklist.push_back(ASC->getOperand(0));
+    } else if (auto *ITP = dyn_cast<IntToPtrInst>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), ITP->getOperand(0), BackInfo))
+        Worklist.push_back(ITP->getOperand(0));
+    } else if (auto *PTI = dyn_cast<PtrToIntInst>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), PTI->getOperand(0), BackInfo))
+        Worklist.push_back(PTI->getOperand(0));
+    } else if (auto *BO = dyn_cast<BinaryOperator>(const_cast<Value *>(V))) {
+      for (Value *Op : BO->operands())
+        if (mergeIntoMap(Res.getMap(), Op, BackInfo))
+          Worklist.push_back(Op);
+    } else if (auto *PHI = dyn_cast<PHINode>(const_cast<Value *>(V))) {
+      for (Value *InV : PHI->incoming_values())
+        if (mergeIntoMap(Res.getMap(), InV, BackInfo))
+          Worklist.push_back(InV);
+    } else if (auto *Sel = dyn_cast<SelectInst>(const_cast<Value *>(V))) {
+      if (mergeIntoMap(Res.getMap(), Sel->getTrueValue(), BackInfo))
+        Worklist.push_back(Sel->getTrueValue());
+      if (mergeIntoMap(Res.getMap(), Sel->getFalseValue(), BackInfo))
+        Worklist.push_back(Sel->getFalseValue());
+    } else if (auto *Arg = dyn_cast<Argument>(const_cast<Value *>(V))) {
+      Function *Fn = Arg->getParent();
+      for (const User *FnU : Fn->users()) {
+        if (auto *CB = dyn_cast<CallBase>(FnU)) {
+          if (Arg->getArgNo() < CB->arg_size()) {
+            Value *CallOp = CB->getArgOperand(Arg->getArgNo());
+            if (mergeIntoMap(Res.getMap(), CallOp, BackInfo))
+              Worklist.push_back(CallOp);
+          }
+        }
+      }
+    } else if (auto *CB = dyn_cast<CallBase>(const_cast<Value *>(V))) {
+      Function *Callee = CB->getCalledFunction();
+      if (!Callee)
+        Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+      if (Callee && !Callee->isDeclaration() && V->getType()->isPointerTy()) {
+        for (BasicBlock &CalleeBB : *Callee) {
+          if (auto *RI = dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
+            if (Value *RetVal = RI->getReturnValue()) {
+              if (mergeIntoMap(Res.getMap(), RetVal, BackInfo))
+                Worklist.push_back(RetVal);
+            }
+          }
+        }
+      }
     }
   }
+
   return Res;
 }
 
-HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeFunction(Function &F) {
-  FunctionAnalysisManager DummyFAM;
+HeapProvenanceAnalysisResult HeapProvenanceAnalysis::analyzeModule(Module &M) {
+  ModuleAnalysisManager DummyMAM;
   ForwardHeapProvenanceAnalysis ForwardHPA;
   BackwardHeapProvenanceAnalysis BackwardHPA;
-  return Result(ForwardHPA.run(F, DummyFAM), BackwardHPA.run(F, DummyFAM));
+  return Result(ForwardHPA.run(M, DummyMAM), BackwardHPA.run(M, DummyMAM));
 }
 
 HeapProvenanceAnalysis::Result
-HeapProvenanceAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
-  auto ForwardRes = FAM.getResult<ForwardHeapProvenanceAnalysis>(F);
-  auto BackwardRes = FAM.getResult<BackwardHeapProvenanceAnalysis>(F);
+HeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+  auto ForwardRes = MAM.getResult<ForwardHeapProvenanceAnalysis>(M);
+  auto BackwardRes = MAM.getResult<BackwardHeapProvenanceAnalysis>(M);
   return Result(std::move(ForwardRes), std::move(BackwardRes));
 }
 
-PreservedAnalyses HeapProvenancePrinterPass::run(Function &F,
-                                                 FunctionAnalysisManager &FAM) {
-  auto &Result = FAM.getResult<HeapProvenanceAnalysis>(F);
-  if (F.isDeclaration())
-    return PreservedAnalyses::all();
-  OS << "Printing analysis 'Heap Provenance Analysis' for function '"
-     << F.getName() << "':\n";
-  for (Argument &Arg : F.args()) {
-    auto Info = Result.getInfo(&Arg);
-    if (Info.isValid()) {
-      OS << "  argument ";
-      Arg.printAsOperand(OS, false);
-      OS << ": "
-         << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
-                 ? "HeapChunkHead"
-                 : "HeapChunkInterim")
-         << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
-    }
-  }
-  for (BasicBlock &BB : F) {
-    for (Instruction &I : BB) {
-      auto Info = Result.getInfo(&I);
+PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
+                                                 ModuleAnalysisManager &MAM) {
+  auto &Result = MAM.getResult<HeapProvenanceAnalysis>(M);
+  OS << "Printing analysis 'Heap Provenance Analysis' for module '"
+     << M.getName() << "':\n";
+  for (Function &F : M) {
+    if (F.isDeclaration()) continue;
+    for (Argument &Arg : F.args()) {
+      auto Info = Result.getInfo(&Arg);
       if (Info.isValid()) {
-        OS << "  ";
-        I.printAsOperand(OS, false);
+        OS << "  argument ";
+        Arg.printAsOperand(OS, false);
         OS << ": "
            << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
                    ? "HeapChunkHead"
-                   : "HeapChunkInterim")
+                   : "HeapChunkInterior")
            << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
       }
     }
+    for (BasicBlock &BB : F) {
+      for (Instruction &I : BB) {
+        auto Info = Result.getInfo(&I);
+        if (Info.isValid()) {
+          OS << "  ";
+          I.printAsOperand(OS, false);
+          OS << ": "
+             << (Info.State == HeapProvenanceLattice::StateKind::HeapChunkHead
+                     ? "HeapChunkHead"
+                     : "HeapChunkInterior")
+             << " (" << Info.getExpr() << ")" << Info.getDirectionStr() << "\n";
+        }
+      }
+    }
   }
   return PreservedAnalyses::all();
 }
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index 68c0f36163b8d..9750cc376b964 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1275,7 +1275,7 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
   }
 
   Value *HeadVal = const_cast<Value *>(Info.getHead());
-  if (!HeadVal)
+  if (!HeadVal || !HeadVal->getType()->isPointerTy() || !V->getType()->isPointerTy())
     return false;
 
   auto getFn = [](const Value *X) -> const Function * {
@@ -1290,41 +1290,53 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
   if (VF && HF && VF != HF)
     return false;
 
-  if (auto *HeadInst = dyn_cast<Instruction>(HeadVal)) {
-    if (VF && HF && VF == HF &&
-        (isa<GetElementPtrInst>(HeadInst) || isa<CastInst>(HeadInst) || isa<BinaryOperator>(HeadInst)) &&
-        HeadInst->getParent() != &VF->getEntryBlock()) {
-      bool CanMoveToEntry = true;
-      for (Value *Op : HeadInst->operands()) {
-        if (isa<Instruction>(Op)) {
-          CanMoveToEntry = false;
-          break;
-        }
+  BasicBlock *CallBlock = Builder.GetInsertBlock();
+  if (CallBlock) {
+    if (auto *HI = dyn_cast<Instruction>(HeadVal)) {
+      if (HI->getParent() == CallBlock) {
+        if (HI != &*Builder.GetInsertPoint() && !HI->comesBefore(&*Builder.GetInsertPoint()))
+          return false;
+      } else if (HI->getParent() != &VF->getEntryBlock()) {
+        return false;
       }
-      if (CanMoveToEntry) {
-        Instruction *Cloned = HeadInst->clone();
-        Cloned->insertBefore(const_cast<Function *>(VF)->getEntryBlock().getFirstInsertionPt());
-        HeadVal = Cloned;
+    }
+    if (auto *VI = dyn_cast<Instruction>(V)) {
+      if (VI->getParent() == CallBlock) {
+        if (VI != &*Builder.GetInsertPoint() && !VI->comesBefore(&*Builder.GetInsertPoint()))
+          return false;
+      } else if (VI->getParent() != &VF->getEntryBlock()) {
+        return false;
       }
     }
   }
 
-  if (auto *FinalHeadInst = dyn_cast<Instruction>(HeadVal)) {
-    if (auto *VI = dyn_cast<Instruction>(V)) {
-      if (FinalHeadInst->getParent() == VI->getParent() &&
-          !FinalHeadInst->comesBefore(VI))
+  Instruction *InsertAfter = nullptr;
+  if (auto *VI = dyn_cast<Instruction>(V)) {
+    if (auto *HI = dyn_cast<Instruction>(HeadVal)) {
+      if (VI->getParent() == HI->getParent()) {
+        InsertAfter = VI->comesBefore(HI) ? HI : VI;
+      } else if (HI->getParent() == &VF->getEntryBlock()) {
+        InsertAfter = VI;
+      } else if (VI->getParent() == &VF->getEntryBlock()) {
+        InsertAfter = HI;
+      } else {
         return false;
-    }
-    if (FinalHeadInst->getParent() == Builder.GetInsertBlock()) {
-      if (!FinalHeadInst->comesBefore(&*Builder.GetInsertPoint())) {
-        if (FinalHeadInst->getNextNode())
-          Builder.SetInsertPoint(FinalHeadInst->getNextNode());
-        else
-          Builder.SetInsertPoint(FinalHeadInst->getParent());
       }
-    } else if (FinalHeadInst->getParent() != &VF->getEntryBlock()) {
-      return false;
+    } else {
+      InsertAfter = VI;
     }
+  } else if (auto *HI = dyn_cast<Instruction>(HeadVal)) {
+    InsertAfter = HI;
+  }
+
+  IRBuilder<>::InsertPointGuard Guard(Builder);
+  if (InsertAfter) {
+    if (auto OptIP = InsertAfter->getInsertionPointAfterDef())
+      Builder.SetInsertPoint(*OptIP);
+    else
+      return false;
+  } else {
+    Builder.SetInsertPoint(const_cast<Function *>(VF)->getEntryBlock().getFirstInsertionPt());
   }
 
   Value *OffsetVal = nullptr;
@@ -1349,10 +1361,12 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     Fn->addFnAttr(Attribute::NoSync);
   }
   Value *ChunkPayloadSize = Builder.CreateCall(GetSizeFC, {HeadVal}, "meta.size");
+  Value *IsHeap = Builder.CreateICmpNE(ChunkPayloadSize, Zero, "is.heap");
   if (HeadVal != V) {
     Value *Cond = Builder.CreateICmpUGE(VInt, HeadInt, "ptr.ge.head");
     ChunkPayloadSize = Builder.CreateSelect(Cond, ChunkPayloadSize, Zero, "chunk.size");
   }
+  ChunkPayloadSize = Builder.CreateSelect(IsHeap, ChunkPayloadSize, ConstantInt::getAllOnesValue(IntTy), "meta.size.checked");
   if (OffsetVal->getType() != IntTy)
     OffsetVal = Builder.CreateZExtOrTrunc(OffsetVal, IntTy);
 
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 08d10cbaa1257..489cf4bd77de8 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -21,6 +21,9 @@
 MODULE_ANALYSIS("callgraph", CallGraphAnalysis())
 MODULE_ANALYSIS("collector-metadata", CollectorMetadataAnalysis())
 MODULE_ANALYSIS("ctx-prof-analysis", CtxProfAnalysis())
+MODULE_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
+MODULE_ANALYSIS("forward-heap-provenance", ForwardHeapProvenanceAnalysis())
+MODULE_ANALYSIS("backward-heap-provenance", BackwardHeapProvenanceAnalysis())
 MODULE_ANALYSIS("dxil-metadata", DXILMetadataAnalysis())
 MODULE_ANALYSIS("dxil-resources", DXILResourceAnalysis())
 MODULE_ANALYSIS("dxil-resource-type", DXILResourceTypeAnalysis())
@@ -147,6 +150,8 @@ MODULE_PASS("print-profile-summary", ProfileSummaryPrinterPass(errs()))
 MODULE_PASS("print-stack-safety", StackSafetyGlobalPrinterPass(errs()))
 MODULE_PASS("print<dxil-metadata>", DXILMetadataAnalysisPrinterPass(errs()))
 MODULE_PASS("print<dxil-resources>", DXILResourcePrinterPass(errs()))
+MODULE_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
+MODULE_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 MODULE_PASS("print<inline-advisor>", InlineAdvisorAnalysisPrinterPass(errs()))
 MODULE_PASS("print<ir2vec>", IR2VecPrinterPass(errs()))
 MODULE_PASS("print<ir2vec-vocab>", IR2VecVocabPrinterPass(errs()))
@@ -370,9 +375,6 @@ FUNCTION_ANALYSIS("ephemerals", EphemeralValuesAnalysis())
 FUNCTION_ANALYSIS("func-properties", FunctionPropertiesAnalysis())
 FUNCTION_ANALYSIS("machine-function-info", MachineFunctionAnalysis(*TM))
 FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis())
-FUNCTION_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
-FUNCTION_ANALYSIS("forward-heap-provenance", ForwardHeapProvenanceAnalysis())
-FUNCTION_ANALYSIS("backward-heap-provenance", BackwardHeapProvenanceAnalysis())
 FUNCTION_ANALYSIS("last-run-tracking", LastRunTrackingAnalysis())
 FUNCTION_ANALYSIS("lazy-value-info", LazyValueAnalysis())
 FUNCTION_ANALYSIS("loops", LoopAnalysis())
@@ -525,8 +527,6 @@ FUNCTION_PASS("print<demanded-bits>", DemandedBitsPrinterPass(errs()))
 FUNCTION_PASS("print<domfrontier>", DominanceFrontierPrinterPass(errs()))
 FUNCTION_PASS("print<domtree>", DominatorTreePrinterPass(errs()))
 FUNCTION_PASS("print<func-properties>", FunctionPropertiesPrinterPass(errs()))
-FUNCTION_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
-FUNCTION_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 FUNCTION_PASS("print<inline-cost>", InlineCostAnnotationPrinterPass(errs()))
 FUNCTION_PASS("print<lazy-value-info>", LazyValueInfoPrinterPass(errs()))
 FUNCTION_PASS("print<loops>", LoopPrinterPass(errs()))
diff --git a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
index 4925fdcc99803..d7bc0b57906b5 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -298,10 +298,15 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
 PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
-  auto &FwdHPA = AM.getResult<ForwardHeapProvenanceAnalysis>(F);
-  auto &BwdHPA = AM.getResult<BackwardHeapProvenanceAnalysis>(F);
+  auto *MAMProxy = AM.getCachedResult<ModuleAnalysisManagerFunctionProxy>(F);
+  const ForwardHeapProvenanceAnalysisResult *FwdHPA = nullptr;
+  const BackwardHeapProvenanceAnalysisResult *BwdHPA = nullptr;
+  if (MAMProxy) {
+    FwdHPA = MAMProxy->getCachedResult<ForwardHeapProvenanceAnalysis>(*F.getParent());
+    BwdHPA = MAMProxy->getCachedResult<BackwardHeapProvenanceAnalysis>(*F.getParent());
+  }
 
-  if (!addBoundsChecking(F, TLI, SE, Opts, &FwdHPA, &BwdHPA))
+  if (!addBoundsChecking(F, TLI, SE, Opts, FwdHPA, BwdHPA))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();

>From bc1d4bf50408385329553877b13accb254765c7c Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 26 Jun 2026 16:34:07 -0700
Subject: [PATCH 15/18] Fix HPA meet bottom property, restrict libcall free
 solver, and reconstruct PHIs in computeFallbackHeapMetadata

Enforce dataflow bottom property when known heap allocation properties meet Unknown or conflicting heads in HPA. Restrict getFreeLibCallOperand to genuine TLI free functions. Stop backward propagation at PHI nodes and Select instructions for soundness. Implement recursive PHI and Select reconstruction in computeFallbackHeapMetadata.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/include/llvm/Analysis/MemoryBuiltins.h   |  2 +
 .../Scalar/LowerConstantIntrinsics.h          |  4 +-
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp  | 58 +++++------
 llvm/lib/Analysis/MemoryBuiltins.cpp          | 97 ++++++++++++++++++-
 llvm/lib/Passes/PassBuilderPipelines.cpp      |  2 +
 .../InstCombine/InstructionCombining.cpp      |  4 +-
 .../Scalar/LowerConstantIntrinsics.cpp        | 11 ++-
 7 files changed, 140 insertions(+), 38 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index d500310bfa71b..06303f643a9f5 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -365,6 +365,8 @@ class ObjectSizeOffsetEvaluator
 
   SizeOffsetValue compute_(Value *V);
   bool computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result);
+  bool computeFallbackHeapMetadata_(Value *V, SizeOffsetValue &Result,
+                                    CacheMapTy &FallbackCacheMap);
 
 public:
   LLVM_ABI ObjectSizeOffsetEvaluator(
diff --git a/llvm/include/llvm/Transforms/Scalar/LowerConstantIntrinsics.h b/llvm/include/llvm/Transforms/Scalar/LowerConstantIntrinsics.h
index 5d0d551bdcbb4..f81dd2edf8d32 100644
--- a/llvm/include/llvm/Transforms/Scalar/LowerConstantIntrinsics.h
+++ b/llvm/include/llvm/Transforms/Scalar/LowerConstantIntrinsics.h
@@ -21,10 +21,12 @@ namespace llvm {
 
 class DominatorTree;
 class Function;
+class HeapProvenanceAnalysisResult;
 class TargetLibraryInfo;
 
 LLVM_ABI bool lowerConstantIntrinsics(Function &F, const TargetLibraryInfo &TLI,
-                                      DominatorTree *DT);
+                                      DominatorTree *DT,
+                                      const HeapProvenanceAnalysisResult *HPA = nullptr);
 
 struct LowerConstantIntrinsicsPass
     : OptionalPassInfoMixin<LowerConstantIntrinsicsPass> {
diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index b27bf30caf3ec..a8ce43da40120 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -34,7 +34,17 @@ static bool isAllocLibCall(const Value *V, const TargetLibraryInfo *TLI) {
 }
 
 static Value *getFreeLibCallOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
-  return getFreedOperand(CB, TLI);
+  if (!CB || !TLI)
+    return nullptr;
+  const Function *Callee = CB->getCalledFunction();
+  if (!Callee)
+    return nullptr;
+  LibFunc TLIFn;
+  if (TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn) &&
+      isLibFreeFunction(Callee, TLIFn)) {
+    return CB->getArgOperand(0);
+  }
+  return nullptr;
 }
 
 static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
@@ -59,11 +69,14 @@ static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
     return false;
   }
 
-  if (Src.State == Lattice::StateKind::Unknown) {
+  if (Src.State == Lattice::StateKind::Unknown ||
+      Dest.State == Lattice::StateKind::Unknown) {
+    bool Changed =
+        (Dest.State != Lattice::StateKind::Unknown) || (Dest.Dir != MergedDir);
     Dest.State = Lattice::StateKind::Unknown;
     Dest.Dir = MergedDir;
     Dest.HeadPayload = {Lattice::Payload::Kind::None, nullptr};
-    return true;
+    return Changed;
   }
 
   bool Changed = false;
@@ -75,29 +88,15 @@ static bool mergeLattice(const Value *Target, HeapProvenanceLattice &Dest,
   if (Dest.State == Src.State && Dest.HeadPayload == Src.HeadPayload)
     return Changed;
 
-  if (isa_and_nonnull<PHINode>(Target)) {
-    Lattice::Payload NewPayload{Lattice::Payload::Kind::Phi, Target};
-    if (Dest.State != Lattice::StateKind::HeapChunkInterior ||
-        Dest.HeadPayload != NewPayload) {
+  if (Dest.isUninit()) {
+    Dest.State = Src.State;
+    if (isa_and_nonnull<PHINode>(Target) || isa_and_nonnull<SelectInst>(Target))
       Dest.State = Lattice::StateKind::HeapChunkInterior;
-      Dest.HeadPayload = NewPayload;
-      Changed = true;
-    }
-    return Changed;
-  }
-
-  if (isa_and_nonnull<SelectInst>(Target)) {
-    Lattice::Payload NewPayload{Lattice::Payload::Kind::Select, Target};
-    if (Dest.State != Lattice::StateKind::HeapChunkInterior ||
-        Dest.HeadPayload != NewPayload) {
-      Dest.State = Lattice::StateKind::HeapChunkInterior;
-      Dest.HeadPayload = NewPayload;
-      Changed = true;
-    }
-    return Changed;
+    Dest.HeadPayload = Src.HeadPayload;
+    return true;
   }
 
-  if (Dest.HeadPayload.Val != Src.HeadPayload.Val) {
+  if (Dest.HeadPayload != Src.HeadPayload) {
     Dest.State = Lattice::StateKind::Unknown;
     Dest.HeadPayload = {Lattice::Payload::Kind::None, nullptr};
     return true;
@@ -250,15 +249,10 @@ BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
       for (Value *Op : BO->operands())
         if (mergeIntoMap(Res.getMap(), Op, BackInfo))
           Worklist.push_back(Op);
-    } else if (auto *PHI = dyn_cast<PHINode>(const_cast<Value *>(V))) {
-      for (Value *InV : PHI->incoming_values())
-        if (mergeIntoMap(Res.getMap(), InV, BackInfo))
-          Worklist.push_back(InV);
-    } else if (auto *Sel = dyn_cast<SelectInst>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), Sel->getTrueValue(), BackInfo))
-        Worklist.push_back(Sel->getTrueValue());
-      if (mergeIntoMap(Res.getMap(), Sel->getFalseValue(), BackInfo))
-        Worklist.push_back(Sel->getFalseValue());
+    } else if (isa<PHINode>(const_cast<Value *>(V)) || isa<SelectInst>(const_cast<Value *>(V))) {
+      // TODO: Propagating backward deallocation provenance across PHI nodes or
+      // Select instructions without path-sensitive join conditions is unsound.
+      // Stop backward propagation here to maintain lattice join soundness.
     } else if (auto *Arg = dyn_cast<Argument>(const_cast<Value *>(V))) {
       Function *Fn = Arg->getParent();
       for (const User *FnU : Fn->users()) {
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index 9750cc376b964..3d1316dc40426 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1246,6 +1246,87 @@ ObjectSizeOffsetEvaluator::ObjectSizeOffsetEvaluator(
 }
 
 bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffsetValue &Result) {
+  CacheMapTy FallbackCacheMap;
+  return computeFallbackHeapMetadata_(V, Result, FallbackCacheMap);
+}
+
+bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata_(
+    Value *V, SizeOffsetValue &Result, CacheMapTy &FallbackCacheMap) {
+  auto CacheIt = FallbackCacheMap.find(V);
+  if (CacheIt != FallbackCacheMap.end()) {
+    Result = CacheIt->second;
+    return Result.bothKnown();
+  }
+
+  if (auto *PHI = dyn_cast<PHINode>(V)) {
+    IRBuilder<>::InsertPointGuard Guard(Builder);
+    Builder.SetInsertPoint(PHI->getParent(), PHI->getParent()->getFirstInsertionPt());
+    PHINode *SizePHI = Builder.CreatePHI(IntTy, PHI->getNumIncomingValues(), "fallback.size");
+    PHINode *OffsetPHI = Builder.CreatePHI(IntTy, PHI->getNumIncomingValues(), "fallback.offset");
+    InsertedInstructions.insert(SizePHI);
+    InsertedInstructions.insert(OffsetPHI);
+    FallbackCacheMap[PHI] = SizeOffsetWeakTrackingVH(SizePHI, OffsetPHI);
+
+    for (unsigned i = 0, e = PHI->getNumIncomingValues(); i != e; ++i) {
+      BasicBlock *IncomingBlock = PHI->getIncomingBlock(i);
+      Builder.SetInsertPoint(IncomingBlock->getTerminator());
+      SizeOffsetValue EdgeResult;
+      if (!computeFallbackHeapMetadata_(PHI->getIncomingValue(i), EdgeResult, FallbackCacheMap)) {
+        SizePHI->replaceAllUsesWith(PoisonValue::get(IntTy));
+        SizePHI->eraseFromParent();
+        InsertedInstructions.erase(SizePHI);
+        OffsetPHI->replaceAllUsesWith(PoisonValue::get(IntTy));
+        OffsetPHI->eraseFromParent();
+        InsertedInstructions.erase(OffsetPHI);
+        FallbackCacheMap.erase(PHI);
+        return false;
+      }
+      SizePHI->addIncoming(EdgeResult.Size, IncomingBlock);
+      OffsetPHI->addIncoming(EdgeResult.Offset, IncomingBlock);
+    }
+
+    Value *Size = SizePHI, *Offset = OffsetPHI;
+    if (Value *Tmp = SizePHI->hasConstantValue()) {
+      Size = Tmp;
+      SizePHI->replaceAllUsesWith(Size);
+      SizePHI->eraseFromParent();
+      InsertedInstructions.erase(SizePHI);
+    }
+    if (Value *Tmp = OffsetPHI->hasConstantValue()) {
+      Offset = Tmp;
+      OffsetPHI->replaceAllUsesWith(Offset);
+      OffsetPHI->eraseFromParent();
+      InsertedInstructions.erase(OffsetPHI);
+    }
+    Result = SizeOffsetValue(Size, Offset);
+    FallbackCacheMap[PHI] = SizeOffsetWeakTrackingVH(Result);
+    return true;
+  }
+
+  if (auto *SI = dyn_cast<SelectInst>(V)) {
+    IRBuilder<>::InsertPointGuard Guard(Builder);
+    if (auto OptIP = SI->getInsertionPointAfterDef())
+      Builder.SetInsertPoint(*OptIP);
+    else
+      return false;
+
+    SizeOffsetValue TrueResult, FalseResult;
+    if (!computeFallbackHeapMetadata_(SI->getTrueValue(), TrueResult, FallbackCacheMap) ||
+        !computeFallbackHeapMetadata_(SI->getFalseValue(), FalseResult, FallbackCacheMap))
+      return false;
+
+    Value *Size = Builder.CreateSelect(SI->getCondition(), TrueResult.Size, FalseResult.Size, "fallback.select.size");
+    Value *Offset = Builder.CreateSelect(SI->getCondition(), TrueResult.Offset, FalseResult.Offset, "fallback.select.offset");
+    if (auto *I = dyn_cast<Instruction>(Size))
+      InsertedInstructions.insert(I);
+    if (auto *I = dyn_cast<Instruction>(Offset))
+      InsertedInstructions.insert(I);
+
+    Result = SizeOffsetValue(Size, Offset);
+    FallbackCacheMap[SI] = SizeOffsetWeakTrackingVH(Result);
+    return true;
+  }
+
   Module *M = nullptr;
   if (auto *I = dyn_cast<Instruction>(V))
     M = I->getModule();
@@ -1260,6 +1341,9 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
   if (!Info.isValid())
     return false;
 
+  if (Info.State != HeapProvenanceLattice::StateKind::HeapChunkHead)
+    return false;
+
   if (auto *I = dyn_cast<Instruction>(V)) {
     if (auto *PHI = dyn_cast<PHINode>(I))
       Builder.SetInsertPoint(PHI->getParent(), PHI->getParent()->getFirstInsertionPt());
@@ -1346,8 +1430,11 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     OffsetVal = Zero;
   } else {
     VInt = Builder.CreatePtrToInt(V, IntTy, "v.int");
+    if (auto *I = dyn_cast<Instruction>(VInt)) InsertedInstructions.insert(I);
     HeadInt = Builder.CreatePtrToInt(HeadVal, IntTy, "head.int");
+    if (auto *I = dyn_cast<Instruction>(HeadInt)) InsertedInstructions.insert(I);
     OffsetVal = Builder.CreateSub(VInt, HeadInt, "meta.offset");
+    if (auto *I = dyn_cast<Instruction>(OffsetVal)) InsertedInstructions.insert(I);
   }
 
   FunctionCallee GetSizeFC = M->getOrInsertFunction(
@@ -1361,16 +1448,24 @@ bool ObjectSizeOffsetEvaluator::computeFallbackHeapMetadata(Value *V, SizeOffset
     Fn->addFnAttr(Attribute::NoSync);
   }
   Value *ChunkPayloadSize = Builder.CreateCall(GetSizeFC, {HeadVal}, "meta.size");
+  if (auto *I = dyn_cast<Instruction>(ChunkPayloadSize)) InsertedInstructions.insert(I);
   Value *IsHeap = Builder.CreateICmpNE(ChunkPayloadSize, Zero, "is.heap");
+  if (auto *I = dyn_cast<Instruction>(IsHeap)) InsertedInstructions.insert(I);
   if (HeadVal != V) {
     Value *Cond = Builder.CreateICmpUGE(VInt, HeadInt, "ptr.ge.head");
+    if (auto *I = dyn_cast<Instruction>(Cond)) InsertedInstructions.insert(I);
     ChunkPayloadSize = Builder.CreateSelect(Cond, ChunkPayloadSize, Zero, "chunk.size");
+    if (auto *I = dyn_cast<Instruction>(ChunkPayloadSize)) InsertedInstructions.insert(I);
   }
   ChunkPayloadSize = Builder.CreateSelect(IsHeap, ChunkPayloadSize, ConstantInt::getAllOnesValue(IntTy), "meta.size.checked");
-  if (OffsetVal->getType() != IntTy)
+  if (auto *I = dyn_cast<Instruction>(ChunkPayloadSize)) InsertedInstructions.insert(I);
+  if (OffsetVal->getType() != IntTy) {
     OffsetVal = Builder.CreateZExtOrTrunc(OffsetVal, IntTy);
+    if (auto *I = dyn_cast<Instruction>(OffsetVal)) InsertedInstructions.insert(I);
+  }
 
   Result = SizeOffsetValue(ChunkPayloadSize, OffsetVal);
+  FallbackCacheMap[V] = SizeOffsetWeakTrackingVH(Result);
   return true;
 }
 
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 9eea552fd263e..6501101fd0eda 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -21,6 +21,7 @@
 #include "llvm/Analysis/CtxProfAnalysis.h"
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/InlineAdvisor.h"
 #include "llvm/Analysis/InstCount.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
@@ -1677,6 +1678,7 @@ PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
                           .hoistLoadsStoresWithCondFaulting(true)));
 
   // Add the core optimizing pipeline.
+  MPM.addPass(RequireAnalysisPass<HeapProvenanceAnalysis, llvm::Module>());
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(OptimizePM),
                                                 PTO.EagerlyInvalidateAnalyses));
 
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 6ed643e138e5f..75cf7986e5323 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3962,7 +3962,9 @@ Instruction *InstCombinerImpl::visitAllocSite(Instruction &MI) {
         if (II->getIntrinsicID() == Intrinsic::objectsize) {
           SmallVector<Instruction *> InsertedInstructions;
           Value *Result = lowerObjectSizeCall(
-              II, DL, &TLI, AA, /*MustSucceed=*/true, &InsertedInstructions);
+              II, DL, &TLI, AA, /*MustSucceed=*/false, &InsertedInstructions);
+          if (!Result)
+            continue;
           for (Instruction *Inserted : InsertedInstructions)
             Worklist.add(Inserted);
           replaceInstUsesWith(*I, Result);
diff --git a/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp b/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
index 0c94ed27880f0..845344e9281d5 100644
--- a/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp
@@ -17,6 +17,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
@@ -97,7 +98,8 @@ static bool replaceConditionalBranchesOnConstant(Instruction *II,
 }
 
 bool llvm::lowerConstantIntrinsics(Function &F, const TargetLibraryInfo &TLI,
-                                   DominatorTree *DT) {
+                                   DominatorTree *DT,
+                                   const HeapProvenanceAnalysisResult *HPA) {
   std::optional<DomTreeUpdater> DTU;
   if (DT)
     DTU.emplace(DT, DomTreeUpdater::UpdateStrategy::Lazy);
@@ -141,7 +143,7 @@ bool llvm::lowerConstantIntrinsics(Function &F, const TargetLibraryInfo &TLI,
       IsConstantIntrinsicsHandled++;
       break;
     case Intrinsic::objectsize:
-      NewValue = lowerObjectSizeCall(II, DL, &TLI, true);
+      NewValue = lowerObjectSizeCall(II, DL, &TLI, true, HPA);
       LLVM_DEBUG(dbgs() << "Folding " << *II << " to " << *NewValue << "\n");
       ObjectSizeIntrinsicsHandled++;
       break;
@@ -156,8 +158,11 @@ bool llvm::lowerConstantIntrinsics(Function &F, const TargetLibraryInfo &TLI,
 
 PreservedAnalyses
 LowerConstantIntrinsicsPass::run(Function &F, FunctionAnalysisManager &AM) {
+  auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+  auto *HPA = MAMProxy.getCachedResult<HeapProvenanceAnalysis>(*F.getParent());
+
   if (lowerConstantIntrinsics(F, AM.getResult<TargetLibraryAnalysis>(F),
-                              AM.getCachedResult<DominatorTreeAnalysis>(F))) {
+                              AM.getCachedResult<DominatorTreeAnalysis>(F), HPA)) {
     PreservedAnalyses PA;
     PA.preserve<DominatorTreeAnalysis>();
     return PA;

>From 9adca5d0bc1a2065785021b1afafd4bfefc833f6 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Fri, 26 Jun 2026 16:46:09 -0700
Subject: [PATCH 16/18] Enforce post-dominance verification in Backward HPA

Guarantee soundness in Backward HPA by requiring deallocation basic blocks to post-dominate the function entry block before marking values as heap chunk heads. This prevents conditional deallocations (such as small vector inline capacity destructors or conditional scalar object deletions) from falsely attributing heap provenance to stack or non-heap pointers on alternative control flow paths.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index a8ce43da40120..b116ab4e3b6b5 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -1,5 +1,6 @@
 #include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
+#include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InstIterator.h"
@@ -202,9 +203,12 @@ BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
     if (F.isDeclaration())
       continue;
     auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
+    auto &PDT = FAM.getResult<PostDominatorTreeAnalysis>(F);
     for (Instruction &I : instructions(F)) {
       if (auto *CB = dyn_cast<CallBase>(&I)) {
         if (Value *ArgVal = getFreeLibCallOperand(CB, &TLI)) {
+          if (!PDT.dominates(CB->getParent(), &F.getEntryBlock()))
+            continue;
           HeapProvenanceLattice Info;
           Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
           Info.Dir = HeapProvenanceLattice::Backward;

>From 8c0c1b2af0a41bc7823212d56c8b874e8f485aff Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Mon, 29 Jun 2026 09:26:13 -0700
Subject: [PATCH 17/18] [BoundsChecking] Implement MLIR-structured sparse
 backward dataflow solver with SCC interprocedural tracking

Replaces flow-insensitive flat map scanning in BackwardHeapProvenanceAnalysis with a sparse basic-block solver and bottom-up SCC traversal. This correctly merges multi-branch deallocations while preventing conditional deallocations from contaminating caller arguments or stack storage.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 321 ++++++++++++++-----
 1 file changed, 236 insertions(+), 85 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index b116ab4e3b6b5..d081f4c5d5277 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -1,7 +1,10 @@
 #include "llvm/Analysis/HeapProvenanceAnalysis.h"
+#include "llvm/ADT/SCCIterator.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
-#include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
@@ -15,18 +18,21 @@ AnalysisKey ForwardHeapProvenanceAnalysis::Key;
 AnalysisKey BackwardHeapProvenanceAnalysis::Key;
 AnalysisKey HeapProvenanceAnalysis::Key;
 
-bool ForwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                                     ModuleAnalysisManager::Invalidator &) {
+bool ForwardHeapProvenanceAnalysisResult::invalidate(
+    Module &, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &) {
   return false;
 }
 
-bool BackwardHeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                                      ModuleAnalysisManager::Invalidator &) {
+bool BackwardHeapProvenanceAnalysisResult::invalidate(
+    Module &, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &) {
   return false;
 }
 
-bool HeapProvenanceAnalysisResult::invalidate(Module &, const PreservedAnalyses &PA,
-                                              ModuleAnalysisManager::Invalidator &) {
+bool HeapProvenanceAnalysisResult::invalidate(
+    Module &, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &) {
   return false;
 }
 
@@ -34,7 +40,8 @@ static bool isAllocLibCall(const Value *V, const TargetLibraryInfo *TLI) {
   return isAllocationFn(V, TLI);
 }
 
-static Value *getFreeLibCallOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
+static Value *getFreeLibCallOperand(const CallBase *CB,
+                                    const TargetLibraryInfo *TLI) {
   if (!CB || !TLI)
     return nullptr;
   const Function *Callee = CB->getCalledFunction();
@@ -169,7 +176,8 @@ ForwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
           if (CB->getArgOperand(i) == V) {
             Function *Callee = CB->getCalledFunction();
             if (!Callee)
-              Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
+              Callee = dyn_cast<Function>(
+                  CB->getCalledOperand()->stripPointerCasts());
             if (Callee && !Callee->isDeclaration() && i < Callee->arg_size()) {
               Argument *Arg = Callee->getArg(i);
               if (mergeIntoMap(Res.getMap(), Arg, Derived))
@@ -195,96 +203,238 @@ ForwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
 BackwardHeapProvenanceAnalysis::Result
 BackwardHeapProvenanceAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
   Result Res;
-  SmallVector<const Value *, 64> Worklist;
   FunctionAnalysisManager &FAM =
       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  CallGraph &CG = MAM.getResult<CallGraphAnalysis>(M);
+
+  for (scc_iterator<CallGraph *> SCCI = scc_begin(&CG); !SCCI.isAtEnd();
+       ++SCCI) {
+    for (CallGraphNode *CGN : *SCCI) {
+      Function *FnPtr = CGN->getFunction();
+      if (!FnPtr || FnPtr->isDeclaration())
+        continue;
+      Function &F = *FnPtr;
+      auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
+
+      // 1. Fast Bailout & Candidate Discovery
+      SmallVector<std::pair<Instruction *, Value *>, 4> Deallocs;
+      for (Instruction &I : instructions(F)) {
+        if (auto *CB = dyn_cast<CallBase>(&I)) {
+          if (Value *ArgVal = getFreeLibCallOperand(CB, &TLI)) {
+            Deallocs.push_back({&I, ArgVal});
+          } else if (Function *Callee = CB->getCalledFunction()) {
+            if (!Callee->isDeclaration()) {
+              for (Argument &Arg : Callee->args()) {
+                if (Arg.getArgNo() < CB->arg_size()) {
+                  HeapProvenanceLattice ArgInfo = Res.getInfo(&Arg);
+                  if (ArgInfo.State ==
+                      HeapProvenanceLattice::StateKind::HeapChunkHead)
+                    Deallocs.push_back({&I, CB->getArgOperand(Arg.getArgNo())});
+                }
+              }
+            }
+          }
+        }
+      }
+      if (Deallocs.empty())
+        continue;
+
+      SmallPtrSet<const Value *, 16> Candidates;
+      SmallVector<const Value *, 16> CandWorklist;
+      for (auto &Pair : Deallocs) {
+        if (Candidates.insert(Pair.second).second)
+          CandWorklist.push_back(Pair.second);
+      }
+      while (!CandWorklist.empty()) {
+        const Value *V = CandWorklist.pop_back_val();
+        auto addCand = [&](const Value *Op) {
+          if (Candidates.insert(Op).second)
+            CandWorklist.push_back(Op);
+        };
+        if (auto *GEP = dyn_cast<GEPOperator>(const_cast<Value *>(V)))
+          addCand(GEP->getPointerOperand());
+        else if (auto *BC = dyn_cast<BitCastInst>(const_cast<Value *>(V)))
+          addCand(BC->getOperand(0));
+        else if (auto *ASC =
+                     dyn_cast<AddrSpaceCastInst>(const_cast<Value *>(V)))
+          addCand(ASC->getOperand(0));
+        else if (auto *ITP = dyn_cast<IntToPtrInst>(const_cast<Value *>(V)))
+          addCand(ITP->getOperand(0));
+        else if (auto *PTI = dyn_cast<PtrToIntInst>(const_cast<Value *>(V)))
+          addCand(PTI->getOperand(0));
+        else if (auto *BO = dyn_cast<BinaryOperator>(const_cast<Value *>(V))) {
+          for (Value *Op : BO->operands())
+            addCand(Op);
+        } else if (auto *PHI = dyn_cast<PHINode>(const_cast<Value *>(V))) {
+          for (Value *InV : PHI->incoming_values())
+            addCand(InV);
+        } else if (auto *Sel = dyn_cast<SelectInst>(const_cast<Value *>(V))) {
+          addCand(Sel->getTrueValue());
+          addCand(Sel->getFalseValue());
+        }
+      }
 
-  for (Function &F : M) {
-    if (F.isDeclaration())
-      continue;
-    auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
-    auto &PDT = FAM.getResult<PostDominatorTreeAnalysis>(F);
-    for (Instruction &I : instructions(F)) {
-      if (auto *CB = dyn_cast<CallBase>(&I)) {
-        if (Value *ArgVal = getFreeLibCallOperand(CB, &TLI)) {
-          if (!PDT.dominates(CB->getParent(), &F.getEntryBlock()))
-            continue;
-          HeapProvenanceLattice Info;
-          Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
-          Info.Dir = HeapProvenanceLattice::Backward;
-          Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref, ArgVal};
-          if (mergeIntoMap(Res.getMap(), ArgVal, Info))
-            Worklist.push_back(ArgVal);
+      // 2. Sparse Block-Level State Initialization
+      DenseMap<const BasicBlock *,
+               DenseMap<const Value *, HeapProvenanceLattice>>
+          BlockEntryState;
+      SmallVector<const BasicBlock *, 16> BBWorklist;
+      for (BasicBlock &BB : F) {
+        if (succ_empty(&BB)) {
+          for (const Value *C : Candidates) {
+            HeapProvenanceLattice Info;
+            Info.State = HeapProvenanceLattice::StateKind::Unknown;
+            Info.Dir = HeapProvenanceLattice::Backward;
+            BlockEntryState[&BB][C] = Info;
+          }
+          BBWorklist.push_back(&BB);
         }
       }
-    }
-  }
+      if (BBWorklist.empty()) {
+        for (BasicBlock &BB : F)
+          BBWorklist.push_back(&BB);
+      }
 
-  if (Worklist.empty())
-    return Res;
+      SmallPtrSet<const BasicBlock *, 16> InWorklist;
+      for (const BasicBlock *BB : BBWorklist)
+        InWorklist.insert(BB);
+
+      // 3. Meet & Backward Transfer Worklist Loop
+      while (!BBWorklist.empty()) {
+        const BasicBlock *BB = BBWorklist.pop_back_val();
+        InWorklist.erase(BB);
+
+        DenseMap<const Value *, HeapProvenanceLattice> ExitState;
+        if (succ_empty(BB)) {
+          ExitState = BlockEntryState[BB];
+        } else {
+          bool FirstSucc = true;
+          for (const BasicBlock *Succ : successors(BB)) {
+            auto It = BlockEntryState.find(Succ);
+            if (It == BlockEntryState.end())
+              continue;
+            const auto &SuccState = It->second;
+            for (const Value *C : Candidates) {
+              auto CIt = SuccState.find(C);
+              if (CIt == SuccState.end() || CIt->second.isUninit())
+                continue;
+              HeapProvenanceLattice SInfo = CIt->second;
+              if (auto *PHI = dyn_cast<PHINode>(const_cast<Value *>(C))) {
+                if (PHI->getParent() == Succ) {
+                  Value *InV = PHI->getIncomingValueForBlock(
+                      const_cast<BasicBlock *>(BB));
+                  auto InIt = SuccState.find(InV);
+                  if (InIt != SuccState.end())
+                    SInfo = InIt->second;
+                }
+              }
+              if (FirstSucc)
+                ExitState[C] = SInfo;
+              else
+                mergeLattice(C, ExitState[C], SInfo);
+            }
+            FirstSucc = false;
+          }
+        }
 
-  while (!Worklist.empty()) {
-    const Value *V = Worklist.pop_back_val();
-    HeapProvenanceLattice Info = Res.getInfo(V);
-    if (!Info.isValid() || !(Info.Dir & HeapProvenanceLattice::Backward))
-      continue;
+        DenseMap<const Value *, HeapProvenanceLattice> CurState = ExitState;
+        for (const Instruction &IConst : llvm::reverse(*BB)) {
+          Instruction *I = const_cast<Instruction *>(&IConst);
+          if (auto *CB = dyn_cast<CallBase>(I)) {
+            if (Value *ArgVal = getFreeLibCallOperand(CB, &TLI)) {
+              if (Candidates.count(ArgVal)) {
+                HeapProvenanceLattice Info;
+                Info.State = HeapProvenanceLattice::StateKind::HeapChunkHead;
+                Info.Dir = HeapProvenanceLattice::Backward;
+                Info.HeadPayload = {HeapProvenanceLattice::Payload::Kind::Ref,
+                                    ArgVal};
+                CurState[ArgVal] = Info;
+              }
+            } else if (Function *Callee = CB->getCalledFunction()) {
+              if (!Callee->isDeclaration()) {
+                for (Argument &Arg : Callee->args()) {
+                  if (Arg.getArgNo() < CB->arg_size()) {
+                    HeapProvenanceLattice ArgInfo = Res.getInfo(&Arg);
+                    if (ArgInfo.State ==
+                        HeapProvenanceLattice::StateKind::HeapChunkHead) {
+                      Value *CallOp = CB->getArgOperand(Arg.getArgNo());
+                      if (Candidates.count(CallOp)) {
+                        HeapProvenanceLattice Info;
+                        Info.State =
+                            HeapProvenanceLattice::StateKind::HeapChunkHead;
+                        Info.Dir = HeapProvenanceLattice::Backward;
+                        Info.HeadPayload = {
+                            HeapProvenanceLattice::Payload::Kind::Ref, CallOp};
+                        CurState[CallOp] = Info;
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+          if (Candidates.count(I)) {
+            auto It = CurState.find(I);
+            if (It != CurState.end() && It->second.isValid()) {
+              HeapProvenanceLattice BackInfo = It->second;
+              if (BackInfo.State ==
+                  HeapProvenanceLattice::StateKind::HeapChunkHead)
+                BackInfo.State =
+                    HeapProvenanceLattice::StateKind::HeapChunkInterior;
+              auto transferOp = [&](Value *Op) {
+                if (Candidates.count(Op))
+                  mergeLattice(Op, CurState[Op], BackInfo);
+              };
+              if (auto *GEP = dyn_cast<GEPOperator>(I))
+                transferOp(GEP->getPointerOperand());
+              else if (auto *BC = dyn_cast<BitCastInst>(I))
+                transferOp(BC->getOperand(0));
+              else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(I))
+                transferOp(ASC->getOperand(0));
+              else if (auto *ITP = dyn_cast<IntToPtrInst>(I))
+                transferOp(ITP->getOperand(0));
+              else if (auto *PTI = dyn_cast<PtrToIntInst>(I))
+                transferOp(PTI->getOperand(0));
+              else if (auto *BO = dyn_cast<BinaryOperator>(I)) {
+                for (Value *Op : BO->operands())
+                  transferOp(Op);
+              } else if (auto *PHI = dyn_cast<PHINode>(I)) {
+                for (Value *InV : PHI->incoming_values())
+                  transferOp(InV);
+              } else if (auto *Sel = dyn_cast<SelectInst>(I)) {
+                transferOp(Sel->getTrueValue());
+                transferOp(Sel->getFalseValue());
+              }
+            }
+          }
+        }
 
-    HeapProvenanceLattice BackInfo = Info;
-    BackInfo.Dir = HeapProvenanceLattice::Backward;
-    if (BackInfo.State == HeapProvenanceLattice::StateKind::HeapChunkHead)
-      BackInfo.State = HeapProvenanceLattice::StateKind::HeapChunkInterior;
-
-    if (auto *GEP = dyn_cast<GEPOperator>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), GEP->getPointerOperand(), BackInfo))
-        Worklist.push_back(GEP->getPointerOperand());
-    } else if (auto *BC = dyn_cast<BitCastInst>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), BC->getOperand(0), BackInfo))
-        Worklist.push_back(BC->getOperand(0));
-    } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), ASC->getOperand(0), BackInfo))
-        Worklist.push_back(ASC->getOperand(0));
-    } else if (auto *ITP = dyn_cast<IntToPtrInst>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), ITP->getOperand(0), BackInfo))
-        Worklist.push_back(ITP->getOperand(0));
-    } else if (auto *PTI = dyn_cast<PtrToIntInst>(const_cast<Value *>(V))) {
-      if (mergeIntoMap(Res.getMap(), PTI->getOperand(0), BackInfo))
-        Worklist.push_back(PTI->getOperand(0));
-    } else if (auto *BO = dyn_cast<BinaryOperator>(const_cast<Value *>(V))) {
-      for (Value *Op : BO->operands())
-        if (mergeIntoMap(Res.getMap(), Op, BackInfo))
-          Worklist.push_back(Op);
-    } else if (isa<PHINode>(const_cast<Value *>(V)) || isa<SelectInst>(const_cast<Value *>(V))) {
-      // TODO: Propagating backward deallocation provenance across PHI nodes or
-      // Select instructions without path-sensitive join conditions is unsound.
-      // Stop backward propagation here to maintain lattice join soundness.
-    } else if (auto *Arg = dyn_cast<Argument>(const_cast<Value *>(V))) {
-      Function *Fn = Arg->getParent();
-      for (const User *FnU : Fn->users()) {
-        if (auto *CB = dyn_cast<CallBase>(FnU)) {
-          if (Arg->getArgNo() < CB->arg_size()) {
-            Value *CallOp = CB->getArgOperand(Arg->getArgNo());
-            if (mergeIntoMap(Res.getMap(), CallOp, BackInfo))
-              Worklist.push_back(CallOp);
+        bool Changed = false;
+        auto &EntryState = BlockEntryState[BB];
+        for (const Value *C : Candidates) {
+          if (EntryState[C] != CurState[C]) {
+            EntryState[C] = CurState[C];
+            Changed = true;
           }
         }
-      }
-    } else if (auto *CB = dyn_cast<CallBase>(const_cast<Value *>(V))) {
-      Function *Callee = CB->getCalledFunction();
-      if (!Callee)
-        Callee = dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
-      if (Callee && !Callee->isDeclaration() && V->getType()->isPointerTy()) {
-        for (BasicBlock &CalleeBB : *Callee) {
-          if (auto *RI = dyn_cast<ReturnInst>(CalleeBB.getTerminator())) {
-            if (Value *RetVal = RI->getReturnValue()) {
-              if (mergeIntoMap(Res.getMap(), RetVal, BackInfo))
-                Worklist.push_back(RetVal);
-            }
+        if (Changed) {
+          for (const BasicBlock *Pred : predecessors(BB)) {
+            if (InWorklist.insert(Pred).second)
+              BBWorklist.push_back(Pred);
           }
         }
       }
+
+      // 4. Populate Res.getMap()
+      for (const Value *C : Candidates) {
+        if (isa<Argument>(const_cast<Value *>(C))) {
+          mergeIntoMap(Res.getMap(), C, BlockEntryState[&F.getEntryBlock()][C]);
+        } else if (auto *I = dyn_cast<Instruction>(const_cast<Value *>(C))) {
+          mergeIntoMap(Res.getMap(), C, BlockEntryState[I->getParent()][C]);
+        }
+      }
     }
   }
-
   return Res;
 }
 
@@ -308,7 +458,8 @@ PreservedAnalyses HeapProvenancePrinterPass::run(Module &M,
   OS << "Printing analysis 'Heap Provenance Analysis' for module '"
      << M.getName() << "':\n";
   for (Function &F : M) {
-    if (F.isDeclaration()) continue;
+    if (F.isDeclaration())
+      continue;
     for (Argument &Arg : F.args()) {
       auto Info = Result.getInfo(&Arg);
       if (Info.isValid()) {

>From d792aabea88f595216e64d5bce5d38f2917984f9 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Mon, 29 Jun 2026 11:28:37 -0700
Subject: [PATCH 18/18] [BoundsChecking] Include realloc and freed operands in
 Backward HPA

Updates getFreeLibCallOperand to check both getFreedOperand and getReallocatedOperand, ensuring realloc calls act as valid deallocation roots for backward heap provenance analysis.

TAG=agy
CONV=aa2b66c1-619b-4d82-aea4-94158ab3bf11
---
 llvm/lib/Analysis/HeapProvenanceAnalysis.cpp | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
index d081f4c5d5277..3afdaf24bdb99 100644
--- a/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
+++ b/llvm/lib/Analysis/HeapProvenanceAnalysis.cpp
@@ -42,16 +42,12 @@ static bool isAllocLibCall(const Value *V, const TargetLibraryInfo *TLI) {
 
 static Value *getFreeLibCallOperand(const CallBase *CB,
                                     const TargetLibraryInfo *TLI) {
-  if (!CB || !TLI)
+  if (!CB)
     return nullptr;
-  const Function *Callee = CB->getCalledFunction();
-  if (!Callee)
-    return nullptr;
-  LibFunc TLIFn;
-  if (TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn) &&
-      isLibFreeFunction(Callee, TLIFn)) {
-    return CB->getArgOperand(0);
-  }
+  if (Value *Freed = getFreedOperand(CB, TLI))
+    return Freed;
+  if (Value *Reallocated = getReallocatedOperand(CB))
+    return Reallocated;
   return nullptr;
 }
 



More information about the cfe-commits mailing list