[clang] [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:59:57 PDT 2026


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

>From c1d6f887e3127d7db6ba66ad1a91be419a98ba12 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 01/16] [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 6dd9ab0eec054..e5f5d10af9ddb 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 64331e04c155b..d8403ba38c674 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())
 
@@ -146,6 +147,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 0a810227ce23b4281265e042fb2bf7c6ea3f390c 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 02/16] [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 a42d662146563..83c596ce3f9f3 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 0364a8bff9e53..15a55e8db256b 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"
@@ -1240,6 +1241,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()));
@@ -1247,6 +1340,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 7ed28de700ce7854aa2ae66f987492c38fe3e086 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 03/16] [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 83c596ce3f9f3..19c811bc873ba 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 15a55e8db256b..37bc85a1c3d48 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -636,15 +636,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!");
 
@@ -675,7 +677,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()) {
@@ -1231,12 +1233,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.
 }
@@ -1247,11 +1249,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;
 
@@ -1269,54 +1270,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 d8403ba38c674..e423c79f45461 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 58bc18b379048..65cb6fb42b7c1 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"
@@ -198,7 +199,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;
 
@@ -206,7 +208,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
@@ -304,8 +306,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 90f65502d50c9e9185c94b993a1dd7a4fc04cd81 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 04/16] [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 3f81fb241b77a77a726795389b86b107d8e029b5 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 05/16] [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 38e1b8d2166d18481f51406619d258daa0d3fc1e 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 06/16] [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 e33480f6c6416..3f0565451d23c 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -70,6 +70,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"
@@ -92,6 +93,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>
@@ -1051,7 +1053,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;
@@ -1078,6 +1087,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 84df1bf2750f9dfba4563ae1f57124c19811b0df 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 07/16] [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 34301334a16f70d9b78c62e9ca175b7d244b6411 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 08/16] [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 37bc85a1c3d48..18d92dcaa7708 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1286,6 +1286,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 7585f1208a48b4fa5f793f9c9eacbd3473c060f2 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 09/16] [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 3f0565451d23c..250c3f3ca7348 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -1054,15 +1054,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) {
@@ -1085,7 +1085,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 19c811bc873ba..9d7c7c2a7983a 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 18d92dcaa7708..fa187ec4961e1 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -677,7 +677,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()) {
@@ -1233,12 +1235,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.
 }
@@ -1249,10 +1252,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;
 
@@ -1287,7 +1292,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)) {
@@ -1295,8 +1302,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 65cb6fb42b7c1..7b1a3ae250fed 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -200,7 +200,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;
 
@@ -208,7 +209,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
@@ -307,9 +308,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 445f3d3c40d0ff1d3ec4d66607a0efc3fafac181 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 10/16] [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 fa187ec4961e1..13056ee69db61 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1310,12 +1310,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");
   }
 
@@ -1329,7 +1349,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 d1977f1739a99110b63e269f2e9c7afa5911bbbf 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 11/16] [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 250c3f3ca7348..e33480f6c6416 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -70,7 +70,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"
@@ -93,7 +92,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>
@@ -1053,16 +1051,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) {
@@ -1085,11 +1076,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 e423c79f45461..07a51f9886062 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())
 
@@ -149,8 +146,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()))
@@ -377,6 +372,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())
@@ -528,6 +526,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 7b1a3ae250fed..80e325d1c9dea 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -307,11 +307,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 816c0583636fcd38dde311fb04eec2dd70d69801 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 12/16] [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 e33480f6c6416..f117a21d170f5 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -68,6 +68,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"
@@ -92,6 +93,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>
@@ -1051,9 +1053,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) {
@@ -1076,8 +1083,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 13056ee69db61..753c71208518e 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1276,7 +1276,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 * {
@@ -1291,41 +1291,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;
@@ -1350,10 +1362,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 07a51f9886062..0412642bc0051 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())
@@ -149,6 +152,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()))
@@ -372,9 +377,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())
@@ -526,8 +528,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 80e325d1c9dea..a151662551177 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -307,10 +307,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 e7828077f76e776ba9a69c173165cb436fe2e7a2 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 13/16] 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 9d7c7c2a7983a..cd2ef5f612547 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 753c71208518e..c77717a870ed1 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -1247,6 +1247,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();
@@ -1261,6 +1342,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());
@@ -1347,8 +1431,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(
@@ -1362,16 +1449,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 6b8f3c8806c86..0c99534d9914c 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"
@@ -1678,6 +1679,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 6d901e835a1bd..3c1b0cd2d2692 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3967,7 +3967,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 07c173c17822c763f77e07a98f2c423c7af491cb 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 14/16] 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 e9b8f16d6be7677f9aa435b75014fcfeede50740 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 15/16] [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 983457d449e23762b18a67fd1d34951686389153 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 16/16] [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