outline call_impl to save on code size

This does cause more move constructions, as shown by the needed update to test_copy_move. Up to reviewers whether they want more code size or more moves.

Author: swolchok

include/pybind11/detail/function_ref.h

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//===- llvm/ADT/STLFunctionalExtras.h - Extras for <functional> -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains some extension to <functional>.
+//
+// No library is required when using these functions.
+//
+//===----------------------------------------------------------------------===//
+//     Extra additions to <functional>
+//===----------------------------------------------------------------------===//
+
+/// An efficient, type-erasing, non-owning reference to a callable. This is
+/// intended for use as the type of a function parameter that is not used
+/// after the function in question returns.
+///
+/// This class does not own the callable, so it is not in general safe to store
+/// a FunctionRef.
+
+// pybind11: modified again from executorch::runtime::FunctionRef
+// - renamed back to function_ref
+// - use pybind11 enable_if_t, remove_cvref_t, and remove_reference_t
+
+// torch::executor: modified from llvm::function_ref
+// - renamed to FunctionRef
+// - removed LLVM_GSL_POINTER and LLVM_LIFETIME_BOUND macro uses
+// - use namespaced internal::remove_cvref_t
+
+#pragma once
+
+#include <pybind11/detail/common.h>
+
+#include <cstdint>
+#include <type_traits>
+#include <utility>
+
+PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
+PYBIND11_NAMESPACE_BEGIN(detail)
+
+//===----------------------------------------------------------------------===//
+//     Features from C++20
+//===----------------------------------------------------------------------===//
+
+template <typename Fn>
+class function_ref;
+
+template <typename Ret, typename... Params>
+class function_ref<Ret(Params...)> {
+    Ret (*callback)(intptr_t callable, Params... params) = nullptr;
+    intptr_t callable;
+
+    template <typename Callable>
+    static Ret callback_fn(intptr_t callable, Params... params) {
+        return (*reinterpret_cast<Callable *>(callable))(std::forward<Params>(params)...);
+    }
+
+public:
+    function_ref() = default;
+    function_ref(std::nullptr_t) {}
+
+    template <typename Callable>
+    function_ref(
+        Callable &&callable,
+        // This is not the copy-constructor.
+        enable_if_t<!std::is_same<remove_cvref_t<Callable>, function_ref>::value> * = nullptr,
+        // Functor must be callable and return a suitable type.
+        enable_if_t<
+            std::is_void<Ret>::value
+            || std::is_convertible<decltype(std::declval<Callable>()(std::declval<Params>()...)),
+                                   Ret>::value> * = nullptr)
+        : callback(callback_fn<remove_reference_t<Callable>>),
+          callable(reinterpret_cast<intptr_t>(&callable)) {}
+
+    Ret operator()(Params... params) const {
+        return callback(callable, std::forward<Params>(params)...);
+    }
+
+    explicit operator bool() const { return callback; }
+
+    bool operator==(const function_ref<Ret(Params...)> &Other) const {
+        return callable == Other.callable;
+    }
+};
+PYBIND11_NAMESPACE_END(detail)
+PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)

include/pybind11/pybind11.h

@@ -13,6 +13,7 @@
 #include "detail/dynamic_raw_ptr_cast_if_possible.h"
 #include "detail/exception_translation.h"
 #include "detail/function_record_pyobject.h"
+#include "detail/function_ref.h"
 #include "detail/init.h"
 #include "detail/native_enum_data.h"
 #include "detail/using_smart_holder.h"
@@ -379,6 +380,40 @@ class cpp_function : public function {
         return unique_function_record(new detail::function_record());
     }
 
+private:
+    // This is outlined from the dispatch lambda in initialize to save
+    // on code size. Crucially, we use function_ref to type-erase the
+    // actual function lambda so that we can get code reuse for
+    // functions with the same Return, Args, and Guard.
+    template <typename Return, typename Guard, typename ArgsConverter, typename... Args>
+    static handle call_impl(detail::function_call &call, detail::function_ref<Return(Args...)> f) {
+        using namespace detail;
+        using cast_out
+            = make_caster<conditional_t<std::is_void<Return>::value, void_type, Return>>;
+
+        ArgsConverter args_converter;
+        if (!args_converter.load_args(call)) {
+            return PYBIND11_TRY_NEXT_OVERLOAD;
+        }
+
+        /* Override policy for rvalues -- usually to enforce rvp::move on an rvalue */
+        return_value_policy policy
+            = return_value_policy_override<Return>::policy(call.func.policy);
+
+        /* Perform the function call */
+        handle result;
+        if (call.func.is_setter) {
+            (void) std::move(args_converter).template call<Return, Guard>(f);
+            result = none().release();
+        } else {
+            result = cast_out::cast(
+                std::move(args_converter).template call<Return, Guard>(f), policy, call.parent);
+        }
+
+        return result;
+    }
+
+protected:
     /// Special internal constructor for functors, lambda functions, etc.
     template <typename Func, typename Return, typename... Args, typename... Extra>
     void initialize(Func &&f, Return (*)(Args...), const Extra &...extra) {
@@ -441,13 +476,6 @@ class cpp_function : public function {
 
         /* Dispatch code which converts function arguments and performs the actual function call */
         rec->impl = [](function_call &call) -> handle {
-            cast_in args_converter;
-
-            /* Try to cast the function arguments into the C++ domain */
-            if (!args_converter.load_args(call)) {
-                return PYBIND11_TRY_NEXT_OVERLOAD;
-            }
-
             /* Invoke call policy pre-call hook */
             process_attributes<Extra...>::precall(call);
 
@@ -456,24 +484,11 @@ class cpp_function : public function {
                                                                           : call.func.data[0]);
             auto *cap = const_cast<capture *>(reinterpret_cast<const capture *>(data));
 
-            /* Override policy for rvalues -- usually to enforce rvp::move on an rvalue */
-            return_value_policy policy
-                = return_value_policy_override<Return>::policy(call.func.policy);
-
-            /* Function scope guard -- defaults to the compile-to-nothing `void_type` */
-            using Guard = extract_guard_t<Extra...>;
-
-            /* Perform the function call */
-            handle result;
-            if (call.func.is_setter) {
-                (void) std::move(args_converter).template call<Return, Guard>(cap->f);
-                result = none().release();
-            } else {
-                result = cast_out::cast(
-                    std::move(args_converter).template call<Return, Guard>(cap->f),
-                    policy,
-                    call.parent);
-            }
+            auto result = call_impl<Return,
+                                    /* Function scope guard -- defaults to the compile-to-nothing
+                                       `void_type` */
+                                    extract_guard_t<Extra...>,
+                                    cast_in>(call, detail::function_ref<Return(Args...)>(cap->f));
 
             /* Invoke call policy post-call hook */
             process_attributes<Extra...>::postcall(call, result);

tests/test_copy_move.py

@@ -70,12 +70,12 @@ def test_move_and_copy_loads():
 
     assert c_m.copy_assignments + c_m.copy_constructions == 0
     assert c_m.move_assignments == 6
-    assert c_m.move_constructions == 9
+    assert c_m.move_constructions == 21
     assert c_mc.copy_assignments + c_mc.copy_constructions == 0
     assert c_mc.move_assignments == 5
-    assert c_mc.move_constructions == 8
+    assert c_mc.move_constructions == 18
     assert c_c.copy_assignments == 4
-    assert c_c.copy_constructions == 6
+    assert c_c.copy_constructions == 14
     assert c_m.alive() + c_mc.alive() + c_c.alive() == 0