Introduce and plumb ProvenanceCaptureRange

This backs out the sleight of hand from the SmallBuddyRange

Co-authored-by: Matthew Parkinson <[email protected]>

src/snmalloc/backend/meta_protected_range.h

@@ -101,6 +101,7 @@ namespace snmalloc
         LocalCacheSizeBits - SubRangeRatioBits,
         bits::BITS - 1,
         Pagemap>,
+      ProvenanceCaptureRange<Pagemap>,
       SmallBuddyRange<>>;
 
     ObjectRange object_range;
@@ -123,7 +124,10 @@ namespace snmalloc
     // Create global range that can service small meta-data requests.
     // Don't want to add the SmallBuddyRange to the CentralMetaRange as that
     // would require committing memory inside the main global lock.
-    using GlobalMetaRange =
-      Pipe<CentralMetaRange, SmallBuddyRange<>, GlobalRange<>>;
+    using GlobalMetaRange = Pipe<
+      CentralMetaRange,
+      ProvenanceCaptureRange<Pagemap>,
+      SmallBuddyRange<>,
+      GlobalRange<>>;
   };
-} // namespace snmalloc
+} // namespace snmalloc

src/snmalloc/backend/standard_range.h

@@ -56,7 +56,10 @@ namespace snmalloc
         page_size_bits>>;
 
   private:
-    using ObjectRange = Pipe<LargeObjectRange, SmallBuddyRange<>>;
+    using ObjectRange = Pipe<
+      LargeObjectRange,
+      ProvenanceCaptureRange<Pagemap>,
+      SmallBuddyRange<>>;
 
     ObjectRange object_range;
 

src/snmalloc/backend_helpers/backend_helpers.h

@@ -10,6 +10,7 @@
 #include "pagemap.h"
 #include "pagemapregisterrange.h"
 #include "palrange.h"
+#include "provenance_capture_range.h"
 #include "range_helpers.h"
 #include "smallbuddyrange.h"
 #include "statsrange.h"

src/snmalloc/backend_helpers/provenance_capture_range.h

@@ -0,0 +1,76 @@
+#pragma once
+
+#include "empty_range.h"
+
+namespace snmalloc
+{
+  /**
+   * This Range converts between Arena-bounded chunks parent-wards and
+   * Chunk-bounded chunks child-wards.  In the alloc path, it stores the
+   * high-authority pointer in the Pagemap, whence it retrieves it on the
+   * dealloc path.
+   */
+  template<
+    SNMALLOC_CONCEPT(IsWritablePagemap) Pagemap,
+    bool strict_provenance = false,
+    typename ParentRange = EmptyRange<>>
+  class ProvenanceCaptureRange : public ContainsParent<ParentRange>
+  {
+    using ContainsParent<ParentRange>::parent;
+
+  public:
+    /**
+     * We use a nested Apply type to enable a Pipe operation.
+     */
+    template<typename ParentRange2>
+    using Apply =
+      ProvenanceCaptureRange<Pagemap, strict_provenance, ParentRange2>;
+
+    constexpr ProvenanceCaptureRange() = default;
+
+    static constexpr bool Aligned = ParentRange::Aligned;
+
+    static constexpr bool ConcurrencySafe = ParentRange::ConcurrencySafe;
+
+    using ChunkBounds = capptr::bounds::Chunk;
+    static_assert(
+      std::is_same_v<typename ParentRange::ChunkBounds, capptr::bounds::Arena>);
+
+    capptr::Chunk<void> alloc_range(size_t size)
+    {
+      SNMALLOC_ASSERT((size & (MIN_CHUNK_SIZE - 1)) == 0);
+
+      auto arena = parent.alloc_range(size);
+
+      if (arena == nullptr)
+        return nullptr;
+
+      if constexpr (strict_provenance)
+      {
+        Pagemap::template get_metaentry_mut<false>(address_cast(arena))
+          .get_backend_word(Pagemap::Entry::Word::One) = arena.unsafe_uintptr();
+      }
+
+      return Aal::capptr_bound<void, capptr::bounds::Chunk>(arena, size);
+    }
+
+    void dealloc_range(capptr::Chunk<void> chunk, size_t size)
+    {
+      capptr::Arena<void> arena;
+      if constexpr (strict_provenance)
+      {
+        arena = capptr::Arena<void>(reinterpret_cast<void*>(
+          Pagemap::template get_metaentry_mut<false>(address_cast(chunk))
+            .get_backend_word(Pagemap::Entry::Word::One)
+            .get()));
+        SNMALLOC_ASSERT(address_cast(arena) == address_cast(chunk));
+      }
+      else
+      {
+        arena = capptr::Arena<void>(chunk.unsafe_ptr());
+      }
+
+      parent.dealloc_range(arena, size);
+    }
+  };
+} // namespace snmalloc

src/snmalloc/backend_helpers/smallbuddyrange.h

@@ -143,7 +143,7 @@ namespace snmalloc
     }
   };
 
-  template<typename ParentRange = EmptyRange<>>
+  template<typename ParentRange = EmptyRange<capptr::bounds::Chunk>>
   class SmallBuddyRange : public ContainsParent<ParentRange>
   {
     using ContainsParent<ParentRange>::parent;
@@ -166,26 +166,21 @@ namespace snmalloc
               .template as_reinterpret<void>();
           if (overflow != nullptr)
           {
-            // XXX: Sleight of hand; on StrictProvenance, we must reload this
-            auto over_arena = capptr::Arena<void>(overflow.unsafe_ptr());
-            parent.dealloc_range(over_arena, bits::one_at_bit(MIN_CHUNK_BITS));
+            parent.dealloc_range(overflow, bits::one_at_bit(MIN_CHUNK_BITS));
           }
         });
     }
 
     capptr::Chunk<void> refill(size_t size)
     {
-      capptr::Arena<void> refill = parent.alloc_range(MIN_CHUNK_SIZE);
+      capptr::Chunk<void> refill = parent.alloc_range(MIN_CHUNK_SIZE);
 
       if (refill == nullptr)
         return nullptr;
 
-      // XXX: Sleight of hand; on StrictProvenance, we must stash the arena
-      auto refill_chunk = capptr::Chunk<void>(refill.unsafe_ptr());
+      add_range(pointer_offset(refill, size), MIN_CHUNK_SIZE - size);
 
-      add_range(pointer_offset(refill_chunk, size), MIN_CHUNK_SIZE - size);
-
-      return refill_chunk;
+      return refill;
     }
 
   public:
@@ -200,9 +195,8 @@ namespace snmalloc
 
     static constexpr bool ConcurrencySafe = false;
 
-    using ChunkBounds = capptr::bounds::Chunk;
-    static_assert(
-      std::is_same_v<typename ParentRange::ChunkBounds, capptr::bounds::Arena>);
+    using ChunkBounds = typename ParentRange::ChunkBounds;
+    static_assert(std::is_same_v<ChunkBounds, capptr::bounds::Chunk>);
 
     constexpr SmallBuddyRange() = default;