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

Schrodinger ZHU Yifan via cfe-commits cfe-commits at lists.llvm.org
Fri Jun 26 09:16:34 PDT 2026


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

>From e7eed7b4e0600b9d62c6c44ebf355eb1bbfc2c3c Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Wed, 10 Jun 2026 13:50:35 -0700
Subject: [PATCH 01/15] Refactor FreeTrie and Implement TLSF FreeStore

- Make FreeTrie a proxy holding root by reference.
- Move root and range storage outside FreeTrie (into FreeStore).
- Add find_min and pop_min implementations.
- Update tests and add pop_min test.
- Replace freestore.h with TLSF implementation.
- Add USE_TRIE_FOR_LARGE_FREELIST option to use FreeTrie in TLSF.
- Remove obsolete set_range interface from FreeStore, FreeListHeap, and tests.

TAG=agy
CONV=fe3b4efa-7a5b-4c74-8257-e53f0d6e4850
---
 libc/src/__support/freelist.h              |   3 +
 libc/src/__support/freelist_heap.h         |   1 -
 libc/src/__support/freestore.h             | 403 +++++++++++++++++----
 libc/src/__support/freetrie.h              |  65 +++-
 libc/test/src/__support/freestore_test.cpp |  27 +-
 libc/test/src/__support/freetrie_test.cpp  |  93 ++++-
 6 files changed, 478 insertions(+), 114 deletions(-)

diff --git a/libc/src/__support/freelist.h b/libc/src/__support/freelist.h
index 48e70c7c29df6..ae2de684b3a24 100644
--- a/libc/src/__support/freelist.h
+++ b/libc/src/__support/freelist.h
@@ -41,6 +41,9 @@ class FreeList {
     /// @returns The inner size of blocks in the list containing this node.
     LIBC_INLINE size_t size() const { return block().inner_size(); }
 
+    /// @returns The next node in the list containing this node.
+    LIBC_INLINE Node *next_node() const { return next; }
+
   private:
     // Circularly linked pointers to adjacent nodes.
     Node *prev;
diff --git a/libc/src/__support/freelist_heap.h b/libc/src/__support/freelist_heap.h
index a591d63ddd4e5..976a6eec4023e 100644
--- a/libc/src/__support/freelist_heap.h
+++ b/libc/src/__support/freelist_heap.h
@@ -84,7 +84,6 @@ LIBC_INLINE void FreeListHeap::init() {
   LIBC_ASSERT(!is_initialized && "duplicate initialization");
   auto result = BlockRef::init(region());
   BlockRef block = *result;
-  free_store.set_range({0, cpp::bit_ceil(block.inner_size())});
   free_store.insert(block);
   is_initialized = true;
 }
diff --git a/libc/src/__support/freestore.h b/libc/src/__support/freestore.h
index adc0e061ace93..4fb90cabf47bd 100644
--- a/libc/src/__support/freestore.h
+++ b/libc/src/__support/freestore.h
@@ -7,114 +7,377 @@
 //===----------------------------------------------------------------------===//
 ///
 /// \file
-/// Interface for freestore.
+/// This file contains a two-level segregated fit free block store.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIBC_SRC___SUPPORT_FREESTORE_H
 #define LLVM_LIBC_SRC___SUPPORT_FREESTORE_H
 
-#include "freetrie.h"
+#include "hdr/stdint_proxy.h"
+#include "hdr/types/size_t.h"
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/bit.h"
+#include "src/__support/CPP/limits.h"
+#include "src/__support/block.h"
+#include "src/__support/freelist.h"
+#include "src/__support/freetrie.h"
+#include "src/__support/macros/config.h"
+#include "src/__support/macros/optimization.h"
 
 namespace LIBC_NAMESPACE_DECL {
 
-/// A best-fit store of variously-sized free blocks. Blocks can be inserted and
-/// removed in logarithmic time.
-class FreeStore {
-  friend class FreeListHeap;
-
-public:
-  FreeStore() = default;
-  FreeStore(const FreeStore &other) = delete;
-  FreeStore &operator=(const FreeStore &other) = delete;
-
-  /// Sets the range of possible block sizes. This can only be called when the
-  /// trie is empty.
-  LIBC_INLINE void set_range(FreeTrie::SizeRange range) {
-    large_trie.set_range(range);
-  }
-
-  /// Insert a free block. If the block is too small to be tracked, nothing
-  /// happens.
-  void insert(BlockRef block);
+/// Configuration for TLSFFreeStore.
+template <size_t UNIT_SIZE_VAL, size_t STEP_SIZE_BITS_VAL,
+          size_t NUM_STEP_BITS_VAL, size_t NUM_TABLE_ENTRIES_VAL,
+          bool USE_TRIE_FOR_LARGE_FREELIST_VAL = false>
+struct TLSFFreeStoreConfig {
+  static constexpr size_t UNIT_SIZE = UNIT_SIZE_VAL;
+  static constexpr size_t STEP_SIZE_BITS = STEP_SIZE_BITS_VAL;
+  static constexpr size_t NUM_STEP_BITS = NUM_STEP_BITS_VAL;
+  static constexpr size_t NUM_TABLE_ENTRIES = NUM_TABLE_ENTRIES_VAL;
+  static constexpr bool USE_TRIE_FOR_LARGE_FREELIST =
+      USE_TRIE_FOR_LARGE_FREELIST_VAL;
+};
 
-  /// Remove a free block. If the block is too small to be tracked, nothing
-  /// happens.
-  void remove(BlockRef block);
+// A two-level segregated fit store for free blocks.
+//
+// The store starts with small lists that grow linearly for small sizes, which
+// covers [0, ... UNIT_SIZE * EXP_BASE]. For larger sizes, the bits are managed
+// in a 2-D table. One can think of each row containing NUM_STEPS lists. Along
+// the row, the size grows by 2 exponentially; along the column, the size
+// increases by STEP_SIZE linearly.
+//
+// Mathematical layout:
+//   STEP_SIZE = 1 << STEP_SIZE_BITS
+//   NUM_STEPS = 1 << NUM_STEP_BITS
+//   EXP_BASE = STEP_SIZE * NUM_STEPS
+//   LARGE_SIZE_THRESHOLD = UNIT_SIZE * EXP_BASE
+//
+// Visual representation with example parameters:
+//   UNIT_SIZE = 32, STEP_SIZE = 8, NUM_STEPS = 4
+//   EXP_BASE = 32, THRESHOLD = 1024 B (1 KiB)
+//
+// 1. Small Sizes (Linear Array):
+//    Covers [0, ... 1024 B] growing directly by UNIT_SIZE = 32 B
+//   +-------+-------+-------+-------+-------+-----------+---------------+
+//   | [0 B] | [32B] | [64B] | [96B] |  ...  | [992 B]   | [1024 B (Th)] |
+//   +-------+-------+-------+-------+-------+-----------+---------------+
+//
+// 2. Large Sizes (2-D Table):
+//    Rows = FL (Exponential growth), Columns = SL (Linear steps)
+//    One can think of each Row containing NUM_STEPS (4) lists.
+//
+//                       LINEAR INCREASE ALONG COLUMN (SL) --->
+//             +---------------+---------------+---------------+---------------+
+//             |    Col = 0    |    Col = 1    |    Col = 2    |    Col = 3    |
+//             |    (Base)     |   (+25% FL)   |   (+50% FL)   |   (+75% FL)   |
+//   +---------+---------------+---------------+---------------+---------------+
+// E | Row = 0 |    1024 B     |    1280 B     |    1536 B     |    1792 B     |
+// X |(Base 1K)| [1024 - 1279] | [1280 - 1535] | [1536 - 1791] | [1792 - 2047] |
+// P +---------+---------------+---------------+---------------+---------------+
+// O | Row = 1 |    2048 B     |    2560 B     |    3072 B     |    3584 B     |
+// N |(Base 2K)| [2048 - 2559] | [2560 - 3071] | [3072 - 3583] | [3584 - 4095] |
+// E +---------+---------------+---------------+---------------+---------------+
+// N | Row = 2 |    4096 B     |    5120 B     |    6144 B     |    7168 B     |
+// T |(Base 4K)| [4096 - 5119] | [5120 - 6143] | [6144 - 7167] | [7168 - 8191] |
+// I +---------+---------------+---------------+---------------+---------------+
+// A | Row = 3 |    8192 B     |   10240 B     |   12288 B     |   14336 B     |
+// L |(Base 8K)|[8192 - 10239]|[10240 - 12287]|[12288 - 14335]|[14336 - 16383]|
+//   +---------+---------------+---------------+---------------+---------------+
+//
+// Note: For the real implementation, we don't actually store the lists in a
+// 2-D structure. Instead, we flatten the entire 2-D layout into a single
+// flat 1-D array of size TOTAL_BITS (free_lists), and map sizes directly to
+// a continuous 1-D index using size_to_bit_index. The allocation state is
+// tracked compactly in the lookup_table bitmask array.
+template <typename CONFIG> class TLSFFreeStoreImpl {
+protected:
+  static_assert(cpp::has_single_bit(CONFIG::UNIT_SIZE),
+                "unit size must be a power of two");
+  static_assert(CONFIG::NUM_TABLE_ENTRIES > 0,
+                "the lookup table must have at least one entry");
 
-  /// Remove a best-fit free block that can contain the given size when
-  /// allocated. Returns nullptr if there is no such block.
-  BlockRef remove_best_fit(size_t size);
+  static constexpr size_t STEP_SIZE = size_t(1) << CONFIG::STEP_SIZE_BITS;
+  static constexpr size_t NUM_STEPS = size_t(1) << CONFIG::NUM_STEP_BITS;
+  static constexpr size_t EXP_BASE = STEP_SIZE * NUM_STEPS;
+  static constexpr int UNIT_SIZE_LOG2 = cpp::bit_width(CONFIG::UNIT_SIZE) - 1;
+  static constexpr int EXP_BASE_LOG2 =
+      CONFIG::STEP_SIZE_BITS + CONFIG::NUM_STEP_BITS;
+  static constexpr size_t BITS_PER_ENTRY =
+      cpp::numeric_limits<uintptr_t>::digits;
+  static constexpr size_t TOTAL_BITS =
+      CONFIG::NUM_TABLE_ENTRIES * BITS_PER_ENTRY;
+  static constexpr bool USE_TRIE = CONFIG::USE_TRIE_FOR_LARGE_FREELIST;
 
-private:
+public:
   static constexpr size_t MIN_OUTER_SIZE = align_up(
       BlockRef::HEADER_SIZE + sizeof(FreeList::Node), BlockRef::MIN_ALIGN);
-  static constexpr size_t MIN_LARGE_OUTER_SIZE = align_up(
-      BlockRef::HEADER_SIZE + sizeof(FreeTrie::Node), BlockRef::MIN_ALIGN);
-  static constexpr size_t NUM_SMALL_SIZES =
-      (MIN_LARGE_OUTER_SIZE - MIN_OUTER_SIZE) / BlockRef::MIN_ALIGN;
 
+  LIBC_INLINE TLSFFreeStoreImpl() = default;
+  LIBC_INLINE TLSFFreeStoreImpl(const TLSFFreeStoreImpl &other) = delete;
+  LIBC_INLINE TLSFFreeStoreImpl &
+  operator=(const TLSFFreeStoreImpl &other) = delete;
+
+  LIBC_INLINE void insert(BlockRef block);
+  LIBC_INLINE void remove(BlockRef block);
+  LIBC_INLINE BlockRef remove_best_fit(size_t size) {
+    return find_and_remove_fit(size);
+  }
+  LIBC_INLINE BlockRef find_and_remove_fit(size_t size);
+
+protected:
   LIBC_INLINE static bool too_small(BlockRef block) {
     return block.outer_size() < MIN_OUTER_SIZE;
   }
-  LIBC_INLINE static bool is_small(BlockRef block) {
-    return block.outer_size() < MIN_LARGE_OUTER_SIZE;
-  }
 
-  FreeList &small_list(BlockRef block);
-  FreeList *find_best_small_fit(size_t size);
+  union ListOrTrie {
+    FreeList list;
+    FreeTrie::Node *trie_root;
+
+    LIBC_INLINE constexpr ListOrTrie() : trie_root(nullptr) {}
+  };
 
-  cpp::array<FreeList, NUM_SMALL_SIZES> small_lists;
-  FreeTrie large_trie;
+  cpp::array<uintptr_t, CONFIG::NUM_TABLE_ENTRIES> lookup_table{};
+  cpp::array<ListOrTrie, TOTAL_BITS> free_lists{};
+
+  LIBC_INLINE static constexpr size_t size_to_bit_index(size_t size);
+  LIBC_INLINE void set_bit(size_t bit_index);
+  LIBC_INLINE void clear_bit(size_t bit_index);
+  LIBC_INLINE bool get_bit(size_t bit_index) const;
+  LIBC_INLINE size_t find_first_bit_set_after(size_t bit_index) const;
+  LIBC_INLINE BlockRef remove_first_fit_in_list(size_t index, size_t size);
+  LIBC_INLINE static constexpr FreeTrie::SizeRange index_to_range(size_t index);
+  LIBC_INLINE FreeTrie get_trie(size_t index);
+  LIBC_INLINE BlockRef find_and_remove_fit_in_trie(size_t index, size_t size);
+  LIBC_INLINE BlockRef pop_min_in_trie(size_t index);
 };
 
-LIBC_INLINE void FreeStore::insert(BlockRef block) {
-  if (too_small(block))
-    return;
-  if (is_small(block))
-    small_list(block).push(block);
-  else
-    large_trie.push(block);
+template <typename CONFIG>
+LIBC_INLINE constexpr size_t
+TLSFFreeStoreImpl<CONFIG>::size_to_bit_index(size_t size) {
+  if (size <= (EXP_BASE << UNIT_SIZE_LOG2))
+    return size >> UNIT_SIZE_LOG2;
+
+  size_t size_ilog2 = static_cast<size_t>(cpp::bit_width(size) - 1);
+  size_t exp_offset = (size_ilog2 - UNIT_SIZE_LOG2 - EXP_BASE_LOG2 - 1)
+                      << CONFIG::NUM_STEP_BITS;
+  size_t step_index = size >> (size_ilog2 - CONFIG::NUM_STEP_BITS);
+  size_t index = EXP_BASE + exp_offset + step_index;
+
+  return index < TOTAL_BITS ? index : TOTAL_BITS - 1;
 }
 
-LIBC_INLINE void FreeStore::remove(BlockRef block) {
-  if (too_small(block))
-    return;
-  if (is_small(block)) {
-    small_list(block).remove(
-        reinterpret_cast<FreeList::Node *>(block.usable_space()));
-  } else {
-    large_trie.remove(reinterpret_cast<FreeTrie::Node *>(block.usable_space()));
+template <typename CONFIG>
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::set_bit(size_t bit_index) {
+  size_t entry_index = bit_index / BITS_PER_ENTRY;
+  size_t bit_offset = bit_index % BITS_PER_ENTRY;
+  lookup_table[entry_index] |= uintptr_t(1) << bit_offset;
+}
+
+template <typename CONFIG>
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::clear_bit(size_t bit_index) {
+  size_t entry_index = bit_index / BITS_PER_ENTRY;
+  size_t bit_offset = bit_index % BITS_PER_ENTRY;
+  lookup_table[entry_index] &= ~(uintptr_t(1) << bit_offset);
+}
+
+template <typename CONFIG>
+LIBC_INLINE bool TLSFFreeStoreImpl<CONFIG>::get_bit(size_t bit_index) const {
+  size_t entry_index = bit_index / BITS_PER_ENTRY;
+  size_t bit_offset = bit_index % BITS_PER_ENTRY;
+  return (lookup_table[entry_index] & (uintptr_t(1) << bit_offset)) != 0;
+}
+
+template <typename CONFIG>
+LIBC_INLINE size_t
+TLSFFreeStoreImpl<CONFIG>::find_first_bit_set_after(size_t bit_index) const {
+  if (bit_index >= TOTAL_BITS - 1)
+    return TOTAL_BITS;
+
+  size_t target_index = bit_index + 1;
+  size_t start_entry = target_index / BITS_PER_ENTRY;
+  size_t bit_offset = target_index % BITS_PER_ENTRY;
+
+  uintptr_t value = lookup_table[start_entry] & (~uintptr_t(0) << bit_offset);
+  if (value != 0)
+    return start_entry * BITS_PER_ENTRY +
+           static_cast<size_t>(cpp::countr_zero(value));
+
+  for (size_t i = start_entry + 1; i < CONFIG::NUM_TABLE_ENTRIES; ++i) {
+    value = lookup_table[i];
+    if (value != 0)
+      return i * BITS_PER_ENTRY + static_cast<size_t>(cpp::countr_zero(value));
   }
+  return TOTAL_BITS;
 }
 
-LIBC_INLINE BlockRef FreeStore::remove_best_fit(size_t size) {
-  if (FreeList *list = find_best_small_fit(size)) {
-    BlockRef block = list->front();
-    list->pop();
-    return block;
+template <typename CONFIG>
+LIBC_INLINE constexpr FreeTrie::SizeRange
+TLSFFreeStoreImpl<CONFIG>::index_to_range(size_t index) {
+  LIBC_ASSERT(index >= EXP_BASE && "only call for large lists");
+  size_t local_index = index - EXP_BASE;
+  size_t exp_index = local_index >> CONFIG::NUM_STEP_BITS;
+  size_t linear_index = local_index & (NUM_STEPS - 1);
+
+  size_t row_base = (EXP_BASE << exp_index) << UNIT_SIZE_LOG2;
+  size_t step_size = (STEP_SIZE << exp_index) << UNIT_SIZE_LOG2;
+  size_t min_size = row_base + linear_index * step_size;
+
+  if (index == TOTAL_BITS - 1) {
+    constexpr size_t width = size_t(1)
+                             << (cpp::numeric_limits<size_t>::digits - 2);
+    LIBC_ASSERT(min_size < width &&
+                "min_size too large for overflow bin width");
+    return FreeTrie::SizeRange(min_size, width);
   }
-  if (FreeTrie::Node *best_fit = large_trie.find_best_fit(size)) {
+
+  return FreeTrie::SizeRange(min_size, step_size);
+}
+
+template <typename CONFIG>
+LIBC_INLINE FreeTrie TLSFFreeStoreImpl<CONFIG>::get_trie(size_t index) {
+  LIBC_ASSERT(index >= EXP_BASE && "only call for large lists");
+  return FreeTrie(free_lists[index].trie_root, index_to_range(index));
+}
+
+template <typename CONFIG>
+LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit_in_trie(
+    size_t index, size_t size) {
+  FreeTrie trie = get_trie(index);
+  if (FreeTrie::Node *best_fit = trie.find_best_fit(size)) {
     BlockRef block = best_fit->block();
-    large_trie.remove(best_fit);
+    trie.remove(best_fit);
+    if (trie.empty())
+      clear_bit(index);
     return block;
   }
   return BlockRef();
 }
 
-LIBC_INLINE FreeList &FreeStore::small_list(BlockRef block) {
-  LIBC_ASSERT(is_small(block) && "only legal for small blocks");
-  return small_lists[(block.outer_size() - MIN_OUTER_SIZE) /
-                     BlockRef::MIN_ALIGN];
+template <typename CONFIG>
+LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::pop_min_in_trie(size_t index) {
+  FreeTrie trie = get_trie(index);
+  FreeTrie::Node *min_node = trie.pop_min();
+  LIBC_ASSERT(min_node && "bit was set but trie is empty");
+  BlockRef block = min_node->block();
+  if (trie.empty())
+    clear_bit(index);
+  return block;
+}
+
+template <typename CONFIG>
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(BlockRef block) {
+  if (too_small(block))
+    return;
+  size_t bit_index = size_to_bit_index(block.inner_size());
+
+  if constexpr (USE_TRIE)
+    if (bit_index > EXP_BASE) {
+      get_trie(bit_index).push(block);
+      set_bit(bit_index);
+      return;
+    }
+
+  free_lists[bit_index].list.push(block);
+  set_bit(bit_index);
+}
+
+template <typename CONFIG>
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(BlockRef block) {
+  if (too_small(block))
+    return;
+  size_t bit_index = size_to_bit_index(block.inner_size());
+
+  if constexpr (USE_TRIE)
+    if (bit_index > EXP_BASE) {
+      FreeTrie trie = get_trie(bit_index);
+      trie.remove(reinterpret_cast<FreeTrie::Node *>(block.usable_space()));
+      if (trie.empty())
+        clear_bit(bit_index);
+      return;
+    }
+
+  free_lists[bit_index].list.remove(
+      reinterpret_cast<FreeList::Node *>(block.usable_space()));
+  if (free_lists[bit_index].list.empty())
+    clear_bit(bit_index);
+}
+
+template <typename CONFIG>
+LIBC_INLINE BlockRef
+TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(size_t index, size_t size) {
+  FreeList::Node *begin_node = free_lists[index].list.begin();
+  if (begin_node == nullptr)
+    return BlockRef();
+
+  FreeList::Node *cur = begin_node;
+  do {
+    if (cur->size() >= size) {
+      free_lists[index].list.remove(cur);
+      if (free_lists[index].list.empty())
+        clear_bit(index);
+      return cur->block();
+    }
+    cur = cur->next_node();
+  } while (cur != begin_node);
+
+  return BlockRef();
 }
 
-LIBC_INLINE FreeList *FreeStore::find_best_small_fit(size_t size) {
-  for (FreeList &list : small_lists)
-    if (!list.empty() && list.size() >= size)
-      return &list;
-  return nullptr;
+template <typename CONFIG>
+LIBC_INLINE BlockRef
+TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
+  size_t bit_index = size_to_bit_index(size);
+
+  if (LIBC_UNLIKELY(bit_index >= TOTAL_BITS - 1)) {
+    if constexpr (USE_TRIE)
+      return find_and_remove_fit_in_trie(TOTAL_BITS - 1, size);
+    else
+      return remove_first_fit_in_list(TOTAL_BITS - 1, size);
+  }
+
+  // 1. Try oversized bins (guaranteed fit, but larger).
+  size_t oversized_bit = find_first_bit_set_after(bit_index);
+  if (LIBC_LIKELY(oversized_bit < TOTAL_BITS)) {
+    if constexpr (USE_TRIE)
+      if (oversized_bit > EXP_BASE)
+        return pop_min_in_trie(oversized_bit);
+
+    BlockRef block = free_lists[oversized_bit].list.front();
+    free_lists[oversized_bit].list.pop();
+    if (free_lists[oversized_bit].list.empty())
+      clear_bit(oversized_bit);
+    return block;
+  }
+
+  // 2. Try exact fit (fallback).
+  if (get_bit(bit_index)) {
+    if constexpr (USE_TRIE) {
+      if (bit_index > EXP_BASE)
+        return find_and_remove_fit_in_trie(bit_index, size);
+      else if (BlockRef block = remove_first_fit_in_list(bit_index, size))
+        return block;
+    } else if (BlockRef block = remove_first_fit_in_list(bit_index, size))
+      return block;
+  }
+
+  return BlockRef();
 }
 
+template <size_t UNIT_SIZE, size_t STEP_SIZE_BITS, size_t NUM_STEP_BITS,
+          size_t NUM_TABLE_ENTRIES, bool USE_TRIE = false>
+using TLSFFreeStore = TLSFFreeStoreImpl<TLSFFreeStoreConfig<
+    UNIT_SIZE, STEP_SIZE_BITS, NUM_STEP_BITS, NUM_TABLE_ENTRIES, USE_TRIE>>;
+
+#ifndef LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST
+#define LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST false
+#endif
+
+using FreeStore =
+    TLSFFreeStore<BlockRef::MIN_ALIGN, 3, 2, (sizeof(uintptr_t) == 8 ? 3 : 6),
+                  LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST>;
+
 } // namespace LIBC_NAMESPACE_DECL
 
 #endif // LLVM_LIBC_SRC___SUPPORT_FREESTORE_H
diff --git a/libc/src/__support/freetrie.h b/libc/src/__support/freetrie.h
index 9e35463462b38..9e07f36ad0a53 100644
--- a/libc/src/__support/freetrie.h
+++ b/libc/src/__support/freetrie.h
@@ -87,15 +87,8 @@ class FreeTrie {
     }
   };
 
-  LIBC_INLINE constexpr FreeTrie() : FreeTrie(SizeRange{0, 0}) {}
-  LIBC_INLINE constexpr FreeTrie(SizeRange range) : range(range) {}
-
-  /// Sets the range of possible block sizes. This can only be called when the
-  /// trie is empty.
-  LIBC_INLINE void set_range(FreeTrie::SizeRange new_range) {
-    LIBC_ASSERT(empty() && "cannot change the range of a preexisting trie");
-    range = new_range;
-  }
+  LIBC_INLINE FreeTrie(Node *&root, SizeRange range)
+      : root(root), range(range) {}
 
   /// @returns Whether the trie contains any blocks.
   LIBC_INLINE bool empty() const { return !root; }
@@ -110,15 +103,24 @@ class FreeTrie {
   /// nullptr.
   Node *find_best_fit(size_t size);
 
+  /// @returns The node with the minimum size in the trie; otherwise nullptr.
+  LIBC_INLINE Node *find_min();
+
+  /// Removes and returns the node with the minimum size in the trie; otherwise
+  /// nullptr.
+  LIBC_INLINE Node *pop_min();
+
 private:
   /// @returns Whether a node is the head of its containing freelist.
-  bool is_head(Node *node) const { return node->parent || node == root; }
+  LIBC_INLINE bool is_head(Node *node) const {
+    return node->parent || node == root;
+  }
 
   /// Replaces references to one node with another (or nullptr) in all adjacent
   /// parent and child nodes.
   void replace_node(Node *node, Node *new_node);
 
-  Node *root = nullptr;
+  Node *&root;
   SizeRange range;
 };
 
@@ -236,6 +238,47 @@ LIBC_INLINE FreeTrie::Node *FreeTrie::find_best_fit(size_t size) {
     }
   }
 }
+LIBC_INLINE FreeTrie::Node *FreeTrie::find_min() {
+  if (empty())
+    return nullptr;
+
+  Node *cur = root;
+  SizeRange cur_range = range;
+  Node *best_min = nullptr;
+
+  while (cur) {
+    if (cur->lower) {
+      if (cur_range.lower().contains(cur->size())) {
+        if (!best_min || cur->size() < best_min->size()) {
+          best_min = cur;
+        }
+      }
+      cur = cur->lower;
+      cur_range = cur_range.lower();
+    } else {
+      if (cur_range.lower().contains(cur->size())) {
+        if (!best_min || cur->size() < best_min->size()) {
+          best_min = cur;
+        }
+        break;
+      } else {
+        if (!best_min || cur->size() < best_min->size()) {
+          best_min = cur;
+        }
+        cur = cur->upper;
+        cur_range = cur_range.upper();
+      }
+    }
+  }
+  return best_min;
+}
+
+LIBC_INLINE FreeTrie::Node *FreeTrie::pop_min() {
+  Node *min_node = find_min();
+  if (min_node)
+    remove(min_node);
+  return min_node;
+}
 
 } // namespace LIBC_NAMESPACE_DECL
 
diff --git a/libc/test/src/__support/freestore_test.cpp b/libc/test/src/__support/freestore_test.cpp
index 61103a9126c08..ce655367da00f 100644
--- a/libc/test/src/__support/freestore_test.cpp
+++ b/libc/test/src/__support/freestore_test.cpp
@@ -38,7 +38,6 @@ TEST(LlvmLibcFreeStore, TooSmall) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.set_range({0, 4096});
   store.insert(too_small);
   store.insert(remainder);
 
@@ -68,21 +67,26 @@ TEST(LlvmLibcFreeStore, RemoveBestFit) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.set_range({0, 4096});
   store.insert(smallest);
   if (largest_small != smallest)
     store.insert(largest_small);
   store.insert(remainder);
 
-  // Find exact match for smallest.
-  ASSERT_EQ(store.remove_best_fit(smallest.inner_size()).addr(),
-            smallest.addr());
-  store.insert(smallest);
-
-  // Find exact match for largest.
-  ASSERT_EQ(store.remove_best_fit(largest_small.inner_size()).addr(),
-            largest_small.addr());
-  store.insert(largest_small);
+  // For TLSF (oversized first), asking for a size will return the block from
+  // the first non-empty oversized bin if one exists, bypassing the exact bin.
+  if (largest_small != smallest) {
+    BlockRef block = store.remove_best_fit(smallest.inner_size());
+    ASSERT_EQ(block.addr(), largest_small.addr());
+    store.insert(block);
+
+    BlockRef block2 = store.remove_best_fit(largest_small.inner_size());
+    ASSERT_EQ(block2.addr(), remainder.addr());
+    store.insert(block2);
+  } else {
+    BlockRef block = store.remove_best_fit(smallest.inner_size());
+    ASSERT_EQ(block.addr(), remainder.addr());
+    store.insert(block);
+  }
 
   // Search small list for best fit.
   BlockRef next_smallest =
@@ -108,7 +112,6 @@ TEST(LlvmLibcFreeStore, Remove) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.set_range({0, 4096});
   store.insert(small);
   store.insert(remainder);
 
diff --git a/libc/test/src/__support/freetrie_test.cpp b/libc/test/src/__support/freetrie_test.cpp
index bf3284b2faf64..7eea72bcdbc37 100644
--- a/libc/test/src/__support/freetrie_test.cpp
+++ b/libc/test/src/__support/freetrie_test.cpp
@@ -22,7 +22,9 @@ using LIBC_NAMESPACE::cpp::byte;
 using LIBC_NAMESPACE::cpp::optional;
 
 TEST(LlvmLibcFreeTrie, FindBestFitRoot) {
-  FreeTrie trie({0, 4096});
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
   EXPECT_EQ(trie.find_best_fit(123), static_cast<FreeTrie::Node *>(nullptr));
 
   byte mem[1024];
@@ -31,10 +33,10 @@ TEST(LlvmLibcFreeTrie, FindBestFitRoot) {
   BlockRef block = *maybeBlock;
   trie.push(block);
 
-  FreeTrie::Node *root = trie.find_best_fit(0);
-  ASSERT_EQ(root->block().addr(), block.addr());
-  EXPECT_EQ(trie.find_best_fit(block.inner_size() - 1), root);
-  EXPECT_EQ(trie.find_best_fit(block.inner_size()), root);
+  FreeTrie::Node *found = trie.find_best_fit(0);
+  ASSERT_EQ(found->block().addr(), block.addr());
+  EXPECT_EQ(trie.find_best_fit(block.inner_size() - 1), found);
+  EXPECT_EQ(trie.find_best_fit(block.inner_size()), found);
   EXPECT_EQ(trie.find_best_fit(block.inner_size() + 1),
             static_cast<FreeTrie::Node *>(nullptr));
   EXPECT_EQ(trie.find_best_fit(4095), static_cast<FreeTrie::Node *>(nullptr));
@@ -47,10 +49,12 @@ TEST(LlvmLibcFreeTrie, FindBestFitLower) {
   BlockRef lower = *maybeBlock;
   maybeBlock = lower.split(512);
   ASSERT_TRUE(maybeBlock.has_value());
-  BlockRef root = *maybeBlock;
+  BlockRef root_block = *maybeBlock;
 
-  FreeTrie trie({0, 4096});
-  trie.push(root);
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
+  trie.push(root_block);
   trie.push(lower);
 
   EXPECT_EQ(trie.find_best_fit(0)->block().addr(), lower.addr());
@@ -60,36 +64,40 @@ TEST(LlvmLibcFreeTrie, FindBestFitUpper) {
   byte mem[4096];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
   ASSERT_TRUE(maybeBlock.has_value());
-  BlockRef root = *maybeBlock;
-  maybeBlock = root.split(512);
+  BlockRef root_block = *maybeBlock;
+  maybeBlock = root_block.split(512);
   ASSERT_TRUE(maybeBlock.has_value());
   BlockRef upper = *maybeBlock;
 
-  FreeTrie trie({0, 4096});
-  trie.push(root);
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
+  trie.push(root_block);
   trie.push(upper);
 
-  EXPECT_EQ(trie.find_best_fit(root.inner_size() + 1)->block().addr(),
+  EXPECT_EQ(trie.find_best_fit(root_block.inner_size() + 1)->block().addr(),
             upper.addr());
   // The upper subtrie should be skipped if it could not contain a better fit.
-  EXPECT_EQ(trie.find_best_fit(root.inner_size() - 1)->block().addr(),
-            root.addr());
+  EXPECT_EQ(trie.find_best_fit(root_block.inner_size() - 1)->block().addr(),
+            root_block.addr());
 }
 
 TEST(LlvmLibcFreeTrie, FindBestFitLowerAndUpper) {
   byte mem[4096];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
   ASSERT_TRUE(maybeBlock.has_value());
-  BlockRef root = *maybeBlock;
-  maybeBlock = root.split(1024);
+  BlockRef root_block = *maybeBlock;
+  maybeBlock = root_block.split(1024);
   ASSERT_TRUE(maybeBlock.has_value());
   BlockRef lower = *maybeBlock;
   maybeBlock = lower.split(128);
   ASSERT_TRUE(maybeBlock.has_value());
   BlockRef upper = *maybeBlock;
 
-  FreeTrie trie({0, 4096});
-  trie.push(root);
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
+  trie.push(root_block);
   trie.push(lower);
   trie.push(upper);
 
@@ -114,7 +122,9 @@ TEST(LlvmLibcFreeTrie, Remove) {
   BlockRef large = *maybeBlock;
 
   // Removing the root empties the trie.
-  FreeTrie trie({0, 4096});
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
   trie.push(large);
   FreeTrie::Node *large_node = trie.find_best_fit(0);
   ASSERT_EQ(large_node->block().addr(), large.addr());
@@ -132,3 +142,46 @@ TEST(LlvmLibcFreeTrie, Remove) {
   EXPECT_EQ(trie.find_best_fit(large.inner_size())->block().addr(),
             large.addr());
 }
+
+TEST(LlvmLibcFreeTrie, PopMin) {
+  alignas(BlockRef::MIN_ALIGN) byte mem[4096];
+  optional<BlockRef> maybe_block = BlockRef::init(mem);
+  ASSERT_TRUE(maybe_block.has_value());
+  BlockRef root_block = *maybe_block;
+  maybe_block = root_block.split(1024);
+  ASSERT_TRUE(maybe_block.has_value());
+  BlockRef lower = *maybe_block;
+  maybe_block = lower.split(128);
+  ASSERT_TRUE(maybe_block.has_value());
+  BlockRef upper = *maybe_block;
+
+  FreeTrie::Node *root = nullptr;
+  FreeTrie::SizeRange range{0, 4096};
+  FreeTrie trie(root, range);
+
+  // Empty pop
+  EXPECT_EQ(trie.pop_min(), static_cast<FreeTrie::Node *>(nullptr));
+
+  trie.push(root_block);
+  trie.push(lower);
+  trie.push(upper);
+
+  // Min should be lower (~128)
+  FreeTrie::Node *min1 = trie.pop_min();
+  ASSERT_NE(min1, static_cast<FreeTrie::Node *>(nullptr));
+  EXPECT_EQ(min1->block().addr(), lower.addr());
+
+  // Next min should be root_block (~1024)
+  FreeTrie::Node *min2 = trie.pop_min();
+  ASSERT_NE(min2, static_cast<FreeTrie::Node *>(nullptr));
+  EXPECT_EQ(min2->block().addr(), root_block.addr());
+
+  // Next min should be upper (~2944)
+  FreeTrie::Node *min3 = trie.pop_min();
+  ASSERT_NE(min3, static_cast<FreeTrie::Node *>(nullptr));
+  EXPECT_EQ(min3->block().addr(), upper.addr());
+
+  // Now empty
+  EXPECT_EQ(trie.pop_min(), static_cast<FreeTrie::Node *>(nullptr));
+  EXPECT_TRUE(trie.empty());
+}

>From a193a91eb59cad755f51d35e4672bea893f136c0 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Thu, 11 Jun 2026 15:23:40 -0700
Subject: [PATCH 02/15] [libc] restrict FreeTrie to the overflow bin in TLSF

- Rename config option to USE_TRIE_FOR_OVERFLOW_BIN.
- Restrict FreeTrie usage to only the last overflow bin (bit_index == TOTAL_BITS - 1).
- Simplify exact fit search to only use FreeList, as large sizes now always go to the overflow bin.
- Simplify away index_to_range by using a fixed [0, INF) range for the overflow Trie.
- Remove index argument from Trie helper functions as they only operate on the overflow bin.

TAG=agy
CONV=fe3b4efa-7a5b-4c74-8257-e53f0d6e4850
---
 libc/src/__support/freestore.h | 90 ++++++++++++----------------------
 1 file changed, 32 insertions(+), 58 deletions(-)

diff --git a/libc/src/__support/freestore.h b/libc/src/__support/freestore.h
index 4fb90cabf47bd..13b79c0debf39 100644
--- a/libc/src/__support/freestore.h
+++ b/libc/src/__support/freestore.h
@@ -30,14 +30,14 @@ namespace LIBC_NAMESPACE_DECL {
 /// Configuration for TLSFFreeStore.
 template <size_t UNIT_SIZE_VAL, size_t STEP_SIZE_BITS_VAL,
           size_t NUM_STEP_BITS_VAL, size_t NUM_TABLE_ENTRIES_VAL,
-          bool USE_TRIE_FOR_LARGE_FREELIST_VAL = false>
+          bool USE_TRIE_FOR_OVERFLOW_BIN_VAL = false>
 struct TLSFFreeStoreConfig {
   static constexpr size_t UNIT_SIZE = UNIT_SIZE_VAL;
   static constexpr size_t STEP_SIZE_BITS = STEP_SIZE_BITS_VAL;
   static constexpr size_t NUM_STEP_BITS = NUM_STEP_BITS_VAL;
   static constexpr size_t NUM_TABLE_ENTRIES = NUM_TABLE_ENTRIES_VAL;
-  static constexpr bool USE_TRIE_FOR_LARGE_FREELIST =
-      USE_TRIE_FOR_LARGE_FREELIST_VAL;
+  static constexpr bool USE_TRIE_FOR_OVERFLOW_BIN =
+      USE_TRIE_FOR_OVERFLOW_BIN_VAL;
 };
 
 // A two-level segregated fit store for free blocks.
@@ -108,7 +108,9 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
       cpp::numeric_limits<uintptr_t>::digits;
   static constexpr size_t TOTAL_BITS =
       CONFIG::NUM_TABLE_ENTRIES * BITS_PER_ENTRY;
-  static constexpr bool USE_TRIE = CONFIG::USE_TRIE_FOR_LARGE_FREELIST;
+  static constexpr bool USE_TRIE = CONFIG::USE_TRIE_FOR_OVERFLOW_BIN;
+  static constexpr size_t OVERFLOW_WIDTH =
+      size_t(1) << (cpp::numeric_limits<size_t>::digits - 2);
 
 public:
   static constexpr size_t MIN_OUTER_SIZE = align_up(
@@ -147,10 +149,9 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
   LIBC_INLINE bool get_bit(size_t bit_index) const;
   LIBC_INLINE size_t find_first_bit_set_after(size_t bit_index) const;
   LIBC_INLINE BlockRef remove_first_fit_in_list(size_t index, size_t size);
-  LIBC_INLINE static constexpr FreeTrie::SizeRange index_to_range(size_t index);
-  LIBC_INLINE FreeTrie get_trie(size_t index);
-  LIBC_INLINE BlockRef find_and_remove_fit_in_trie(size_t index, size_t size);
-  LIBC_INLINE BlockRef pop_min_in_trie(size_t index);
+  LIBC_INLINE FreeTrie get_trie();
+  LIBC_INLINE BlockRef find_and_remove_fit_in_trie(size_t size);
+  LIBC_INLINE BlockRef pop_min_in_trie();
 };
 
 template <typename CONFIG>
@@ -213,56 +214,33 @@ TLSFFreeStoreImpl<CONFIG>::find_first_bit_set_after(size_t bit_index) const {
 }
 
 template <typename CONFIG>
-LIBC_INLINE constexpr FreeTrie::SizeRange
-TLSFFreeStoreImpl<CONFIG>::index_to_range(size_t index) {
-  LIBC_ASSERT(index >= EXP_BASE && "only call for large lists");
-  size_t local_index = index - EXP_BASE;
-  size_t exp_index = local_index >> CONFIG::NUM_STEP_BITS;
-  size_t linear_index = local_index & (NUM_STEPS - 1);
-
-  size_t row_base = (EXP_BASE << exp_index) << UNIT_SIZE_LOG2;
-  size_t step_size = (STEP_SIZE << exp_index) << UNIT_SIZE_LOG2;
-  size_t min_size = row_base + linear_index * step_size;
-
-  if (index == TOTAL_BITS - 1) {
-    constexpr size_t width = size_t(1)
-                             << (cpp::numeric_limits<size_t>::digits - 2);
-    LIBC_ASSERT(min_size < width &&
-                "min_size too large for overflow bin width");
-    return FreeTrie::SizeRange(min_size, width);
-  }
-
-  return FreeTrie::SizeRange(min_size, step_size);
-}
-
-template <typename CONFIG>
-LIBC_INLINE FreeTrie TLSFFreeStoreImpl<CONFIG>::get_trie(size_t index) {
-  LIBC_ASSERT(index >= EXP_BASE && "only call for large lists");
-  return FreeTrie(free_lists[index].trie_root, index_to_range(index));
+LIBC_INLINE FreeTrie TLSFFreeStoreImpl<CONFIG>::get_trie() {
+  return FreeTrie(free_lists[TOTAL_BITS - 1].trie_root,
+                  FreeTrie::SizeRange(0, OVERFLOW_WIDTH));
 }
 
 template <typename CONFIG>
-LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit_in_trie(
-    size_t index, size_t size) {
-  FreeTrie trie = get_trie(index);
+LIBC_INLINE BlockRef
+TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit_in_trie(size_t size) {
+  FreeTrie trie = get_trie();
   if (FreeTrie::Node *best_fit = trie.find_best_fit(size)) {
     BlockRef block = best_fit->block();
     trie.remove(best_fit);
     if (trie.empty())
-      clear_bit(index);
+      clear_bit(TOTAL_BITS - 1);
     return block;
   }
   return BlockRef();
 }
 
 template <typename CONFIG>
-LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::pop_min_in_trie(size_t index) {
-  FreeTrie trie = get_trie(index);
+LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::pop_min_in_trie() {
+  FreeTrie trie = get_trie();
   FreeTrie::Node *min_node = trie.pop_min();
   LIBC_ASSERT(min_node && "bit was set but trie is empty");
   BlockRef block = min_node->block();
   if (trie.empty())
-    clear_bit(index);
+    clear_bit(TOTAL_BITS - 1);
   return block;
 }
 
@@ -273,8 +251,8 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(BlockRef block) {
   size_t bit_index = size_to_bit_index(block.inner_size());
 
   if constexpr (USE_TRIE)
-    if (bit_index > EXP_BASE) {
-      get_trie(bit_index).push(block);
+    if (bit_index == TOTAL_BITS - 1) {
+      get_trie().push(block);
       set_bit(bit_index);
       return;
     }
@@ -290,8 +268,8 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(BlockRef block) {
   size_t bit_index = size_to_bit_index(block.inner_size());
 
   if constexpr (USE_TRIE)
-    if (bit_index > EXP_BASE) {
-      FreeTrie trie = get_trie(bit_index);
+    if (bit_index == TOTAL_BITS - 1) {
+      FreeTrie trie = get_trie();
       trie.remove(reinterpret_cast<FreeTrie::Node *>(block.usable_space()));
       if (trie.empty())
         clear_bit(bit_index);
@@ -332,7 +310,7 @@ TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
 
   if (LIBC_UNLIKELY(bit_index >= TOTAL_BITS - 1)) {
     if constexpr (USE_TRIE)
-      return find_and_remove_fit_in_trie(TOTAL_BITS - 1, size);
+      return find_and_remove_fit_in_trie(size);
     else
       return remove_first_fit_in_list(TOTAL_BITS - 1, size);
   }
@@ -340,9 +318,10 @@ TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
   // 1. Try oversized bins (guaranteed fit, but larger).
   size_t oversized_bit = find_first_bit_set_after(bit_index);
   if (LIBC_LIKELY(oversized_bit < TOTAL_BITS)) {
-    if constexpr (USE_TRIE)
-      if (oversized_bit > EXP_BASE)
-        return pop_min_in_trie(oversized_bit);
+    if constexpr (USE_TRIE) {
+      if (oversized_bit == TOTAL_BITS - 1)
+        return pop_min_in_trie();
+    }
 
     BlockRef block = free_lists[oversized_bit].list.front();
     free_lists[oversized_bit].list.pop();
@@ -353,12 +332,7 @@ TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
 
   // 2. Try exact fit (fallback).
   if (get_bit(bit_index)) {
-    if constexpr (USE_TRIE) {
-      if (bit_index > EXP_BASE)
-        return find_and_remove_fit_in_trie(bit_index, size);
-      else if (BlockRef block = remove_first_fit_in_list(bit_index, size))
-        return block;
-    } else if (BlockRef block = remove_first_fit_in_list(bit_index, size))
+    if (BlockRef block = remove_first_fit_in_list(bit_index, size))
       return block;
   }
 
@@ -370,13 +344,13 @@ template <size_t UNIT_SIZE, size_t STEP_SIZE_BITS, size_t NUM_STEP_BITS,
 using TLSFFreeStore = TLSFFreeStoreImpl<TLSFFreeStoreConfig<
     UNIT_SIZE, STEP_SIZE_BITS, NUM_STEP_BITS, NUM_TABLE_ENTRIES, USE_TRIE>>;
 
-#ifndef LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST
-#define LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST false
+#ifndef LIBC_COPT_USE_TRIE_FOR_OVERFLOW_BIN
+#define LIBC_COPT_USE_TRIE_FOR_OVERFLOW_BIN false
 #endif
 
 using FreeStore =
     TLSFFreeStore<BlockRef::MIN_ALIGN, 3, 2, (sizeof(uintptr_t) == 8 ? 3 : 6),
-                  LIBC_COPT_USE_TRIE_FOR_LARGE_FREELIST>;
+                  LIBC_COPT_USE_TRIE_FOR_OVERFLOW_BIN>;
 
 } // namespace LIBC_NAMESPACE_DECL
 

>From a4697e8f113033c1151b3260a1c0323148edc7aa Mon Sep 17 00:00:00 2001
From: yfzhu <yfzhu at google.com>
Date: Tue, 23 Jun 2026 09:15:24 -0700
Subject: [PATCH 03/15] [libc] make baremetal freelist heap headers-only

This change moves the implementations of FreeList, FreeTrie and FreeListHeap
to their respective header files, making the baremetal freelist heap library
headers-only. This simplifies integration and avoids compilation issues.

TAG=agy
CONV=4118a58d-30c6-4fd4-8e45-defe707d8bba
---
 libc/src/__support/CMakeLists.txt             | 15 +--
 libc/src/__support/freelist.cpp               | 47 ----------
 libc/src/__support/freelist.h                 | 30 +++++-
 libc/src/__support/freelist_heap.cpp          | 19 ----
 libc/src/__support/freelist_heap.h            | 26 ++++--
 libc/src/__support/freetrie.cpp               | 64 -------------
 libc/src/__support/freetrie.h                 | 53 ++++++++++-
 .../llvm-project-overlay/libc/BUILD.bazel     | 92 +++++++++++++++++++
 8 files changed, 190 insertions(+), 156 deletions(-)
 delete mode 100644 libc/src/__support/freelist.cpp
 delete mode 100644 libc/src/__support/freelist_heap.cpp
 delete mode 100644 libc/src/__support/freetrie.cpp

diff --git a/libc/src/__support/CMakeLists.txt b/libc/src/__support/CMakeLists.txt
index 9d08b4bcf7303..2c5f38a8d2f72 100644
--- a/libc/src/__support/CMakeLists.txt
+++ b/libc/src/__support/CMakeLists.txt
@@ -25,12 +25,10 @@ add_header_library(
     libc.src.__support.math_extras
 )
 
-add_object_library(
+add_header_library(
   freelist
   HDRS
     freelist.h
-  SRCS
-    freelist.cpp
   DEPENDS
     .block
     libc.src.__support.fixedvector
@@ -40,12 +38,10 @@ add_object_library(
     libc.src.__support.CPP.span
 )
 
-add_object_library(
+add_header_library(
   freetrie
   HDRS
     freetrie.h
-  SRCS
-    freetrie.cpp
   DEPENDS
     .block
     .freelist
@@ -59,15 +55,10 @@ add_header_library(
     .freetrie
 )
 
-libc_set_definition(libc_freelist_malloc_size "LIBC_FREELIST_MALLOC_SIZE=${LIBC_CONF_FREELIST_MALLOC_BUFFER_SIZE}")
-add_object_library(
+add_header_library(
   freelist_heap
-  SRCS
-    freelist_heap.cpp
   HDRS
     freelist_heap.h
-  COMPILE_OPTIONS
-    ${libc_freelist_malloc_size}
   DEPENDS
     .block
     .freelist
diff --git a/libc/src/__support/freelist.cpp b/libc/src/__support/freelist.cpp
deleted file mode 100644
index 28f73b2726d7d..0000000000000
--- a/libc/src/__support/freelist.cpp
+++ /dev/null
@@ -1,47 +0,0 @@
-//===----------------------------------------------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-///
-/// \file
-/// Implementation for freelist.
-///
-//===----------------------------------------------------------------------===//
-
-#include "freelist.h"
-
-namespace LIBC_NAMESPACE_DECL {
-
-void FreeList::push(Node *node) {
-  if (begin_) {
-    LIBC_ASSERT(BlockRef::from_usable_space(node).outer_size() ==
-                    begin_->block().outer_size() &&
-                "freelist entries must have the same size");
-    // Since the list is circular, insert the node immediately before begin_.
-    node->prev = begin_->prev;
-    node->next = begin_;
-    begin_->prev->next = node;
-    begin_->prev = node;
-  } else {
-    begin_ = node->prev = node->next = node;
-  }
-}
-
-void FreeList::remove(Node *node) {
-  LIBC_ASSERT(begin_ && "cannot remove from empty list");
-  if (node == node->next) {
-    LIBC_ASSERT(node == begin_ &&
-                "a self-referential node must be the only element");
-    begin_ = nullptr;
-  } else {
-    node->prev->next = node->next;
-    node->next->prev = node->prev;
-    if (begin_ == node)
-      begin_ = node->next;
-  }
-}
-
-} // namespace LIBC_NAMESPACE_DECL
diff --git a/libc/src/__support/freelist.h b/libc/src/__support/freelist.h
index ae2de684b3a24..55d1be9115c7d 100644
--- a/libc/src/__support/freelist.h
+++ b/libc/src/__support/freelist.h
@@ -15,6 +15,8 @@
 #define LLVM_LIBC_SRC___SUPPORT_FREELIST_H
 
 #include "block.h"
+#include "src/__support/libc_assert.h"
+#include "src/__support/macros/config.h"
 
 namespace LIBC_NAMESPACE_DECL {
 
@@ -80,13 +82,37 @@ class FreeList {
 
   /// Push an already-constructed node to the back of the list.
   /// This allows pushing derived node types with additional data.
-  void push(Node *node);
+  LIBC_INLINE void push(Node* node) {
+    if (begin_) {
+      LIBC_ASSERT(BlockRef::from_usable_space(node).outer_size() ==
+                      begin_->block().outer_size() &&
+                  "freelist entries must have the same size");
+      // Since the list is circular, insert the node immediately before begin_.
+      node->prev = begin_->prev;
+      node->next = begin_;
+      begin_->prev->next = node;
+      begin_->prev = node;
+    } else {
+      begin_ = node->prev = node->next = node;
+    }
+  }
 
   /// Pop the first node from the list.
   LIBC_INLINE void pop() { remove(begin_); }
 
   /// Remove an arbitrary node from the list.
-  void remove(Node *node);
+  LIBC_INLINE void remove(Node* node) {
+    LIBC_ASSERT(begin_ && "cannot remove from empty list");
+    if (node == node->next) {
+      LIBC_ASSERT(node == begin_ &&
+                  "a self-referential node must be the only element");
+      begin_ = nullptr;
+    } else {
+      node->prev->next = node->next;
+      node->next->prev = node->prev;
+      if (begin_ == node) begin_ = node->next;
+    }
+  }
 
 private:
   Node *begin_;
diff --git a/libc/src/__support/freelist_heap.cpp b/libc/src/__support/freelist_heap.cpp
deleted file mode 100644
index 4deb0e0f09e22..0000000000000
--- a/libc/src/__support/freelist_heap.cpp
+++ /dev/null
@@ -1,19 +0,0 @@
-//===-- Implementation for freelist_heap ----------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "src/__support/freelist_heap.h"
-#include "src/__support/macros/config.h"
-
-#include <stddef.h>
-
-namespace LIBC_NAMESPACE_DECL {
-
-static LIBC_CONSTINIT FreeListHeap freelist_heap_symbols;
-FreeListHeap *freelist_heap = &freelist_heap_symbols;
-
-} // namespace LIBC_NAMESPACE_DECL
diff --git a/libc/src/__support/freelist_heap.h b/libc/src/__support/freelist_heap.h
index 976a6eec4023e..4f1cf1dac16b2 100644
--- a/libc/src/__support/freelist_heap.h
+++ b/libc/src/__support/freelist_heap.h
@@ -21,6 +21,7 @@
 #include "src/__support/CPP/optional.h"
 #include "src/__support/CPP/span.h"
 #include "src/__support/libc_assert.h"
+#include "src/__support/macros/attributes.h"
 #include "src/__support/macros/config.h"
 #include "src/__support/math_extras.h"
 #include "src/string/memory_utils/inline_memcpy.h"
@@ -80,7 +81,7 @@ template <size_t BUFF_SIZE> class FreeListHeapBuffer : public FreeListHeap {
   cpp::byte buffer[BUFF_SIZE];
 };
 
-LIBC_INLINE void FreeListHeap::init() {
+[[gnu::noinline]] LIBC_INLINE void FreeListHeap::init() {
   LIBC_ASSERT(!is_initialized && "duplicate initialization");
   auto result = BlockRef::init(region());
   BlockRef block = *result;
@@ -88,7 +89,8 @@ LIBC_INLINE void FreeListHeap::init() {
   is_initialized = true;
 }
 
-LIBC_INLINE void *FreeListHeap::allocate_impl(size_t alignment, size_t size) {
+[[gnu::noinline]] LIBC_INLINE void* FreeListHeap::allocate_impl(
+    size_t alignment, size_t size) {
   if (size == 0)
     return nullptr;
 
@@ -113,12 +115,12 @@ LIBC_INLINE void *FreeListHeap::allocate_impl(size_t alignment, size_t size) {
   return block_info.block.usable_space();
 }
 
-LIBC_INLINE void *FreeListHeap::allocate(size_t size) {
+[[gnu::noinline]] LIBC_INLINE void* FreeListHeap::allocate(size_t size) {
   return allocate_impl(BlockRef::MIN_ALIGN, size);
 }
 
-LIBC_INLINE void *FreeListHeap::aligned_allocate(size_t alignment,
-                                                 size_t size) {
+[[gnu::noinline]] LIBC_INLINE void* FreeListHeap::aligned_allocate(
+    size_t alignment, size_t size) {
   // The alignment must be an integral power of two.
   if (!IsPow2(alignment))
     return nullptr;
@@ -133,7 +135,7 @@ LIBC_INLINE void *FreeListHeap::aligned_allocate(size_t alignment,
   return allocate_impl(alignment, size);
 }
 
-LIBC_INLINE void FreeListHeap::free(void *ptr) {
+[[gnu::noinline]] LIBC_INLINE void FreeListHeap::free(void* ptr) {
   if (ptr == nullptr)
     return;
 
@@ -164,7 +166,8 @@ LIBC_INLINE void FreeListHeap::free(void *ptr) {
   free_store.insert(block);
 }
 
-LIBC_INLINE bool FreeListHeap::shrink_in_place(BlockRef block, size_t size) {
+[[gnu::noinline]] LIBC_INLINE bool FreeListHeap::shrink_in_place(BlockRef block,
+                                                                 size_t size) {
   size_t min_outer_size = BlockRef::outer_size(cpp::max(size, sizeof(size_t)));
   uintptr_t next_block_start = BlockRef::next_possible_block_start(
       block.addr() + min_outer_size, BlockRef::MIN_ALIGN);
@@ -193,7 +196,8 @@ LIBC_INLINE bool FreeListHeap::shrink_in_place(BlockRef block, size_t size) {
 
 // Follows constract of the C standard realloc() function
 // If ptr is free'd, will return nullptr.
-LIBC_INLINE void *FreeListHeap::realloc(void *ptr, size_t size) {
+[[gnu::noinline]] LIBC_INLINE void* FreeListHeap::realloc(void* ptr,
+                                                          size_t size) {
   if (size == 0) {
     free(ptr);
     return nullptr;
@@ -228,7 +232,8 @@ LIBC_INLINE void *FreeListHeap::realloc(void *ptr, size_t size) {
   return new_ptr;
 }
 
-LIBC_INLINE void *FreeListHeap::calloc(size_t num, size_t size) {
+[[gnu::noinline]] LIBC_INLINE void* FreeListHeap::calloc(size_t num,
+                                                         size_t size) {
   size_t bytes;
   if (__builtin_mul_overflow(num, size, &bytes))
     return nullptr;
@@ -238,7 +243,8 @@ LIBC_INLINE void *FreeListHeap::calloc(size_t num, size_t size) {
   return ptr;
 }
 
-extern FreeListHeap *freelist_heap;
+LIBC_INLINE_VAR LIBC_CONSTINIT FreeListHeap freelist_heap_symbols;
+LIBC_INLINE_VAR FreeListHeap* freelist_heap = &freelist_heap_symbols;
 
 } // namespace LIBC_NAMESPACE_DECL
 
diff --git a/libc/src/__support/freetrie.cpp b/libc/src/__support/freetrie.cpp
deleted file mode 100644
index e76efe717f215..0000000000000
--- a/libc/src/__support/freetrie.cpp
+++ /dev/null
@@ -1,64 +0,0 @@
-//===-- Implementation for freetrie ---------------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "freetrie.h"
-
-namespace LIBC_NAMESPACE_DECL {
-
-void FreeTrie::remove(Node *node) {
-  LIBC_ASSERT(!empty() && "cannot remove from empty trie");
-  FreeList list = node;
-  list.pop();
-  Node *new_node = static_cast<Node *>(list.begin());
-  if (!new_node) {
-    // The freelist is empty. Replace the subtrie root with an arbitrary leaf.
-    // This is legal because there is no relationship between the size of the
-    // root and its children.
-    Node *leaf = node;
-    while (leaf->lower || leaf->upper)
-      leaf = leaf->lower ? leaf->lower : leaf->upper;
-    if (leaf == node) {
-      // If the root is a leaf, then removing it empties the subtrie.
-      replace_node(node, nullptr);
-      return;
-    }
-
-    replace_node(leaf, nullptr);
-    new_node = leaf;
-  }
-
-  if (!is_head(node))
-    return;
-
-  // Copy the trie links to the new head.
-  new_node->lower = node->lower;
-  new_node->upper = node->upper;
-  new_node->parent = node->parent;
-  replace_node(node, new_node);
-}
-
-void FreeTrie::replace_node(Node *node, Node *new_node) {
-  LIBC_ASSERT(is_head(node) && "only head nodes contain trie links");
-
-  if (node->parent) {
-    Node *&parent_child =
-        node->parent->lower == node ? node->parent->lower : node->parent->upper;
-    LIBC_ASSERT(parent_child == node &&
-                "no reference to child node found in parent");
-    parent_child = new_node;
-  } else {
-    LIBC_ASSERT(root == node && "non-root node had no parent");
-    root = new_node;
-  }
-  if (node->lower)
-    node->lower->parent = new_node;
-  if (node->upper)
-    node->upper->parent = new_node;
-}
-
-} // namespace LIBC_NAMESPACE_DECL
diff --git a/libc/src/__support/freetrie.h b/libc/src/__support/freetrie.h
index 9e07f36ad0a53..f238a6b441cd5 100644
--- a/libc/src/__support/freetrie.h
+++ b/libc/src/__support/freetrie.h
@@ -15,6 +15,7 @@
 #define LLVM_LIBC_SRC___SUPPORT_FREETRIE_H
 
 #include "freelist.h"
+#include "src/__support/libc_assert.h"
 
 namespace LIBC_NAMESPACE_DECL {
 
@@ -97,7 +98,7 @@ class FreeTrie {
   void push(BlockRef block);
 
   /// Remove a node from this trie node's free list.
-  void remove(Node *node);
+  LIBC_INLINE void remove(Node* node);
 
   /// @returns A smallest node that can allocate the given size; otherwise
   /// nullptr.
@@ -118,7 +119,7 @@ class FreeTrie {
 
   /// Replaces references to one node with another (or nullptr) in all adjacent
   /// parent and child nodes.
-  void replace_node(Node *node, Node *new_node);
+  LIBC_INLINE void replace_node(Node* node, Node* new_node);
 
   Node *&root;
   SizeRange range;
@@ -280,6 +281,54 @@ LIBC_INLINE FreeTrie::Node *FreeTrie::pop_min() {
   return min_node;
 }
 
+LIBC_INLINE void FreeTrie::remove(Node* node) {
+  LIBC_ASSERT(!empty() && "cannot remove from empty trie");
+  FreeList list = node;
+  list.pop();
+  Node* new_node = static_cast<Node*>(list.begin());
+  if (!new_node) {
+    // The freelist is empty. Replace the subtrie root with an arbitrary leaf.
+    // This is legal because there is no relationship between the size of the
+    // root and its children.
+    Node* leaf = node;
+    while (leaf->lower || leaf->upper)
+      leaf = leaf->lower ? leaf->lower : leaf->upper;
+    if (leaf == node) {
+      // If the root is a leaf, then removing it empties the subtrie.
+      replace_node(node, nullptr);
+      return;
+    }
+
+    replace_node(leaf, nullptr);
+    new_node = leaf;
+  }
+
+  if (!is_head(node)) return;
+
+  // Copy the trie links to the new head.
+  new_node->lower = node->lower;
+  new_node->upper = node->upper;
+  new_node->parent = node->parent;
+  replace_node(node, new_node);
+}
+
+LIBC_INLINE void FreeTrie::replace_node(Node* node, Node* new_node) {
+  LIBC_ASSERT(is_head(node) && "only head nodes contain trie links");
+
+  if (node->parent) {
+    Node*& parent_child =
+        node->parent->lower == node ? node->parent->lower : node->parent->upper;
+    LIBC_ASSERT(parent_child == node &&
+                "no reference to child node found in parent");
+    parent_child = new_node;
+  } else {
+    LIBC_ASSERT(root == node && "non-root node had no parent");
+    root = new_node;
+  }
+  if (node->lower) node->lower->parent = new_node;
+  if (node->upper) node->upper->parent = new_node;
+}
+
 } // namespace LIBC_NAMESPACE_DECL
 
 #endif // LLVM_LIBC_SRC___SUPPORT_FREETRIE_H
diff --git a/utils/bazel/llvm-project-overlay/libc/BUILD.bazel b/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
index 7c0a9faff678d..623d9f0858eac 100644
--- a/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
@@ -17,6 +17,7 @@ load(
     "libc_math_function",
     "libc_support_library",
 )
+load(":libc_namespace.bzl", "LIBC_NAMESPACE")
 load(":platforms.bzl", "PLATFORM_CPU_ARM64", "PLATFORM_CPU_X86_64")
 
 package(
@@ -16250,3 +16251,94 @@ libc_function(
         ":types_wchar_t",
     ],
 )
+
+libc_support_library(
+    name = "__support_block",
+    hdrs = ["src/__support/block.h"],
+    deps = [
+        ":__support_common",
+        ":__support_cpp_algorithm",
+        ":__support_cpp_cstddef",
+        ":__support_cpp_limits",
+        ":__support_cpp_new",
+        ":__support_cpp_optional",
+        ":__support_cpp_span",
+        ":__support_libc_assert",
+        ":__support_macros_attributes",
+        ":__support_macros_config",
+        ":__support_math_extras",
+        ":hdr_stdint_proxy",
+        ":string_memory_utils",
+        ":types_size_t",
+    ],
+)
+
+libc_support_library(
+    name = "freelist",
+    hdrs = ["src/__support/freelist.h"],
+    deps = [
+        ":__support_block",
+        ":__support_common",
+        ":__support_libc_assert",
+        ":__support_macros_config",
+    ],
+)
+
+libc_support_library(
+    name = "freetrie",
+    hdrs = ["src/__support/freetrie.h"],
+    deps = [
+        ":__support_common",
+        ":__support_libc_assert",
+        ":__support_macros_config",
+        ":freelist",
+    ],
+)
+
+libc_support_library(
+    name = "freestore",
+    hdrs = ["src/__support/freestore.h"],
+    deps = [
+        ":__support_block",
+        ":__support_common",
+        ":__support_cpp_array",
+        ":__support_cpp_bit",
+        ":__support_cpp_limits",
+        ":__support_macros_config",
+        ":__support_macros_optimization",
+        ":freelist",
+        ":freetrie",
+        ":hdr_stdint_proxy",
+        ":types_size_t",
+    ],
+)
+
+libc_support_library(
+    name = "freelist_heap",
+    hdrs = ["src/__support/freelist_heap.h"],
+    deps = [
+        ":__support_block",
+        ":__support_common",
+        ":__support_cpp_optional",
+        ":__support_cpp_span",
+        ":__support_libc_assert",
+        ":__support_macros_attributes",
+        ":__support_macros_config",
+        ":__support_math_extras",
+        ":freestore",
+        ":string_memory_utils",
+    ],
+)
+
+cc_library(
+    name = "freelist_heap_lib",
+    defines = [
+        "LLVM_LIBC_SRC___SUPPORT_CPP_NEW_H",
+        "LIBC_NAMESPACE=" + LIBC_NAMESPACE,
+    ],
+    includes = ["."],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":freelist_heap",
+    ],
+)

>From e6b3c88d7bbef93b628b090b8f6c9a9bf871e5e8 Mon Sep 17 00:00:00 2001
From: yfzhu <yfzhu at google.com>
Date: Tue, 23 Jun 2026 09:17:39 -0700
Subject: [PATCH 04/15] [libc] add option for hardened freelist

This change adds the LIBC_COPT_HARDEN_FREELIST option which enables
encryption of forward (next) and backward (prev) pointers in the baremetal
freelist heap. When enabled, the FreeListHeap class stores three const
uintptr_t keys, which are initialized in its constructors. These keys are
passed down to FreeStore and FreeList operations to perform encoding and
corruption verification. If corruption is detected, the allocator will trap.

TAG=agy
CONV=4118a58d-30c6-4fd4-8e45-defe707d8bba
---
 libc/src/__support/freelist.h                 | 132 +++++++++++++++---
 libc/src/__support/freelist_heap.h            |  91 ++++++++----
 libc/src/__support/freestore.h                |  45 +++---
 libc/src/__support/freetrie.h                 |   4 +-
 .../test/src/__support/freelist_heap_test.cpp |  21 ++-
 libc/test/src/__support/freelist_test.cpp     |  83 +++++++++--
 libc/test/src/__support/freestore_test.cpp    |  52 ++++---
 .../llvm-project-overlay/libc/BUILD.bazel     |   1 +
 .../libc/test/src/__support/BUILD.bazel       |  50 +++++++
 9 files changed, 371 insertions(+), 108 deletions(-)

diff --git a/libc/src/__support/freelist.h b/libc/src/__support/freelist.h
index 55d1be9115c7d..46b1893f038a7 100644
--- a/libc/src/__support/freelist.h
+++ b/libc/src/__support/freelist.h
@@ -15,11 +15,75 @@
 #define LLVM_LIBC_SRC___SUPPORT_FREELIST_H
 
 #include "block.h"
+#include "hdr/stdint_proxy.h"
 #include "src/__support/libc_assert.h"
 #include "src/__support/macros/config.h"
 
+#ifndef LIBC_COPT_HARDEN_FREELIST
+#define LIBC_COPT_HARDEN_FREELIST false
+#endif
+
+#if LIBC_COPT_HARDEN_FREELIST
+#define LIBC_HARDENING_ASSERT(cond) \
+  do {                              \
+    if (LIBC_UNLIKELY(!(cond))) {   \
+      __builtin_trap();             \
+    }                               \
+  } while (0)
+#else
+#define LIBC_HARDENING_ASSERT(cond) LIBC_ASSERT(cond)
+#endif
+
 namespace LIBC_NAMESPACE_DECL {
 
+struct FreeListSecrets {
+#if LIBC_COPT_HARDEN_FREELIST
+  uintptr_t k0;
+  uintptr_t k1;
+  uintptr_t k2;
+#endif
+
+  template <typename T>
+  LIBC_INLINE T* decrypt_next(T* next_val) const {
+#if LIBC_COPT_HARDEN_FREELIST
+    return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(next_val) ^ k0);
+#else
+    return next_val;
+#endif
+  }
+
+  template <typename T>
+  LIBC_INLINE T* decrypt_prev([[maybe_unused]] const void* node,
+                              T* prev_val) const {
+#if LIBC_COPT_HARDEN_FREELIST
+    return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(prev_val) ^ k1 ^
+                                reinterpret_cast<uintptr_t>(node) ^ k2);
+#else
+    return prev_val;
+#endif
+  }
+
+  template <typename T>
+  LIBC_INLINE T* encrypt_next(T* next_val) const {
+#if LIBC_COPT_HARDEN_FREELIST
+    return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(next_val) ^ k0);
+#else
+    return next_val;
+#endif
+  }
+
+  template <typename T>
+  LIBC_INLINE T* encrypt_prev([[maybe_unused]] const void* node,
+                              T* prev_val) const {
+#if LIBC_COPT_HARDEN_FREELIST
+    return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(prev_val) ^ k1 ^
+                                reinterpret_cast<uintptr_t>(node) ^ k2);
+#else
+    return prev_val;
+#endif
+  }
+};
+
 /// A circularly-linked FIFO list storing free Blocks. All Blocks on a list
 /// are the same size. The blocks are referenced by Nodes in the list; the list
 /// refers to these, but it does not own them.
@@ -43,9 +107,6 @@ class FreeList {
     /// @returns The inner size of blocks in the list containing this node.
     LIBC_INLINE size_t size() const { return block().inner_size(); }
 
-    /// @returns The next node in the list containing this node.
-    LIBC_INLINE Node *next_node() const { return next; }
-
   private:
     // Circularly linked pointers to adjacent nodes.
     Node *prev;
@@ -53,7 +114,7 @@ class FreeList {
     friend class FreeList;
   };
 
-  LIBC_INLINE constexpr FreeList() : FreeList(nullptr) {}
+  LIBC_INLINE constexpr FreeList() : begin_(nullptr) {}
   LIBC_INLINE constexpr FreeList(Node *begin) : begin_(begin) {}
 
   LIBC_INLINE bool empty() const { return !begin_; }
@@ -70,47 +131,72 @@ class FreeList {
   /// @returns The first block in the list.
   LIBC_INLINE BlockRef front() { return begin_->block(); }
 
+  LIBC_INLINE Node* next_node(const Node* node,
+                              const FreeListSecrets& secrets) const {
+    return node ? secrets.decrypt_next(node->next) : nullptr;
+  }
+
+  LIBC_INLINE Node* prev_node(const Node* node,
+                              const FreeListSecrets& secrets) const {
+    return node ? secrets.decrypt_prev(node, node->prev) : nullptr;
+  }
+
   /// Push a block to the back of the list.
   /// The block must be large enough to contain a node.
-  LIBC_INLINE void push(BlockRef block) {
+  LIBC_INLINE void push(BlockRef block, const FreeListSecrets& secrets) {
     LIBC_ASSERT(!block.used() &&
                 "only free blocks can be placed on free lists");
-    LIBC_ASSERT(block.inner_size_free() >= sizeof(FreeList) &&
+    LIBC_ASSERT(block.inner_size_free() >= sizeof(Node) &&
                 "block too small to accomodate free list node");
-    push(new (block.usable_space()) Node);
+    push(new (block.usable_space()) Node, secrets);
   }
 
   /// Push an already-constructed node to the back of the list.
   /// This allows pushing derived node types with additional data.
-  LIBC_INLINE void push(Node* node) {
+  LIBC_INLINE void push(Node* node, const FreeListSecrets& secrets) {
     if (begin_) {
-      LIBC_ASSERT(BlockRef::from_usable_space(node).outer_size() ==
-                      begin_->block().outer_size() &&
-                  "freelist entries must have the same size");
-      // Since the list is circular, insert the node immediately before begin_.
-      node->prev = begin_->prev;
-      node->next = begin_;
-      begin_->prev->next = node;
-      begin_->prev = node;
+      Node* begin_prev = secrets.decrypt_prev(begin_, begin_->prev);
+
+      LIBC_HARDENING_ASSERT(secrets.decrypt_next(begin_prev->next) == begin_ &&
+                            "Corrupted free list links (push check)");
+
+      node->prev = secrets.encrypt_prev(node, begin_prev);
+      node->next = secrets.encrypt_next(begin_);
+      begin_prev->next = secrets.encrypt_next(node);
+      begin_->prev = secrets.encrypt_prev(begin_, node);
     } else {
-      begin_ = node->prev = node->next = node;
+      begin_ = node;
+      node->next = secrets.encrypt_next(node);
+      node->prev = secrets.encrypt_prev(node, node);
     }
   }
 
   /// Pop the first node from the list.
-  LIBC_INLINE void pop() { remove(begin_); }
+  LIBC_INLINE void pop(const FreeListSecrets& secrets) {
+    remove(begin_, secrets);
+  }
 
   /// Remove an arbitrary node from the list.
-  LIBC_INLINE void remove(Node* node) {
+  LIBC_INLINE void remove(Node* node, const FreeListSecrets& secrets) {
     LIBC_ASSERT(begin_ && "cannot remove from empty list");
-    if (node == node->next) {
+    Node* node_next = secrets.decrypt_next(node->next);
+    if (node == node_next) {
       LIBC_ASSERT(node == begin_ &&
                   "a self-referential node must be the only element");
       begin_ = nullptr;
     } else {
-      node->prev->next = node->next;
-      node->next->prev = node->prev;
-      if (begin_ == node) begin_ = node->next;
+      Node* node_prev = secrets.decrypt_prev(node, node->prev);
+
+      LIBC_HARDENING_ASSERT(
+          secrets.decrypt_next(node_prev->next) == node &&
+          "Corrupted free list links (remove check prev->next)");
+      LIBC_HARDENING_ASSERT(
+          secrets.decrypt_prev(node_next, node_next->prev) == node &&
+          "Corrupted free list links (remove check next->prev)");
+
+      node_prev->next = secrets.encrypt_next(node_next);
+      node_next->prev = secrets.encrypt_prev(node_next, node_prev);
+      if (begin_ == node) begin_ = node_next;
     }
   }
 
diff --git a/libc/src/__support/freelist_heap.h b/libc/src/__support/freelist_heap.h
index 4f1cf1dac16b2..e82822825cd9a 100644
--- a/libc/src/__support/freelist_heap.h
+++ b/libc/src/__support/freelist_heap.h
@@ -37,22 +37,45 @@ using cpp::span;
 
 LIBC_INLINE constexpr bool IsPow2(size_t x) { return x && (x & (x - 1)) == 0; }
 
+#ifndef LIBC_COPT_FREELIST_KEY0
+#define LIBC_COPT_FREELIST_KEY0 0
+#endif
+#ifndef LIBC_COPT_FREELIST_KEY1
+#define LIBC_COPT_FREELIST_KEY1 0
+#endif
+#ifndef LIBC_COPT_FREELIST_KEY2
+#define LIBC_COPT_FREELIST_KEY2 0
+#endif
+
 class FreeListHeap {
 public:
-  constexpr FreeListHeap() : begin(&_end), end(&__llvm_libc_heap_limit) {}
-
-  constexpr FreeListHeap(span<cpp::byte> region)
-      : begin(region.begin()), end(region.end()) {}
-
-  void *allocate(size_t size);
-  void *aligned_allocate(size_t alignment, size_t size);
-  // NOTE: All pointers passed to free must come from one of the other
-  // allocation functions: `allocate`, `aligned_allocate`, `realloc`, `calloc`.
-  void free(void *ptr);
-  void *realloc(void *ptr, size_t size);
-  void *calloc(size_t num, size_t size);
-
-  cpp::span<cpp::byte> region() const { return {begin, end}; }
+#if LIBC_COPT_HARDEN_FREELIST
+ constexpr FreeListHeap(const FreeListSecrets& secrets =
+                            {LIBC_COPT_FREELIST_KEY0, LIBC_COPT_FREELIST_KEY1,
+                             LIBC_COPT_FREELIST_KEY2})
+     : begin(&_end), end(&__llvm_libc_heap_limit), secrets(secrets) {}
+
+ constexpr FreeListHeap(span<cpp::byte> region,
+                        const FreeListSecrets& secrets =
+                            {LIBC_COPT_FREELIST_KEY0, LIBC_COPT_FREELIST_KEY1,
+                             LIBC_COPT_FREELIST_KEY2})
+     : begin(region.begin()), end(region.end()), secrets(secrets) {}
+#else
+ constexpr FreeListHeap() : begin(&_end), end(&__llvm_libc_heap_limit) {}
+
+ constexpr FreeListHeap(span<cpp::byte> region)
+     : begin(region.begin()), end(region.end()) {}
+#endif
+
+ void* allocate(size_t size);
+ void* aligned_allocate(size_t alignment, size_t size);
+ // NOTE: All pointers passed to free must come from one of the other
+ // allocation functions: `allocate`, `aligned_allocate`, `realloc`, `calloc`.
+ void free(void* ptr);
+ void* realloc(void* ptr, size_t size);
+ void* calloc(size_t num, size_t size);
+
+ cpp::span<cpp::byte> region() const { return {begin, end}; }
 
 private:
   void init();
@@ -71,21 +94,39 @@ class FreeListHeap {
   cpp::byte *end;
   bool is_initialized = false;
   FreeStore free_store;
+#if LIBC_COPT_HARDEN_FREELIST
+  const FreeListSecrets secrets;
+#endif
 };
 
 template <size_t BUFF_SIZE> class FreeListHeapBuffer : public FreeListHeap {
 public:
-  constexpr FreeListHeapBuffer() : FreeListHeap{buffer}, buffer{} {}
+#if LIBC_COPT_HARDEN_FREELIST
+ constexpr FreeListHeapBuffer(
+     const FreeListSecrets& secrets = {LIBC_COPT_FREELIST_KEY0,
+                                       LIBC_COPT_FREELIST_KEY1,
+                                       LIBC_COPT_FREELIST_KEY2})
+     : FreeListHeap{buffer, secrets}, buffer{} {}
+#else
+ constexpr FreeListHeapBuffer() : FreeListHeap{buffer}, buffer{} {}
+#endif
 
 private:
   cpp::byte buffer[BUFF_SIZE];
 };
 
+#if LIBC_COPT_HARDEN_FREELIST
+#define HEAP_SECRETS secrets
+#else
+#define HEAP_SECRETS \
+  FreeListSecrets {}
+#endif
+
 [[gnu::noinline]] LIBC_INLINE void FreeListHeap::init() {
   LIBC_ASSERT(!is_initialized && "duplicate initialization");
   auto result = BlockRef::init(region());
   BlockRef block = *result;
-  free_store.insert(block);
+  free_store.insert(block, HEAP_SECRETS);
   is_initialized = true;
 }
 
@@ -101,15 +142,13 @@ template <size_t BUFF_SIZE> class FreeListHeapBuffer : public FreeListHeap {
   if (!request_size)
     return nullptr;
 
-  BlockRef block = free_store.remove_best_fit(request_size);
+  BlockRef block = free_store.remove_best_fit(request_size, HEAP_SECRETS);
   if (!block)
     return nullptr;
 
   auto block_info = BlockRef::allocate(block, alignment, size);
-  if (block_info.next)
-    free_store.insert(block_info.next);
-  if (block_info.prev)
-    free_store.insert(block_info.prev);
+  if (block_info.next) free_store.insert(block_info.next, HEAP_SECRETS);
+  if (block_info.prev) free_store.insert(block_info.prev, HEAP_SECRETS);
 
   block_info.block.mark_used();
   return block_info.block.usable_space();
@@ -154,16 +193,16 @@ template <size_t BUFF_SIZE> class FreeListHeapBuffer : public FreeListHeap {
 
   if (prev_free) {
     // Remove from free store and merge.
-    free_store.remove(prev_free);
+    free_store.remove(prev_free, HEAP_SECRETS);
     block = prev_free;
     block.merge_next();
   }
   if (!next.used()) {
-    free_store.remove(next);
+    free_store.remove(next, HEAP_SECRETS);
     block.merge_next();
   }
   // Add back to the freelist
-  free_store.insert(block);
+  free_store.insert(block, HEAP_SECRETS);
 }
 
 [[gnu::noinline]] LIBC_INLINE bool FreeListHeap::shrink_in_place(BlockRef block,
@@ -184,10 +223,10 @@ template <size_t BUFF_SIZE> class FreeListHeapBuffer : public FreeListHeap {
       // to be non-null.
       LIBC_ASSERT(right && "right block must be non-null");
       if (!right.used()) {
-        free_store.remove(right);
+        free_store.remove(right, HEAP_SECRETS);
         next_block.merge_next();
       }
-      free_store.insert(next_block);
+      free_store.insert(next_block, HEAP_SECRETS);
     }
     return true;
   }
diff --git a/libc/src/__support/freestore.h b/libc/src/__support/freestore.h
index 13b79c0debf39..ac152e587a112 100644
--- a/libc/src/__support/freestore.h
+++ b/libc/src/__support/freestore.h
@@ -121,14 +121,16 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
   LIBC_INLINE TLSFFreeStoreImpl &
   operator=(const TLSFFreeStoreImpl &other) = delete;
 
-  LIBC_INLINE void insert(BlockRef block);
-  LIBC_INLINE void remove(BlockRef block);
-  LIBC_INLINE BlockRef remove_best_fit(size_t size) {
-    return find_and_remove_fit(size);
+  LIBC_INLINE void insert(BlockRef block, const FreeListSecrets& secrets);
+  LIBC_INLINE void remove(BlockRef block, const FreeListSecrets& secrets);
+  LIBC_INLINE BlockRef remove_best_fit(size_t size,
+                                       const FreeListSecrets& secrets) {
+    return find_and_remove_fit(size, secrets);
   }
-  LIBC_INLINE BlockRef find_and_remove_fit(size_t size);
+  LIBC_INLINE BlockRef find_and_remove_fit(size_t size,
+                                           const FreeListSecrets& secrets);
 
-protected:
+ protected:
   LIBC_INLINE static bool too_small(BlockRef block) {
     return block.outer_size() < MIN_OUTER_SIZE;
   }
@@ -148,7 +150,8 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
   LIBC_INLINE void clear_bit(size_t bit_index);
   LIBC_INLINE bool get_bit(size_t bit_index) const;
   LIBC_INLINE size_t find_first_bit_set_after(size_t bit_index) const;
-  LIBC_INLINE BlockRef remove_first_fit_in_list(size_t index, size_t size);
+  LIBC_INLINE BlockRef remove_first_fit_in_list(size_t index, size_t size,
+                                                const FreeListSecrets& secrets);
   LIBC_INLINE FreeTrie get_trie();
   LIBC_INLINE BlockRef find_and_remove_fit_in_trie(size_t size);
   LIBC_INLINE BlockRef pop_min_in_trie();
@@ -245,7 +248,8 @@ LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::pop_min_in_trie() {
 }
 
 template <typename CONFIG>
-LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(BlockRef block) {
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(
+    BlockRef block, const FreeListSecrets& secrets) {
   if (too_small(block))
     return;
   size_t bit_index = size_to_bit_index(block.inner_size());
@@ -257,12 +261,13 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(BlockRef block) {
       return;
     }
 
-  free_lists[bit_index].list.push(block);
+  free_lists[bit_index].list.push(block, secrets);
   set_bit(bit_index);
 }
 
 template <typename CONFIG>
-LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(BlockRef block) {
+LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(
+    BlockRef block, const FreeListSecrets& secrets) {
   if (too_small(block))
     return;
   size_t bit_index = size_to_bit_index(block.inner_size());
@@ -277,14 +282,14 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(BlockRef block) {
     }
 
   free_lists[bit_index].list.remove(
-      reinterpret_cast<FreeList::Node *>(block.usable_space()));
+      reinterpret_cast<FreeList::Node*>(block.usable_space()), secrets);
   if (free_lists[bit_index].list.empty())
     clear_bit(bit_index);
 }
 
 template <typename CONFIG>
-LIBC_INLINE BlockRef
-TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(size_t index, size_t size) {
+LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(
+    size_t index, size_t size, const FreeListSecrets& secrets) {
   FreeList::Node *begin_node = free_lists[index].list.begin();
   if (begin_node == nullptr)
     return BlockRef();
@@ -292,27 +297,27 @@ TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(size_t index, size_t size) {
   FreeList::Node *cur = begin_node;
   do {
     if (cur->size() >= size) {
-      free_lists[index].list.remove(cur);
+      free_lists[index].list.remove(cur, secrets);
       if (free_lists[index].list.empty())
         clear_bit(index);
       return cur->block();
     }
-    cur = cur->next_node();
+    cur = free_lists[index].list.next_node(cur, secrets);
   } while (cur != begin_node);
 
   return BlockRef();
 }
 
 template <typename CONFIG>
-LIBC_INLINE BlockRef
-TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
+LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(
+    size_t size, const FreeListSecrets& secrets) {
   size_t bit_index = size_to_bit_index(size);
 
   if (LIBC_UNLIKELY(bit_index >= TOTAL_BITS - 1)) {
     if constexpr (USE_TRIE)
       return find_and_remove_fit_in_trie(size);
     else
-      return remove_first_fit_in_list(TOTAL_BITS - 1, size);
+      return remove_first_fit_in_list(TOTAL_BITS - 1, size, secrets);
   }
 
   // 1. Try oversized bins (guaranteed fit, but larger).
@@ -324,7 +329,7 @@ TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
     }
 
     BlockRef block = free_lists[oversized_bit].list.front();
-    free_lists[oversized_bit].list.pop();
+    free_lists[oversized_bit].list.pop(secrets);
     if (free_lists[oversized_bit].list.empty())
       clear_bit(oversized_bit);
     return block;
@@ -332,7 +337,7 @@ TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(size_t size) {
 
   // 2. Try exact fit (fallback).
   if (get_bit(bit_index)) {
-    if (BlockRef block = remove_first_fit_in_list(bit_index, size))
+    if (BlockRef block = remove_first_fit_in_list(bit_index, size, secrets))
       return block;
   }
 
diff --git a/libc/src/__support/freetrie.h b/libc/src/__support/freetrie.h
index f238a6b441cd5..bc2981d22c676 100644
--- a/libc/src/__support/freetrie.h
+++ b/libc/src/__support/freetrie.h
@@ -155,7 +155,7 @@ LIBC_INLINE void FreeTrie::push(BlockRef block) {
   } else {
     node->parent = nullptr;
   }
-  list.push(node);
+  list.push(node, FreeListSecrets{});
   *cur = static_cast<Node *>(list.begin());
 }
 
@@ -284,7 +284,7 @@ LIBC_INLINE FreeTrie::Node *FreeTrie::pop_min() {
 LIBC_INLINE void FreeTrie::remove(Node* node) {
   LIBC_ASSERT(!empty() && "cannot remove from empty trie");
   FreeList list = node;
-  list.pop();
+  list.pop(FreeListSecrets{});
   Node* new_node = static_cast<Node*>(list.begin());
   if (!new_node) {
     // The freelist is empty. Replace the subtrie root with an arbitrary leaf.
diff --git a/libc/test/src/__support/freelist_heap_test.cpp b/libc/test/src/__support/freelist_heap_test.cpp
index 68cc30152cd6d..34f96cc5fac33 100644
--- a/libc/test/src/__support/freelist_heap_test.cpp
+++ b/libc/test/src/__support/freelist_heap_test.cpp
@@ -32,9 +32,18 @@ using LIBC_NAMESPACE::BlockRef;
 using LIBC_NAMESPACE::freelist_heap;
 using LIBC_NAMESPACE::FreeListHeap;
 using LIBC_NAMESPACE::FreeListHeapBuffer;
+using LIBC_NAMESPACE::FreeListSecrets;
 using LIBC_NAMESPACE::cpp::byte;
 using LIBC_NAMESPACE::cpp::span;
 
+#if LIBC_COPT_HARDEN_FREELIST
+#define TEST_SECRETS_INIT FreeListSecrets{0x123, 0x456, 0x789}
+#define TEST_SECRETS_ARG , FreeListSecrets{0x123, 0x456, 0x789}
+#else
+#define TEST_SECRETS_INIT
+#define TEST_SECRETS_ARG
+#endif
+
 // Similar to `LlvmLibcBlockTest` in block_test.cpp, we'd like to run the same
 // tests independently for different parameters. In this case, we'd like to test
 // functionality for a `FreeListHeap` and the global `freelist_heap` which was
@@ -49,16 +58,16 @@ using LIBC_NAMESPACE::cpp::span;
   class LlvmLibcFreeListHeapTest##TestCase                                     \
       : public LIBC_NAMESPACE::testing::Test {                                 \
   public:                                                                      \
-    FreeListHeapBuffer<BufferSize> fake_global_buffer;                         \
+    FreeListHeapBuffer<BufferSize> fake_global_buffer{TEST_SECRETS_INIT};      \
     void SetUp() override {                                                    \
-      freelist_heap =                                                          \
-          new (&fake_global_buffer) FreeListHeapBuffer<BufferSize>;            \
+      freelist_heap = new (&fake_global_buffer)                                \
+          FreeListHeapBuffer<BufferSize>{TEST_SECRETS_INIT};                   \
     }                                                                          \
     void RunTest(FreeListHeap &allocator, [[maybe_unused]] size_t N);          \
   };                                                                           \
   TEST_F(LlvmLibcFreeListHeapTest##TestCase, TestCase) {                       \
     byte buf[BufferSize] = {byte(0)};                                          \
-    FreeListHeap allocator(buf);                                               \
+    FreeListHeap allocator(buf TEST_SECRETS_ARG);                              \
     RunTest(allocator, BufferSize);                                            \
     RunTest(*freelist_heap, freelist_heap->region().size());                   \
   }                                                                            \
@@ -112,7 +121,7 @@ TEST(LlvmLibcFreeListHeap, ReturnsNullWhenFull) {
   constexpr size_t N = 2048;
   byte buf[N];
 
-  FreeListHeap allocator(buf);
+  FreeListHeap allocator(buf TEST_SECRETS_ARG);
 
   bool went_null = false;
   for (size_t i = 0; i < N; i++) {
@@ -308,7 +317,7 @@ TEST(LlvmLibcFreeListHeap, AlignedAllocUnalignedBuffer) {
   byte buf[4096] = {byte(0)};
 
   // Ensure the underlying buffer is poorly aligned.
-  FreeListHeap allocator(span<byte>(buf).subspan(1));
+  FreeListHeap allocator(span<byte>(buf).subspan(1) TEST_SECRETS_ARG);
 
   constexpr size_t ALIGNMENTS[] = {1, 2, 4, 8, 16, 32, 64, 128, 256};
   constexpr size_t SIZE_SCALES[] = {1, 2, 3, 4, 5};
diff --git a/libc/test/src/__support/freelist_test.cpp b/libc/test/src/__support/freelist_test.cpp
index 580f59ae62bd8..17b96b4d50380 100644
--- a/libc/test/src/__support/freelist_test.cpp
+++ b/libc/test/src/__support/freelist_test.cpp
@@ -18,9 +18,17 @@
 
 using LIBC_NAMESPACE::BlockRef;
 using LIBC_NAMESPACE::FreeList;
+using LIBC_NAMESPACE::FreeListSecrets;
 using LIBC_NAMESPACE::cpp::byte;
 using LIBC_NAMESPACE::cpp::optional;
 
+#if LIBC_COPT_HARDEN_FREELIST
+#define TEST_SECRETS FreeListSecrets{0x1234, 0x5678, 0x9abc}
+#else
+#define TEST_SECRETS                                                           \
+  FreeListSecrets {}
+#endif
+
 TEST(LlvmLibcFreeList, FreeList) {
   byte mem[1024];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
@@ -35,24 +43,83 @@ TEST(LlvmLibcFreeList, FreeList) {
   ASSERT_TRUE(maybeBlock.has_value());
 
   FreeList list;
-  list.push(block1);
+  list.push(block1, TEST_SECRETS);
   ASSERT_FALSE(list.empty());
   EXPECT_EQ(list.front().addr(), block1.addr());
 
-  list.push(block2);
+  list.push(block2, TEST_SECRETS);
   EXPECT_EQ(list.front().addr(), block1.addr());
 
-  list.pop();
+  list.pop(TEST_SECRETS);
   ASSERT_FALSE(list.empty());
   EXPECT_EQ(list.front().addr(), block2.addr());
 
-  list.pop();
+  list.pop(TEST_SECRETS);
   ASSERT_TRUE(list.empty());
 
-  list.push(block1);
-  list.push(block2);
-  list.remove(reinterpret_cast<FreeList::Node *>(block2.usable_space()));
+  list.push(block1, TEST_SECRETS);
+  list.push(block2, TEST_SECRETS);
+  list.remove(reinterpret_cast<FreeList::Node *>(block2.usable_space()),
+              TEST_SECRETS);
   EXPECT_EQ(list.front().addr(), block1.addr());
-  list.pop();
+  list.pop(TEST_SECRETS);
   ASSERT_TRUE(list.empty());
 }
+
+#if LIBC_COPT_HARDEN_FREELIST
+TEST(LlvmLibcFreeList, HardenedCorruptNext) {
+  byte mem[1024];
+  optional<BlockRef> maybeBlock = BlockRef::init(mem);
+  ASSERT_TRUE(maybeBlock.has_value());
+  BlockRef block1 = *maybeBlock;
+
+  maybeBlock = block1.split(128);
+  ASSERT_TRUE(maybeBlock.has_value());
+  BlockRef block2 = *maybeBlock;
+
+  FreeList list;
+  list.push(block1, TEST_SECRETS);
+  list.push(block2, TEST_SECRETS);
+
+  struct RawNode {
+    void *prev;
+    void *next;
+  };
+  RawNode *raw_node2 = reinterpret_cast<RawNode *>(block2.usable_space());
+  raw_node2->next = reinterpret_cast<void *>(0xDEADBEEF); // Corrupt next
+
+  EXPECT_DEATH(
+      [&] {
+        list.pop(TEST_SECRETS); // Should trap due to corrupted block2->next
+      },
+      WITH_SIGNAL(-1));
+}
+
+TEST(LlvmLibcFreeList, HardenedCorruptPrev) {
+  byte mem[1024];
+  optional<BlockRef> maybeBlock = BlockRef::init(mem);
+  ASSERT_TRUE(maybeBlock.has_value());
+  BlockRef block1 = *maybeBlock;
+
+  maybeBlock = block1.split(128);
+  ASSERT_TRUE(maybeBlock.has_value());
+  BlockRef block2 = *maybeBlock;
+
+  FreeList list;
+  list.push(block1, TEST_SECRETS);
+  list.push(block2, TEST_SECRETS);
+
+  struct RawNode {
+    void *prev;
+    void *next;
+  };
+  RawNode *raw_node2 = reinterpret_cast<RawNode *>(block2.usable_space());
+  raw_node2->prev = reinterpret_cast<void *>(0xDEADBEEF); // Corrupt prev
+
+  EXPECT_DEATH(
+      [&] {
+        list.pop(TEST_SECRETS); // Should trap due to corrupted block2->prev
+      },
+      WITH_SIGNAL(-1));
+}
+#endif
diff --git a/libc/test/src/__support/freestore_test.cpp b/libc/test/src/__support/freestore_test.cpp
index ce655367da00f..d0433dba9437c 100644
--- a/libc/test/src/__support/freestore_test.cpp
+++ b/libc/test/src/__support/freestore_test.cpp
@@ -18,6 +18,7 @@
 
 using LIBC_NAMESPACE::BlockRef;
 using LIBC_NAMESPACE::FreeList;
+using LIBC_NAMESPACE::FreeListSecrets;
 using LIBC_NAMESPACE::FreeStore;
 using LIBC_NAMESPACE::FreeTrie;
 using LIBC_NAMESPACE::cpp::byte;
@@ -25,6 +26,7 @@ using LIBC_NAMESPACE::cpp::optional;
 
 // Inserting or removing blocks too small to be tracked does nothing.
 TEST(LlvmLibcFreeStore, TooSmall) {
+  FreeListSecrets secrets{};
   byte mem[1024];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
   ASSERT_TRUE(maybeBlock.has_value());
@@ -38,15 +40,16 @@ TEST(LlvmLibcFreeStore, TooSmall) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.insert(too_small);
-  store.insert(remainder);
+  store.insert(too_small, secrets);
+  store.insert(remainder, secrets);
 
-  EXPECT_EQ(store.remove_best_fit(too_small.inner_size()).addr(),
+  EXPECT_EQ(store.remove_best_fit(too_small.inner_size(), secrets).addr(),
             remainder.addr());
-  store.remove(too_small);
+  store.remove(too_small, secrets);
 }
 
 TEST(LlvmLibcFreeStore, RemoveBestFit) {
+  FreeListSecrets secrets{};
   byte mem[1024];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
   ASSERT_TRUE(maybeBlock.has_value());
@@ -67,40 +70,43 @@ TEST(LlvmLibcFreeStore, RemoveBestFit) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.insert(smallest);
+  store.insert(smallest, secrets);
   if (largest_small != smallest)
-    store.insert(largest_small);
-  store.insert(remainder);
+    store.insert(largest_small, secrets);
+  store.insert(remainder, secrets);
 
   // For TLSF (oversized first), asking for a size will return the block from
   // the first non-empty oversized bin if one exists, bypassing the exact bin.
   if (largest_small != smallest) {
-    BlockRef block = store.remove_best_fit(smallest.inner_size());
+    BlockRef block = store.remove_best_fit(smallest.inner_size(), secrets);
     ASSERT_EQ(block.addr(), largest_small.addr());
-    store.insert(block);
+    store.insert(block, secrets);
 
-    BlockRef block2 = store.remove_best_fit(largest_small.inner_size());
+    BlockRef block2 =
+        store.remove_best_fit(largest_small.inner_size(), secrets);
     ASSERT_EQ(block2.addr(), remainder.addr());
-    store.insert(block2);
+    store.insert(block2, secrets);
   } else {
-    BlockRef block = store.remove_best_fit(smallest.inner_size());
+    BlockRef block = store.remove_best_fit(smallest.inner_size(), secrets);
     ASSERT_EQ(block.addr(), remainder.addr());
-    store.insert(block);
+    store.insert(block, secrets);
   }
 
   // Search small list for best fit.
   BlockRef next_smallest =
       largest_small == smallest ? remainder : largest_small;
-  ASSERT_EQ(store.remove_best_fit(smallest.inner_size() + 1).addr(),
+  ASSERT_EQ(store.remove_best_fit(smallest.inner_size() + 1, secrets).addr(),
             next_smallest.addr());
-  store.insert(next_smallest);
+  store.insert(next_smallest, secrets);
 
   // Continue search for best fit to large blocks.
-  EXPECT_EQ(store.remove_best_fit(largest_small.inner_size() + 1).addr(),
-            remainder.addr());
+  EXPECT_EQ(
+      store.remove_best_fit(largest_small.inner_size() + 1, secrets).addr(),
+      remainder.addr());
 }
 
 TEST(LlvmLibcFreeStore, Remove) {
+  FreeListSecrets secrets{};
   byte mem[1024];
   optional<BlockRef> maybeBlock = BlockRef::init(mem);
   ASSERT_TRUE(maybeBlock.has_value());
@@ -112,13 +118,13 @@ TEST(LlvmLibcFreeStore, Remove) {
   BlockRef remainder = *maybeBlock;
 
   FreeStore store;
-  store.insert(small);
-  store.insert(remainder);
+  store.insert(small, secrets);
+  store.insert(remainder, secrets);
 
-  store.remove(remainder);
-  ASSERT_EQ(store.remove_best_fit(remainder.inner_size()).addr(),
+  store.remove(remainder, secrets);
+  ASSERT_EQ(store.remove_best_fit(remainder.inner_size(), secrets).addr(),
             BlockRef().addr());
-  store.remove(small);
-  ASSERT_EQ(store.remove_best_fit(small.inner_size()).addr(),
+  store.remove(small, secrets);
+  ASSERT_EQ(store.remove_best_fit(small.inner_size(), secrets).addr(),
             BlockRef().addr());
 }
diff --git a/utils/bazel/llvm-project-overlay/libc/BUILD.bazel b/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
index 623d9f0858eac..83115da15f4e4 100644
--- a/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/libc/BUILD.bazel
@@ -16281,6 +16281,7 @@ libc_support_library(
         ":__support_common",
         ":__support_libc_assert",
         ":__support_macros_config",
+        ":hdr_stdint_proxy",
     ],
 )
 
diff --git a/utils/bazel/llvm-project-overlay/libc/test/src/__support/BUILD.bazel b/utils/bazel/llvm-project-overlay/libc/test/src/__support/BUILD.bazel
index 0698f45ceabc9..7a221e62bfee3 100644
--- a/utils/bazel/llvm-project-overlay/libc/test/src/__support/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/libc/test/src/__support/BUILD.bazel
@@ -121,3 +121,53 @@ libc_test(
         "//libc:__support_macros_properties_types",
     ],
 )
+
+libc_test(
+    name = "freelist_test",
+    srcs = ["freelist_test.cpp"],
+    deps = [
+        "//libc:__support_block",
+        "//libc:__support_cpp_optional",
+        "//libc:freelist",
+    ],
+)
+
+libc_test(
+    name = "freelist_hardened_test",
+    srcs = ["freelist_test.cpp"],
+    defines = ["LIBC_COPT_HARDEN_FREELIST=1"],
+    deps = [
+        "//libc:__support_block",
+        "//libc:__support_cpp_optional",
+        "//libc:freelist",
+    ],
+)
+
+libc_test(
+    name = "freelist_heap_test",
+    srcs = ["freelist_heap_test.cpp"],
+    deps = [
+        "//libc:__support_block",
+        "//libc:__support_cpp_optional",
+        "//libc:__support_cpp_span",
+        "//libc:freelist_heap",
+        "//libc:memcmp",
+        "//libc:memcpy",
+        "//libc:string_memory_utils",
+    ],
+)
+
+libc_test(
+    name = "freelist_heap_hardened_test",
+    srcs = ["freelist_heap_test.cpp"],
+    defines = ["LIBC_COPT_HARDEN_FREELIST=1"],
+    deps = [
+        "//libc:__support_block",
+        "//libc:__support_cpp_optional",
+        "//libc:__support_cpp_span",
+        "//libc:freelist_heap",
+        "//libc:memcmp",
+        "//libc:memcpy",
+        "//libc:string_memory_utils",
+    ],
+)

>From 70abce682c122b870c0134eacd7195241f3e632b Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yfzhu at google.com>
Date: Wed, 24 Jun 2026 11:04:26 -0700
Subject: [PATCH 05/15] [libc] Disable loop unrolling in FreeStore only for
 32-bit targets

Loop unrolling in list search and table scanning in FreeStore increases
code size with negligible performance benefit for 32-bit targets (like Cortex-M33)
where heap sizes are small and lists are short.

This change introduces LIBC_LOOP_NOUNROLL_32 which disables unrolling
only when __SIZEOF_POINTER__ == 4.

Also, default FreeStore config for 32-bit is adjusted to (8, 2, 2, 2)
to optimize lookup table size and reduce fragmentation for small heaps
(e.g., 336 KiB).
---
 libc/src/__support/freestore.h | 41 ++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/libc/src/__support/freestore.h b/libc/src/__support/freestore.h
index ac152e587a112..7f0427854890e 100644
--- a/libc/src/__support/freestore.h
+++ b/libc/src/__support/freestore.h
@@ -25,6 +25,12 @@
 #include "src/__support/macros/config.h"
 #include "src/__support/macros/optimization.h"
 
+#if defined(__SIZEOF_POINTER__) && __SIZEOF_POINTER__ == 4
+#define LIBC_LOOP_NOUNROLL_32 _Pragma("clang loop unroll(disable)")
+#else
+#define LIBC_LOOP_NOUNROLL_32
+#endif
+
 namespace LIBC_NAMESPACE_DECL {
 
 /// Configuration for TLSFFreeStore.
@@ -121,16 +127,16 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
   LIBC_INLINE TLSFFreeStoreImpl &
   operator=(const TLSFFreeStoreImpl &other) = delete;
 
-  LIBC_INLINE void insert(BlockRef block, const FreeListSecrets& secrets);
-  LIBC_INLINE void remove(BlockRef block, const FreeListSecrets& secrets);
+  LIBC_INLINE void insert(BlockRef block, const FreeListSecrets &secrets);
+  LIBC_INLINE void remove(BlockRef block, const FreeListSecrets &secrets);
   LIBC_INLINE BlockRef remove_best_fit(size_t size,
-                                       const FreeListSecrets& secrets) {
+                                       const FreeListSecrets &secrets) {
     return find_and_remove_fit(size, secrets);
   }
   LIBC_INLINE BlockRef find_and_remove_fit(size_t size,
-                                           const FreeListSecrets& secrets);
+                                           const FreeListSecrets &secrets);
 
- protected:
+protected:
   LIBC_INLINE static bool too_small(BlockRef block) {
     return block.outer_size() < MIN_OUTER_SIZE;
   }
@@ -151,7 +157,7 @@ template <typename CONFIG> class TLSFFreeStoreImpl {
   LIBC_INLINE bool get_bit(size_t bit_index) const;
   LIBC_INLINE size_t find_first_bit_set_after(size_t bit_index) const;
   LIBC_INLINE BlockRef remove_first_fit_in_list(size_t index, size_t size,
-                                                const FreeListSecrets& secrets);
+                                                const FreeListSecrets &secrets);
   LIBC_INLINE FreeTrie get_trie();
   LIBC_INLINE BlockRef find_and_remove_fit_in_trie(size_t size);
   LIBC_INLINE BlockRef pop_min_in_trie();
@@ -208,6 +214,7 @@ TLSFFreeStoreImpl<CONFIG>::find_first_bit_set_after(size_t bit_index) const {
     return start_entry * BITS_PER_ENTRY +
            static_cast<size_t>(cpp::countr_zero(value));
 
+  LIBC_LOOP_NOUNROLL_32
   for (size_t i = start_entry + 1; i < CONFIG::NUM_TABLE_ENTRIES; ++i) {
     value = lookup_table[i];
     if (value != 0)
@@ -248,8 +255,9 @@ LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::pop_min_in_trie() {
 }
 
 template <typename CONFIG>
-LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(
-    BlockRef block, const FreeListSecrets& secrets) {
+LIBC_INLINE void
+TLSFFreeStoreImpl<CONFIG>::insert(BlockRef block,
+                                  const FreeListSecrets &secrets) {
   if (too_small(block))
     return;
   size_t bit_index = size_to_bit_index(block.inner_size());
@@ -266,8 +274,9 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::insert(
 }
 
 template <typename CONFIG>
-LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(
-    BlockRef block, const FreeListSecrets& secrets) {
+LIBC_INLINE void
+TLSFFreeStoreImpl<CONFIG>::remove(BlockRef block,
+                                  const FreeListSecrets &secrets) {
   if (too_small(block))
     return;
   size_t bit_index = size_to_bit_index(block.inner_size());
@@ -282,19 +291,20 @@ LIBC_INLINE void TLSFFreeStoreImpl<CONFIG>::remove(
     }
 
   free_lists[bit_index].list.remove(
-      reinterpret_cast<FreeList::Node*>(block.usable_space()), secrets);
+      reinterpret_cast<FreeList::Node *>(block.usable_space()), secrets);
   if (free_lists[bit_index].list.empty())
     clear_bit(bit_index);
 }
 
 template <typename CONFIG>
 LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(
-    size_t index, size_t size, const FreeListSecrets& secrets) {
+    size_t index, size_t size, const FreeListSecrets &secrets) {
   FreeList::Node *begin_node = free_lists[index].list.begin();
   if (begin_node == nullptr)
     return BlockRef();
 
   FreeList::Node *cur = begin_node;
+  LIBC_LOOP_NOUNROLL_32
   do {
     if (cur->size() >= size) {
       free_lists[index].list.remove(cur, secrets);
@@ -310,7 +320,7 @@ LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::remove_first_fit_in_list(
 
 template <typename CONFIG>
 LIBC_INLINE BlockRef TLSFFreeStoreImpl<CONFIG>::find_and_remove_fit(
-    size_t size, const FreeListSecrets& secrets) {
+    size_t size, const FreeListSecrets &secrets) {
   size_t bit_index = size_to_bit_index(size);
 
   if (LIBC_UNLIKELY(bit_index >= TOTAL_BITS - 1)) {
@@ -354,7 +364,10 @@ using TLSFFreeStore = TLSFFreeStoreImpl<TLSFFreeStoreConfig<
 #endif
 
 using FreeStore =
-    TLSFFreeStore<BlockRef::MIN_ALIGN, 3, 2, (sizeof(uintptr_t) == 8 ? 3 : 6),
+    TLSFFreeStore<BlockRef::MIN_ALIGN,
+                  (sizeof(uintptr_t) == 8 ? 3 : 2), // STEP_SIZE_BITS
+                  2,                                // NUM_STEP_BITS
+                  (sizeof(uintptr_t) == 8 ? 3 : 2), // NUM_TABLE_ENTRIES
                   LIBC_COPT_USE_TRIE_FOR_OVERFLOW_BIN>;
 
 } // namespace LIBC_NAMESPACE_DECL

>From 14846a71558f638ee012c202edf0c73f85c2118e 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 06/15] [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 f7db63ef8bf74..046969f98620e 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -71,6 +71,7 @@
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
 #include "llvm/Analysis/ScalarEvolutionDivision.h"
 #include "llvm/Analysis/ScopedNoAliasAA.h"
+#include "llvm/Analysis/HeapProvenanceAnalysis.h"
 #include "llvm/Analysis/StackLifetime.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/StructuralHash.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 8f70c7cefa408..2962394549056 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -37,6 +37,7 @@ MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
 MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
 MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
 MODULE_ANALYSIS("runtime-libcall-info", RuntimeLibraryAnalysis())
+MODULE_ANALYSIS("heap-provenance", HeapProvenanceAnalysis())
 MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
 MODULE_ANALYSIS("verify", VerifierAnalysis())
 
@@ -144,6 +145,8 @@ MODULE_PASS("print-lcg-dot", LazyCallGraphDOTPrinterPass(errs()))
 MODULE_PASS("print-must-be-executed-contexts",
             MustBeExecutedContextPrinterPass(errs()))
 MODULE_PASS("print-profile-summary", ProfileSummaryPrinterPass(errs()))
+MODULE_PASS("print-heap-provenance", HeapProvenancePrinterPass(errs()))
+MODULE_PASS("print<heap-provenance>", HeapProvenancePrinterPass(errs()))
 MODULE_PASS("print-stack-safety", StackSafetyGlobalPrinterPass(errs()))
 MODULE_PASS("print<dxil-metadata>", DXILMetadataAnalysisPrinterPass(errs()))
 MODULE_PASS("print<dxil-resources>", DXILResourcePrinterPass(errs()))

>From ecf9ddd925205ed720fb22406e8f7185635935dd 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 07/15] [Analysis] add a fallback fortification method for
 baremetal heap

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

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

>From c05453831aa26c19ef402307062109221727cc1b 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 08/15] [Analysis] Refactor HeapProvenanceAnalysis to use LLVM
 SparseSolver

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

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

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

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

>From 9ac6d37697a530983e8d011766e0d71edc5202fb 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 09/15] [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 d2c1d101e0d3287c6dd1c93e1ce2d1c3e228fdfa 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 10/15] [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 ae1003b90b50928b66ec908fe2461b84b35edde9 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 11/15] [Clang][Sanitizer] Split bounds checking pipeline to
 run HPA after mem2reg

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

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

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

>From 385c587b532a96eb0f9a9128c98aae54f8a333bf 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 12/15] [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 062e76751930e26e2a65f1dab0773542584eb2dc 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 13/15] [Analysis][HeapProvenance] Fix backward provenance
 propagation and dominance insertion

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

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

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

>From 166e710fe6e61913d04d1359bfc4478703d417e4 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 14/15] [Clang][Sanitizer][HeapProvenance] Group bounds check
 pipeline and separate HPA passes

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

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

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

>From 9368c2aec84f2ad9dfc3028803950af1d78bc1a1 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 15/15] [Clang][Sanitizer][HeapProvenance] Check interim
 pointer bounds and verify use list dominance

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

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

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



More information about the cfe-commits mailing list