SP: introduce CapPtr<> wrapper type & bounds taxonomy

This commit just fills out ds/ptrwrap.h with the new types and adds utility
methods to ds/address.h.

src/ds/address.h

@@ -1,6 +1,7 @@
 #pragma once
 #include "../pal/pal_consts.h"
 #include "bits.h"
+#include "ptrwrap.h"
 
 #include <cstdint>
 
@@ -24,6 +25,13 @@ namespace snmalloc
     return reinterpret_cast<U*>(reinterpret_cast<char*>(base) + diff);
   }
 
+  template<enum capptr_bounds bounds, typename T>
+  inline CapPtr<void, bounds>
+  pointer_offset(CapPtr<T, bounds> base, size_t diff)
+  {
+    return CapPtr<void, bounds>(pointer_offset(base.unsafe_capptr, diff));
+  }
+
   /**
    * Perform pointer arithmetic and return the adjusted pointer.
    */
@@ -33,6 +41,14 @@ namespace snmalloc
     return reinterpret_cast<U*>(reinterpret_cast<char*>(base) + diff);
   }
 
+  template<enum capptr_bounds bounds, typename T>
+  inline CapPtr<void, bounds>
+  pointer_offset_signed(CapPtr<T, bounds> base, ptrdiff_t diff)
+  {
+    return CapPtr<void, bounds>(
+      pointer_offset_signed(base.unsafe_capptr, diff));
+  }
+
   /**
    * Cast from a pointer type to an address.
    */
@@ -42,16 +58,36 @@ namespace snmalloc
     return reinterpret_cast<address_t>(ptr);
   }
 
+  /*
+   * Provide address_cast methods for the provenance-hinting pointer wrapper
+   * types as well.  While we'd prefer that these be methods on the wrapper
+   * type, they have to be defined later, because the AAL both define address_t,
+   * as per above, and uses the wrapper types in its own definition, e.g., of
+   * capptr_bound.
+   */
+
+  template<typename T, enum capptr_bounds bounds>
+  inline address_t address_cast(CapPtr<T, bounds> a)
+  {
+    return address_cast(a.unsafe_capptr);
+  }
+
   /**
    * Test if a pointer is aligned to a given size, which must be a power of
    * two.
    */
   template<size_t alignment>
-  static inline bool is_aligned_block(void* p, size_t size)
+  static inline bool is_aligned_block(address_t p, size_t size)
   {
     static_assert(bits::is_pow2(alignment));
 
-    return ((address_cast(p) | size) & (alignment - 1)) == 0;
+    return ((p | size) & (alignment - 1)) == 0;
+  }
+
+  template<size_t alignment>
+  static inline bool is_aligned_block(void* p, size_t size)
+  {
+    return is_aligned_block<alignment>(address_cast(p), size);
   }
 
   /**
@@ -76,6 +112,12 @@ namespace snmalloc
     }
   }
 
+  template<size_t alignment, typename T, capptr_bounds bounds>
+  inline CapPtr<T, bounds> pointer_align_down(CapPtr<void, bounds> p)
+  {
+    return CapPtr<T, bounds>(pointer_align_down<alignment, T>(p.unsafe_capptr));
+  }
+
   template<size_t alignment>
   inline address_t address_align_down(address_t p)
   {
@@ -104,6 +146,12 @@ namespace snmalloc
     }
   }
 
+  template<size_t alignment, typename T = void, enum capptr_bounds bounds>
+  inline CapPtr<T, bounds> pointer_align_up(CapPtr<void, bounds> p)
+  {
+    return CapPtr<T, bounds>(pointer_align_up<alignment, T>(p.unsafe_capptr));
+  }
+
   template<size_t alignment>
   inline address_t address_align_up(address_t p)
   {
@@ -144,6 +192,13 @@ namespace snmalloc
 #endif
   }
 
+  template<typename T = void, enum capptr_bounds bounds>
+  inline CapPtr<T, bounds>
+  pointer_align_up(CapPtr<void, bounds> p, size_t alignment)
+  {
+    return CapPtr<T, bounds>(pointer_align_up<T>(p.unsafe_capptr, alignment));
+  }
+
   /**
    * Compute the difference in pointers in units of char.  base is
    * expected to point to the base of some (sub)allocation into which cursor
@@ -156,6 +211,16 @@ namespace snmalloc
       static_cast<char*>(cursor) - static_cast<char*>(base));
   }
 
+  template<
+    typename T = void,
+    typename U = void,
+    enum capptr_bounds Tbounds,
+    enum capptr_bounds Ubounds>
+  inline size_t pointer_diff(CapPtr<T, Tbounds> base, CapPtr<U, Ubounds> cursor)
+  {
+    return pointer_diff(base.unsafe_capptr, cursor.unsafe_capptr);
+  }
+
   /**
    * Compute the difference in pointers in units of char. This can be used
    * across allocations.
@@ -166,4 +231,15 @@ namespace snmalloc
       static_cast<char*>(cursor) - static_cast<char*>(base));
   }
 
+  template<
+    typename T = void,
+    typename U = void,
+    enum capptr_bounds Tbounds,
+    enum capptr_bounds Ubounds>
+  inline ptrdiff_t
+  pointer_diff_signed(CapPtr<T, Tbounds> base, CapPtr<U, Ubounds> cursor)
+  {
+    return pointer_diff_signed(base.unsafe_capptr, cursor.unsafe_capptr);
+  }
+
 } // namespace snmalloc

src/ds/ptrwrap.h

@@ -15,4 +15,272 @@ namespace snmalloc
 
   template<typename T>
   using AtomicPointer = std::atomic<T*>;
+
+  /**
+   * Summaries of StrictProvenance metadata.  We abstract away the particular
+   * size and any offset into the bounds.
+   *
+   * CBArena is as powerful as our pointers get: they're results from mmap(),
+   * and so confer as much authority as the kernel has given us.
+   *
+   * CBChunk is restricted to either a single chunk (SUPERSLAB_SIZE) or perhaps
+   * to several if we've requesed a large allocation (see capptr_chunk_is_alloc
+   * and its uses).
+   *
+   * CBChunkD is curious: we often use CBArena-bounded pointers to derive
+   * pointers to Allocslab metadata, and on most fast paths these pointers end
+   * up being ephemeral.  As such, on NDEBUG builds, we elide the capptr_bounds
+   * that would bound these to chunks and instead just unsafely inherit the
+   * CBArena bounds.  The use of CBChunkD thus helps to ensure that we
+   * eventually do invoke capptr_bounds when these pointers end up being longer
+   * lived!
+   *
+   * *E forms are "exported" and have had platform constraints applied.  That
+   * means, for example, on CheriBSD, that they have had their VMMAP permission
+   * stripped.
+   *
+   * Yes, I wish the start-of-comment characters were aligned below as well.
+   * I blame clang format.
+   */
+  enum capptr_bounds
+  {
+    /*                Spatial  Notes                                      */
+    CBArena, /*        Arena                                              */
+    CBChunkD, /*       Arena   Chunk-bounded in debug; internal use only! */
+    CBChunk, /*        Chunk                                              */
+    CBChunkE, /*       Chunk   (+ platform constraints)                   */
+    CBAlloc, /*        Alloc                                              */
+    CBAllocE /*        Alloc   (+ platform constraints)                   */
+  };
+
+  /**
+   * Compute the "exported" variant of a capptr_bounds annotation.  This is
+   * used by the PAL's capptr_export function to compute its return value's
+   * annotation.
+   */
+  template<capptr_bounds B>
+  constexpr capptr_bounds capptr_export_type()
+  {
+    static_assert(
+      (B == CBChunk) || (B == CBAlloc), "capptr_export_type of bad type");
+
+    switch (B)
+    {
+      case CBChunk:
+        return CBChunkE;
+      case CBAlloc:
+        return CBAllocE;
+    }
+  }
+
+  template<capptr_bounds BI, capptr_bounds BO>
+  constexpr bool capptr_is_bounds_refinement()
+  {
+    switch (BI)
+    {
+      case CBAllocE:
+        return BO == CBAllocE;
+      case CBAlloc:
+        return BO == CBAlloc;
+      case CBChunkE:
+        return BO == CBAllocE || BO == CBChunkE;
+      case CBChunk:
+        return BO == CBAlloc || BO == CBChunk || BO == CBChunkD;
+      case CBChunkD:
+        return BO == CBAlloc || BO == CBChunk || BO == CBChunkD;
+      case CBArena:
+        return BO == CBAlloc || BO == CBChunk || BO == CBChunkD ||
+          BO == CBArena;
+    }
+  }
+
+  /**
+   * A pointer annotated with a "phantom type parameter" carrying a static
+   * summary of its StrictProvenance metadata.
+   */
+  template<typename T, capptr_bounds bounds>
+  struct CapPtr
+  {
+    T* unsafe_capptr;
+
+    /**
+     * nullptr is implicitly constructable at any bounds type
+     */
+    CapPtr(const std::nullptr_t n) : unsafe_capptr(n) {}
+
+    CapPtr() : CapPtr(nullptr) {}
+
+    /**
+     * all other constructions must be explicit
+     */
+    explicit CapPtr(T* p) : unsafe_capptr(p) {}
+
+    /**
+     * Allow static_cast<>-s that preserve bounds but vary the target type.
+     */
+    template<typename U>
+    SNMALLOC_FAST_PATH CapPtr<U, bounds> as_static()
+    {
+      return CapPtr<U, bounds>(static_cast<U*>(this->unsafe_capptr));
+    }
+
+    SNMALLOC_FAST_PATH CapPtr<void, bounds> as_void()
+    {
+      return this->as_static<void>();
+    }
+
+    /**
+     * A more aggressive bounds-preserving cast, using reinterpret_cast
+     */
+    template<typename U>
+    SNMALLOC_FAST_PATH CapPtr<U, bounds> as_reinterpret()
+    {
+      return CapPtr<U, bounds>(reinterpret_cast<U*>(this->unsafe_capptr));
+    }
+
+    SNMALLOC_FAST_PATH bool operator==(const CapPtr& rhs) const
+    {
+      return this->unsafe_capptr == rhs.unsafe_capptr;
+    }
+
+    SNMALLOC_FAST_PATH bool operator!=(const CapPtr& rhs) const
+    {
+      return this->unsafe_capptr != rhs.unsafe_capptr;
+    }
+
+    SNMALLOC_FAST_PATH bool operator<(const CapPtr& rhs) const
+    {
+      return this->unsafe_capptr < rhs.unsafe_capptr;
+    }
+
+    SNMALLOC_FAST_PATH T* operator->() const
+    {
+      /*
+       * CBAllocE bounds are associated with objects coming from or going to the
+       * client; we should be doing nothing with them.
+       */
+      static_assert(bounds != CBAllocE);
+      return this->unsafe_capptr;
+    }
+  };
+
+  static_assert(sizeof(CapPtr<void, CBArena>) == sizeof(void*));
+  static_assert(alignof(CapPtr<void, CBArena>) == alignof(void*));
+
+  template<typename T>
+  using CapPtrCBArena = CapPtr<T, CBArena>;
+
+  template<typename T>
+  using CapPtrCBChunk = CapPtr<T, CBChunk>;
+
+  template<typename T>
+  using CapPtrCBChunkE = CapPtr<T, CBChunkE>;
+
+  template<typename T>
+  using CapPtrCBAlloc = CapPtr<T, CBAlloc>;
+
+  /**
+   * Sometimes (with large allocations) we really mean the entire chunk (or even
+   * several chunks) to be the allocation.
+   */
+  template<typename T>
+  SNMALLOC_FAST_PATH CapPtr<T, CBAllocE>
+  capptr_chunk_is_alloc(CapPtr<T, CBChunkE> p)
+  {
+    return CapPtr<T, CBAlloc>(p.unsafe_capptr);
+  }
+
+  /**
+   * With all the bounds and constraints in place, it's safe to extract a void
+   * pointer (to reveal to the client).
+   */
+  SNMALLOC_FAST_PATH void* capptr_reveal(CapPtr<void, CBAllocE> p)
+  {
+    return p.unsafe_capptr;
+  }
+
+  /**
+   *
+   * Wrap a std::atomic<T*> with bounds annotation and speak in terms of
+   * bounds-annotated pointers at the interface.
+   *
+   * Note the membranous sleight of hand being pulled here: this class puts
+   * annotations around an un-annotated std::atomic<T*>, to appease C++, yet
+   * will expose or consume only CapPtr<T> with the same bounds annotation.
+   */
+  template<typename T, capptr_bounds bounds>
+  struct AtomicCapPtr
+  {
+    std::atomic<T*> unsafe_capptr;
+
+    /**
+     * nullptr is constructable at any bounds type
+     */
+    AtomicCapPtr(const std::nullptr_t n) : unsafe_capptr(n) {}
+
+    /**
+     * Interconversion with CapPtr
+     */
+    AtomicCapPtr(CapPtr<T, bounds> p) : unsafe_capptr(p.unsafe_capptr) {}
+
+    operator CapPtr<T, bounds>() const noexcept
+    {
+      return CapPtr<T, bounds>(this->unsafe_capptr);
+    }
+
+    // Our copy-assignment operator follows std::atomic and returns a copy of
+    // the RHS.  Clang finds this surprising; we suppress the warning.
+    // NOLINTNEXTLINE(misc-unconventional-assign-operator)
+    CapPtr<T, bounds> operator=(CapPtr<T, bounds> p) noexcept
+    {
+      this->store(p);
+      return p;
+    }
+
+    SNMALLOC_FAST_PATH CapPtr<T, bounds>
+    load(std::memory_order order = std::memory_order_seq_cst) noexcept
+    {
+      return CapPtr<T, bounds>(this->unsafe_capptr.load(order));
+    }
+
+    SNMALLOC_FAST_PATH void store(
+      CapPtr<T, bounds> desired,
+      std::memory_order order = std::memory_order_seq_cst) noexcept
+    {
+      this->unsafe_capptr.store(desired.unsafe_capptr, order);
+    }
+
+    SNMALLOC_FAST_PATH CapPtr<T, bounds> exchange(
+      CapPtr<T, bounds> desired,
+      std::memory_order order = std::memory_order_seq_cst) noexcept
+    {
+      return CapPtr<T, bounds>(
+        this->unsafe_capptr.exchange(desired.unsafe_capptr, order));
+    }
+
+    SNMALLOC_FAST_PATH bool operator==(const AtomicCapPtr& rhs) const
+    {
+      return this->unsafe_capptr == rhs.unsafe_capptr;
+    }
+
+    SNMALLOC_FAST_PATH bool operator!=(const AtomicCapPtr& rhs) const
+    {
+      return this->unsafe_capptr != rhs.unsafe_capptr;
+    }
+
+    SNMALLOC_FAST_PATH bool operator<(const AtomicCapPtr& rhs) const
+    {
+      return this->unsafe_capptr < rhs.unsafe_capptr;
+    }
+  };
+
+  template<typename T>
+  using AtomicCapPtrCBArena = AtomicCapPtr<T, CBArena>;
+
+  template<typename T>
+  using AtomicCapPtrCBChunk = AtomicCapPtr<T, CBChunk>;
+
+  template<typename T>
+  using AtomicCapPtrCBAlloc = AtomicCapPtr<T, CBAlloc>;
+
 } // namespace snmalloc