Refactor check_bounds (#756)

This changes the shape of check_bounds to take a continuation to call if the bounds check succeeds.  This is designed to allow for easily wrapping existing code with a bounds check, e.g.

```
void* memcpy(void* dest, const void* src, size_t n) {
  return check_bounds(dest, n, [&] {
    return memcpy_impl(dest, src, n);
  });
}
```

src/snmalloc/global/bounds_checks.h

@@ -81,26 +81,33 @@ namespace snmalloc
   /**
    * Check whether a pointer + length is in the same object as the pointer.
    *
-   * Returns true if the checks succeeds.
+   * Returns the result of the supplied continuation
    *
    * The template parameter indicates whether the check should be performed.  It
-   * defaults to true. If it is false, the check will always succeed.
+   * defaults to true. If it is false, it short cuts to calling the continuation
+   * directly.
    */
-  template<bool PerformCheck = true, SNMALLOC_CONCEPT(IsConfig) Config = Config>
-  SNMALLOC_FAST_PATH_INLINE bool check_bounds(const void* ptr, size_t len)
+  template<
+    bool PerformCheck = true,
+    typename F,
+    SNMALLOC_CONCEPT(IsConfig) Config = Config>
+  SNMALLOC_FAST_PATH_INLINE auto check_bound(
+    const void* ptr, size_t len, const char* msg, F f = []() {})
   {
     if constexpr (PerformCheck)
     {
-      if (SNMALLOC_LIKELY(Config::is_initialised()))
+      if (SNMALLOC_LIKELY(len != 0))
       {
-        return remaining_bytes(address_cast(ptr)) >= len;
+        if (SNMALLOC_UNLIKELY(remaining_bytes<Config>(address_cast(ptr)) < len))
+        {
+          return report_fatal_bounds_error(ptr, len, msg);
+        }
       }
-      return true;
     }
     else
     {
-      UNUSED(ptr, len);
-      return true;
+      UNUSED(ptr, len, msg);
     }
+    return f();
   }
 } // namespace snmalloc

src/snmalloc/global/globalalloc.h

@@ -135,7 +135,7 @@ namespace snmalloc
    * remaining_bytes(p + n) == size - n     provided n < size
    */
   template<SNMALLOC_CONCEPT(IsConfig) Config_ = Config>
-  size_t remaining_bytes(address_t p)
+  size_t SNMALLOC_FAST_PATH_INLINE remaining_bytes(address_t p)
   {
     const auto& entry = Config_::Backend::template get_metaentry<true>(p);
 

src/snmalloc/global/memcpy.h

@@ -161,9 +161,9 @@ namespace snmalloc
     /**
      * Hook for architecture-specific optimisations.
      */
-    static SNMALLOC_FAST_PATH_INLINE void
-    copy(void* dst, const void* src, size_t len)
+    static void* copy(void* dst, const void* src, size_t len)
     {
+      auto orig_dst = dst;
       // If this is a small size, use a jump table for small sizes.
       if (len <= LargestRegisterSize)
       {
@@ -175,6 +175,7 @@ namespace snmalloc
         block_copy<LargestRegisterSize>(dst, src, len);
         copy_end<LargestRegisterSize>(dst, src, len);
       }
+      return orig_dst;
     }
   };
 
@@ -197,9 +198,9 @@ namespace snmalloc
 
     static constexpr size_t LargestRegisterSize = 16;
 
-    static SNMALLOC_FAST_PATH_INLINE void
-    copy(void* dst, const void* src, size_t len)
+    static void* copy(void* dst, const void* src, size_t len)
     {
+      auto orig_dst = dst;
       /*
        * As a function of misalignment relative to pointers, how big do we need
        * to be such that the span could contain an aligned pointer?  We'd need
@@ -306,6 +307,7 @@ namespace snmalloc
         block_copy<LargestRegisterSize>(dst, src, len);
         copy_end<LargestRegisterSize>(dst, src, len);
       }
+      return orig_dst;
     }
   };
 
@@ -331,9 +333,9 @@ namespace snmalloc
     /**
      * Platform-specific copy hook.  For large copies, use `rep movsb`.
      */
-    static SNMALLOC_FAST_PATH_INLINE void
-    copy(void* dst, const void* src, size_t len)
+    static inline void* copy(void* dst, const void* src, size_t len)
     {
+      auto orig_dst = dst;
       // If this is a small size, use a jump table for small sizes, like on the
       // generic architecture case above.
       if (len <= LargestRegisterSize)
@@ -370,6 +372,7 @@ namespace snmalloc
         block_copy<LargestRegisterSize>(dst, src, len);
         copy_end<LargestRegisterSize>(dst, src, len);
       }
+      return orig_dst;
     }
   };
 #endif
@@ -386,9 +389,9 @@ namespace snmalloc
      * For large copies (128 bytes or above), use a copy loop that moves up to
      * 128 bytes at once with pre-loop alignment up to 64 bytes.
      */
-    static SNMALLOC_FAST_PATH_INLINE void
-    copy(void* dst, const void* src, size_t len)
+    static void* copy(void* dst, const void* src, size_t len)
     {
+      auto orig_dst = dst;
       if (len < LargestRegisterSize)
       {
         block_copy<1>(dst, src, len);
@@ -408,6 +411,7 @@ namespace snmalloc
         block_copy<LargestRegisterSize>(dst, src, len);
         copy_end<LargestRegisterSize>(dst, src, len);
       }
+      return orig_dst;
     }
   };
 #endif
@@ -441,24 +445,13 @@ namespace snmalloc
     typename Arch = DefaultArch>
   SNMALLOC_FAST_PATH_INLINE void* memcpy(void* dst, const void* src, size_t len)
   {
-    auto orig_dst = dst;
-    // 0 is a very common size for memcpy and we don't need to do external
-    // pointer checks if we hit it.  It's also the fastest case, to encourage
-    // the compiler to favour the other cases.
-    if (SNMALLOC_UNLIKELY(len == 0))
-    {
-      return dst;
-    }
-
-    // Check the bounds of the arguments.
-    if (SNMALLOC_UNLIKELY(!check_bounds<(Checked && ReadsChecked)>(src, len)))
-      return report_fatal_bounds_error(
-        src, len, "memcpy with source out of bounds of heap allocation");
-    if (SNMALLOC_UNLIKELY(!check_bounds<Checked>(dst, len)))
-      return report_fatal_bounds_error(
-        dst, len, "memcpy with destination out of bounds of heap allocation");
-
-    Arch::copy(dst, src, len);
-    return orig_dst;
+    return check_bound<(Checked && ReadsChecked)>(
+      src, len, "memcpy with source out of bounds of heap allocation", [&]() {
+        return check_bound<Checked>(
+          dst,
+          len,
+          "memcpy with destination out of bounds of heap allocation",
+          [&]() { return Arch::copy(dst, src, len); });
+      });
   }
 } // namespace snmalloc

src/test/perf/memcpy/memcpy.cc

@@ -94,15 +94,13 @@ void memcpy_unchecked(void* dst, const void* src, size_t size)
 }
 
 NOINLINE
-void memcpy_platform_checked(void* dst, const void* src, size_t size)
+void* memcpy_platform_checked(void* dst, const void* src, size_t size)
 {
-  if (SNMALLOC_UNLIKELY(!check_bounds(dst, size)))
-  {
-    report_fatal_bounds_error(dst, size, "");
-    return;
-  }
-
-  memcpy(dst, src, size);
+  return check_bound(
+    dst,
+    size,
+    "memcpy with destination out of bounds of heap allocation",
+    [&]() { return memcpy(dst, src, size); });
 }
 
 int main(int argc, char** argv)