folly/experimental/fibers/WhenN-inl.h

   1 /*
   2  * Copyright 2015 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #include <folly/Optional.h>
  17
  18 #include <folly/experimental/fibers/FiberManager.h>
  19 #include <folly/experimental/fibers/ForEach.h>
  20
  21 namespace folly { namespace fibers {
  22
  23 template <class InputIterator>
  24 typename std::vector<
  25   typename std::enable_if<
  26     !std::is_same<
  27       typename std::result_of<
  28         typename std::iterator_traits<InputIterator>::value_type()>::type, void
  29       >::value,
  30     typename std::pair<
  31       size_t,
  32       typename std::result_of<
  33         typename std::iterator_traits<InputIterator>::value_type()>::type>
  34     >::type
  35   >
  36 collectN(InputIterator first, InputIterator last, size_t n) {
  37   typedef typename std::result_of<
  38     typename std::iterator_traits<InputIterator>::value_type()>::type Result;
  39   assert(n > 0);
  40   assert(n <= std::distance(first, last));
  41
  42   struct Context {
  43     std::vector<std::pair<size_t, Result>> results;
  44     size_t tasksTodo;
  45     std::exception_ptr e;
  46     folly::Optional<Promise<void>> promise;
  47
  48     Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
  49       this->results.reserve(tasksTodo_);
  50     }
  51   };
  52   auto context = std::make_shared<Context>(n);
  53
  54   await(
  55     [first, last, context](Promise<void> promise) mutable {
  56       context->promise = std::move(promise);
  57       for (size_t i = 0; first != last; ++i, ++first) {
  58 #ifdef __clang__
  59 #pragma clang diagnostic push // ignore generalized lambda capture warning
  60 #pragma clang diagnostic ignored "-Wc++1y-extensions"
  61 #endif
  62         addTask(
  63           [i, context, f = std::move(*first)]() {
  64             try {
  65               auto result = f();
  66               if (context->tasksTodo == 0) {
  67                 return;
  68               }
  69               context->results.emplace_back(i, std::move(result));
  70             } catch (...) {
  71               if (context->tasksTodo == 0) {
  72                 return;
  73               }
  74               context->e = std::current_exception();
  75             }
  76             if (--context->tasksTodo == 0) {
  77               context->promise->setValue();
  78             }
  79           });
  80 #ifdef __clang__
  81 #pragma clang diagnostic pop
  82 #endif
  83       }
  84     });
  85
  86   if (context->e != std::exception_ptr()) {
  87     std::rethrow_exception(context->e);
  88   }
  89
  90   return std::move(context->results);
  91 }
  92
  93 template <class InputIterator>
  94 typename std::enable_if<
  95   std::is_same<
  96     typename std::result_of<
  97       typename std::iterator_traits<InputIterator>::value_type()>::type, void
  98     >::value, std::vector<size_t>>::type
  99 collectN(InputIterator first, InputIterator last, size_t n) {
 100   assert(n > 0);
 101   assert(n <= std::distance(first, last));
 102
 103   struct Context {
 104     std::vector<size_t> taskIndices;
 105     std::exception_ptr e;
 106     size_t tasksTodo;
 107     folly::Optional<Promise<void>> promise;
 108
 109     Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
 110       this->taskIndices.reserve(tasksTodo_);
 111     }
 112   };
 113   auto context = std::make_shared<Context>(n);
 114
 115   await(
 116     [first, last, context](Promise<void> promise) mutable {
 117       context->promise = std::move(promise);
 118       for (size_t i = 0; first != last; ++i, ++first) {
 119 #ifdef __clang__
 120 #pragma clang diagnostic push // ignore generalized lambda capture warning
 121 #pragma clang diagnostic ignored "-Wc++1y-extensions"
 122 #endif
 123         addTask(
 124           [i, context, f = std::move(*first)]() {
 125             try {
 126               f();
 127               if (context->tasksTodo == 0) {
 128                 return;
 129               }
 130               context->taskIndices.push_back(i);
 131             } catch (...) {
 132               if (context->tasksTodo == 0) {
 133                 return;
 134               }
 135               context->e = std::current_exception();
 136             }
 137             if (--context->tasksTodo == 0) {
 138               context->promise->setValue();
 139             }
 140           });
 141 #ifdef __clang__
 142 #pragma clang diagnostic pop
 143 #endif
 144       }
 145     });
 146
 147   if (context->e != std::exception_ptr()) {
 148     std::rethrow_exception(context->e);
 149   }
 150
 151   return context->taskIndices;
 152 }
 153
 154 template <class InputIterator>
 155 typename std::vector<
 156   typename std::enable_if<
 157     !std::is_same<
 158       typename std::result_of<
 159         typename std::iterator_traits<InputIterator>::value_type()>::type, void
 160       >::value,
 161     typename std::result_of<
 162       typename std::iterator_traits<InputIterator>::value_type()>::type>::type>
 163 inline collectAll(InputIterator first, InputIterator last) {
 164   typedef typename std::result_of<
 165     typename std::iterator_traits<InputIterator>::value_type()>::type Result;
 166   size_t n = std::distance(first, last);
 167   std::vector<Result> results;
 168   std::vector<size_t> order(n);
 169   results.reserve(n);
 170
 171   forEach(first, last,
 172     [&results, &order] (size_t id, Result result) {
 173       order[id] = results.size();
 174       results.emplace_back(std::move(result));
 175     });
 176   assert(results.size() == n);
 177
 178   std::vector<Result> orderedResults;
 179   orderedResults.reserve(n);
 180
 181   for (size_t i = 0; i < n; ++i) {
 182     orderedResults.emplace_back(std::move(results[order[i]]));
 183   }
 184
 185   return orderedResults;
 186 }
 187
 188 template <class InputIterator>
 189 typename std::enable_if<
 190   std::is_same<
 191     typename std::result_of<
 192       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 193     >::value, void>::type
 194 inline collectAll(InputIterator first, InputIterator last) {
 195   forEach(first, last, [] (size_t id) {});
 196 }
 197
 198 template <class InputIterator>
 199 typename std::enable_if<
 200   !std::is_same<
 201     typename std::result_of<
 202       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 203     >::value,
 204   typename std::pair<
 205     size_t,
 206     typename std::result_of<
 207       typename std::iterator_traits<InputIterator>::value_type()>::type>
 208   >::type
 209 inline collectAny(InputIterator first, InputIterator last) {
 210   auto result = collectN(first, last, 1);
 211   assert(result.size() == 1);
 212   return std::move(result[0]);
 213 }
 214
 215 template <class InputIterator>
 216 typename std::enable_if<
 217   std::is_same<
 218     typename std::result_of<
 219       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 220     >::value, size_t>::type
 221 inline collectAny(InputIterator first, InputIterator last) {
 222   auto result = collectN(first, last, 1);
 223   assert(result.size() == 1);
 224   return std::move(result[0]);
 225 }
 226
 227 }}