如何实现能够返回生成器的模板驱动的嵌套 for 循环？

Question

我正在尝试实现一个在多维数组上提供下标的生成器。具体来说，我正在使用 CPP Chrono 的生成器实现。由于维数由输入

shape

数组的大小驱动，因此我需要一个递归函数来生成嵌套

for

循环。 到目前为止，我已经实现了大小扣除嵌套 for 循环的逻辑：

#include <iostream>
#include <array>

using Size_t = size_t;
using Int64_t = int;

/// @brief Nested for loop helper
/// @details Recursively generates NLoops nested for loops
template <size_t Index = 0, Size_t NLoops>
void nested_for_loop(const std::array<Size_t, NLoops>& shape, std::array<Int64_t, NLoops>& outSubscripts) {
    if constexpr (Index == NLoops) {
        // Base case: All nested loops have run.
        // Operation is performed here
        for (size_t i = 0; i < NLoops; ++i) {
            std::cout << outSubscripts[i] << ' ';
        }
        std::cout << std::endl;

    } else {
        // Recursive case: Loop from 0 to array[Index]
        for (size_t i = 0; i <= shape[Index]; ++i) {
            outSubscripts[Index] = i; // Store the current subscript
            nested_for_loop<Index + 1>(shape, outSubscripts); // Recursively call for the next loop
        }
    }
}

int main() {
    std::array<Size_t, 3> shape = {2, 30, 6};
    std::array<int, 3> outSubscripts{};

    Int64_t outIndex = 0;
    nested_for_loop(shape, outSubscripts);

    return 0;
}

这样就可以打印出我想要的下标了。如果这只是一个常规的旧 for 循环，我可以在操作逻辑所在的位置添加一个

co_yield

，将函数返回类型更新为 

generator<MyDesiredArrayType>

，然后就到此为止了。然而，似乎不可能将这种递归方法转换为生成器函数。

是否有更好的方法来展开这个

for

 循环，以便我可以正确地从协程中屈服？我能想到的唯一方法是

这个答案中描述的方法。我宁愿不必诉诸这样的方法，因为我认为没有一种 constexpr

 方法可以将线性循环索引转换为 

N

 所需的下标，而且我不想招致运行时额外加法/乘法的成本。

我总是可以选择手动写出一系列

if constexpr

语句。我可以添加 

M

 条件块，因此我将支持从 0 到 

M

 嵌套 for 循环。然而，这将是非常乏味和限制性的，以至于它是愚蠢的。肯定有更好的方法！

Answer 1

要实现可返回生成器的模板驱动的嵌套 for 循环，您可以使用 C++20 的协程和 '

co_yield 语句' 。以下是如何修改现有代码以实现此目的的示例：

#include <iostream>
#include <array>
#include <coroutine>

using Size_t = size_t;
using Int64_t = int;

/// @brief Nested for loop generator
/// @details Generates subscripts over a multi-dimensional array
template <Size_t NLoops>
struct NestedForLoopGenerator {
    std::array<Size_t, NLoops> shape;
    std::array<Int64_t, NLoops> subscripts;

    NestedForLoopGenerator(const std::array<Size_t, NLoops>& shape) : shape(shape) {}

    struct iterator {
        NestedForLoopGenerator& parent;
        size_t index;

        bool operator!=(const iterator& other) const {
            return index < other.parent.shape[0];
        }

        iterator& operator++() {
            ++index;
            return *this;
        }

        std::array<Int64_t, NLoops>& operator*() {
            parent.subscripts[0] = index;
            for (size_t i = 1; i < NLoops; ++i) {
                parent.subscripts[i] = 0;
            }
            return parent.subscripts;
        }

        void operator++(int) {
            operator++();
        }
    };

    iterator begin() {
        return { *this, 0 };
    }

    iterator end() {
        return { *this, shape[0] };
    }
};

int main() {
    std::array<Size_t, 3> shape = {2, 3, 4};

    NestedForLoopGenerator<3> generator(shape);
    for (auto subscripts : generator) {
        for (const auto& val : subscripts) {
            std::cout << val << ' ';
        }
        std::cout << std::endl;
    }

    return 0;
}

在此代码中，我们定义了一个“

NestedForLoopGenerator”结构，该结构将形状作为输入。它有一个迭代器，可以为嵌套的 for 循环生成下标。 'begin' 和 'end' 方法返回允许您循环下标的迭代器。

您可以通过更改 NLoops 模板参数轻松修改此代码以支持不同数量的维度。

通过使用C++20协程，您可以直接在循环中产生下标，而不需要递归函数调用。该代码应该用作生成器并为多维数组生成所需的下标。