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

[kuhar]

Summary

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.

[dblaikie]

Would std::distance suffice?

[kuhar]

In D146231#4199846, @dblaikie wrote:

Would std::distance suffice?

It does work and this is what we use in the assertions now. The issue is that this degrades to a linear scan over some iterator types (e.g., with sets/maps)

[dblaikie]

llvm/include/llvm/ADT/STLExtras.h

1770–1784

Rather than detecting both member and non-member, could we standardize on non-member, like adl_begin/adl_end works? Could expose adl_size to match adl_begin/end too, in addition to this wrapper that has the fallback to distance.

1793–1794

I'd guess this can be trivially optimized from the generic std::distance call? Maybe not worth the special case?

[kuhar]

llvm/include/llvm/ADT/STLExtras.h

1770–1784

Good idea. I wonder why std::ranges::size doesn't do it too.

1793–1794

Once we go through std::size it should cover it too. This will make this diverge from std::ranges::size but I don't think anyone will notice the difference, and our implementation will become simpler.

[kuhar]

Use adl_size to simplify the implementation and reduce branching.

[zero9178]

LGTM with tiny note

llvm/include/llvm/ADT/STLExtras.h

2420

Guessing you still need to rebase this ontop of https://reviews.llvm.org/D146340?
Does that happen to also fix the compiler bug from that revision or will we still require the workaround using the intializer list?

[kuhar]

llvm/include/llvm/ADT/STLExtras.h

2420

Yes, this needs rebasing. I've just checked and the workaround is still necessary.

[kuhar]

Rebase

[dblaikie]

Looks good, thanks!

llvm/include/llvm/ADT/STLExtras.h

 template <typename R, typename UnaryPredicate>
 bool all_of(R &&range, UnaryPredicate P);
 
 template <typename R, typename UnaryPredicate>
 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;
 
 // We have to alias this since inlining the actual type at the usage site
 // in the parameter list of iterator_facade_base<> below ICEs MSVC 2017.
 template<typename... Iters> struct ZipTupleType {
   using type = std::tuple<decltype(*declval<Iters>())...>;

 }
 
 /// zip iterator that assumes that all iteratees have the same length.
 /// In builds with assertions on, this assumption is checked before the
 /// iteration starts.
 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)...);
 }
 
 /// zip iterator that, for the sake of efficiency, assumes the first iteratee to
 /// be the shortest. Iteration continues until the end of the first iteratee is
 /// reached. In builds with assertions on, we check that the assumption about
 /// the first iteratee being the shortest holds.
 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...>(
       std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
 }
 
 namespace detail {
 template <typename Iter>
auto size(R &&Range,
           std::enable_if_t<
               std::is_base_of<std::random_access_iterator_tag,
                               typename std::iterator_traits<decltype(
                                   Range.begin())>::iterator_category>::value,
               void> * = nullptr) {
   return std::distance(Range.begin(), Range.end());
 }
 
+namespace detail {
+template <typename Range>
+using check_has_member_size = decltype(std::declval<Range &>().size());
+
+template <typename Range>
+static constexpr bool HasMemberSize =
+    is_detected<check_has_member_size, Range>::value;
+
+template <typename Range>
+using check_has_free_function_size = decltype(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 follows `std::ranges::size` from C++20 and delegates the size
+/// check to `std::extent`, `Range.size()`, `size(Range)`, or `std::distance`,
+/// as available, 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 (std::is_array_v<R>)
+    return std::extent<R>::value;
+  else if constexpr (detail::HasMemberSize<R>)
+    return Range.size();
+  else if constexpr (detail::HasFreeFunctionSize<R>)
+    return size(Range);
+  else
+    return 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>
 UnaryFunction for_each(R &&Range, UnaryFunction F) {
   return std::for_each(adl_begin(Range), adl_end(Range), F);
 }
 
 /// Provide wrappers to std::all_of which take ranges instead of having to pass

 ///   Item 0: 20
 ///   Item 1: 21
 ///   Item 2: 22
 ///   Item 3: 23
 ///
 template <typename FirstRange, typename... RestRanges>
 auto enumerate(FirstRange &&First, RestRanges &&...Rest) {
   assert((sizeof...(Rest) == 0 ||
-          all_equal({std::distance(adl_begin(First), adl_end(First)),
-                     std::distance(adl_begin(Rest), adl_end(Rest))...})) &&
+          all_equal({range_size(First), range_size(Rest)...})) &&
          "Ranges have different length");
   using enumerator = detail::zippy<detail::zip_enumerator, detail::index_stream,
                                    FirstRange, RestRanges...>;
   return enumerator(detail::index_stream{}, std::forward<FirstRange>(First),
                     std::forward<RestRanges>(Rest)...);
 }
 
 namespace detail {

llvm/unittests/ADT/IteratorTest.cpp

 TEST(RangeTest, Distance) {
   std::vector<int> v1;
   std::vector<int> v2{1, 2, 3};
 
   EXPECT_EQ(std::distance(v1.begin(), v1.end()), size(v1));
   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; }
+  friend size_t size(const FooWithMemberSize &) { return 13; }
+  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);
+  EXPECT_EQ(size(container), 13u); // Silence unused function warning.
+}
+
+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