Add folly::coro::concat algorithm

Summary: The `concat` algorithm will concatenate a series of streams known at compile time into a single stream of the same type. Reviewed By: kirkshoop Differential Revision: D18695001 fbshipit-source-id: 2da7f53cbc835333ffd9551ceae1aa5db085936d

Add folly::coro::concat algorithm
Summary: The `concat` algorithm will concatenate a series of streams known at compile time into a single stream of the same type. Reviewed By: kirkshoop Differential Revision: D18695001 fbshipit-source-id: 2da7f53cbc835333ffd9551ceae1aa5db085936d
522fc79c · Taylor Hopper · Facebook Github Bot · dc6c17ca · 522fc79c · 522fc79c
Commit 522fc79c authored Nov 27, 2019 by Taylor Hopper Committed by Facebook Github Bot Nov 27, 2019
4 changed files
--- a/folly/experimental/coro/AsyncGenerator.h
+++ b/folly/experimental/coro/AsyncGenerator.h
@@ -496,6 +496,10 @@ template <typename Reference, typename Value>
 inline constexpr bool is_async_generator_v<AsyncGenerator<Reference, Value>> =
    true;

+template <typename T>
+using enable_if_async_generator_t =
+    std::enable_if_t<detail::is_async_generator_v<T>, T>;
+
 } // namespace detail

 // Helper for immediately invoking a lambda with captures that returns an

--- a/folly/experimental/coro/Concat-inl.h
+++ b/folly/experimental/coro/Concat-inl.h
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/experimental/coro/Concat.h>
+
+namespace folly {
+namespace coro {
+
+template <typename Head>
+detail::enable_if_async_generator_t<Head> concat(Head head) {
+  return std::move(head);
+}
+
+template <typename Head, typename... Tail>
+detail::enable_if_async_generator_t<Head> concat(Head head, Tail... tail) {
+  using Reference = typename Head::reference;
+  while (auto val = co_await head.next()) {
+    co_yield std::move(val).value();
+  }
+
+  auto tailGen = concat(std::move(tail)...);
+  while (auto val = co_await tailGen.next()) {
+    co_yield std::move(val).value();
+  }
+}
+
+} // namespace coro
+} // namespace folly
--- a/folly/experimental/coro/Concat.h
+++ b/folly/experimental/coro/Concat.h
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/experimental/coro/AsyncGenerator.h>
+
+namespace folly {
+namespace coro {
+
+// Concatenate the values from multiple streams into a single stream such
+// that each stream is exhausted before the next one begins.
+//
+// The input is a variadic list of AsyncGenerators, where each input has the
+// same Reference and Value types.
+//
+// The output is a single AsyncGenerator over all of the input generators.
+//
+// Example:
+//  AsyncGenerator<int> stream();
+//
+//  Task<int> consumer() {
+//    auto values = concat(stream(), stream(), stream());
+//
+//    int result = 0;
+//    while (auto item = co_await values.next()) {
+//      result += *item;
+//    }
+//
+//    return result;
+//  }
+template <typename Head>
+detail::enable_if_async_generator_t<Head> concat(Head head);
+
+template <typename Head, typename... Tail>
+detail::enable_if_async_generator_t<Head> concat(Head head, Tail... tail);
+
+} // namespace coro
+} // namespace folly
+
+#include <folly/experimental/coro/Concat-inl.h>
--- a/folly/experimental/coro/test/ConcatTest.cpp
+++ b/folly/experimental/coro/test/ConcatTest.cpp
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/Portability.h>
+
+#if FOLLY_HAS_COROUTINES
+
+#include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Concat.h>
+#include <folly/experimental/coro/Task.h>
+
+#include <folly/portability/GTest.h>
+
+using namespace ::testing;
+using namespace folly::coro;
+
+AsyncGenerator<int> generateInts(int begin, int end) {
+  for (int i = begin; i < end; i++) {
+    co_await co_reschedule_on_current_executor;
+    co_yield i;
+  }
+}
+
+Task<std::vector<int>> toVector(AsyncGenerator<int> generator) {
+  std::vector<int> result;
+  while (auto x = co_await generator.next()) {
+    result.push_back(*x);
+  }
+  co_return result;
+}
+
+TEST(ConcatTest, ConcatSingle) {
+  auto gen = concat(generateInts(0, 5));
+  auto result = blockingWait(toVector(std::move(gen)));
+  std::vector<int> expected{0, 1, 2, 3, 4};
+
+  EXPECT_EQ(result, expected);
+}
+
+TEST(ConcatTest, ConcatMultiple) {
+  auto gen =
+      concat(generateInts(0, 5), generateInts(7, 10), generateInts(12, 15));
+  auto result = blockingWait(toVector(std::move(gen)));
+  std::vector<int> expected{0, 1, 2, 3, 4, 7, 8, 9, 12, 13, 14};
+
+  EXPECT_EQ(result, expected);
+}
+
+#endif