[ADT] Add llvm::range_size function for generic ranges

This function follows `std::ranges::size` from C++20. It is intended
mainly for generic code that does not know the exact range type.
I did not modify the existing `llvm::size` function because it has a strict
guarantee of O(1) running time, and we cannot guarantee that when we delegate
size check to user-defined size functions.

Use `range_size` to optimize size checks in `zip`* and `enumerate`
functions. Before that, we would have to perform linear scans for ranges
without random access iterators.

This is the last change I have planned in the series that overhauls
`zip`* and `enumerate`.

Reviewed By: dblaikie, zero9178

Differential Revision: https://reviews.llvm.org/D146231

llvm/include/llvm/ADT/STLExtras.h

@@ -77,6 +77,14 @@ constexpr void swap_impl(T &&lhs,
   swap(std::forward<T>(lhs), std::forward<T>(rhs));
 }
 
+using std::size;
+
+template <typename RangeT>
+constexpr auto size_impl(RangeT &&range)
+    -> decltype(size(std::forward<RangeT>(range))) {
+  return size(std::forward<RangeT>(range));
+}
+
 } // end namespace adl_detail
 
 /// Returns the begin iterator to \p range using `std::begin` and
@@ -103,6 +111,14 @@ constexpr void adl_swap(T &&lhs, T &&rhs) noexcept(
   adl_detail::swap_impl(std::forward<T>(lhs), std::forward<T>(rhs));
 }
 
+/// Returns the size of \p range using `std::size` and functions found through
+/// Argument-Dependent Lookup (ADL).
+template <typename RangeT>
+constexpr auto adl_size(RangeT &&range)
+    -> decltype(adl_detail::size_impl(std::forward<RangeT>(range))) {
+  return adl_detail::size_impl(std::forward<RangeT>(range));
+}
+
 namespace detail {
 
 template <typename RangeT>
@@ -745,6 +761,8 @@ bool any_of(R &&range, UnaryPredicate P);
 
 template <typename T> bool all_equal(std::initializer_list<T> Values);
 
+template <typename R> constexpr size_t range_size(R &&Range);
+
 namespace detail {
 
 using std::declval;
@@ -936,9 +954,7 @@ detail::zippy<detail::zip_shortest, T, U, Args...> zip(T &&t, U &&u,
 template <typename T, typename U, typename... Args>
 detail::zippy<detail::zip_first, T, U, Args...> zip_equal(T &&t, U &&u,
                                                           Args &&...args) {
-  assert(all_equal({std::distance(adl_begin(t), adl_end(t)),
-                    std::distance(adl_begin(u), adl_end(u)),
-                    std::distance(adl_begin(args), adl_end(args))...}) &&
+  assert(all_equal({range_size(t), range_size(u), range_size(args)...}) &&
          "Iteratees do not have equal length");
   return detail::zippy<detail::zip_first, T, U, Args...>(
       std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
@@ -951,9 +967,7 @@ detail::zippy<detail::zip_first, T, U, Args...> zip_equal(T &&t, U &&u,
 template <typename T, typename U, typename... Args>
 detail::zippy<detail::zip_first, T, U, Args...> zip_first(T &&t, U &&u,
                                                           Args &&...args) {
-  assert(std::distance(adl_begin(t), adl_end(t)) <=
-             std::min({std::distance(adl_begin(u), adl_end(u)),
-                       std::distance(adl_begin(args), adl_end(args))...}) &&
+  assert(range_size(t) <= std::min({range_size(u), range_size(args)...}) &&
          "First iteratee is not the shortest");
 
   return detail::zippy<detail::zip_first, T, U, Args...>(
@@ -1769,6 +1783,29 @@ auto size(R &&Range,
   return std::distance(Range.begin(), Range.end());
 }
 
+namespace detail {
+template <typename Range>
+using check_has_free_function_size =
+    decltype(adl_size(std::declval<Range &>()));
+
+template <typename Range>
+static constexpr bool HasFreeFunctionSize =
+    is_detected<check_has_free_function_size, Range>::value;
+} // namespace detail
+
+/// Returns the size of the \p Range, i.e., the number of elements. This
+/// implementation takes inspiration from `std::ranges::size` from C++20 and
+/// delegates the size check to `adl_size` or `std::distance`, in this order of
+/// preference. Unlike `llvm::size`, this function does *not* guarantee O(1)
+/// running time, and is intended to be used in generic code that does not know
+/// the exact range type.
+template <typename R> constexpr size_t range_size(R &&Range) {
+  if constexpr (detail::HasFreeFunctionSize<R>)
+    return adl_size(Range);
+  else
+    return static_cast<size_t>(std::distance(adl_begin(Range), adl_end(Range)));
+}
+
 /// Provide wrappers to std::for_each which take ranges instead of having to
 /// pass begin/end explicitly.
 template <typename R, typename UnaryFunction>
@@ -2389,9 +2426,7 @@ auto enumerate(FirstRange &&First, RestRanges &&...Rest) {
 #ifndef NDEBUG
     // Note: Create an array instead of an initializer list to work around an
     // Apple clang 14 compiler bug.
-    size_t sizes[] = {
-        static_cast<size_t>(std::distance(adl_begin(First), adl_end(First))),
-        static_cast<size_t>(std::distance(adl_begin(Rest), adl_end(Rest)))...};
+    size_t sizes[] = {range_size(First), range_size(Rest)...};
     assert(all_equal(sizes) && "Ranges have different length");
 #endif
   }

llvm/unittests/ADT/IteratorTest.cpp

@@ -743,4 +743,65 @@ TEST(RangeTest, Distance) {
   EXPECT_EQ(std::distance(v2.begin(), v2.end()), size(v2));
 }
 
+TEST(RangeSizeTest, CommonRangeTypes) {
+  SmallVector<int> v1 = {1, 2, 3};
+  EXPECT_EQ(range_size(v1), 3u);
+
+  std::map<int, int> m1 = {{1, 1}, {2, 2}};
+  EXPECT_EQ(range_size(m1), 2u);
+
+  auto it_range = llvm::make_range(m1.begin(), m1.end());
+  EXPECT_EQ(range_size(it_range), 2u);
+
+  static constexpr int c_arr[5] = {};
+  static_assert(range_size(c_arr) == 5u);
+
+  static constexpr std::array<int, 6> cpp_arr = {};
+  static_assert(range_size(cpp_arr) == 6u);
+}
+
+struct FooWithMemberSize {
+  size_t size() const { return 42; }
+  auto begin() { return Data.begin(); }
+  auto end() { return Data.end(); }
+
+  std::set<int> Data;
+};
+
+TEST(RangeSizeTest, MemberSize) {
+  // Make sure that member `.size()` is preferred over the free fuction and
+  // `std::distance`.
+  FooWithMemberSize container;
+  EXPECT_EQ(range_size(container), 42u);
+}
+
+struct FooWithFreeSize {
+  friend size_t size(const FooWithFreeSize &) { return 13; }
+  auto begin() { return Data.begin(); }
+  auto end() { return Data.end(); }
+
+  std::set<int> Data;
+};
+
+TEST(RangeSizeTest, FreeSize) {
+  // Make sure that `size(x)` is preferred over `std::distance`.
+  FooWithFreeSize container;
+  EXPECT_EQ(range_size(container), 13u);
+}
+
+struct FooWithDistance {
+  auto begin() { return Data.begin(); }
+  auto end() { return Data.end(); }
+
+  std::set<int> Data;
+};
+
+TEST(RangeSizeTest, Distance) {
+  // Make sure that we can fall back to `std::distance` even the iterator is not
+  // random-access.
+  FooWithDistance container;
+  EXPECT_EQ(range_size(container), 0u);
+  container.Data = {1, 2, 3, 4};
+  EXPECT_EQ(range_size(container), 4u);
+}
 } // anonymous namespace

llvm/unittests/ADT/STLExtrasTest.cpp

@@ -575,6 +575,39 @@ TEST(STLExtrasTest, ADLTestConstexpr) {
   SUCCEED();
 }
 
+struct FooWithMemberSize {
+  size_t size() const { return 42; }
+  auto begin() { return Data.begin(); }
+  auto end() { return Data.end(); }
+
+  std::set<int> Data;
+};
+
+namespace some_namespace {
+struct FooWithFreeSize {
+  auto begin() { return Data.begin(); }
+  auto end() { return Data.end(); }
+
+  std::set<int> Data;
+};
+
+size_t size(const FooWithFreeSize &) { return 13; }
+} // namespace some_namespace
+
+TEST(STLExtrasTest, ADLSizeTest) {
+  FooWithMemberSize foo1;
+  EXPECT_EQ(adl_size(foo1), 42u);
+
+  some_namespace::FooWithFreeSize foo2;
+  EXPECT_EQ(adl_size(foo2), 13u);
+
+  static constexpr int c_arr[] = {1, 2, 3};
+  static_assert(adl_size(c_arr) == 3u);
+
+  static constexpr std::array<int, 4> cpp_arr = {};
+  static_assert(adl_size(cpp_arr) == 4u);
+}
+
 TEST(STLExtrasTest, DropBeginTest) {
   SmallVector<int, 5> vec{0, 1, 2, 3, 4};