clang/test/CodeGenCoroutines/coro-halo.cpp

   1 // This tests that the coroutine heap allocation elision optimization could happen succesfully.
   2 // RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -O2 -emit-llvm %s -o - | FileCheck %s
   3
   4 #include "Inputs/coroutine.h"
   5 #include "Inputs/numeric.h"
   6
   7 template <typename T> struct generator {
   8   struct promise_type {
   9     T current_value;
  10     std::suspend_always yield_value(T value) {
  11       this->current_value = value;
  12       return {};
  13     }
  14     std::suspend_always initial_suspend() { return {}; }
  15     std::suspend_always final_suspend() noexcept { return {}; }
  16     generator get_return_object() { return generator{this}; };
  17     void unhandled_exception() {}
  18     void return_void() {}
  19   };
  20
  21   struct iterator {
  22     std::coroutine_handle<promise_type> _Coro;
  23     bool _Done;
  24
  25     iterator(std::coroutine_handle<promise_type> Coro, bool Done)
  26         : _Coro(Coro), _Done(Done) {}
  27
  28     iterator &operator++() {
  29       _Coro.resume();
  30       _Done = _Coro.done();
  31       return *this;
  32     }
  33
  34     bool operator==(iterator const &_Right) const {
  35       return _Done == _Right._Done;
  36     }
  37
  38     bool operator!=(iterator const &_Right) const { return !(*this == _Right); }
  39     T const &operator*() const { return _Coro.promise().current_value; }
  40     T const *operator->() const { return &(operator*()); }
  41   };
  42
  43   iterator begin() {
  44     p.resume();
  45     return {p, p.done()};
  46   }
  47
  48   iterator end() { return {p, true}; }
  49
  50   generator(generator const &) = delete;
  51   generator(generator &&rhs) : p(rhs.p) { rhs.p = nullptr; }
  52
  53   ~generator() {
  54     if (p)
  55       p.destroy();
  56   }
  57
  58 private:
  59   explicit generator(promise_type *p)
  60       : p(std::coroutine_handle<promise_type>::from_promise(*p)) {}
  61
  62   std::coroutine_handle<promise_type> p;
  63 };
  64
  65 template <typename T>
  66 generator<T> seq() {
  67   for (T i = {};; ++i)
  68     co_yield i;
  69 }
  70
  71 template <typename T>
  72 generator<T> take_until(generator<T> &g, T limit) {
  73   for (auto &&v : g)
  74     if (v < limit)
  75       co_yield v;
  76     else
  77       break;
  78 }
  79
  80 template <typename T>
  81 generator<T> multiply(generator<T> &g, T factor) {
  82   for (auto &&v : g)
  83     co_yield v *factor;
  84 }
  85
  86 template <typename T>
  87 generator<T> add(generator<T> &g, T adder) {
  88   for (auto &&v : g)
  89     co_yield v + adder;
  90 }
  91
  92 int main() {
  93   auto s = seq<int>();
  94   auto t = take_until(s, 10);
  95   auto m = multiply(t, 2);
  96   auto a = add(m, 110);
  97   return std::accumulate(a.begin(), a.end(), 0);
  98 }
  99
 100 // CHECK-LABEL: define{{.*}} i32 @main(
 101 //   CHECK: ret i32 1190
 102 //   CHECK-NOT: call{{.*}}_Znwm