customization for folly::index_iterator (#2259)

Summary:
Pull Request resolved: #2259

Lack of customization for `folly::index_iterator` results in a lot of boilerplate.
`operator[]` having a different meaning is an important case for maps.

Reviewed By: yfeldblum

Differential Revision: D59811691

fbshipit-source-id: 7da2db4fcf30c448080adc4f6ff280c29db347e1

folly/container/BUCK

@@ -41,6 +41,7 @@ cpp_library(
         ":access",
         "//folly:traits",
         "//folly:utility",
+        "//folly/functional:invoke",
         "//folly/lang:rvalue_reference_wrapper",
     ],
 )

folly/container/Iterator.h

@@ -26,6 +26,7 @@
 #include <folly/Traits.h>
 #include <folly/Utility.h>
 #include <folly/container/Access.h>
+#include <folly/functional/Invoke.h>
 #include <folly/lang/RValueReferenceWrapper.h>
 
 namespace folly {
@@ -597,8 +598,24 @@ struct arrow_proxy {
   explicit arrow_proxy(Ref* ref) : res(*ref) {}
 };
 
+struct index_iterator_access_at {
+  template <typename Container, typename Index>
+  constexpr decltype(auto) operator()(Container& container, Index index) const {
+    if constexpr (folly::is_tag_invocable_v<
+                      index_iterator_access_at,
+                      Container&,
+                      Index>) {
+      return folly::tag_invoke(*this, container, index);
+    } else {
+      return container[index];
+    }
+  }
+};
+
 } // namespace detail
 
+FOLLY_DEFINE_CPO(detail::index_iterator_access_at, index_iterator_access_at)
+
 /**
  * index_iterator
  *
@@ -625,8 +642,20 @@ struct arrow_proxy {
  *  Note that `some_ref_type` can be any proxy reference, as long as the
  *  algorithms support that (for example from range-v3).
  *
- * NOTE: there is no way to override `operator[]`, if that's needed
- *       we recommend to wrap your data in a struct with `operator[]`.
+ * NOTE: if `operator[]` doesn't work for you for some reason
+ *       you can specify:
+ *
+ * ```
+ *  friend some_ref_type tag_invoke(
+ *    folly::cpo_t<index_iterator_access_at>,
+ *    Container& c,
+ *    std::size_t index);
+ *
+ *  friend some_cref_type tag_invoke(
+ *    folly::cpo_t<index_iterator_access_at>,
+ *    const Container& c,
+ *    std::size_t index);
+ * ```
  **/
 
 template <typename Container>
@@ -637,6 +666,12 @@ class index_iterator {
   template <typename T>
   using get_difference_type_t = typename std::remove_cv_t<T>::difference_type;
 
+  template <typename IndexType>
+  constexpr static decltype(auto) get_reference_by_index(
+      Container& container, IndexType index) {
+    return index_iterator_access_at(container, index);
+  }
+
  public:
   // index iterator specific types
 
@@ -647,7 +682,8 @@ class index_iterator {
 
   using value_type = typename std::remove_const_t<container_type>::value_type;
   using iterator_category = std::random_access_iterator_tag;
-  using reference = decltype(FOLLY_DECLVAL(container_type&)[size_type{}]);
+  using reference = decltype(get_reference_by_index(
+      FOLLY_DECLVAL(container_type&), size_type{}));
   using difference_type =
       detected_or_t<std::ptrdiff_t, get_difference_type_t, Container>;
 
@@ -683,7 +719,9 @@ class index_iterator {
 
   // access ---
 
-  constexpr reference operator*() const { return (*container_)[index_]; }
+  constexpr reference operator*() const {
+    return get_reference_by_index(*container_, index_);
+  }
 
   pointer operator->() const {
     // It's equivalent to pointer{&**this} but compiler stops

folly/container/test/IteratorTest.cpp

@@ -681,6 +681,42 @@ struct has_both_size_and_difference_type : just_operator {
   using difference_type = std::int32_t;
 };
 
+struct no_operator {
+  using value_type = int;
+
+  friend ref_type tag_invoke(
+      folly::cpo_t<index_iterator_access_at>, no_operator&, std::size_t) {
+    return {};
+  }
+
+  friend cref_type tag_invoke(
+      folly::cpo_t<index_iterator_access_at>, const no_operator&, std::size_t) {
+    return {};
+  }
+
+  [[maybe_unused]] int operator[](std::size_t);
+  [[maybe_unused]] int operator[](std::size_t) const;
+};
+
+struct no_operator_private {
+  using value_type = int;
+
+ private:
+  friend ref_type tag_invoke(
+      folly::cpo_t<index_iterator_access_at>,
+      no_operator_private&,
+      std::size_t) {
+    return {};
+  }
+
+  friend cref_type tag_invoke(
+      folly::cpo_t<index_iterator_access_at>,
+      const no_operator_private&,
+      std::size_t) {
+    return {};
+  }
+};
+
 template <typename T>
 using all_types = std::tuple<
     typename std::iterator_traits<folly::index_iterator<T>>::value_type,
@@ -719,6 +755,18 @@ TEST(IndexIterator, Types) {
   is_same_test(
       all_types<const has_both_size_and_difference_type>{},
       std::tuple<int, cref_type, std::int32_t, std::uint32_t>{});
+  is_same_test(
+      all_types<no_operator>{},
+      std::tuple<int, ref_type, std::ptrdiff_t, std::size_t>{});
+  is_same_test(
+      all_types<const no_operator>{},
+      std::tuple<int, cref_type, std::ptrdiff_t, std::size_t>{});
+  is_same_test(
+      all_types<no_operator_private>{},
+      std::tuple<int, ref_type, std::ptrdiff_t, std::size_t>{});
+  is_same_test(
+      all_types<const no_operator_private>{},
+      std::tuple<int, cref_type, std::ptrdiff_t, std::size_t>{});
 }
 
 } // namespace index_iterator_type_tests
@@ -832,6 +880,31 @@ struct IndexedVector {
   const_iterator cend() const { return {*this, v_.size()}; }
 };
 
+struct CustomMapWithIndexedIterator {
+ public:
+  using size_type = std::uint32_t;
+  using value_type = std::pair<int, int>;
+
+ public:
+  using iterator = folly::index_iterator<CustomMapWithIndexedIterator>;
+  using const_iterator =
+      folly::index_iterator<const CustomMapWithIndexedIterator>;
+
+  iterator begin() { return {*this, 0}; }
+  const_iterator cbegin() const { return {*this, 0}; }
+
+  std::vector<int> keys;
+  std::vector<int> values;
+
+ private:
+  friend std::pair<int, int> tag_invoke(
+      folly::cpo_t<index_iterator_access_at>,
+      const CustomMapWithIndexedIterator& self,
+      size_type idx) {
+    return {self.keys[idx], self.values[idx]};
+  }
+};
+
 } // namespace
 
 TEST(IndexIterator, UseProxyReferences) {
@@ -900,3 +973,17 @@ TEST(IndexIterator, OperatorArrowForNonProxies) {
 
   ASSERT_EQ(expected, v);
 }
+
+TEST(IndexIterator, NoOperatorIntegration) {
+  CustomMapWithIndexedIterator m{.keys{1, 2, 3}, .values{4, 5, 6}};
+
+  auto it = m.begin();
+  ASSERT_EQ((std::pair{1, 4}), it[0]);
+  ASSERT_EQ((std::pair{2, 5}), it[1]);
+  ASSERT_EQ((std::pair{3, 6}), it[2]);
+
+  auto cit = m.cbegin();
+  ASSERT_EQ((std::pair{1, 4}), cit[0]);
+  ASSERT_EQ((std::pair{2, 5}), cit[1]);
+  ASSERT_EQ((std::pair{3, 6}), cit[2]);
+}