[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

1786–1800

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.

1809–1810

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

1786–1800

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

1809–1810

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

2426

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

2426

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 T>
 constexpr void swap_impl(T &&lhs,
                          T &&rhs) noexcept(noexcept(swap(std::declval<T>(),
                                                          std::declval<T>()))) {
   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
 /// function found through Argument-Dependent Lookup (ADL).
 template <typename RangeT>
 constexpr auto adl_begin(RangeT &&range)
     -> decltype(adl_detail::begin_impl(std::forward<RangeT>(range))) {
   return adl_detail::begin_impl(std::forward<RangeT>(range));

 /// Swaps \p lhs with \p rhs using `std::swap` and functions found through
 /// Argument-Dependent Lookup (ADL).
 template <typename T>
 constexpr void adl_swap(T &&lhs, T &&rhs) noexcept(
     noexcept(adl_detail::swap_impl(std::declval<T>(), std::declval<T>()))) {
   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>
 using IterOfRange = decltype(adl_begin(std::declval<RangeT &>()));
 
 template <typename RangeT>
 using ValueOfRange =
     std::remove_reference_t<decltype(*adl_begin(std::declval<RangeT &>()))>;

 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_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 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
 /// begin/end explicitly.
 template <typename R, typename UnaryPredicate>

 ///   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; }
+  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

TEST(STLExtrasTest, ADLTestConstexpr) {
   static_assert(adl_end(c_arr) == c_arr + 3);
 
   static constexpr std::array<int, 2> std_arr = {1, 2};
   static_assert(adl_begin(std_arr) == std_arr.begin());
   static_assert(adl_end(std_arr) == std_arr.end());
   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};
 
   for (int n = 0; n < 5; ++n) {
     int i = n;
     for (auto &v : drop_begin(vec, n)) {
       EXPECT_EQ(v, i);
       i += 1;