Add Iterator::try_collect()

sus/CMakeLists.txt

@@ -115,6 +115,7 @@ target_sources(subspace PUBLIC
     "iter/iterator.h"
     "iter/iterator_concept.h"
     "iter/iterator_defn.h"
+    "iter/iterator_impl.h"
     "iter/iterator_ref.h"
     "iter/once.h"
     "iter/product.h"
@@ -123,6 +124,7 @@ target_sources(subspace PUBLIC
     "iter/size_hint.h"
     "iter/size_hint_impl.h"
     "iter/successors.h"
+    "iter/try_from_iterator.h"
     "iter/zip.h"
     "macros/__private/compiler_bugs.h"
     "macros/assume.h"

sus/iter/from_iterator.h

@@ -50,9 +50,10 @@ concept FromIterator =
 /// typically called through calling `collect()` on an iterator. However this
 /// function can be preferrable for some readers, especially in generic template
 /// code.
-template <class ToType, ::sus::iter::IntoIteratorAny IntoIter>
-  requires(
-      FromIterator<ToType, typename IntoIteratorOutputType<IntoIter>::Item>)
+template <class ToType, IntoIteratorAny IntoIter>
+  requires(FromIterator<ToType,
+                        typename IntoIteratorOutputType<IntoIter>::Item> &&  //
+           ::sus::mem::IsMoveRef<IntoIter &&>)
 constexpr inline ToType from_iter(IntoIter&& into_iter) noexcept {
   return FromIteratorImpl<ToType>::from_iter(
       ::sus::move(into_iter).into_iter());

sus/iter/iterator.h

@@ -16,6 +16,7 @@
 
 // IWYU pragma: begin_exports
 #include "sus/iter/iterator_defn.h"
+#include "sus/iter/iterator_impl.h"
 
 // The usize formatting is in unsigned_integer_impl.h which has an include cycle
 // with usize->Array->Array iterators->SizeHint->usize. So the SizeHint

sus/iter/iterator_defn.h

@@ -927,6 +927,57 @@ class IteratorBase {
       ::sus::fn::FnMut<bool(const std::remove_reference_t<Item>&)> auto
           pred) && noexcept;
 
+  /// Fallibly transforms an iterator into a collection, short circuiting if a
+  /// failure is encountered.
+  ///
+  /// `try_collect()` is a variation of `collect()` that allows fallible
+  /// conversions during collection. Its main use case is simplifying
+  /// conversions from iterators yielding `Option<T>` into
+  /// `Option<Collection<T>>`, or similarly for other Try types (e.g. `Result`
+  /// or `std::optional`).
+  ///
+  /// Importantly, `try_collect()` doesn’t require that the outer `Try` type
+  /// also implements FromIterator; only the `Try` type's `Output` type must
+  /// implement it. Concretely, this means that collecting into
+  /// `TryThing<Vec<i32>, _>` can be valid because `Vec<i32>` implements
+  /// FromIterator, even if `TryThing` doesn’t.
+  ///
+  /// Also, if a failure is encountered during `try_collect()`, the iterator is
+  /// still valid and may continue to be used, in which case it will continue
+  /// iterating starting after the element that triggered the failure. See the
+  /// last example below for an example of how this works.
+  ///
+  /// # Examples
+  /// Successfully collecting an iterator of `Option<i32>` into
+  /// `Option<Vec<i32>>`:
+  /// ```
+  /// auto u = Vec<Option<i32>>::with(some(1), some(2), some(3));
+  /// auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+  /// sus::check(v == some(Vec<i32>::with(1, 2, 3 )));
+  /// ```
+  /// Failing to collect in the same way:
+  /// ```
+  /// auto u = Vec<Option<i32>>::with(some(1), some(2), none(), some(3));
+  /// auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+  /// sus::check(v == none());
+  /// ```
+  /// A similar example, but with [`Result`](sus::result::Result):
+  /// ```
+  /// enum Error { ERROR };
+  /// auto u = Vec<Result<i32, Error>>::with(ok(1), ok(2), ok(3));
+  /// auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+  /// sus::check(v == ok(Vec<i32>::with(1, 2, 3)));
+  /// auto w = Vec<Result<i32, Error>>::with(ok(1), ok(2), err(ERROR), ok(3));
+  /// auto x = sus::move(w).into_iter().try_collect<Vec<i32>>();
+  /// sus::check(x == err(ERROR));
+  /// ```
+  template <class C>
+    requires(::sus::ops::Try<ItemT> &&  //
+             FromIterator<C, ::sus::ops::TryOutputType<ItemT>> &&
+             // Void can not be collected from.
+             !std::is_void_v<::sus::ops::TryOutputType<ItemT>>)
+  constexpr auto try_collect() noexcept;
+
   /// This function acts like `fold()` but the closure returns a type that
   /// satisfies `sus::ops::Try` and which converts to the accumulator type on
   /// success through the Try concept. If the closure ever returns failure, the
@@ -1059,7 +1110,7 @@ class IteratorBase {
   /// sus::move(iter).collect<MyContainer<i32>>()
   /// ```
   template <FromIterator<ItemT> C>
-  constexpr FromIterator<ItemT> auto collect() && noexcept;
+  constexpr C collect() && noexcept;
 
   /// Transforms an iterator into a Vec.
   ///
@@ -1433,7 +1484,8 @@ template <class Iter, class Item>
 template <
     ::sus::fn::FnMut<::sus::fn::NonVoid(const std::remove_reference_t<Item>&)>
         KeyFn,
-    int&..., class Key>
+    int&...,
+    class Key>
   requires(::sus::ops::Ord<Key> &&  //
            !std::is_reference_v<Key>)
 constexpr Option<Item> IteratorBase<Iter, Item>::max_by_key(
@@ -1490,7 +1542,8 @@ template <class Iter, class Item>
 template <
     ::sus::fn::FnMut<::sus::fn::NonVoid(const std::remove_reference_t<Item>&)>
         KeyFn,
-    int&..., class Key>
+    int&...,
+    class Key>
   requires(::sus::ops::Ord<Key> &&  //
            !std::is_reference_v<Key>)
 constexpr Option<Item> IteratorBase<Iter, Item>::min_by_key(
@@ -1699,7 +1752,8 @@ constexpr Option<usize> IteratorBase<Iter, Item>::rposition(
 
 template <class Iter, class Item>
 template <class State, ::sus::fn::FnMut<::sus::fn::NonVoid(State&, Item&&)> F,
-          int&..., class R, class InnerR>
+          int&..., class R,
+          class InnerR>
   requires(::sus::option::__private::IsOptionType<R>::value &&  //
            !std::is_reference_v<State>)
 constexpr Iterator<InnerR> auto IteratorBase<Iter, Item>::scan(
@@ -1859,14 +1913,13 @@ constexpr std::weak_ordering IteratorBase<Iter, Item>::weak_cmp_by(
 
 template <class Iter, class Item>
 template <FromIterator<Item> C>
-constexpr FromIterator<Item> auto
-IteratorBase<Iter, Item>::collect() && noexcept {
+constexpr C IteratorBase<Iter, Item>::collect() && noexcept {
   return from_iter<C>(static_cast<Iter&&>(*this));
 }
 
 template <class Iter, class Item>
-::sus::containers::Vec<Item> constexpr IteratorBase<
-    Iter, Item>::collect_vec() && noexcept {
+constexpr ::sus::containers::Vec<Item>
+IteratorBase<Iter, Item>::collect_vec() && noexcept {
   return from_iter<::sus::containers::Vec<Item>>(static_cast<Iter&&>(*this));
 }
 

sus/iter/iterator_impl.h

@@ -0,0 +1,38 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// IWYU pragma: private, include "sus/iter/iterator.h"
+// IWYU pragma: friend "sus/.*"
+#pragma once
+
+#include <type_traits>
+
+#include "sus/iter/try_from_iterator.h"
+
+namespace sus::iter {
+
+// Implementation of IteratorBase::try_collect has to live here to avoid
+// cyclical includes, as the impl of try_from_iter() needs to see IteratorBase
+// in order to make something that can be consumed by FromIterator.
+template <class Iter, class Item>
+template <class C>
+  requires(::sus::ops::Try<Item> &&  //
+           FromIterator<C, ::sus::ops::TryOutputType<Item>> &&
+           // Void can not be collected from.
+           !std::is_void_v<::sus::ops::TryOutputType<Item>>)
+constexpr auto IteratorBase<Iter, Item>::try_collect() noexcept {
+  return try_from_iter<C>(static_cast<Iter&&>(*this));
+}
+
+}  // namespace sus::iter

sus/iter/iterator_unittest.cc

@@ -33,10 +33,13 @@
 #include "sus/num/overflow_integer.h"
 #include "sus/ops/eq.h"
 #include "sus/prelude.h"
+#include "sus/test/no_copy_move.h"
 
 using sus::containers::Array;
 using sus::iter::IteratorBase;
 using sus::option::Option;
+using sus::result::Result;
+using sus::test::NoCopyMove;
 
 namespace sus::test::iter {
 
@@ -45,6 +48,10 @@ struct CollectSum {
   T sum;
 };
 
+struct CollectRefs {
+  sus::Vec<const NoCopyMove*> vec;
+};
+
 }  // namespace sus::test::iter
 
 template <class T>
@@ -57,6 +64,16 @@ struct sus::iter::FromIteratorImpl<sus::test::iter::CollectSum<T>> {
   }
 };
 
+template <>
+struct sus::iter::FromIteratorImpl<sus::test::iter::CollectRefs> {
+  static sus::test::iter::CollectRefs from_iter(
+      sus::iter::IntoIterator<const NoCopyMove&> auto iter) noexcept {
+    auto v = sus::Vec<const NoCopyMove*>();
+    for (const NoCopyMove& t : sus::move(iter).into_iter()) v.push(&t);
+    return sus::test::iter::CollectRefs(sus::move(v));
+  }
+};
+
 namespace {
 using namespace sus::test::iter;
 
@@ -460,6 +477,10 @@ TEST(Iterator, Collect) {
       ArrayIterator<i32, 5>::with_array(nums).collect<CollectSum<i32>>();
   EXPECT_EQ(collected.sum, 1 + 2 + 3 + 4 + 5);
 
+  auto from = sus::iter::from_iter<CollectSum<i32>>(
+      ArrayIterator<i32, 5>::with_array(nums));
+  EXPECT_EQ(from.sum, 1 + 2 + 3 + 4 + 5);
+
   static_assert(sus::Array<i32, 5>::with(1, 2, 3, 4, 5)
                     .into_iter()
                     .collect<CollectSum<i32>>()
@@ -481,6 +502,97 @@ TEST(Iterator, CollectVec) {
       sus::Vec<i32>::with(1, 2, 3, 4, 5));
 }
 
+TEST(Iterator, TryCollect) {
+  // Option.
+  {
+    auto collected = sus::Array<Option<i32>, 3>::with(
+                         ::sus::some(1), ::sus::some(2), ::sus::some(3))
+                         .into_iter()
+                         .try_collect<Vec<i32>>();
+    static_assert(std::same_as<decltype(collected), Option<Vec<i32>>>);
+    EXPECT_EQ(collected.is_some(), true);
+    EXPECT_EQ(collected.as_value(), sus::Vec<i32>::with(1, 2, 3));
+
+    auto it = sus::Array<Option<i32>, 3>::with(::sus::some(1), ::sus::none(),
+                                               ::sus::some(3))
+                  .into_iter();
+    auto up_to_none = it.try_collect<Vec<i32>>();
+    EXPECT_EQ(up_to_none, sus::none());
+    auto after_none = it.try_collect<Vec<i32>>();
+    EXPECT_EQ(after_none.as_value(), sus::Vec<i32>::with(3));
+
+    NoCopyMove n[3];
+    auto refs = sus::Array<Option<const NoCopyMove&>, 3>::with(
+                    ::sus::some(n[0]), ::sus::some(n[1]), ::sus::some(n[2]))
+                    .into_iter()
+                    .try_collect<CollectRefs>();
+    EXPECT_EQ(refs.is_some(), true);
+    EXPECT_EQ(refs.as_value().vec[0u], &n[0]);
+    EXPECT_EQ(refs.as_value().vec[1u], &n[1]);
+    EXPECT_EQ(refs.as_value().vec[2u], &n[2]);
+  }
+  // Result.
+  enum Error { ERROR };
+  {
+    auto collected = sus::Array<Result<i32, Error>, 3>::with(
+                         ::sus::ok(1), ::sus::ok(2), ::sus::ok(3))
+                         .into_iter()
+                         .try_collect<Vec<i32>>();
+    static_assert(std::same_as<decltype(collected), Result<Vec<i32>, Error>>);
+    EXPECT_EQ(collected.as_ok(), sus::Vec<i32>::with(1, 2, 3));
+
+    auto it = sus::Array<Result<i32, Error>, 3>::with(
+                  ::sus::ok(1), ::sus::err(ERROR), ::sus::ok(3))
+                  .into_iter();
+    auto up_to_err = it.try_collect<Vec<i32>>();
+    EXPECT_EQ(up_to_err, sus::err(ERROR));
+    auto after_err = it.try_collect<Vec<i32>>();
+    EXPECT_EQ(after_err.as_ok(), sus::Vec<i32>::with(3));
+  }
+
+  auto from =
+      sus::iter::try_from_iter<Vec<i32>>(sus::Array<Option<i32>, 3>::with(
+          ::sus::some(1), ::sus::some(2), ::sus::some(3)));
+  EXPECT_EQ(from.as_value(), sus::Vec<i32>::with(1, 2, 3));
+
+  static_assert(
+      sus::Array<Option<i32>, 3>::with(sus::some(1), sus::some(2), sus::some(3))
+          .into_iter()
+          .try_collect<Vec<i32>>()
+          .unwrap()
+          .into_iter()
+          .sum() == 1 + 2 + 3);
+}
+
+TEST(Iterator, TryCollect_Example) {
+  using sus::none;
+  using sus::ok;
+  using sus::err;
+  using sus::Option;
+  using sus::result::Result;
+  using sus::some;
+  using sus::Vec;
+  {
+    auto u = Vec<Option<i32>>::with(some(1), some(2), some(3));
+    auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+    sus::check(v == some(Vec<i32>::with(1, 2, 3)));
+  }
+  {
+    auto u = Vec<Option<i32>>::with(some(1), some(2), none(), some(3));
+    auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+    sus::check(v == none());
+  }
+  {
+    enum Error { ERROR };
+    auto u = Vec<Result<i32, Error>>::with(ok(1), ok(2), ok(3));
+    auto v = sus::move(u).into_iter().try_collect<Vec<i32>>();
+    sus::check(v == ok(Vec<i32>::with(1, 2, 3)));
+    auto w = Vec<Result<i32, Error>>::with(ok(1), ok(2), err(ERROR), ok(3));
+    auto x = sus::move(w).into_iter().try_collect<Vec<i32>>();
+    sus::check(x == err(ERROR));
+  }
+}
+
 TEST(Iterator, Rev) {
   i32 nums[5] = {1, 2, 3, 4, 5};