Make array iterable

Author: danakj

CMakeLists.txt

@@ -29,6 +29,7 @@ add_library(subspace STATIC
     "assertions/unreachable.h"
     "assertions/windows.h"
     "containers/array.h"
+    "containers/array_iter.h"
     "construct/from.h"
     "construct/into.h"
     "construct/make_default.h"
@@ -39,9 +40,10 @@ add_library(subspace STATIC
     "fn/fn_bind.h"
     "fn/fn_defn.h"
     "fn/fn_impl.h"
+    "iter/__private/iterator_end.h"
+    "iter/__private/iterator_loop.h"
     "iter/filter.h"
     "iter/iterator.h"
-    "iter/iterator.cc"
     "iter/iterator_defn.h"
     "iter/iterator_impl.h"
     "iter/once.h"

containers/array.h

@@ -22,6 +22,7 @@
 
 #include "assertions/check.h"
 #include "construct/make_default.h"
+#include "containers/array_iter.h"
 #include "fn/callable.h"
 #include "marker/unsafe.h"
 #include "mem/move.h"
@@ -125,6 +126,19 @@ class Array final {
     return storage_.data_;
   }
 
+  ::sus::iter::Iterator<ArrayIter<T, N>> iter() const& {
+    return ArrayIter<T, N>::with(*this);
+  }
+  ::sus::iter::Iterator<ArrayIter<T, N>> iter() && = delete;
+
+  ::sus::iter::Iterator<ArrayIterMut<T, N>> iter_mut() & {
+    return ArrayIterMut<T, N>::with(*this);
+  }
+
+  ::sus::iter::Iterator<ArrayIntoIter<T, N>> into_iter() && {
+    return ArrayIntoIter<T, N>::with(static_cast<Array&&>(*this));
+  }
+
   /// sus::num::Eq<Array<U, N>> trait.
   template <::sus::num::Eq<T> U>
   constexpr bool operator==(const Array<U, N>& r) const& noexcept {
@@ -222,6 +236,10 @@ constexpr inline auto operator<=>(const Array<T, N>& l,
                               std::make_index_sequence<N>());
 }
 
+// Implicit for-ranged loop iteration via `Array::iter()`.
+using ::sus::iter::__private::begin;
+using ::sus::iter::__private::end;
+
 }  // namespace sus::containers
 
 namespace sus {

containers/array_iter.h

@@ -0,0 +1,97 @@
+// 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.
+
+#pragma once
+
+#include <stdint.h>
+
+#include <type_traits>
+
+#include "iter/iterator_defn.h"
+#include "mem/move.h"
+#include "num/unsigned_integer.h"
+
+namespace sus::containers {
+
+template <class T, size_t N>
+  requires(N <= PTRDIFF_MAX)
+class Array;
+
+template <class Item, size_t N>
+  requires(std::is_move_constructible_v<Item>)
+struct ArrayIter : public ::sus::iter::IteratorBase<const Item&> {
+ public:
+  static constexpr auto with(const Array<Item, N>& array) noexcept {
+    return ::sus::iter::Iterator<ArrayIter>(array);
+  }
+
+  Option<const Item&> next() noexcept final {
+    if (next_index.primitive_value == N) return Option<const Item&>::none();
+    return Option<const Item&>::some(
+        array.get(::sus::mem::replace(mref(next_index), next_index + 1_usize)));
+  }
+
+ protected:
+  ArrayIter(const Array<Item, N>& array) noexcept : array(array) {}
+
+ private:
+  usize next_index = 0_usize;
+  const Array<Item, N>& array;
+};
+
+template <class Item, size_t N>
+  requires(std::is_move_constructible_v<Item>)
+struct ArrayIterMut : public ::sus::iter::IteratorBase<Item&> {
+ public:
+  static constexpr auto with(Array<Item, N>& array) noexcept {
+    return ::sus::iter::Iterator<ArrayIterMut>(array);
+  }
+
+  Option<Item&> next() noexcept final {
+    if (next_index.primitive_value == N) return Option<Item&>::none();
+    return Option<Item&>::some(mref(array.get_mut(
+        ::sus::mem::replace(mref(next_index), next_index + 1_usize))));
+  }
+
+ protected:
+  ArrayIterMut(Array<Item, N>& array) noexcept : array(array) {}
+
+ private:
+  usize next_index = 0_usize;
+  Array<Item, N>& array;
+};
+
+template <class Item, size_t N>
+  requires(std::is_move_constructible_v<Item>)
+struct ArrayIntoIter : public ::sus::iter::IteratorBase<Item> {
+ public:
+  static constexpr auto with(Array<Item, N>&& array) noexcept {
+    return ::sus::iter::Iterator<ArrayIntoIter>(move(array));
+  }
+
+  Option<Item> next() noexcept final {
+    if (next_index.primitive_value == N) return Option<Item>::none();
+    return Option<Item>::some(move(array.get_mut(
+        ::sus::mem::replace(mref(next_index), next_index + 1_usize))));
+  }
+
+ protected:
+  ArrayIntoIter(Array<Item, N>&& array) noexcept : array(move(array)) {}
+
+ private:
+  usize next_index = 0_usize;
+  Array<Item, N> array;
+};
+
+}  // namespace sus::containers

containers/array_unittest.cc

@@ -17,7 +17,9 @@
 #include <type_traits>
 
 #include "construct/into.h"
+#include "iter/iterator.h"
 #include "marker/unsafe.h"
+#include "mem/move.h"
 #include "mem/relocate.h"
 #include "num/types.h"
 #include "third_party/googletest/googletest/include/gtest/gtest.h"
@@ -255,4 +257,47 @@ TEST(Array, PartialOrder) {
   EXPECT_EQ(std::partial_order(a, b), std::partial_ordering::less);
 }
 
+TEST(Array, Iter) {
+  const auto a = Array<usize, 5>::with_value(3u);
+  auto sum = 0_usize;
+  for (const usize& i : a.iter()) {
+    sum += i;
+  }
+  EXPECT_EQ(sum, 15_usize);
+}
+
+TEST(Array, IterMut) {
+  auto a = Array<usize, 5>::with_value(3u);
+  auto sum = 0_usize;
+  for (usize& i : a.iter_mut()) {
+    sum += i;
+    i += 1_usize;
+  }
+  EXPECT_EQ(sum, 15_usize);
+
+  sum = 0_usize;
+  for (const usize& i : a.iter()) {
+    sum += i;
+  }
+  EXPECT_EQ(sum, 20_usize);
+}
+
+TEST(Array, IntoIter) {
+  auto a = Array<usize, 5>::with_value(3u);
+  auto sum = 0_usize;
+  for (usize i : sus::move(a).into_iter()) {
+    sum += i;
+  }
+  EXPECT_EQ(sum, 15_usize);
+}
+
+TEST(Array, ImplicitIter) {
+  const auto a = Array<usize, 5>::with_value(3u);
+  auto sum = 0_usize;
+  for (const usize& i : a) {
+    sum += i;
+  }
+  EXPECT_EQ(sum, 15_usize);
+}
+
 }  // namespace

iter/iterator.cc renamed to iter/__private/iterator_end.h

@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "iter/iterator.h"
+#pragma once
 
-namespace sus::iter {
-const IteratorEnd iterator_end;
-}
+namespace sus::iter::__private {
+
+struct IteratorEnd {};
+
+}  // namespace sus::iter::__private

iter/__private/iterator_loop.h

@@ -0,0 +1,57 @@
+// 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.
+
+#pragma once
+
+#include "option/option.h"
+
+namespace sus::iter::__private {
+
+/// An adaptor for range-based for loops.
+template <class Iterator>
+class IteratorLoop final {
+  using Item = typename std::remove_reference_t<Iterator>::Item;
+
+ public:
+  IteratorLoop(Iterator iter)
+      : iter_(static_cast<Iterator&&>(iter)), item_(iter_.next()) {}
+
+  inline bool operator==(const __private::IteratorEnd&) const {
+    return item_.is_nome();
+  }
+  inline bool operator!=(const __private::IteratorEnd&) const {
+    return item_.is_some();
+  }
+  inline void operator++() & { item_ = iter_.next(); }
+  inline Item operator*() & { return item_.take().unwrap(); }
+
+ private:
+  /* TODO: NonNull<IteratorBase<Item>> */ Iterator iter_;
+  Option<Item> item_;
+};
+
+// ADL helpers to call T::iter() in a range-based for loop, which will call
+// `begin(T)`.
+
+template <class T>
+constexpr auto begin(const T& t) noexcept {
+  return IteratorLoop(t.iter());
+}
+
+template <class T>
+constexpr auto end(const T&) noexcept {
+  return ::sus::iter::__private::IteratorEnd();
+}
+
+}  // namespace sus::iter::__private

iter/iterator_defn.h

@@ -15,8 +15,8 @@
 #pragma once
 
 #include "fn/fn_defn.h"
-#include "option/option.h"
 #include "num/unsigned_integer.h"
+#include "option/option.h"
 
 namespace sus::iter {
 
@@ -26,10 +26,20 @@ using ::sus::option::Option;
 template <class Item, size_t InnerIterSize, size_t InnerIterAlign>
 class Filter;
 
+namespace __private {
 template <class IteratorBase>
 class IteratorLoop;
+
+template <class Iterator>
+class IteratorImplicitLoop;
 struct IteratorEnd;
 
+template <class T>
+constexpr auto begin(const T& t) noexcept;
+template <class T>
+constexpr auto end(const T& t) noexcept;
+}  // namespace __private
+
 // TODO: Do we want to be able to pass IteratorBase& as a "generic" iterator?
 // Then it needs access to the adapator methods of Iterator<T>, so make them
 // virtual methods on IteratorBase?
@@ -89,9 +99,9 @@ class IteratorBase {
   virtual usize count() noexcept;
 
   /// Adaptor for use in ranged for loops.
-  IteratorLoop<IteratorBase<Item>> begin() & noexcept;
+  __private::IteratorLoop<IteratorBase<Item>&> begin() & noexcept;
   /// Adaptor for use in ranged for loops.
-  IteratorEnd end() & noexcept;
+  __private::IteratorEnd end() & noexcept;
 
  protected:
   IteratorBase() = default;
@@ -101,7 +111,9 @@ template <class I>
 class Iterator final : public I {
  private:
   using typename I::Item;
-  friend I;  // I::foo() can construct Iterator<I>.
+  template <class T>
+  friend class __private::IteratorLoop;  // Can see Item.
+  friend I;                              // I::foo() can construct Iterator<I>.
 
   template <class J>
   friend class Iterator;  // Iterator<J>::foo() can construct Iterator<I>.

iter/iterator_impl.h

@@ -15,6 +15,8 @@
 #pragma once
 
 #include "fn/fn.h"
+#include "iter/__private/iterator_end.h"
+#include "iter/__private/iterator_loop.h"
 #include "iter/iterator_defn.h"
 #include "iter/sized_iterator.h"
 #include "mem/move.h"
@@ -25,35 +27,15 @@
 
 namespace sus::iter {
 
-struct IteratorEnd {};
-extern const IteratorEnd iterator_end;
-
-/// An adaptor for range-based for loops.
-template <class IteratorBase>
-class IteratorLoop final {
-  using Item = typename IteratorBase::Item;
-
- public:
-  IteratorLoop(IteratorBase& iter) : iter_(iter), item_(iter_.next()) {}
-
-  inline bool operator==(const IteratorEnd&) const { return item_.is_nome(); }
-  inline bool operator!=(const IteratorEnd&) const { return item_.is_some(); }
-  inline void operator++() & { item_ = iter_.next(); }
-  inline Item operator*() & { return item_.take().unwrap(); }
-
- private:
-  /* TODO: NonNull<IteratorBase<Item>> */ IteratorBase& iter_;
-  Option<Item> item_;
-};
-
 template <class Item>
-IteratorLoop<IteratorBase<Item>> IteratorBase<Item>::begin() & noexcept {
+__private::IteratorLoop<IteratorBase<Item>&>
+IteratorBase<Item>::begin() & noexcept {
   return {*this};
 }
 
 template <class Item>
-IteratorEnd IteratorBase<Item>::end() & noexcept {
-  return iterator_end;
+__private::IteratorEnd IteratorBase<Item>::end() & noexcept {
+  return __private::IteratorEnd();
 }
 
 template <class Item>

mem/mref.h

@@ -86,7 +86,7 @@ struct Mref<T&> final {
   constexpr Mref& operator=(Mref&&) noexcept = default;
 
   // Prevent constructing an Mref argument without writing mref().
-  Mref(T& t) = delete;
+  consteval Mref(T& t) = delete;
   // Prevent passing an Mref argument along without writing mref() again.
   Mref(Mref&) = delete;