Fixing Until, Take

Summary:
`apply()` returns false if and only if the handler it was //directly//
passed returned false. `Take` didn't do this right, and `Until` was just broken.

Test Plan: More thorough unit tests.

Reviewed By: kittipat@fb.com

FB internal diff: D913185

folly/experimental/Gen-inl.h

@@ -159,11 +159,11 @@ class GenImpl : public FBounded<Self> {
   typedef typename std::decay<Value>::type StorageType;
 
   /**
-   * apply() - Send all values produced by this generator to given
-   * handler until the handler returns false. Returns false if and only if the
-   * handler returns false. Note: It should return true even if it completes
-   * (without the handler returning false), as 'Chain' uses the return value of
-   * apply to determine if it should process the second object in its chain.
+   * apply() - Send all values produced by this generator to given handler until
+   * the handler returns false. Returns false if and only if the handler passed
+   * in returns false. Note: It should return true even if it completes (without
+   * the handler returning false), as 'Chain' uses the return value of apply to
+   * determine if it should process the second object in its chain.
    */
   template<class Handler>
   bool apply(Handler&& handler) const;
@@ -695,10 +695,8 @@ class Until : public Operator<Until<Predicate>> {
   {}
 
   template<class Value,
-           class Source,
-           class Result = typename std::result_of<Predicate(Value)>::type>
-  class Generator :
-      public GenImpl<Result, Generator<Value, Source, Result>> {
+           class Source>
+  class Generator : public GenImpl<Value, Generator<Value, Source>> {
     Source source_;
     Predicate pred_;
    public:
@@ -707,10 +705,18 @@ class Until : public Operator<Until<Predicate>> {
 
     template<class Handler>
     bool apply(Handler&& handler) const {
-      return source_.apply([&](Value value) -> bool {
-        return !pred_(std::forward<Value>(value))
-            && handler(std::forward<Value>(value));
+      bool cancelled = false;
+      source_.apply([&](Value value) -> bool {
+        if (pred_(value)) { // un-forwarded to disable move
+          return false;
+        }
+        if (!handler(std::forward<Value>(value))) {
+          cancelled = true;
+          return false;
+        }
+        return true;
       });
+      return !cancelled;
     }
   };
 
@@ -761,12 +767,15 @@ class Take : public Operator<Take> {
     bool apply(Handler&& handler) const {
       if (count_ == 0) { return false; }
       size_t n = count_;
-      return source_.apply([&](Value value) -> bool {
-          if (!handler(std::forward<Value>(value))) {
-            return false;
-          }
-          return --n;
-        });
+      bool cancelled = false;
+      source_.apply([&](Value value) -> bool {
+        if (!handler(std::forward<Value>(value))) {
+          cancelled = true;
+          return false;
+        }
+        return --n;
+      });
+      return !cancelled;
     }
   };
 
@@ -907,7 +916,7 @@ class Skip : public Operator<Skip> {
     template<class Handler>
     bool apply(Handler&& handler) const {
       if (count_ == 0) {
-        return source_.apply(handler);
+        return source_.apply(std::forward<Handler>(handler));
       }
       size_t n = 0;
       return source_.apply([&](Value value) -> bool {

folly/experimental/test/GenTest.cpp

@@ -241,13 +241,37 @@ TEST(Gen, Contains) {
 }
 
 TEST(Gen, Take) {
-  auto expected = vector<int>{1, 4, 9, 16};
-  auto actual =
+  {
+    auto expected = vector<int>{1, 4, 9, 16};
+    auto actual =
       seq(1, 1000)
-    | mapped([](int x) { return x * x; })
-    | take(4)
-    | as<vector<int>>();
-  EXPECT_EQ(expected, actual);
+      | mapped([](int x) { return x * x; })
+      | take(4)
+      | as<vector<int>>();
+    EXPECT_EQ(expected, actual);
+  }
+  {
+    auto expected = vector<int>{ 0, 1, 4, 5, 8 };
+    auto actual
+      = ((seq(0) | take(2)) +
+         (seq(4) | take(2)) +
+         (seq(8) | take(2)))
+      | take(5)
+      | as<vector>();
+    EXPECT_EQ(expected, actual);
+  }
+  {
+    auto expected = vector<int>{ 0, 1, 4, 5, 8 };
+    auto actual
+      = seq(0)
+      | mapped([](int i) {
+          return seq(i * 4) | take(2);
+        })
+      | concat
+      | take(5)
+      | as<vector>();
+    EXPECT_EQ(expected, actual);
+  }
 }
 
 TEST(Gen, Sample) {
@@ -294,11 +318,39 @@ TEST(Gen, Skip) {
 }
 
 TEST(Gen, Until) {
-  auto gen =
-      seq(1) //infinite
-    | mapped([](int x) { return x * x; })
-    | until([](int x) { return x >= 1000; });
-  EXPECT_EQ(31, gen | count);
+  {
+    auto expected = vector<int>{1, 4, 9, 16};
+    auto actual
+      = seq(1, 1000)
+      | mapped([](int x) { return x * x; })
+      | until([](int x) { return x > 20; })
+      | as<vector<int>>();
+    EXPECT_EQ(expected, actual);
+  }
+  {
+    auto expected = vector<int>{ 0, 1, 4, 5, 8 };
+    auto actual
+      = ((seq(0) | until([](int i) { return i > 1; })) +
+         (seq(4) | until([](int i) { return i > 5; })) +
+         (seq(8) | until([](int i) { return i > 9; })))
+      | until([](int i) { return i > 8; })
+      | as<vector<int>>();
+    EXPECT_EQ(expected, actual);
+  }
+  /*
+  {
+    auto expected = vector<int>{ 0, 1, 5, 6, 10 };
+    auto actual
+      = seq(0)
+      | mapped([](int i) {
+          return seq(i * 5) | until([=](int j) { return j > i * 5 + 1; });
+        })
+      | concat
+      | until([](int i) { return i > 10; })
+      | as<vector<int>>();
+    EXPECT_EQ(expected, actual);
+  }
+  */
 }
 
 auto even = [](int i) -> bool { return i % 2 == 0; };