[ADT] Support const-qualified unique_functions

[sammccall]

Summary

This technique should extend to rvalue-qualified etc, but I didn't add any.
I removed "volatile" from the future plans, which seems... speculative at best.

While here I moved the callbacks object out of the constructor into a
variable template, which I believe addresses the fixme there about unused
objects.

(I'm not a template guru, so it's always possible the old version was designed
for compile-time performance in a way I'm missing)

[kadircet]

llvm/include/llvm/ADT/FunctionExtras.h

131

nit: let's keep this and make use of it in getCallPtr instead of inlining.

162

nit: drop this (also the one before getNonTriivialCallbacks())?

175

nit: again drop this

192–193

nit: rename this to CalledAsT ?

215

s/to an/an/

218

move and destroy pointers don't really depend on CallableT maybe change this to only take CalledAsT ?

225

again this also doesn't depend on CallableT

228

maybe CalledAs ?

233

s/CallAsT/CalledAsT/

(or rename in the other way for the rest of the file)

361

call Base(nullptr_t), also in the mutable version?

363

drop this ?

373

drop this ?

llvm/unittests/ADT/FunctionExtrasTest.cpp

230

only-moveable captures are already tested elsewhere. is there a specific reason for this capture ?

[sammccall]

llvm/include/llvm/ADT/FunctionExtras.h

218

I don't think using e.g. MoveAs<CalledAsT> will work: you almost always can't move from a const T.

You could do something like MoveAs<remove_cv_t<CalledAsT>> or have MoveAs const_cast. But that's just a different design, and I don't think particularly a better one: having both CallableT and CalledAsT explicitly specified (with the base class responsible for ensuring they're compatible) seems more robust than trying to make assumptions about how to calculate one from the other.

For example, when rvalue-ref qualified functions get added, I think it's nice not to have to worry about updating the metaprogramming to calculate the underlying type.

225

This is a specialization of the template above, so I don't think there's anything we can change here if the above stays the same.

361

This is "just" syntax sugar for the default constructor, and that's hardly obscure.

I think it's unneccesarily confusing to implement the public API in the base class where it doesn't save any complexity, so I've removed the nullptr_t from the base instead.

(I'm torn on operator bool as it's the last thing publicly exposed from the base, but I don't think it matters much)

363

it's required here to call a method from a dependent base.
(That's where all the other extra thises came from - code that I had in one of these subclasses and then moved)

llvm/unittests/ADT/FunctionExtrasTest.cpp

230

Have to make the machinery sweat a little!

More seriously, if the underlying functor is copyable, then MoveImpl<const lambda> *would* compile... but does a copy.

So there are reasonable-looking implementations where move works and const works but move const doesn't work. (but I admit I only worked this out after the fact).

[sammccall]

Address comments

[kadircet]

Thanks for all the explanations, LGTM!

llvm/include/llvm/ADT/FunctionExtras.h

162

return isTrivialCallback() ? getTrivialCallback()->CallPtr : getNonTrivialCallbacks()->CallPtr;

llvm/include/llvm/ADT/FunctionExtras.h

 //===- FunctionExtras.h - Function type erasure utilities -------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file provides a collection of function (or more generally, callable)
 /// type erasure utilities supplementing those provided by the standard library
 /// in `<function>`.
 ///
 /// It provides `unique_function`, which works like `std::function` but supports
-/// move-only callable objects.
+/// move-only callable objects and const-qualification.
 ///
 /// Future plans:
-/// - Add a `function` that provides const, volatile, and ref-qualified support,
-///   which doesn't work with `std::function`.
+/// - Add a `function` that provides 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_EXTRAS_H
 #define LLVM_ADT_FUNCTION_EXTRAS_H
 
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/Support/MemAlloc.h"
 #include "llvm/Support/type_traits.h"
 #include <memory>
+#include <type_traits>
 
 namespace llvm {
 
+/// unique_function is a type-erasing functor similar to std::function.
+///
+/// It can hold move-only function objects, like lambdas capturing unique_ptrs.
+/// Accordingly, it is movable but not copyable.
+///
+/// It supports const-qualification:
+/// - unique_function<int() const> has a const operator().
+///   It can only hold functions which themselves have a const operator().
+/// - unique_function<int()> has a non-const operator().
+///   It can hold functions with a non-const operator(), like mutable lambdas.
 template <typename FunctionT> class unique_function;
 
-template <typename ReturnT, typename... ParamTs>
-class unique_function<ReturnT(ParamTs...)> {
+namespace detail {
+
+template <typename T>
+using EnableIfTrivial =
+    std::enable_if_t<llvm::is_trivially_move_constructible<T>::value &&
+                     std::is_trivially_destructible<T>::value>;
+
+template <typename ReturnT, typename... ParamTs> class UniqueFunctionBase {
+protected:
   static constexpr size_t InlineStorageSize = sizeof(void *) * 3;
 
   // MSVC has a bug and ICEs if we give it a particular dependent value
   // expression as part of the `std::conditional` below. To work around this,
   // we build that into a template struct's constexpr bool.
   template <typename T> struct IsSizeLessThanThresholdT {
     static constexpr bool value = sizeof(T) <= (2 * sizeof(void *));
   };
struct OutOfLineStorageT {
       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 is mutable as an inlined `const unique_function<void() const>` may
+    // still modify its own mutable members.
+    mutable
         typename std::aligned_storage<InlineStorageSize, alignof(void *)>::type
             InlineStorage;
   } StorageUnion;
 
   // 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<TrivialCallback *>();
   }
 
   CallPtrT getTrivialCallback() const {
     return CallbackAndInlineFlag.getPointer().template get<TrivialCallback *>()->CallPtr;
   }
 
   NonTrivialCallbacks *getNonTrivialCallbacks() const {
     return CallbackAndInlineFlag.getPointer()
         .template get<NonTrivialCallbacks *>();
   }
 
-  void *getInlineStorage() { return &StorageUnion.InlineStorage; }
+  CallPtrT getCallPtr() const {
+    return isTrivialCallback() ? CallbackAndInlineFlag.getPointer()
+                                     .template get<TrivialCallback *>()
+                                     ->CallPtr
+                               : getNonTrivialCallbacks()->CallPtr;
+  }
 
-  void *getOutOfLineStorage() {
+  // These three functions are only const in the narrow sense. They return
+  // mutable pointers to function state.
+  // This allows unique_function<T const>::operator() to be const, even if the
+  // underlying functor may be internally mutable.
+  //
+  // const callers must ensure they're only used in const-correct ways.
+  void *getCalleePtr() const {
+    return isInlineStorage() ? getInlineStorage() : getOutOfLineStorage();
+  }
+  void *getInlineStorage() const { return &StorageUnion.InlineStorage; }
+  void *getOutOfLineStorage() const {
     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, AdjustedParamT<ParamTs>... Params) {
-    return (*reinterpret_cast<CallableT *>(CallableAddr))(
-        std::forward<ParamTs>(Params)...);
+  template <typename CalledAsT>
+  static ReturnT CallImpl(void *CallableAddr,
+                          AdjustedParamT<ParamTs>... Params) {
+    auto &Func = *reinterpret_cast<CalledAsT *>(CallableAddr);
+    return Func(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*/) {}
+  // The pointers to call/move/destroy functions are determined for each
+  // callable type (and called-as type, which determines the overload chosen).
+  // (definitions are out-of-line).
 
-  ~unique_function() {
+  // By default, we need an object that contains all the different
+  // type erased behaviors needed. Create a static instance of the struct type
+  // here and each instance will contain a pointer to it.
+  template <typename CallableT, typename CalledAs, typename Enable = void>
+  static NonTrivialCallbacks Callbacks;
+  // See if we can create a trivial callback. We need the callable to be
+  // trivially moved and trivially destroyed so that we don't have to store
+  // type erased callbacks for those operations.
+  template <typename CallableT, typename CalledAs>
+  static TrivialCallback
+      Callbacks<CallableT, CalledAs, EnableIfTrivial<CallableT>>;
+
+  // A simple tag type so the call-as type to be passed to the constructor.
+  template <typename T> struct CalledAs {};
+
+  // Essentially the "main" unique_function constructor, but subclasses
+  // provide the qualified type to be used for the call.
+  // (We always store a T, even if the call will use a pointer to const T).
+  template <typename CallableT, typename CalledAsT>
+  UniqueFunctionBase(CallableT Callable, CalledAs<CalledAsT>) {
+    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));
+    CallbackAndInlineFlag = {&Callbacks<CallableT, CalledAsT>, IsInlineStorage};
+  }
+
+  ~UniqueFunctionBase() {
     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 {
+  UniqueFunctionBase(UniqueFunctionBase &&RHS) noexcept {
     // Copy the callback and inline flag.
     CallbackAndInlineFlag = RHS.CallbackAndInlineFlag;
 
     // If the RHS is empty, just copying the above is sufficient.
     if (!RHS)
       return;
 
     if (!isInlineStorage()) {
UniqueFunctionBase(UniqueFunctionBase &&RHS) noexcept {
     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 {
+  UniqueFunctionBase &operator=(UniqueFunctionBase &&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));
+    this->~UniqueFunctionBase();
+    new (this) UniqueFunctionBase(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);
-    }
+  UniqueFunctionBase() = default;
+
+public:
+  explicit operator bool() const {
+    return (bool)CallbackAndInlineFlag.getPointer();
+};
 
-    // Now move into the storage.
-    new (CallableAddr) CallableT(std::move(Callable));
+template <typename R, typename... P>
+template <typename CallableT, typename CalledAsT, typename Enable>
+typename UniqueFunctionBase<R, P...>::NonTrivialCallbacks
+    UniqueFunctionBase<R, P...>::Callbacks = {
+        &CallImpl<CalledAsT>, &MoveImpl<CallableT>, &DestroyImpl<CallableT>};
+
+template <typename R, typename... P>
+template <typename CallableT, typename CalledAsT>
+typename UniqueFunctionBase<R, P...>::TrivialCallback UniqueFunctionBase<
+    R, P...>::Callbacks<CallableT, CalledAsT, EnableIfTrivial<CallableT>>{
+    &CallImpl<CalledAsT>};
+
+} // namespace detail
+
+template <typename R, typename... P>
+class unique_function<R(P...)> : public detail::UniqueFunctionBase<R, P...> {
+  using Base = detail::UniqueFunctionBase<R, P...>;
 
-    // See if we can create a trivial callback. We need the callable to be
-    // trivially moved and trivially destroyed so that we don't have to store
-    // type erased callbacks for those operations.
-    //
-    // 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 (llvm::is_trivially_move_constructible<CallableT>::value &&
-        std::is_trivially_destructible<CallableT>::value) {
-      // We need to create a nicely aligned object. We use a static variable
-      // for this because it is a trivial struct.
-      static TrivialCallback Callback = { &CallImpl<CallableT> };
-
-      CallbackAndInlineFlag = {&Callback, IsInlineStorage};
-      return;
-    }
+public:
+  unique_function() = default;
+  unique_function(std::nullptr_t) {}
+  unique_function(unique_function &&) = default;
+  unique_function(const unique_function &) = delete;
+  unique_function &operator=(unique_function &&) = default;
+  unique_function &operator=(const unique_function &) = delete;
 
-    // Otherwise, we need to point at an object that contains all the different
-    // type erased behaviors needed. Create a static instance of the struct type
-    // here and then use a pointer to that.
-    static NonTrivialCallbacks Callbacks = {
-        &CallImpl<CallableT>, &MoveImpl<CallableT>, &DestroyImpl<CallableT>};
+  template <typename CallableT>
+  unique_function(CallableT Callable)
+      : Base(std::forward<CallableT>(Callable),
+             typename Base::template CalledAs<CallableT>{}) {}
 
-    CallbackAndInlineFlag = {&Callbacks, IsInlineStorage};
+  R operator()(P... Params) {
+    return this->getCallPtr()(this->getCalleePtr(), Params...);
   }
+};
 
-  ReturnT operator()(ParamTs... Params) {
-    void *CallableAddr =
-        isInlineStorage() ? getInlineStorage() : getOutOfLineStorage();
+template <typename R, typename... P>
+class unique_function<R(P...) const>
+    : public detail::UniqueFunctionBase<R, P...> {
+  using Base = detail::UniqueFunctionBase<R, P...>;
 
-    return (isTrivialCallback()
-                ? getTrivialCallback()
-                : getNonTrivialCallbacks()->CallPtr)(CallableAddr, Params...);
-  }
+public:
+  unique_function() = default;
+  unique_function(std::nullptr_t) {}
+  unique_function(unique_function &&) = default;
+  unique_function(const unique_function &) = delete;
+  unique_function &operator=(unique_function &&) = default;
+  unique_function &operator=(const unique_function &) = delete;
 
-  explicit operator bool() const {
-    return (bool)CallbackAndInlineFlag.getPointer();
+  template <typename CallableT>
+  unique_function(CallableT Callable)
+      : Base(std::forward<CallableT>(Callable),
+             typename Base::template CalledAs<const CallableT>{}) {}
+
+  R operator()(P... Params) const {
+    return this->getCallPtr()(this->getCalleePtr(), Params...);
   }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_ADT_FUNCTION_H

llvm/unittests/ADT/FunctionExtrasTest.cpp

 //===- FunctionExtrasTest.cpp - Unit tests for function type erasure ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/FunctionExtras.h"
 #include "gtest/gtest.h"
 
 #include <memory>
+#include <type_traits>
 
 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);
TEST(UniqueFunctionTest, CountForwardingMoves) {
   };
   unique_function<void(const Unmovable &)> UnmovableF = [](const Unmovable &) {
   };
   UnmovableF(Unmovable());
   Unmovable X;
   UnmovableF(X);
 }
 
+TEST(UniqueFunctionTest, Const) {
+  // Can assign from const lambda.
+  unique_function<int(int) const> Plus2 = [X(std::make_unique<int>(2))](int Y) {
+    return *X + Y;
+  };
+  EXPECT_EQ(5, Plus2(3));
+
+  // Can call through a const ref.
+  const auto &Plus2Ref = Plus2;
+  EXPECT_EQ(5, Plus2Ref(3));
+
+  // Can move-construct and assign.
+  unique_function<int(int) const> Plus2A = std::move(Plus2);
+  EXPECT_EQ(5, Plus2A(3));
+  unique_function<int(int) const> Plus2B;
+  Plus2B = std::move(Plus2A);
+  EXPECT_EQ(5, Plus2B(3));
+
+  // Can convert to non-const function type, but not back.
+  unique_function<int(int)> Plus2C = std::move(Plus2B);
+  EXPECT_EQ(5, Plus2C(3));
+
+  // Overloaded call operator correctly resolved.
+  struct ChooseCorrectOverload {
+    StringRef operator()() { return "non-const"; }
+    StringRef operator()() const { return "const"; }
+  };
+  unique_function<StringRef()> ChooseMutable = ChooseCorrectOverload();
+  ChooseCorrectOverload A;
+  EXPECT_EQ("non-const", ChooseMutable());
+  EXPECT_EQ("non-const", A());
+  unique_function<StringRef() const> ChooseConst = ChooseCorrectOverload();
+  const ChooseCorrectOverload &X = A;
+  EXPECT_EQ("const", ChooseConst());
+  EXPECT_EQ("const", X());
+}
+
 } // anonymous namespace