(Folly/Gen) Make ranges and sequences take a stepping functor

Summary: The functor allows to create ranges and sequences that advance the iterator by more than 1 per iteration. Test Plan: Unit test. Reviewed By: tjackson@fb.com FB internal diff: D1228065
facebook · Apr 18, 2014 · 8cd97e5 · 8cd97e5
1 parent d456b58
commit 8cd97e5
Show file tree

Hide file tree

Showing 3 changed files with 117 additions and 48 deletions.
diff --git a/folly/gen/Base-inl.h b/folly/gen/Base-inl.h
@@ -183,47 +183,29 @@ class RangeSource : public GenImpl<typename Range<Iterator>::reference,
 
 /**
  * Sequence - For generating values from beginning value, incremented along the
- * way with the ++ and += operators. Iteration may continue indefinitely by
- * setting the 'endless' template parameter to true. If set to false, iteration
- * will stop when value reaches 'end', either inclusively or exclusively,
- * depending on the template parameter 'endInclusive'. Value type specified
- * explicitly.
+ * way with the ++ and += operators. Iteration may continue indefinitely.
+ * Value type specified explicitly.
  *
  * This type is primarily used through the 'seq' and 'range' function, like:
  *
  *   int total = seq(1, 10) | sum;
  *   auto indexes = range(0, 10);
+ *   auto endless = seq(0); // 0, 1, 2, 3, ...
  */
-template<class Value,
-         bool endless,
-         bool endInclusive>
-class Sequence : public GenImpl<const Value&,
-                                Sequence<Value, endless, endInclusive>> {
+template<class Value, class SequenceImpl>
+class Sequence : public GenImpl<const Value&, Sequence<Value, SequenceImpl>> {
   static_assert(!std::is_reference<Value>::value &&
                 !std::is_const<Value>::value, "Value mustn't be const or ref.");
-  Value bounds_[endless ? 1 : 2];
+  Value start_;
+  SequenceImpl impl_;
 public:
-  explicit Sequence(Value begin)
-      : bounds_{std::move(begin)} {
-    static_assert(endless, "Must supply 'end'");
-  }
-
-  Sequence(Value begin,
-           Value end)
-    : bounds_{std::move(begin), std::move(end)} {}
+  explicit Sequence(Value start, SequenceImpl impl)
+      : start_(std::move(start)), impl_(std::move(impl)) { }
 
   template<class Handler>
   bool apply(Handler&& handler) const {
-    Value value = bounds_[0];
-    for (;endless || value < bounds_[1]; ++value) {
-      const Value& arg = value;
-      if (!handler(arg)) {
-        return false;
-      }
-    }
-    if (endInclusive && value == bounds_[1]) {
-      const Value& arg = value;
-      if (!handler(arg)) {
+    for (Value current = start_; impl_.test(current); impl_.step(current)) {
+      if (!handler(current)) {
         return false;
       }
     }
@@ -232,18 +214,60 @@ class Sequence : public GenImpl<const Value&,
 
   template<class Body>
   void foreach(Body&& body) const {
-    Value value = bounds_[0];
-    for (;endless || value < bounds_[1]; ++value) {
-      const Value& arg = value;
-      body(arg);
-    }
-    if (endInclusive && value == bounds_[1]) {
-      const Value& arg = value;
-      body(arg);
+    for (Value current = start_; impl_.test(current); impl_.step(current)) {
+      body(current);
     }
   }
+};
+
+/**
+ * Sequence implementations (range, sequence, infinite, with/without step)
+ **/
+template<class Value>
+class RangeImpl {
+  Value end_;
+ public:
+  explicit RangeImpl(Value end) : end_(std::move(end)) { }
+  bool test(const Value& current) const { return current < end_; }
+  void step(Value& current) const { ++current; }
+};
 
-  static constexpr bool infinite = endless;
+template<class Value, class Distance>
+class RangeWithStepImpl {
+  Value end_;
+  Distance step_;
+ public:
+  explicit RangeWithStepImpl(Value end, Distance step)
+    : end_(std::move(end)), step_(std::move(step)) { }
+  bool test(const Value& current) const { return current < end_; }
+  void step(Value& current) const { current += step_; }
+};
+
+template<class Value>
+class SeqImpl {
+  Value end_;
+ public:
+  explicit SeqImpl(Value end) : end_(std::move(end)) { }
+  bool test(const Value& current) const { return current <= end_; }
+  void step(Value& current) const { ++current; }
+};
+
+template<class Value, class Distance>
+class SeqWithStepImpl {
+  Value end_;
+  Distance step_;
+ public:
+  explicit SeqWithStepImpl(Value end, Distance step)
+    : end_(std::move(end)), step_(std::move(step)) { }
+  bool test(const Value& current) const { return current <= end_; }
+  void step(Value& current) const { current += step_; }
+};
+
+template<class Value>
+class InfiniteImpl {
+ public:
+  bool test(const Value& current) const { return true; }
+  void step(Value& current) const { ++current; }
 };
 
 /**

diff --git a/folly/gen/Base.h b/folly/gen/Base.h
@@ -256,9 +256,24 @@ template<class Value,
          class Container = std::vector<typename std::decay<Value>::type>>
 class CopiedSource;
 
-template<class Value, bool endless = false, bool endInclusive = false>
+template<class Value, class SequenceImpl>
 class Sequence;
 
+template <class Value>
+class RangeImpl;
+
+template <class Value, class Distance>
+class RangeWithStepImpl;
+
+template <class Value>
+class SeqImpl;
+
+template <class Value, class Distance>
+class SeqWithStepImpl;
+
+template <class Value>
+class InfiniteImpl;
+
 template<class Value, class Source>
 class Yield;
 
@@ -396,21 +411,36 @@ From from(Container&& source) {
   return From(std::move(source));
 }
 
-template<class Value, class Gen = detail::Sequence<Value, false, false>>
+template<class Value, class Impl = detail::RangeImpl<Value>,
+         class Gen = detail::Sequence<Value, Impl>>
 Gen range(Value begin, Value end) {
-  return Gen(begin, end);
+  return Gen{std::move(begin), Impl{std::move(end)}};
 }
 
-template<class Value,
-         class Gen = detail::Sequence<Value, false, true>>
+template<class Value, class Distance,
+         class Impl = detail::RangeWithStepImpl<Value, Distance>,
+         class Gen = detail::Sequence<Value, Impl>>
+Gen range(Value begin, Value end, Distance step) {
+  return Gen{std::move(begin), Impl{std::move(end), std::move(step)}};
+}
+
+template<class Value, class Impl = detail::SeqImpl<Value>,
+         class Gen = detail::Sequence<Value, Impl>>
 Gen seq(Value first, Value last) {
-  return Gen(first, last);
+  return Gen{std::move(first), Impl{std::move(last)}};
 }
 
-template<class Value,
-         class Gen = detail::Sequence<Value, true>>
-Gen seq(Value begin) {
-  return Gen(begin);
+template<class Value, class Distance,
+         class Impl = detail::SeqWithStepImpl<Value, Distance>,
+         class Gen = detail::Sequence<Value, Impl>>
+Gen seq(Value first, Value last, Distance step) {
+  return Gen{std::move(first), Impl{std::move(last), std::move(step)}};
+}
+
+template<class Value, class Impl = detail::InfiniteImpl<Value>,
+         class Gen = detail::Sequence<Value, Impl>>
+Gen seq(Value first) {
+  return Gen{std::move(first), Impl{}};
 }
 
 template<class Value,

diff --git a/folly/gen/test/BaseTest.cpp b/folly/gen/test/BaseTest.cpp
@@ -186,13 +186,28 @@ TEST(Gen, Seq) {
   }
 }
 
+TEST(Gen, SeqWithStep) {
+  EXPECT_EQ(75, seq(5, 25, 5) | sum);
+}
+
+TEST(Gen, SeqWithStepArray) {
+  const std::array<int, 6> arr{{1, 2, 3, 4, 5, 6}};
+  EXPECT_EQ(9, seq(&arr[0], &arr[5], 2)
+             | map([](const int *i) { return *i; })
+             | sum);
+}
+
 TEST(Gen, Range) {
   // cover the fenceposts of the loop unrolling
   for (int n = 1; n < 100; ++n) {
     EXPECT_EQ(gen::range(0, n) | count, n);
   }
 }
 
+TEST(Gen, RangeWithStep) {
+  EXPECT_EQ(50, range(5, 25, 5) | sum);
+}
+
 TEST(Gen, FromIterators) {
   vector<int> source {2, 3, 5, 7, 11};
   auto gen = from(folly::range(source.begin() + 1, source.end() - 1));