[libcxx-commits] [libcxx] [libc++] `std::ranges::advance`: avoid unneeded bounds checks when advancing iterator (PR #84126)

[Jan Kokemüller via libcxx-commits]

https://github.com/jiixyj updated https://github.com/llvm/llvm-project/pull/84126

>From dc039616cee1453a38729a4f2f6adc03fd93549c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Kokem=C3=BCller?= <jan.kokemueller at gmail.com>
Date: Wed, 6 Mar 2024 07:49:25 +0100
Subject: [PATCH 1/2] skip unneeded checks against '__bound_sentinel' when
 advancing iterator

---
 libcxx/include/__iterator/advance.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/libcxx/include/__iterator/advance.h b/libcxx/include/__iterator/advance.h
index 7959bdeae32643..296db1aaab6526 100644
--- a/libcxx/include/__iterator/advance.h
+++ b/libcxx/include/__iterator/advance.h
@@ -170,14 +170,14 @@ struct __fn {
     } else {
       // Otherwise, if `n` is non-negative, while `bool(i != bound_sentinel)` is true, increments `i` but at
       // most `n` times.
-      while (__i != __bound_sentinel && __n > 0) {
+      while (__n > 0 && __i != __bound_sentinel) {
         ++__i;
         --__n;
       }
 
       // Otherwise, while `bool(i != bound_sentinel)` is true, decrements `i` but at most `-n` times.
       if constexpr (bidirectional_iterator<_Ip> && same_as<_Ip, _Sp>) {
-        while (__i != __bound_sentinel && __n < 0) {
+        while (__n < 0 && __i != __bound_sentinel) {
           --__i;
           ++__n;
         }

>From b23e7fe87d6eaa65cf684749e2eb1ceee51176d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Kokem=C3=BCller?= <jan.kokemueller at gmail.com>
Date: Sat, 9 Mar 2024 11:22:59 +0100
Subject: [PATCH 2/2] add tests

---
 .../iterator_count_sentinel.pass.cpp          | 22 ++++++++++++++++++
 libcxx/test/support/test_iterators.h          | 23 +++++++++++++++++--
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/libcxx/test/std/iterators/iterator.primitives/range.iter.ops/range.iter.ops.advance/iterator_count_sentinel.pass.cpp b/libcxx/test/std/iterators/iterator.primitives/range.iter.ops/range.iter.ops.advance/iterator_count_sentinel.pass.cpp
index a1c15640182162..72584e60d9f89f 100644
--- a/libcxx/test/std/iterators/iterator.primitives/range.iter.ops/range.iter.ops.advance/iterator_count_sentinel.pass.cpp
+++ b/libcxx/test/std/iterators/iterator.primitives/range.iter.ops/range.iter.ops.advance/iterator_count_sentinel.pass.cpp
@@ -42,6 +42,13 @@ constexpr void check_forward(int* first, int* last, std::iter_difference_t<It> n
     // regardless of the iterator category.
     assert(it.stride_count() == M);
     assert(it.stride_displacement() == M);
+    if (n == 0) {
+      assert(it.equals_count() == 0);
+    } else {
+      assert(it.equals_count() > 0);
+      assert(it.equals_count() == M || it.equals_count() == M + 1);
+      assert(it.equals_count() <= n);
+    }
   }
 }
 
@@ -95,6 +102,7 @@ constexpr void check_backward(int* first, int* last, std::iter_difference_t<It>
     (void)std::ranges::advance(it, n, sent);
     assert(it.stride_count() <= 1);
     assert(it.stride_displacement() <= 1);
+    assert(it.equals_count() == 0);
   }
 }
 
@@ -213,6 +221,20 @@ constexpr bool test() {
     assert(i == iota_iterator{INT_MIN+1});
   }
 
+  // Check that we don't do an unneeded bounds check when decrementing a
+  // `bidirectional_iterator` that doesn't model `sized_sentinel_for`.
+  {
+    static_assert(std::bidirectional_iterator<bidirectional_iterator<iota_iterator>>);
+    static_assert(!std::sized_sentinel_for<bidirectional_iterator<iota_iterator>,
+                                           bidirectional_iterator<iota_iterator>>);
+
+    auto it = stride_counting_iterator(bidirectional_iterator(iota_iterator{+1}));
+    auto sent = stride_counting_iterator(bidirectional_iterator(iota_iterator{-2}));
+    assert(std::ranges::advance(it, -3, sent) == 0);
+    assert(base(base(it)) == iota_iterator{-2});
+    assert(it.equals_count() == 3);
+  }
+
   return true;
 }
 
diff --git a/libcxx/test/support/test_iterators.h b/libcxx/test/support/test_iterators.h
index c92ce375348ff4..e551ab5a62d08b 100644
--- a/libcxx/test/support/test_iterators.h
+++ b/libcxx/test/support/test_iterators.h
@@ -725,11 +725,14 @@ struct common_input_iterator {
 #  endif // TEST_STD_VER >= 20
 
 // Iterator adaptor that counts the number of times the iterator has had a successor/predecessor
-// operation called. Has two recorders:
+// operation or an equality comparison operation called. Has three recorders:
 // * `stride_count`, which records the total number of calls to an op++, op--, op+=, or op-=.
 // * `stride_displacement`, which records the displacement of the calls. This means that both
 //   op++/op+= will increase the displacement counter by 1, and op--/op-= will decrease the
 //   displacement counter by 1.
+// * `equals_count`, which records the total number of calls to an op== or op!=. If compared
+//   against a sentinel object, that sentinel object must call the `record_equality_comparison`
+//   function so that the comparison is counted correctly.
 template <class It>
 class stride_counting_iterator {
 public:
@@ -754,6 +757,8 @@ class stride_counting_iterator {
 
     constexpr difference_type stride_displacement() const { return stride_displacement_; }
 
+    constexpr difference_type equals_count() const { return equals_count_; }
+
     constexpr decltype(auto) operator*() const { return *It(base_); }
 
     constexpr decltype(auto) operator[](difference_type n) const { return It(base_)[n]; }
@@ -838,9 +843,15 @@ class stride_counting_iterator {
         return base(x) - base(y);
     }
 
+    constexpr void record_equality_comparison() const
+    {
+        ++equals_count_;
+    }
+
     constexpr bool operator==(stride_counting_iterator const& other) const
         requires std::sentinel_for<It, It>
     {
+        record_equality_comparison();
         return It(base_) == It(other.base_);
     }
 
@@ -875,6 +886,7 @@ class stride_counting_iterator {
     decltype(base(std::declval<It>())) base_;
     difference_type stride_count_ = 0;
     difference_type stride_displacement_ = 0;
+    mutable difference_type equals_count_ = 0;
 };
 template <class It>
 stride_counting_iterator(It) -> stride_counting_iterator<It>;
@@ -887,7 +899,14 @@ class sentinel_wrapper {
 public:
     explicit sentinel_wrapper() = default;
     constexpr explicit sentinel_wrapper(const It& it) : base_(base(it)) {}
-    constexpr bool operator==(const It& other) const { return base_ == base(other); }
+    constexpr bool operator==(const It& other) const {
+      // If supported, record statistics about the equality operator call
+      // inside `other`.
+      if constexpr (requires { other.record_equality_comparison(); }) {
+        other.record_equality_comparison();
+      }
+      return base_ == base(other);
+    }
     friend constexpr It base(const sentinel_wrapper& s) { return It(s.base_); }
 private:
     decltype(base(std::declval<It>())) base_;