folly/fibers/WhenN-inl.h

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