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(std::distance(first, last) >= 0);
  41   assert(n <= static_cast<size_t>(std::distance(first, last)));
  42
  43   struct Context {
  44     std::vector<std::pair<size_t, Result>> results;
  45     size_t tasksTodo;
  46     std::exception_ptr e;
  47     folly::Optional<Promise<void>> promise;
  48
  49     Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
  50       this->results.reserve(tasksTodo_);
  51     }
  52   };
  53   auto context = std::make_shared<Context>(n);
  54
  55   await(
  56     [first, last, context](Promise<void> promise) mutable {
  57       context->promise = std::move(promise);
  58       for (size_t i = 0; first != last; ++i, ++first) {
  59 #ifdef __clang__
  60 #pragma clang diagnostic push // ignore generalized lambda capture warning
  61 #pragma clang diagnostic ignored "-Wc++1y-extensions"
  62 #endif
  63         addTask(
  64           [i, context, f = std::move(*first)]() {
  65             try {
  66               auto result = f();
  67               if (context->tasksTodo == 0) {
  68                 return;
  69               }
  70               context->results.emplace_back(i, std::move(result));
  71             } catch (...) {
  72               if (context->tasksTodo == 0) {
  73                 return;
  74               }
  75               context->e = std::current_exception();
  76             }
  77             if (--context->tasksTodo == 0) {
  78               context->promise->setValue();
  79             }
  80           });
  81 #ifdef __clang__
  82 #pragma clang diagnostic pop
  83 #endif
  84       }
  85     });
  86
  87   if (context->e != std::exception_ptr()) {
  88     std::rethrow_exception(context->e);
  89   }
  90
  91   return std::move(context->results);
  92 }
  93
  94 template <class InputIterator>
  95 typename std::enable_if<
  96   std::is_same<
  97     typename std::result_of<
  98       typename std::iterator_traits<InputIterator>::value_type()>::type, void
  99     >::value, std::vector<size_t>>::type
 100 collectN(InputIterator first, InputIterator last, size_t n) {
 101   assert(n > 0);
 102   assert(std::distance(first, last) >= 0);
 103   assert(n <= static_cast<size_t>(std::distance(first, last)));
 104
 105   struct Context {
 106     std::vector<size_t> taskIndices;
 107     std::exception_ptr e;
 108     size_t tasksTodo;
 109     folly::Optional<Promise<void>> promise;
 110
 111     Context(size_t tasksTodo_) : tasksTodo(tasksTodo_) {
 112       this->taskIndices.reserve(tasksTodo_);
 113     }
 114   };
 115   auto context = std::make_shared<Context>(n);
 116
 117   await(
 118     [first, last, context](Promise<void> promise) mutable {
 119       context->promise = std::move(promise);
 120       for (size_t i = 0; first != last; ++i, ++first) {
 121 #ifdef __clang__
 122 #pragma clang diagnostic push // ignore generalized lambda capture warning
 123 #pragma clang diagnostic ignored "-Wc++1y-extensions"
 124 #endif
 125         addTask(
 126           [i, context, f = std::move(*first)]() {
 127             try {
 128               f();
 129               if (context->tasksTodo == 0) {
 130                 return;
 131               }
 132               context->taskIndices.push_back(i);
 133             } catch (...) {
 134               if (context->tasksTodo == 0) {
 135                 return;
 136               }
 137               context->e = std::current_exception();
 138             }
 139             if (--context->tasksTodo == 0) {
 140               context->promise->setValue();
 141             }
 142           });
 143 #ifdef __clang__
 144 #pragma clang diagnostic pop
 145 #endif
 146       }
 147     });
 148
 149   if (context->e != std::exception_ptr()) {
 150     std::rethrow_exception(context->e);
 151   }
 152
 153   return context->taskIndices;
 154 }
 155
 156 template <class InputIterator>
 157 typename std::vector<
 158   typename std::enable_if<
 159     !std::is_same<
 160       typename std::result_of<
 161         typename std::iterator_traits<InputIterator>::value_type()>::type, void
 162       >::value,
 163     typename std::result_of<
 164       typename std::iterator_traits<InputIterator>::value_type()>::type>::type>
 165 inline collectAll(InputIterator first, InputIterator last) {
 166   typedef typename std::result_of<
 167     typename std::iterator_traits<InputIterator>::value_type()>::type Result;
 168   size_t n = std::distance(first, last);
 169   std::vector<Result> results;
 170   std::vector<size_t> order(n);
 171   results.reserve(n);
 172
 173   forEach(first, last,
 174     [&results, &order] (size_t id, Result result) {
 175       order[id] = results.size();
 176       results.emplace_back(std::move(result));
 177     });
 178   assert(results.size() == n);
 179
 180   std::vector<Result> orderedResults;
 181   orderedResults.reserve(n);
 182
 183   for (size_t i = 0; i < n; ++i) {
 184     orderedResults.emplace_back(std::move(results[order[i]]));
 185   }
 186
 187   return orderedResults;
 188 }
 189
 190 template <class InputIterator>
 191 typename std::enable_if<
 192   std::is_same<
 193     typename std::result_of<
 194       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 195     >::value, void>::type
 196 inline collectAll(InputIterator first, InputIterator last) {
 197   forEach(first, last, [](size_t /* id */) {});
 198 }
 199
 200 template <class InputIterator>
 201 typename std::enable_if<
 202   !std::is_same<
 203     typename std::result_of<
 204       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 205     >::value,
 206   typename std::pair<
 207     size_t,
 208     typename std::result_of<
 209       typename std::iterator_traits<InputIterator>::value_type()>::type>
 210   >::type
 211 inline collectAny(InputIterator first, InputIterator last) {
 212   auto result = collectN(first, last, 1);
 213   assert(result.size() == 1);
 214   return std::move(result[0]);
 215 }
 216
 217 template <class InputIterator>
 218 typename std::enable_if<
 219   std::is_same<
 220     typename std::result_of<
 221       typename std::iterator_traits<InputIterator>::value_type()>::type, void
 222     >::value, size_t>::type
 223 inline collectAny(InputIterator first, InputIterator last) {
 224   auto result = collectN(first, last, 1);
 225   assert(result.size() == 1);
 226   return std::move(result[0]);
 227 }
 228
 229 }}