[ADT] Fix const-correctness issues in `zippy`

[kuhar]

Summary

This defines the iterator tuple based on the storage type of zippy,
instead of its type arguments. This way, we can support temporaries that
gets passed in and allow for them to be modified during iteration.

Because the iterator types to the tuple storage can have different types
when the storage is and isn't const, this defines a const iterator type
and non-const begin/end functions. This way we avoid unintentional
casts, e.g., trying to cast vector<bool>::reference to
vector<bool>::const_reference, which may be unrelated types that are
not convertible.

This patch is a general and free-standing improvement but my primary use
is in the implemention a version of enumerate that accepts multiple ranges:
D144583.

[kuhar]

Drop unnecessary headers

[zero9178]

LGTM, but please give it a day or two to give other reviewers a chance to comment as well

[dblaikie]

llvm/unittests/ADT/IteratorTest.cpp

454–457

These could be static_asserts instead, perhaps? (similarly below)

503–504

Does this code work? How? Wouldn't the lifetimes of the temporary ranges expire before the body of the loop runs?

[kuhar]

Use static_assert to test for constness. Rebase.

llvm/unittests/ADT/IteratorTest.cpp

503–504

These ranges get moved into the tuple<...> storage; inside zippy. From then on, we can use references obtained from this storage to access them. So this should not rely on any lifetime extensions on the temporaries passed to zip_equal.

[kuhar]

Add comment explaining lifetimes of temporaries in the ZipEqualTemporaries test.

[kuhar]

Fix the definition of MakeConst to return values instead of references, to avoid dangling references.

[kuhar]

llvm/unittests/ADT/IteratorTest.cpp

503–504

To confirm, I ran this under ASan and UBSan and they don't report anything once I fixed the definition of MakeConst

[dblaikie]

Might've been nice to separate the adl_* support from the other work in this patch, but at least at a vague glance, it all looks good - thanks!

llvm/unittests/ADT/IteratorTest.cpp

503–504

ah, yeah, MakeConst especially made me uncertain this was correct - but yeah, returning by value, and zip carrying its arguments by value when passed temporaries makes sense that it all works together now, thanks!

May be worth a comment on MakeConst about its return value being const intentionally (though it's obvious enough from the name I suppose) - since const on by-value return types is usually a mistake & I could imagine a linter or something trying to clean that up.

[kuhar]

Clarify the purpose of MakeConst

llvm/include/llvm/ADT/STLExtras.h

public:
 
   zip_shortest(Iters &&... ts) : Base(std::forward<Iters>(ts)...) {}
 
   bool operator==(const zip_shortest<Iters...> &other) const {
     return !test(other, std::index_sequence_for<Iters...>{});
   }
 };
 
+/// Helper to obtain the iterator types for the tuple storage within `zippy`.
+template <template <typename...> class ItType, typename TupleStorageType,
+          typename IndexSequence>
+struct ZippyIteratorTuple;
+
+/// Partial specialization for non-const tuple storage.
+template <template <typename...> class ItType, typename... Args,
+          std::size_t... Ns>
+struct ZippyIteratorTuple<ItType, std::tuple<Args...>,
+                          std::index_sequence<Ns...>> {
+  using type = ItType<decltype(adl_begin(
+      std::get<Ns>(declval<std::tuple<Args...> &>())))...>;
+};
+
+/// Partial specialization for const tuple storage.
+template <template <typename...> class ItType, typename... Args,
+          std::size_t... Ns>
+struct ZippyIteratorTuple<ItType, const std::tuple<Args...>,
+                          std::index_sequence<Ns...>> {
+  using type = ItType<decltype(adl_begin(
+      std::get<Ns>(declval<const std::tuple<Args...> &>())))...>;
+};
+
 template <template <typename...> class ItType, typename... Args> class zippy {
+private:
+  std::tuple<Args...> storage;
+  using IndexSequence = std::index_sequence_for<Args...>;
+
 public:
-  using iterator = ItType<decltype(std::begin(std::declval<Args>()))...>;
+  using iterator = typename ZippyIteratorTuple<ItType, decltype(storage),
+                                               IndexSequence>::type;
+  using const_iterator =
+      typename ZippyIteratorTuple<ItType, const decltype(storage),
+                                  IndexSequence>::type;
   using iterator_category = typename iterator::iterator_category;
   using value_type = typename iterator::value_type;
   using difference_type = typename iterator::difference_type;
   using pointer = typename iterator::pointer;
   using reference = typename iterator::reference;
+  using const_reference = typename const_iterator::reference;
 
-private:
-  std::tuple<Args...> ts;
+  zippy(Args &&...args) : storage(std::forward<Args>(args)...) {}
+
+  const_iterator begin() const { return begin_impl(IndexSequence{}); }
+  iterator begin() { return begin_impl(IndexSequence{}); }
+  const_iterator end() const { return end_impl(IndexSequence{}); }
+  iterator end() { return end_impl(IndexSequence{}); }
 
+private:
   template <size_t... Ns>
-  iterator begin_impl(std::index_sequence<Ns...>) const {
-    return iterator(std::begin(std::get<Ns>(ts))...);
+  const_iterator begin_impl(std::index_sequence<Ns...>) const {
+    return const_iterator(adl_begin(std::get<Ns>(storage))...);
   }
-  template <size_t... Ns> iterator end_impl(std::index_sequence<Ns...>) const {
-    return iterator(std::end(std::get<Ns>(ts))...);
+  template <size_t... Ns> iterator begin_impl(std::index_sequence<Ns...>) {
+    return iterator(adl_begin(std::get<Ns>(storage))...);
   }
 
-public:
-  zippy(Args &&... ts_) : ts(std::forward<Args>(ts_)...) {}
-
-  iterator begin() const {
-    return begin_impl(std::index_sequence_for<Args...>{});
+  template <size_t... Ns>
+  const_iterator end_impl(std::index_sequence<Ns...>) const {
+    return const_iterator(adl_end(std::get<Ns>(storage))...);
+  }
+  template <size_t... Ns> iterator end_impl(std::index_sequence<Ns...>) {
+    return iterator(adl_end(std::get<Ns>(storage))...);
   }
-  iterator end() const { return end_impl(std::index_sequence_for<Args...>{}); }
 };
 
 } // end namespace detail
 
 /// zip iterator for two or more iteratable types. Iteration continues until the
 /// end of the *shortest* iteratee is reached.
 template <typename T, typename U, typename... Args>
 detail::zippy<detail::zip_shortest, T, U, Args...> zip(T &&t, U &&u,

llvm/unittests/ADT/IteratorTest.cpp

 //===- IteratorTest.cpp - Unit tests for iterator utilities ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ADT/ilist.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/ilist.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <optional>
+#include <type_traits>
+#include <vector>
 
 using namespace llvm;
+using testing::ElementsAre;
 
 namespace {
 
 template <int> struct Shadow;
 
 struct WeirdIter
     : llvm::iterator_facade_base<WeirdIter, std::input_iterator_tag, Shadow<0>,
                                  Shadow<1>, Shadow<2>, Shadow<3>> {};
TEST(ZipIteratorTest, ZipEqualBasic) {
   for (auto [lhs, rhs] : zip_equal(vals, pi)) {
     EXPECT_EQ(lhs, rhs & 0x01);
     ++iters;
   }
 
   EXPECT_EQ(iters, 6u);
 }
 
+template <typename T>
+constexpr bool IsConstRef =
+    std::is_reference_v<T> && std::is_const_v<std::remove_reference_t<T>>;
+
+template <typename T>
+constexpr bool IsBoolConstRef =
+    std::is_same_v<llvm::remove_cvref_t<T>, std::vector<bool>::const_reference>;
+
+/// Returns a `const` copy of the passed value. The `const` on the returned
+/// value is intentional here so that `MakeConst` can be used in range-for
+/// loops.
+template <typename T> const T MakeConst(T &&value) {
+  return std::forward<T>(value);
+}
+
+TEST(ZipIteratorTest, ZipEqualConstCorrectness) {
+  const std::vector<unsigned> c_first = {3, 1, 4};
+  std::vector<unsigned> first = c_first;
+  const SmallVector<bool> c_second = {1, 1, 0};
+  SmallVector<bool> second = c_second;
+
+  for (auto [a, b, c, d] : zip_equal(c_first, first, c_second, second)) {
+    b = 0;
+    d = true;
+    static_assert(IsConstRef<decltype(a)>);
+    static_assert(!IsConstRef<decltype(b)>);
+    static_assert(IsConstRef<decltype(c)>);
+    static_assert(!IsConstRef<decltype(d)>);
+  }
+
+  EXPECT_THAT(first, ElementsAre(0, 0, 0));
+  EXPECT_THAT(second, ElementsAre(true, true, true));
+
+  std::vector<bool> nemesis = {true, false, true};
+  const std::vector<bool> c_nemesis = nemesis;
+
+  for (auto &&[a, b, c, d] : zip_equal(first, c_first, nemesis, c_nemesis)) {
+    a = 2;
+    c = true;
+    static_assert(!IsConstRef<decltype(a)>);
+    static_assert(IsConstRef<decltype(b)>);
+    static_assert(!IsBoolConstRef<decltype(c)>);
+    static_assert(IsBoolConstRef<decltype(d)>);
+  }
+
+  EXPECT_THAT(first, ElementsAre(2, 2, 2));
+  EXPECT_THAT(nemesis, ElementsAre(true, true, true));
+
+  unsigned iters = 0;
+  for (const auto &[a, b, c, d] :
+       zip_equal(first, c_first, nemesis, c_nemesis)) {
+    static_assert(!IsConstRef<decltype(a)>);
+    static_assert(IsConstRef<decltype(b)>);
+    static_assert(!IsBoolConstRef<decltype(c)>);
+    static_assert(IsBoolConstRef<decltype(d)>);
+    ++iters;
+  }
+  EXPECT_EQ(iters, 3u);
+  iters = 0;
+
+  for (const auto &[a, b, c, d] :
+       MakeConst(zip_equal(first, c_first, nemesis, c_nemesis))) {
+    static_assert(!IsConstRef<decltype(a)>);
+    static_assert(IsConstRef<decltype(b)>);
+    static_assert(!IsBoolConstRef<decltype(c)>);
+    static_assert(IsBoolConstRef<decltype(d)>);
+    ++iters;
+  }
+  EXPECT_EQ(iters, 3u);
+}
+
+TEST(ZipIteratorTest, ZipEqualTemporaries) {
+  unsigned iters = 0;
+
+  // These temporary ranges get moved into the `tuple<...> storage;` inside
+  // `zippy`. From then on, we can use references obtained from this storage to
+  // access them. This does not rely on any lifetime extensions on the
+  // temporaries passed to `zip_equal`.
+  for (auto [a, b, c] : zip_equal(SmallVector<int>{1, 2, 3}, std::string("abc"),
+                                  std::vector<bool>{true, false, true})) {
+    a = 3;
+    b = 'c';
+    c = false;
+    static_assert(!IsConstRef<decltype(a)>);
+    static_assert(!IsConstRef<decltype(b)>);
+    static_assert(!IsBoolConstRef<decltype(c)>);
+    ++iters;
+  }
+  EXPECT_EQ(iters, 3u);
+  iters = 0;
+
+  for (auto [a, b, c] :
+       MakeConst(zip_equal(SmallVector<int>{1, 2, 3}, std::string("abc"),
+                           std::vector<bool>{true, false, true}))) {
+    static_assert(IsConstRef<decltype(a)>);
+    static_assert(IsConstRef<decltype(b)>);
+    static_assert(IsBoolConstRef<decltype(c)>);
+    ++iters;
+  }
+  EXPECT_EQ(iters, 3u);
+}
+
 #if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
 // Check that an assertion is triggered when ranges passed to `zip_equal` differ
 // in length.
 TEST(ZipIteratorTest, ZipEqualNotEqual) {
   const SmallVector<unsigned, 6> pi = {3, 1, 4, 1, 5, 8};
   const SmallVector<bool, 2> vals = {1, 1};
 
   EXPECT_DEATH(zip_equal(pi, vals), "Iteratees do not have equal length");