Introduce new functions to return pageflags/residency bitmaps in terms of TCMalloc pages instead of native pages.

tcmalloc/internal/pageflags.h

@@ -70,6 +70,17 @@ class PageFlagsBase {
   virtual std::optional<PageStats> Get(const void* addr, size_t size) = 0;
   virtual absl::StatusCode GetSinglePageBitmaps(const void* addr,
                                                 ResidencyBitmap& stale) = 0;
+  template <size_t N>
+  absl::StatusCode GetSinglePageBitmapsWrapper(const void* addr,
+                                               Bitmap<N>& stale) {
+    ResidencyBitmap temp;
+    auto res = GetSinglePageBitmaps(addr, temp);
+    if (res != absl::StatusCode::kOk) {
+      return res;
+    }
+    stale = Contract<N>(temp, kHugePageSize / GetPageSize());
+    return res;
+  }
 };
 
 // PageFlags offers a look at kernel page flags to identify pieces of memory as

tcmalloc/internal/pageflags_test.cc

@@ -33,6 +33,7 @@
 #include <optional>
 #include <ostream>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -86,6 +87,12 @@ class PageFlagsFriend {
     return r_.GetSinglePageBitmaps(addr, stale);
   }
 
+  template <size_t N>
+  absl::StatusCode GetSinglePageBitmapsWrapper(const void* addr,
+                                               Bitmap<N>& stale) {
+    return r_.GetSinglePageBitmapsWrapper(addr, stale);
+  }
+
   void SetCachedScanSeconds(
       decltype(PageFlags::cached_scan_seconds_) scan_seconds) {
     r_.cached_scan_seconds_ = scan_seconds;
@@ -663,6 +670,59 @@ TEST(PageFlagsTest, GetSinglePageBitmapsSuccess) {
   EXPECT_EQ(stale.CountBits(), kMaxResidencyBits / 2);
 }
 
+template <typename T>
+class GetSinglePageBitmapsWrapperTest : public ::testing::Test {};
+
+using WrapperTestSizes =
+    ::testing::Types<std::integral_constant<size_t, kMaxResidencyBits>,
+                     std::integral_constant<size_t, kMaxResidencyBits / 2>,
+                     std::integral_constant<size_t, kMaxResidencyBits * 2>>;
+
+TYPED_TEST_SUITE(GetSinglePageBitmapsWrapperTest, WrapperTestSizes);
+
+TYPED_TEST(GetSinglePageBitmapsWrapperTest, AlternatingStale) {
+  constexpr size_t N = TypeParam::value;
+
+  std::string fake_pageflags =
+      absl::StrCat(testing::TempDir(), "/fake_pageflags_wrapper_", N);
+  const size_t kHardwarePageSize = getpagesize();
+  const size_t kNativePagesInHugePage = kHugePageSize / kHardwarePageSize;
+  std::vector<uint64_t> data(kNativePagesInHugePage, 0);
+
+  data[0] |= kPageStale;
+  data[1] |= kPageStale;
+  data[2] |= kPageStale;
+
+  std::string content(reinterpret_cast<const char*>(data.data()),
+                      data.size() * sizeof(uint64_t));
+  SetContents(fake_pageflags, content);
+
+  PageFlagsFriend s(fake_pageflags);
+  Bitmap<N> stale;
+  stale.Clear();
+
+  EXPECT_EQ(s.GetSinglePageBitmapsWrapper<N>(nullptr, stale),
+            absl::StatusCode::kOk);
+
+  if constexpr (N == kMaxResidencyBits) {
+    EXPECT_EQ(stale.CountBits(), 3);
+    EXPECT_TRUE(stale.GetBit(0));
+    EXPECT_TRUE(stale.GetBit(1));
+    EXPECT_TRUE(stale.GetBit(2));
+    EXPECT_FALSE(stale.GetBit(3));
+  } else if constexpr (N == kMaxResidencyBits / 2) {
+    EXPECT_EQ(stale.CountBits(), 1);
+    EXPECT_TRUE(stale.GetBit(0));
+    EXPECT_FALSE(stale.GetBit(1));
+  } else if constexpr (N == kMaxResidencyBits * 2) {
+    EXPECT_EQ(stale.CountBits(), 6);
+    for (size_t i = 0; i < 6; ++i) {
+      EXPECT_TRUE(stale.GetBit(i)) << i;
+    }
+    EXPECT_FALSE(stale.GetBit(6));
+  }
+}
+
 }  // namespace
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc

tcmalloc/internal/range_tracker.h

@@ -118,6 +118,31 @@ class Bitmap {
   ssize_t FindValueBackwards(size_t index) const;
 };
 
+template <size_t M, size_t N>
+Bitmap<M> Contract(const Bitmap<N>& src, size_t src_len) {
+  TC_ASSERT_LE(src_len, N);
+  TC_ASSERT(src_len % M == 0 || M % src_len == 0);
+
+  Bitmap<M> res;
+  const size_t src_per_dst = std::max<size_t>(1, src_len / M);
+  const size_t dst_per_src = std::max<size_t>(1, M / src_len);
+
+  size_t dst_idx = 0;
+  for (size_t src_idx = 0; src_idx < src_len; src_idx += src_per_dst) {
+    size_t count = 0;
+    for (size_t j = 0; j < src_per_dst; ++j) {
+      if (src.GetBit(src_idx + j)) {
+        count++;
+      }
+    }
+    if (count == src_per_dst) {
+      res.SetRange(dst_idx, dst_per_src);
+    }
+    dst_idx += dst_per_src;
+  }
+  return res;
+}
+
 constexpr size_t kMaxResidencyBits = 512;
 using ResidencyBitmap = Bitmap<kMaxResidencyBits>;
 // Tracks allocations in a range of items of fixed size.  Supports

tcmalloc/internal/range_tracker_test.cc

@@ -347,6 +347,74 @@ TEST_F(BitmapTest, PopBatch) {
   EXPECT_TRUE(map.IsZero());
 }
 
+#ifndef NDEBUG
+TEST_F(BitmapTest, ContractAssertionFailures) {
+  Bitmap<64> map1;
+  // 1. src_len (64) doesn't divide M (10) and vice versa
+  EXPECT_DEATH(Contract<10>(map1, 64), "");
+  // 2. src_len > N (64)
+  EXPECT_DEATH(Contract<256>(map1, 128), "");
+}
+#endif
+
+TEST_F(BitmapTest, Contract) {
+  // Test case 1: N == M (equal)
+  {
+    Bitmap<64> map1;
+    map1.SetBit(3);
+    map1.SetBit(17);
+    map1.SetBit(63);
+
+    auto map2 = Contract<64>(map1, 64);
+    for (size_t i = 0; i < 64; ++i) {
+      EXPECT_EQ(map1.GetBit(i), map2.GetBit(i));
+    }
+  }
+
+  // Test case 2: N > M (Contracting)
+  {
+    Bitmap<64> map1;
+    // We group by 2 bits.
+    // Group 0: bits 0, 1 -> both must be set
+    map1.SetBit(0);
+    map1.SetBit(1);
+
+    // Group 1: bits 2, 3 -> only one set -> should not be set
+    map1.SetBit(2);
+
+    // Group 7: bits 14, 15 -> both set
+    map1.SetBit(14);
+    map1.SetBit(15);
+
+    auto map2 = Contract<8>(map1, 16);
+    EXPECT_TRUE(map2.GetBit(0));
+    EXPECT_FALSE(map2.GetBit(1));
+    EXPECT_FALSE(map2.GetBit(2));
+    EXPECT_FALSE(map2.GetBit(3));
+    EXPECT_FALSE(map2.GetBit(4));
+    EXPECT_FALSE(map2.GetBit(5));
+    EXPECT_FALSE(map2.GetBit(6));
+    EXPECT_TRUE(map2.GetBit(7));
+  }
+
+  // Test case 3: N < M (expanding)
+  {
+    Bitmap<64> map1;
+    // We expand by 2 bits.
+    map1.SetBit(0);
+    map1.SetBit(7);
+
+    auto map2 = Contract<16>(map1, 8);
+    EXPECT_TRUE(map2.GetBit(0));
+    EXPECT_TRUE(map2.GetBit(1));
+    EXPECT_FALSE(map2.GetBit(2));
+    EXPECT_FALSE(map2.GetBit(3));
+
+    EXPECT_TRUE(map2.GetBit(14));
+    EXPECT_TRUE(map2.GetBit(15));
+  }
+}
+
 class RangeTrackerTest : public ::testing::Test {
  protected:
   std::vector<std::pair<size_t, size_t>> FreeRanges() {

tcmalloc/internal/residency.h

@@ -47,18 +47,31 @@ class Residency {
   virtual inline size_t GetNativePagesInHugePage() const = 0;
 
   // Struct is ordered with bitmaps first to optimize cacheline usage.
-  struct SinglePageBitmaps {
-    Bitmap<kMaxResidencyBits> unbacked;
-    Bitmap<kMaxResidencyBits> swapped;
+  template <size_t N>
+  struct SinglePageBitmapsWrapper {
+    Bitmap<N> unbacked;
+    Bitmap<N> swapped;
     absl::StatusCode status;
   };
 
+  using SinglePageBitmaps = SinglePageBitmapsWrapper<kMaxResidencyBits>;
+
   // Using a hugepage-aligned address, parse through /proc/self/pagemap
   // to output two bitmaps - one for pages that are unbacked and one for pages
   // that are swapped. Hugepage-sized regions are assumed to be 2MiB in size. A
   // SinglePageBitmaps struct is returned with the status, the page_unbacked
   // bitmap, and the page_swapped bitmap.
   virtual SinglePageBitmaps GetUnbackedAndSwappedBitmaps(const void* addr) = 0;
+
+  template <size_t N>
+  SinglePageBitmapsWrapper<N> GetUnbackedAndSwappedBitmapsWrapper(
+      const void* addr) {
+    SinglePageBitmaps res = GetUnbackedAndSwappedBitmaps(addr);
+    return SinglePageBitmapsWrapper<N>{
+        .unbacked = Contract<N>(res.unbacked, GetNativePagesInHugePage()),
+        .swapped = Contract<N>(res.swapped, GetNativePagesInHugePage()),
+        .status = res.status};
+  }
 };
 
 // Residency offers information about memory residency: whether or not specific

tcmalloc/internal/residency_test.cc

@@ -25,6 +25,7 @@
 #include <cstdint>
 #include <optional>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -369,6 +370,80 @@ TEST(PageMapIntegrationTest, WorksOnActualData) {
   EXPECT_TRUE(res.swapped.IsZero());
 }
 
+template <typename T>
+class PageMapIntegrationTest : public ::testing::Test {};
+
+using WorksOnActualDataTestSizes =
+    ::testing::Types<std::integral_constant<size_t, kMaxResidencyBits>,
+                     std::integral_constant<size_t, kMaxResidencyBits / 2>,
+                     std::integral_constant<size_t, kMaxResidencyBits * 2>>;
+
+TYPED_TEST_SUITE(PageMapIntegrationTest, WorksOnActualDataTestSizes);
+
+TYPED_TEST(PageMapIntegrationTest, WorksOnActualDataWrapper) {
+  constexpr size_t N = TypeParam::value;
+  std::optional<AllocationGuard> g;
+  void* addr = mmap(nullptr, 4 << 20, PROT_WRITE,
+                    MAP_ANONYMOUS | MAP_POPULATE | MAP_PRIVATE, -1, 0);
+  ASSERT_NE(addr, MAP_FAILED) << errno;
+  auto position = reinterpret_cast<uintptr_t>(addr);
+  if ((position & (kHugePageSize - 1)) != 0) {
+    position |= kHugePageSize - 1;
+    position++;
+    addr = reinterpret_cast<void*>(position);
+  }
+  g.emplace();
+  ResidencyPageMap r;
+  auto res = r.GetUnbackedAndSwappedBitmapsWrapper<N>(addr);
+  g.reset();
+  ASSERT_EQ(res.status, absl::StatusCode::kOk);
+  EXPECT_TRUE(res.unbacked.IsZero());
+  EXPECT_TRUE(res.swapped.IsZero());
+  // Unmap consecutive pages starting at page 2, and page 16
+  ASSERT_EQ(munmap(reinterpret_cast<uint8_t*>(addr) + 2 * 4096, 2 * 4096), 0)
+      << errno;
+  ASSERT_EQ(munmap(reinterpret_cast<uint8_t*>(addr) + 16 * 4096, 2 * 4096), 0)
+      << errno;
+
+  g.emplace();
+  res = r.GetUnbackedAndSwappedBitmapsWrapper<N>(addr);
+  g.reset();
+  ASSERT_EQ(res.status, absl::StatusCode::kOk);
+  EXPECT_FALSE(res.unbacked.IsZero());
+
+  if constexpr (N == kMaxResidencyBits) {
+    ASSERT_TRUE(res.unbacked.GetBit(2));
+    ASSERT_TRUE(res.unbacked.GetBit(3));
+    ASSERT_TRUE(res.unbacked.GetBit(16));
+    ASSERT_TRUE(res.unbacked.GetBit(17));
+    res.unbacked.ClearBit(2);
+    res.unbacked.ClearBit(3);
+    res.unbacked.ClearBit(16);
+    res.unbacked.ClearBit(17);
+  } else if constexpr (N == kMaxResidencyBits / 2) {
+    ASSERT_TRUE(res.unbacked.GetBit(1));
+    ASSERT_TRUE(res.unbacked.GetBit(8));
+    res.unbacked.ClearBit(1);
+    res.unbacked.ClearBit(8);
+  } else if constexpr (N == kMaxResidencyBits * 2) {
+    ASSERT_TRUE(res.unbacked.GetBit(4));
+    ASSERT_TRUE(res.unbacked.GetBit(5));
+    ASSERT_TRUE(res.unbacked.GetBit(32));
+    ASSERT_TRUE(res.unbacked.GetBit(33));
+    res.unbacked.ClearBit(4);
+    res.unbacked.ClearBit(5);
+    res.unbacked.ClearBit(6);
+    res.unbacked.ClearBit(7);
+    res.unbacked.ClearBit(32);
+    res.unbacked.ClearBit(33);
+    res.unbacked.ClearBit(34);
+    res.unbacked.ClearBit(35);
+  }
+
+  EXPECT_TRUE(res.unbacked.IsZero());
+  EXPECT_TRUE(res.swapped.IsZero());
+}
+
 }  // namespace
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc