Make take_while() work with iterables

Author: tcbrindle

include/flux/adaptor/take_while.hpp

@@ -12,7 +12,7 @@ namespace flux {
 
 namespace detail {
 
-template <sequence Base, typename Pred>
+template <iterable Base, typename Pred>
 struct take_while_adaptor : inline_sequence_base<take_while_adaptor<Base, Pred>> {
 private:
     Base base_;
@@ -34,19 +34,43 @@ struct take_while_adaptor : inline_sequence_base<take_while_adaptor<Base, Pred>>
 };
 
 struct take_while_fn {
-    template <adaptable_sequence Seq, std::move_constructible Pred>
-        requires std::predicate<Pred&, element_t<Seq>>
+    template <sink_iterable It, std::move_constructible Pred>
+        requires std::predicate<Pred&, element_t<It>>
     [[nodiscard]]
-    constexpr auto operator()(Seq&& seq, Pred pred) const
+    constexpr auto operator()(It&& it, Pred pred) const
     {
-        return take_while_adaptor<std::decay_t<Seq>, Pred>(
-                    FLUX_FWD(seq), std::move(pred));
+        return take_while_adaptor<std::decay_t<It>, Pred>(
+                    FLUX_FWD(it), std::move(pred));
     }
 };
 
 } // namespace detail
 
-template <typename Base, typename Pred>
+template <iterable Base, typename Pred>
+struct sequence_traits<detail::take_while_adaptor<Base, Pred>>
+    : default_sequence_traits {
+
+    static consteval auto element_type(auto& self)
+        -> element_t<decltype((self.base_))>;
+
+    static constexpr auto iterate(auto& self, auto&& iter_pred) -> bool
+    {
+        bool done = false;
+        bool res = flux::iterate(self.base_, [&](auto&& elem) {
+            if (!std::invoke(self.pred_, elem)) {
+                done = true;
+                return false; // break
+            } else {
+                return std::invoke(iter_pred, FLUX_FWD(elem));
+            }
+        });
+        // Return true if take_while was exhausted
+        return res ? res : done;
+    }
+
+};
+
+template <sequence Base, typename Pred>
 struct sequence_traits<detail::take_while_adaptor<Base, Pred>>
     : detail::passthrough_traits_base
 {
@@ -57,7 +81,21 @@ struct sequence_traits<detail::take_while_adaptor<Base, Pred>>
     static constexpr bool is_infinite = false;
 
     using default_sequence_traits::element_type;
-    using default_sequence_traits::iterate;
+
+    static constexpr auto iterate(auto& self, auto&& iter_pred) -> bool
+    {
+        bool done = false;
+        bool res = flux::iterate(self.base_, [&](auto&& elem) {
+            if (!std::invoke(self.pred_, elem)) {
+                done = true;
+                return false; // break
+            } else {
+                return std::invoke(iter_pred, FLUX_FWD(elem));
+            }
+        });
+        // Return true if take_while was exhausted
+        return res ? res : done;
+    }
 
     template <typename Self>
     static constexpr bool is_last(Self& self, cursor_t<Self> const& cur)

test/test_take_while.cpp

@@ -4,6 +4,7 @@
 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
 
 #include <array>
+#include <ranges>
 
 #include "test_utils.hpp"
 
@@ -70,6 +71,23 @@ constexpr bool test_take_while()
         STATIC_CHECK(check_equal(seq, {4, 16, 36, 64, 100}));
     }
 
+    // Test with non-sequence iterables
+    {
+        auto view = std::views::filter(std::array{1, 2, 3, 4, 5}, flux::pred::true_);
+
+        auto taken = flux::take_while(std::move(view), flux::pred::lt(4));
+
+        using T = decltype(taken);
+
+        static_assert(flux::iterable<T>);
+        static_assert(not flux::sized_iterable<T>);
+        static_assert(not flux::sequence<T>);
+
+        STATIC_CHECK(flux::contains(taken, 2));
+        STATIC_CHECK(not flux::contains(taken, 4));
+        STATIC_CHECK(check_equal(taken, {1, 2, 3}));
+    }
+
     return true;
 }
 static_assert(test_take_while());

test/test_utils.hpp

@@ -47,9 +47,10 @@ inline constexpr struct {
                               std::initializer_list<T> ilist) const
     {
         auto iter = ilist.begin();
-        return flux::iterate(it, [&iter](auto&& elem) {
+        bool r =  flux::iterate(it, [&iter](auto&& elem) {
             return *iter++ == elem;
         });
+        return r ? iter == ilist.end() : r;
     }
 
     template <typename T>
@@ -63,11 +64,12 @@ inline constexpr struct {
                               flux::sequence auto&& seq) const
     {
         auto cur = flux::first(seq);
-        return flux::iterate(it, [&](auto&& elem) {
+        bool k = flux::iterate(it, [&](auto&& elem) {
             bool r = (elem == flux::read_at(seq, cur));
             flux::inc(seq, cur);
             return r;
         });
+        return k ? flux::is_last(seq, cur) : k;
     }
 
     constexpr bool operator()(flux::sequence auto&& seq1,