C++多模式可变参数模板实现向量逐元素apply函数

要实现支持向量与标量混合传入的apply函数，核心是区分向量和标量参数、自动广播标量，同时解决模板推导和友元声明的问题。下面是完整的可行方案，基于C17（兼容C14只需少量调整）：

第一步：定义类型 trait 识别向量类型 #

首先我们需要一个工具来判断某个类型是否是你的vec类，并提取它的元素类型和维度N：

#include <type_traits>
#include <utility> // 用于std::index_sequence

// 前置声明你的vec类
template <typename T, size_t N>
class vec;

// 类型trait：判断是否为vec类型
template <typename T>
struct is_vec : std::false_type {};

template <typename T, size_t N>
struct is_vec<vec<T, N>> : std::true_type {
    using value_type = T;
    static constexpr size_t size = N;
};

// 便捷别名和常量
template <typename T>
constexpr bool is_vec_v = is_vec<T>::value;

template <typename T>
using vec_value_type_t = typename is_vec<T>::value_type;

template <typename T>
constexpr size_t vec_size_v = is_vec<T>::size;

第二步：实现元素获取的辅助函数 #

这个函数负责统一处理向量和标量：对于向量返回指定索引的元素，对于标量直接返回自身（实现广播）：

// 获取第i个元素：标量直接返回
template <typename Arg>
auto get_element(const Arg& arg, size_t /*i*/) -> std::enable_if_t<!is_vec_v<Arg>, const Arg&> {
    return arg;
}

// 获取第i个元素：向量返回对应索引的元素
template <typename Arg>
auto get_element(const Arg& arg, size_t i) -> std::enable_if_t<is_vec_v<Arg>, const vec_value_type_t<Arg>&> {
    return arg[i]; // 假设你的vec类实现了operator[]
}

第三步：实现apply函数及其友元声明 #

首先在vec类内部声明友元，然后实现apply函数。我们需要用std::index_sequence来生成编译期的索引，逐个元素调用目标函数：

template <typename T, size_t N>
class vec {
public:
    // 假设你的vec有这样的构造函数和存储
    vec(std::initializer_list<T> init) {
        size_t i = 0;
        for (auto val : init) v[i++] = val;
    }

    T v[N]; // 存储元素的数组
    T& operator[](size_t i) { return v[i]; }
    const T& operator[](size_t i) const { return v[i]; }

    // 声明apply为友元：注意模板参数的写法
    template <typename F, typename... Args>
    friend auto apply(F&& f, const Args&... args) -> vec<
        std::invoke_result_t<F, decltype(get_element(std::declval<const Args&>(), 0))...>,
        // 从第一个向量参数推导N；如果没有向量参数，这里会编译错误（合理，因为标量不需要apply）
        std::enable_if_t<(is_vec_v<Args> || ...), size_t>
    >;
};

// 辅助函数：实际生成每个元素的调用
template <typename F, typename... Args, size_t... Is>
auto apply_impl(F&& f, const Args&... args, std::index_sequence<Is...>) {
    using ResultType = std::invoke_result_t<F, decltype(get_element(std::declval<const Args&>(), 0))...>;
    static constexpr size_t N = (vec_size_v<Args> + ... + 0) / (is_vec_v<Args> + ... + 0); // 计算统一的N

    // 检查所有向量参数的N是否一致
    static_assert(((is_vec_v<Args> ? (vec_size_v<Args> == N) : true) && ...),
                  "All vector arguments must have the same size");

    return vec<ResultType, N>{std::invoke(std::forward<F>(f), get_element(args, Is)...)...};
}

// 对外暴露的apply函数
template <typename F, typename... Args>
auto apply(F&& f, const Args&... args) -> vec<
    std::invoke_result_t<F, decltype(get_element(std::declval<const Args&>(), 0))...>,
    std::enable_if_t<(is_vec_v<Args> || ...), size_t>
> {
    // 找到第一个向量参数的N，生成索引序列
    static constexpr size_t N = (vec_size_v<Args> + ... + 0) / (is_vec_v<Args> + ... + 0);
    return apply_impl(std::forward<F>(f), args..., std::make_index_sequence<N>{});
}

第四步：测试示例 #

#include <cmath>
#include <iostream>

int main() {
    vec<float, 3> v{1.0f, 2.0f, 3.0f};
    auto result = apply(pow, v, 2); // 广播标量2到每个元素

    for (size_t i = 0; i < 3; ++i) {
        std::cout << result[i] << " "; // 输出1 4 9
    }

    // 也支持多个向量混合标量
    vec<int, 2> a{10, 20};
    vec<int, 2> b{5, 7};
    auto sum = apply([](int x, int y, int z) { return x + y + z; }, a, b, 3);
    // 结果是10+5+3=18，20+7+3=30
    for (size_t i = 0; i < 2; ++i) {
        std::cout << sum[i] << " "; // 输出18 30
    }

    return 0;
}

关键细节解释 #

1. 类型 trait：is_vec用来区分向量和标量，确保我们能正确提取向量的元素类型和维度。
2. 元素广播：get_element函数对不同类型做重载，标量直接返回自身，实现了自动广播到每个元素的效果。
3. 维度一致性检查：用static_assert确保所有传入的向量参数维度相同，避免运行时错误。
4. 模板推导：通过std::invoke_result_t自动推导返回向量的元素类型，从第一个向量参数推导维度N（如果有多个向量，会检查维度一致）。
5. 友元声明：在vec类内部正确声明模板友元，避免访问权限问题。

如果需要兼容C14，只需将std::invoke_result_t替换为std::result_of_t<F(decltype(get_element(std::declval<const Args&>(), 0))...)>，并确保编译器支持C14特性。

内容的提问来源于stack exchange，提问作者Glen Fletcher