Merge pull request #166 from microsoft/function_ref

Replace uses of std::function by function_ref.

src/ds/helpers.h

@@ -101,4 +101,46 @@ namespace snmalloc
       f();
     }
   };
+
+  /**
+   * Non-owning version of std::function. Wraps a reference to a callable object
+   * (eg. a lambda) and allows calling it through dynamic dispatch, with no
+   * allocation. This is useful in the allocator code paths, where we can't
+   * safely use std::function.
+   *
+   * Inspired by the C++ proposal:
+   * http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2018/p0792r2.html
+   */
+  template<typename Fn>
+  struct function_ref;
+  template<typename R, typename... Args>
+  struct function_ref<R(Args...)>
+  {
+    // The enable_if is used to stop this constructor from shadowing the default
+    // copy / move constructors.
+    template<
+      typename Fn,
+      typename =
+        std::enable_if_t<!std::is_same_v<std::decay_t<Fn>, function_ref>>>
+    function_ref(Fn&& fn)
+    {
+      data_ = static_cast<void*>(&fn);
+      fn_ = execute<Fn>;
+    }
+
+    R operator()(Args... args) const
+    {
+      return fn_(data_, args...);
+    }
+
+  private:
+    void* data_;
+    R (*fn_)(void*, Args...);
+
+    template<typename Fn>
+    static R execute(void* p, Args... args)
+    {
+      return (*static_cast<std::add_pointer_t<Fn>>(p))(args...);
+    };
+  };
 } // namespace snmalloc

src/mem/alloc.h

@@ -83,7 +83,7 @@ namespace snmalloc
    */
   template<
     bool (*NeedsInitialisation)(void*),
-    void* (*InitThreadAllocator)(std::function<void*(void*)>&),
+    void* (*InitThreadAllocator)(function_ref<void*(void*)>),
     class MemoryProvider = GlobalVirtual,
     class ChunkMap = SNMALLOC_DEFAULT_CHUNKMAP,
     bool IsQueueInline = true>
@@ -1141,10 +1141,9 @@ namespace snmalloc
     template<ZeroMem zero_mem, AllowReserve allow_reserve>
     SNMALLOC_SLOW_PATH void* small_alloc_first_alloc(size_t size)
     {
-      std::function<void*(void*)> f = [size](void* alloc) {
+      return InitThreadAllocator([size](void* alloc) {
         return reinterpret_cast<Allocator*>(alloc)->alloc(size);
-      };
-      return InitThreadAllocator(f);
+      });
     }
 
     /**
@@ -1321,10 +1320,9 @@ namespace snmalloc
       {
         if (NeedsInitialisation(this))
         {
-          std::function<void*(void*)> f = [size](void* alloc) {
+          return InitThreadAllocator([size](void* alloc) {
             return reinterpret_cast<Allocator*>(alloc)->alloc(size);
-          };
-          return InitThreadAllocator(f);
+          });
         }
         slab = reinterpret_cast<Mediumslab*>(
           large_allocator.template alloc<NoZero, allow_reserve>(
@@ -1395,10 +1393,9 @@ namespace snmalloc
 
       if (NeedsInitialisation(this))
       {
-        std::function<void*(void*)> f = [size](void* alloc) {
+        return InitThreadAllocator([size](void* alloc) {
           return reinterpret_cast<Allocator*>(alloc)->alloc(size);
-        };
-        return InitThreadAllocator(f);
+        });
       }
 
       size_t size_bits = bits::next_pow2_bits(size);
@@ -1423,11 +1420,10 @@ namespace snmalloc
 
       if (NeedsInitialisation(this))
       {
-        std::function<void*(void*)> f = [p](void* alloc) {
+        InitThreadAllocator([p](void* alloc) {
           reinterpret_cast<Allocator*>(alloc)->dealloc(p);
           return nullptr;
-        };
-        InitThreadAllocator(f);
+        });
         return;
       }
 
@@ -1479,11 +1475,10 @@ namespace snmalloc
       // a real allocator and construct one if we aren't.
       if (NeedsInitialisation(this))
       {
-        std::function<void*(void*)> f = [p](void* alloc) {
+        InitThreadAllocator([p](void* alloc) {
           reinterpret_cast<Allocator*>(alloc)->dealloc(p);
           return nullptr;
-        };
-        InitThreadAllocator(f);
+        });
         return;
       }
 

src/mem/globalalloc.h

@@ -7,7 +7,7 @@
 namespace snmalloc
 {
   inline bool needs_initialisation(void*);
-  void* init_thread_allocator(std::function<void*(void*)>&);
+  void* init_thread_allocator(function_ref<void*(void*)>);
 
   template<class MemoryProvider, class Alloc>
   class AllocPool : Pool<Alloc, MemoryProvider>

src/mem/threadalloc.h

@@ -62,7 +62,7 @@ namespace snmalloc
 #    pragma warning(push)
 #    pragma warning(disable : 4702)
 #  endif
-  SNMALLOC_FAST_PATH void* init_thread_allocator(std::function<void*(void*)>& f)
+  SNMALLOC_FAST_PATH void* init_thread_allocator(function_ref<void*(void*)> f)
   {
     error("Critical Error: This should never be called.");
     return f(nullptr);
@@ -101,7 +101,7 @@ namespace snmalloc
    */
   class ThreadAllocCommon
   {
-    friend void* init_thread_allocator(std::function<void*(void*)>&);
+    friend void* init_thread_allocator(function_ref<void*(void*)>);
 
   protected:
     /**
@@ -195,12 +195,11 @@ namespace snmalloc
       auto*& alloc = get_reference();
       if (unlikely(needs_initialisation(alloc)) && !destructor_has_run)
       {
-        std::function<void*(void*)> f = [](void*) { return nullptr; };
         // Call `init_thread_allocator` to perform down call in case
         // register_clean_up does more.
         // During teardown for the destructor based ThreadAlloc this will set
         // alloc to GlobalPlaceHolder;
-        init_thread_allocator(f);
+        init_thread_allocator([](void*) { return nullptr; });
       }
       return alloc;
 #  endif
@@ -277,7 +276,7 @@ namespace snmalloc
    * path.
    * The second component of the return indicates if this TLS is being torndown.
    */
-  SNMALLOC_FAST_PATH void* init_thread_allocator(std::function<void*(void*)>& f)
+  SNMALLOC_FAST_PATH void* init_thread_allocator(function_ref<void*(void*)> f)
   {
     auto*& local_alloc = ThreadAlloc::get_reference();
     // If someone reuses a noncachable call, then we can end up here