[ADT] Add llvm::unique_function which is like std::function but supporting move-only closures.

[chandlerc]

Summary

Most of the core optimizations for std::function are here plus
a potentially novel one that detects trivially movable and destroyable
functors and implements those with fewer indirections.

This is especially useful as we start trying to add concurrency
primitives as those often end up with move-only types (futures,
promises, etc) and wanting them to work through lambdas.

As further work, we could add better support for things like const-qualified
operator()s to support more algorithms, and r-value ref qualified operator()s
to model call-once. None of that is here though.

We can also provide our own llvm::function that has some of the optimizations
used in this class, but with copy semantics instead of move semantics.

This is motivated by increasing usage of things like executors and the task
queue where it is useful to embed move-only types like a std::promise within
a type erased function. That isn't possible without this version of a type
erased function.

[chandlerc]

BTW, one question I have is what the best name for the file is here. function.h seems kinda lame, but not sure what else to call it. Suggestions very welcome here.

[zturner]

Not to suggest the obvious, but UniqueFunction.h?

[chandlerc]

In D48349#1138032, @zturner wrote:

Not to suggest the obvious, but UniqueFunction.h?

If people prefer that, totally fine with me.

I went with this because we might want to merge llvm::function_ref into this file. And we might want to add llvm::function rather than relying on std::function given the amount of churn in that part of the standard library.

[zturner]

In that case maybe FunctionExtras.h in keeping with the llvm naming
spirit

[timshen]

llvm/include/llvm/ADT/function.h

10

Can we add a bit of future plan? Is it going to merge with the future-standard unique_function, if exist at all?

55

teh duh!

81

What's the rationale of choosing this specific data layout? Currently it's

tagged_union { CallPtrT, struct { CallPtrT, MovePtrT, DestroyPtrT }* }

What makes us want it instead of, for example:

struct { CallPtrT, /* nullable */ struct { MovePtrT, DestroyPtrT}* }

Or even:

// MoveOrDestroyPtrT does either Move or Destroy depending on an extra parameter.
struct { CallPtrT, /* nulltable */ MoveOrDestroyPtrT }

?

I'm not saying the examples are better, but just want more details on the decision.

228

Do we want to perfect-forward the parameters?

[dberris]

llvm/include/llvm/ADT/function.h

188

Do you still need the std::move(...) here?

[chandlerc]

Thanks for the comments! Some responses below...

llvm/include/llvm/ADT/function.h

81

The goal is to minimize the footprint by only using a single pointer within the object. Improved the comment to include this and not reference the prior design of a base class.

188

Yeah, within the function its an L-value reference and so this will copy otherwise. Same reason you need std::forward w/ a forwarding reference.

[chandlerc]

Address comments from review, beefing up tests along the way to cover this
stuff much more carefully.

Also fixes several bugs uncovered with better testing and adds a missing
interface (bool conversion).

[chandlerc]

In D48349#1138090, @zturner wrote:

In that case maybe FunctionExtras.h in keeping with the llvm naming
spirit

I'm fine with the idea of this, but I'm always frustrated by the generic "extras" name component. Feels about as useless at "Utils". But I lack a significantly better idea....

Closest thing to a better idea I have is FunctionTypeErasure.h or CallableTypeErasuer.h. Thoughts? Other folks with naming thoughts here?

[chandlerc]

Actually include some of the new tests and fix minor bugs they found.

[timshen]

llvm/include/llvm/ADT/function.h

236

I think I understand the perfect-forwarding problem better now. Suppose we have the following user code:

unique_function<void(string)> f(...);
string a;
f(a);

Is the string copy-constructed once or twice? I figured that it'd be twice, (1) when calling operator(), (2) when calling CallImpl().

Regardless of the solution, can we add a test case for that?

[chandlerc]

Add even more testing and rework how we forward parameters through the various
level of call wrapping and type erasure. This minimizes the copies and moves
when those do occur. We also use an optimization to keep things pass-by-value
where it will end up especially beneficial and can't be observed.

[chandlerc]

Thanks so much for the careful review Tim!

llvm/include/llvm/ADT/function.h

236

Ah, there is an interesting case here. I've added tests that much more carefully cover all the variations of both functionality and the number of moves and copies triggered.

There is a fair amount of logic that goes into getting this right however, so please do look through the updated code.

[timshen]

llvm/include/llvm/ADT/FunctionExtras.h

161 ↗

(On Diff #153607)

The default ctor doesn't initialize CallbackAndInlineFlag to a valid empty state, does it? We probably want to zero-initialize it.

std::function's default ctor is not trivial either.

181 ↗

(On Diff #153607)

if (!(bool)RHS) { ... }

[chandlerc]

Minor fix from code review. Also adjusted the order of things in the private section of the class to be more readable by putting the two data members together.

llvm/include/llvm/ADT/FunctionExtras.h

161 ↗

(On Diff #153607)

Yep, the underlying value is set to zero by an in-class initializer.

This isn't about making something trivial, its just the simplest implementation.

181 ↗

(On Diff #153607)

Not 100% sure what you're going for here, but yes I should bail out early when RHS is null. Let me know if you intended something else.

[yroux]

Hi Chandler,

The added test cases are broken on ARM bots, logs are available at:

http://lab.llvm.org:8011/builders/clang-cmake-thumbv7-full-sh/builds/464

Thanks
Yvan

[chandlerc]

I've (finally) landed r336337 which just completely moves away from using a
function pointer packed into a pointerunion. =/

[yroux]

ARM bots are back, Thanks Chandler.

llvm/include/llvm/ADT/function.h

+//===- function.h - Function type erasure utilities -------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+/// \file
+/// This file provides a collection of function (or more generally, callable)
+/// type erasure utilities supplementing those provided by the standard library.
+///
+/// It provides `unique_function`, which works like `std::function` but supports
+/// move-only callable objects.
+///
+/// Future plans:
+/// - Add a `function` that provides const, volatile, and ref-qualified support,
+///   which doesn't work with `std::function`.
+/// - Provide support for specifying multiple signatures to type erase callable
+///   objects with an overload set, such as those produced by generic lambdas.
+/// - Expand to include a copyable utility that directly replaces std::function
+///   but brings the above improvements.
+///
+/// Note that LLVM's utilities are greatly simplified by not supporting
+/// allocators.
+///
+/// If the standard library ever begins to provide comparable facilities we can
+/// consider switching to those.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ADT_FUNCTION_H
+#define LLVM_ADT_FUNCTION_H
+
+#include "llvm/ADT/PointerIntPair.h"
+#include "llvm/ADT/PointerUnion.h"
+#include <memory>
+#include <type_traits>
+
+namespace llvm {
+
+namespace detail {} // end namespace detail
+
+template <typename FunctionT> class unique_function;
+
+template <typename ReturnT, typename... ParamTs>
+class unique_function<ReturnT(ParamTs...)> {
+  static constexpr int InlineStorageSize = sizeof(void *) * 3;
+
+  union StorageUnionT {
+    // For out-of-line storage we keep a pointer to the underlying storage and
+    // the size. This is enough to deallocate the memory.
+    struct OutOfLineStorageT {
+      void *StoragePtr;
+      size_t Size;
+      size_t Alignment;
+    } OutOfLineStorage;
+    static_assert(
+        sizeof(OutOfLineStorageT) <= InlineStorageSize,
+        "Should always use all of the out-of-line storage for inline storage!");
+
+    // For in-line storage, we just provide an aligned character buffer. We
+    // provide three pointers worth of storage here. This wastes one pointer
+    // worth of memory when in the out-of-line form.
+    typename std::aligned_storage<InlineStorageSize, alignof(void *)>::type
+        InlineStorage;
+  } StorageUnion;
+
+  // The type of the erased function pointer we use as a callback to dispatch to
+  // the stored callable when it is trivial to move and destroy.
+  using CallPtrT = ReturnT (*)(void *CallableAddr, ParamTs... Params);
+  using MovePtrT = void (*)(void *LHSCallableAddr, void *RHSCallableAddr);
+  using DestroyPtrT = void (*)(void *CallableAddr);
+
+  /// A struct we use to aggregate three callbacks when we need full set of
+  /// operations.
+  struct NonTrivialCallbacks {
+    CallPtrT CallPtr;
+    MovePtrT MovePtr;
+    DestroyPtrT DestroyPtr;
+  };
+
+  // Now we can create a pointer union between either a direct, trivial call
+  // pointer and a pointer to a static struct of the call, move, and destroy
+  // pointers. We do this to keep the footprint in this object a single pointer
+  // while supporting all the necessary type-erased operation.
+  using CallbackPointerUnionT = PointerUnion<CallPtrT, NonTrivialCallbacks *>;
+
+  // A compressed pointer to either our dispatching callback or our table of
+  // dispatching callbacks and the flag for whether the callable itself is
+  // stored inline or not.
+  PointerIntPair<CallbackPointerUnionT, 1, bool> CallbackAndInlineFlag;
+
+  bool isInlineStorage() const { return CallbackAndInlineFlag.getInt(); }
+
+  bool isTrivialCallback() const {
+    return CallbackAndInlineFlag.getPointer().template is<CallPtrT>();
+  }
+
+  CallPtrT getTrivialCallback() const {
+    return CallbackAndInlineFlag.getPointer().template get<CallPtrT>();
+  }
+
+  NonTrivialCallbacks *getNonTrivialCallbacks() const {
+    return CallbackAndInlineFlag.getPointer()
+        .template get<NonTrivialCallbacks *>();
+  }
+
+  void *getInlineStorage() { return &StorageUnion.InlineStorage; }
+
+  void *getOutOfLineStorage() {
+    return StorageUnion.OutOfLineStorage.StoragePtr;
+  }
+  size_t getOutOfLineStorageSize() const {
+    return StorageUnion.OutOfLineStorage.Size;
+  }
+  size_t getOutOfLineStorageAlignment() const {
+    return StorageUnion.OutOfLineStorage.Alignment;
+  }
+
+  void setOutOfLineStorage(void *Ptr, size_t Size, size_t Alignment) {
+    StorageUnion.OutOfLineStorage = {Ptr, Size, Alignment};
+  }
+
+  template <typename CallableT>
+  static ReturnT CallImpl(void *CallableAddr, ParamTs... Params) {
+    return (*reinterpret_cast<CallableT *>(CallableAddr))(
+        std::forward<ParamTs>(Params)...);
+  }
+
+  template <typename CallableT>
+  static void MoveImpl(void *LHSCallableAddr, void *RHSCallableAddr) noexcept {
+    new (LHSCallableAddr)
+        CallableT(std::move(*reinterpret_cast<CallableT *>(RHSCallableAddr)));
+  }
+
+  template <typename CallableT>
+  static void DestroyImpl(void *CallableAddr) noexcept {
+    reinterpret_cast<CallableT *>(CallableAddr)->~CallableT();
+  }
+
+public:
+  unique_function() = default;
+  unique_function(std::nullptr_t /*null_callable*/) {}
+
+  ~unique_function() {
+    if (!CallbackAndInlineFlag.getPointer())
+      return;
+
+    // Cache this value so we don't re-check it after type-erased operations.
+    bool IsInlineStorage = isInlineStorage();
+
+    if (!isTrivialCallback())
+      getNonTrivialCallbacks()->DestroyPtr(
+          IsInlineStorage ? getInlineStorage() : getOutOfLineStorage());
+
+    if (!IsInlineStorage)
+      deallocate_buffer(getOutOfLineStorage(), getOutOfLineStorageSize(),
+                        getOutOfLineStorageAlignment());
+  }
+
+  unique_function(unique_function &&RHS) noexcept {
+    // Copy the callback and inline flag.
+    CallbackAndInlineFlag = RHS.CallbackAndInlineFlag;
+
+    if (!isInlineStorage()) {
+      // The out-of-line case is easiest to move.
+      StorageUnion.OutOfLineStorage = RHS.StorageUnion.OutOfLineStorage;
+
+    } else if (isTrivialCallback()) {
+      // Move is trivial, just memcpy the bytes across.
+      memcpy(getInlineStorage(), RHS.getInlineStorage(), InlineStorageSize);
+    } else {
+      // Non-trivial move, so dispatch to a type-erased implementation.
+      getNonTrivialCallbacks()->MovePtr(getInlineStorage(),
+                                        RHS.getInlineStorage());
+    }
+
+    // Clear the old callback and inline flag to get back to as-if-null.
+    RHS.CallbackAndInlineFlag = {};
+
+#ifndef NDEBUG
+    // In debug builds, we also scribble across the rest of the storage.
+    memset(RHS.getInlineStorage(), 0xAD, InlineStorageSize);
+#endif
+  }
+
+  unique_function &operator=(unique_function &&RHS) noexcept {
+    if (this == &RHS)
+      return *this;
+
+    // Because we don't try to provide any exception safety guarantees we can
+    // implement move assignment very simply by first destroying the current
+    // object and then move-constructing over top of it.
+    this->~unique_function();
+    new (this) unique_function(std::move(RHS));
+    return *this;
+  }
+
+  template <typename CallableT> unique_function(CallableT Callable) {
+    bool IsInlineStorage = true;
+    void *CallableAddr = getInlineStorage();
+    if (sizeof(CallableT) > InlineStorageSize ||
+        alignof(CallableT) > alignof(decltype(StorageUnion.InlineStorage))) {
+      IsInlineStorage = false;
+      // Allocate out-of-line storage. FIXME: Use an explicit alignment
+      // parameter in C++17 mode.
+      auto Size = sizeof(CallableT);
+      auto Alignment = alignof(CallableT);
+      CallableAddr = allocate_buffer(Size, Alignment);
+      setOutOfLineStorage(CallableAddr, Size, Alignment);
+    }
+
+    // Now move into the storage.
+    new (CallableAddr) CallableT(std::move(Callable));
+
+    // See if we can create a trivial callback.
+    // FIXME: we should use constexpr if here and below to avoid instantiating
+    // the non-trivial static objects when unnecessary. While the linker should
+    // remove them, it is still wasteful.
+    if (std::is_trivially_move_constructible<CallableT>::value &&
+        std::is_trivially_destructible<CallableT>::value) {
+      CallbackAndInlineFlag = {&CallImpl<CallableT>, IsInlineStorage};
+      return;
+    }
+
+    // Otherwise, we need to point at an object with a vtable that contains all
+    // the different type erased behaviors needed. Create a static instance of
+    // the derived type here and then use a pointer to that.
+    static NonTrivialCallbacks Callbacks = {
+        &CallImpl<CallableT>, &MoveImpl<CallableT>, &DestroyImpl<CallableT>};
+
+    CallbackAndInlineFlag = {&Callbacks, IsInlineStorage};
+  }
+
+  ReturnT operator()(ParamTs... Params) {
+    void *CallableAddr =
+        isInlineStorage() ? getInlineStorage() : getOutOfLineStorage();
+    if (isTrivialCallback())
+      return getTrivialCallback()(CallableAddr,
+                                  std::forward<ParamTs>(Params)...);
+
+    return getNonTrivialCallbacks()->CallPtr(CallableAddr,
+                                             std::forward<ParamTs>(Params)...);
+  }
+
+  explicit operator bool() const {
+    return (bool)CallbackAndInlineFlag.getPointer();
+  }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_ADT_FUNCTION_H

llvm/include/llvm/Support/Compiler.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_COMPILER_H
 #define LLVM_SUPPORT_COMPILER_H
 
 #include "llvm/Config/llvm-config.h"
 
+#include <new>
+#include <stddef.h>
+
 #if defined(_MSC_VER)
 #include <sal.h>
 #endif
 
 #ifndef __has_feature
 # define __has_feature(x) 0
 #endif
 

 #if __has_feature(cxx_exceptions)
 #define LLVM_ENABLE_EXCEPTIONS 1
 #elif defined(__GNUC__) && defined(__EXCEPTIONS)
 #define LLVM_ENABLE_EXCEPTIONS 1
 #elif defined(_MSC_VER) && defined(_CPPUNWIND)
 #define LLVM_ENABLE_EXCEPTIONS 1
 #endif
 
+namespace llvm {
+
+/// Allocate a buffer of memory with the given size and alignment.
+///
+/// When the compiler supports aligned operator new, this will use it to to
+/// handle even over-aligned allocations.
+///
+/// However, this doesn't make any attempt to leverage the fancier techniques
+/// like posix_memalign due to portability. It is mostly intended to allow
+/// compatibility with platforms that, after aligned allocation was added, use
+/// reduced default alignment.
+inline void *allocate_buffer(size_t Size, size_t Alignment) {
+  return ::operator new(Size
+#if __cpp_aligned_new
+                        ,
+                        std::align_val_t(Alignment)
+#endif
+  );
+}
+
+/// Deallocate a buffer of memory with the given size and alignment.
+///
+/// If supported, this will used the sized delete operator. Also if supported,
+/// this will pass the alignment to the delete operator.
+///
+/// The pointer must have been allocated with the corresponding new operator,
+/// most likely using the above helper.
+inline void deallocate_buffer(void *Ptr, size_t Size, size_t Alignment) {
+  ::operator delete(Ptr
+#if __cpp_sized_deallocation
+                    ,
+                    Size
+#endif
+#if __cpp_aligned_new
+                    ,
+                    std::align_val_t(Alignment)
+#endif
+  );
+}
+
+} // End namespace llvm
+
 #endif

llvm/include/llvm/Support/PointerLikeTypeTraits.h

 // structures to reason about pointers and other things that are pointer sized.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_POINTERLIKETYPETRAITS_H
 #define LLVM_SUPPORT_POINTERLIKETYPETRAITS_H
 
 #include "llvm/Support/DataTypes.h"
+#include <assert.h>
 #include <type_traits>
 
 namespace llvm {
 
 /// A traits type that is used to handle pointer types and things that are just
 /// wrappers for pointers as a uniform entity.
 template <typename T> struct PointerLikeTypeTraits;
 
template <> struct PointerLikeTypeTraits<uintptr_t> {
   }
   static inline uintptr_t getFromVoidPointer(void *P) {
     return reinterpret_cast<uintptr_t>(P);
   }
   // No bits are available!
   enum { NumLowBitsAvailable = 0 };
 };
 
+/// Provide suitable custom traits struct for function pointers.
+///
+/// Function pointers can't be directly given these traits as functions can't
+/// have their alignment computed with `alignof` and we need different casting.
+///
+/// To rely on higher alignment for a specialized use, you can provide a
+/// customized form of this template explicitly with higher alignment, and
+/// potentially use alignment attributes on functions to satisfy that.
+template <int Alignment, typename FunctionPointerT>
+struct FunctionPointerLikeTypeTraits {
+  enum { NumLowBitsAvailable = detail::ConstantLog2<Alignment>::value };
+  static inline void *getAsVoidPointer(FunctionPointerT P) {
+    assert((reinterpret_cast<uintptr_t>(P) &
+            ~((uintptr_t)-1 << NumLowBitsAvailable)) == 0 &&
+           "Alignment not satisfied for an actual function pointer!");
+    return reinterpret_cast<void *>(P);
+  }
+  static inline FunctionPointerT getFromVoidPointer(void *P) {
+    return reinterpret_cast<FunctionPointerT>(P);
+  }
+};
+
+/// Provide a default specialization for function pointers that assumes 4-byte
+/// alignment.
+///
+/// We assume here that functions used with this are always at least 4-byte
+/// aligned. This means that, for example, thumb functions won't work or systems
+/// with weird unaligned function pointers won't work. But all practical systems
+/// we support satisfy this requirement.
+template <typename ReturnT, typename... ParamTs>
+struct PointerLikeTypeTraits<ReturnT (*)(ParamTs...)>
+    : FunctionPointerLikeTypeTraits<4, ReturnT (*)(ParamTs...)> {};
+
 } // end namespace llvm
 
 #endif

llvm/unittests/ADT/CMakeLists.txt

add_llvm_unittest(ADTTests
   BumpPtrListTest.cpp
   DAGDeltaAlgorithmTest.cpp
   DeltaAlgorithmTest.cpp
   DenseMapTest.cpp
   DenseSetTest.cpp
   DepthFirstIteratorTest.cpp
   EquivalenceClassesTest.cpp
   FoldingSet.cpp
+  FunctionTest.cpp
   FunctionRefTest.cpp
   HashingTest.cpp
   IListBaseTest.cpp
   IListIteratorTest.cpp
   IListNodeBaseTest.cpp
   IListNodeTest.cpp
   IListSentinelTest.cpp
   IListTest.cpp

llvm/unittests/ADT/FunctionTest.cpp

+//===- FunctionTest.cpp - Unit tests for function type erasure ------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/function.h"
+#include "gtest/gtest.h"
+
+#include <memory>
+
+using namespace llvm;
+
+namespace {
+
+TEST(UniqueFunctionTest, Basic) {
+  unique_function<int(int, int)> Sum = [](int A, int B) { return A + B; };
+  EXPECT_EQ(Sum(1, 2), 3);
+
+  unique_function<int(int, int)> Sum2 = std::move(Sum);
+  EXPECT_EQ(Sum2(1, 2), 3);
+
+  unique_function<int(int, int)> Sum3 = [](int A, int B) { return A + B; };
+  Sum2 = std::move(Sum3);
+  EXPECT_EQ(Sum2(1, 2), 3);
+
+  Sum2 = unique_function<int(int, int)>([](int A, int B) { return A + B; });
+  EXPECT_EQ(Sum2(1, 2), 3);
+
+  // Explicit self-move test.
+  *&Sum2 = std::move(Sum2);
+  EXPECT_EQ(Sum2(1, 2), 3);
+
+  Sum2 = unique_function<int(int, int)>();
+  EXPECT_FALSE(Sum2);
+}
+
+TEST(UniqueFunctionTest, Captures) {
+  long A = 1, B = 2, C = 3, D = 4, E = 5;
+
+  unique_function<long()> Tmp;
+
+  unique_function<long()> C1 = [A]() { return A; };
+  EXPECT_EQ(C1(), 1);
+  Tmp = std::move(C1);
+  EXPECT_EQ(Tmp(), 1);
+
+  unique_function<long()> C2 = [A, B]() { return A + B; };
+  EXPECT_EQ(C2(), 3);
+  Tmp = std::move(C2);
+  EXPECT_EQ(Tmp(), 3);
+
+  unique_function<long()> C3 = [A, B, C]() { return A + B + C; };
+  EXPECT_EQ(C3(), 6);
+  Tmp = std::move(C3);
+  EXPECT_EQ(Tmp(), 6);
+
+  unique_function<long()> C4 = [A, B, C, D]() { return A + B + C + D; };
+  EXPECT_EQ(C4(), 10);
+  Tmp = std::move(C4);
+  EXPECT_EQ(Tmp(), 10);
+
+  unique_function<long()> C5 = [A, B, C, D, E]() { return A + B + C + D + E; };
+  EXPECT_EQ(C5(), 15);
+  Tmp = std::move(C5);
+  EXPECT_EQ(Tmp(), 15);
+}
+
+TEST(UniqueFunctionTest, MoveOnly) {
+  struct SmallCallable {
+    std::unique_ptr<int> A{new int(1)};
+
+    int operator()(int B) { return *A + B; }
+  };
+  unique_function<int(int)> Small = SmallCallable();
+  EXPECT_EQ(Small(2), 3);
+  unique_function<int(int)> Small2 = std::move(Small);
+  EXPECT_EQ(Small2(2), 3);
+
+  struct LargeCallable {
+    std::unique_ptr<int> A{new int(1)};
+    std::unique_ptr<int> B{new int(2)};
+    std::unique_ptr<int> C{new int(3)};
+    std::unique_ptr<int> D{new int(4)};
+    std::unique_ptr<int> E{new int(5)};
+
+    int operator()() { return *A + *B + *C + *D + *E; }
+  };
+  unique_function<int()> Large = LargeCallable();
+  EXPECT_EQ(Large(), 15);
+  unique_function<int()> Large2 = std::move(Large);
+  EXPECT_EQ(Large2(), 15);
+}
+
+} // anonymous namespace
 No newline at end of file