[ADT] Rewrite mapped_iterator in terms of iterator_adaptor_base.

Summary:
This eliminates the boilerplate implementation of the iterator interface in
mapped_iterator.

This patch also adds unit tests that verify that the mapped function is applied
by operator* and operator->, and that references returned by the map function
are returned via operator*.

Reviewers: dblaikie, chandlerc

Subscribers: llvm-commits, mgorny

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

llvm-svn: 317902

3 files changed, 68 insertions, 75 deletions

diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index 1be5bf91385..6e20a70c581 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/llvm/include/llvm/ADT/STLExtras.h
@@ -132,91 +132,32 @@ inline void deleter(T *Ptr) {
 
 // mapped_iterator - This is a simple iterator adapter that causes a function to
 // be applied whenever operator* is invoked on the iterator.
-template <class RootIt, class UnaryFunc>
-class mapped_iterator {
-  RootIt current;
-  UnaryFunc Fn;
 
+template <typename ItTy, typename FuncTy,
+          typename FuncReturnTy =
+            decltype(std::declval<FuncTy>()(*std::declval<ItTy>()))>
+class mapped_iterator
+    : public iterator_adaptor_base<
+             mapped_iterator<ItTy, FuncTy>, ItTy,
+             typename std::iterator_traits<ItTy>::iterator_category,
+             typename std::remove_reference<FuncReturnTy>::type> {
 public:
-  using iterator_category =
-      typename std::iterator_traits<RootIt>::iterator_category;
-  using difference_type =
-      typename std::iterator_traits<RootIt>::difference_type;
-  using value_type =
-      decltype(std::declval<UnaryFunc>()(*std::declval<RootIt>()));
-
-  using pointer = void;
-  using reference = void; // Can't modify value returned by fn
-
-  using iterator_type = RootIt;
-
-  inline explicit mapped_iterator(const RootIt &I, UnaryFunc F)
-    : current(I), Fn(F) {}
-
-  inline value_type operator*() const {   // All this work to do this
-    return Fn(*current);         // little change
-  }
-
-  mapped_iterator &operator++() {
-    ++current;
-    return *this;
-  }
-  mapped_iterator &operator--() {
-    --current;
-    return *this;
-  }
-  mapped_iterator operator++(int) {
-    mapped_iterator __tmp = *this;
-    ++current;
-    return __tmp;
-  }
-  mapped_iterator operator--(int) {
-    mapped_iterator __tmp = *this;
-    --current;
-    return __tmp;
-  }
-  mapped_iterator operator+(difference_type n) const {
-    return mapped_iterator(current + n, Fn);
-  }
-  mapped_iterator &operator+=(difference_type n) {
-    current += n;
-    return *this;
-  }
-  mapped_iterator operator-(difference_type n) const {
-    return mapped_iterator(current - n, Fn);
-  }
-  mapped_iterator &operator-=(difference_type n) {
-    current -= n;
-    return *this;
-  }
-  reference operator[](difference_type n) const { return *(*this + n); }
+  mapped_iterator(ItTy U, FuncTy F)
+    : mapped_iterator::iterator_adaptor_base(std::move(U)), F(std::move(F)) {}
 
-  bool operator!=(const mapped_iterator &X) const { return !operator==(X); }
-  bool operator==(const mapped_iterator &X) const {
-    return current == X.current;
-  }
-  bool operator<(const mapped_iterator &X) const { return current < X.current; }
+  ItTy getCurrent() { return this->I; }
 
-  difference_type operator-(const mapped_iterator &X) const {
-    return current - X.current;
-  }
+  FuncReturnTy operator*() { return F(*this->I); }
 
-  inline const RootIt &getCurrent() const { return current; }
-  inline const UnaryFunc &getFunc() const { return Fn; }
+private:
+  FuncTy F;
 };
 
-template <class Iterator, class Func>
-inline mapped_iterator<Iterator, Func>
-operator+(typename mapped_iterator<Iterator, Func>::difference_type N,
-          const mapped_iterator<Iterator, Func> &X) {
-  return mapped_iterator<Iterator, Func>(X.getCurrent() - N, X.getFunc());
-}
-
 // map_iterator - Provide a convenient way to create mapped_iterators, just like
 // make_pair is useful for creating pairs...
 template <class ItTy, class FuncTy>
-inline mapped_iterator<ItTy, FuncTy> map_iterator(const ItTy &I, FuncTy F) {
-  return mapped_iterator<ItTy, FuncTy>(I, F);
+inline mapped_iterator<ItTy, FuncTy> map_iterator(ItTy I, FuncTy F) {
+  return mapped_iterator<ItTy, FuncTy>(std::move(I), std::move(F));
 }
 
 /// Helper to determine if type T has a member called rbegin().
diff --git a/llvm/unittests/ADT/CMakeLists.txt b/llvm/unittests/ADT/CMakeLists.txt
index d0313951dfa..e5b6bbf2a23 100644
--- a/llvm/unittests/ADT/CMakeLists.txt
+++ b/llvm/unittests/ADT/CMakeLists.txt
@@ -32,6 +32,7 @@ set(ADTSources
   IntrusiveRefCntPtrTest.cpp
   IteratorTest.cpp
   MakeUniqueTest.cpp
+  MappedIteratorTest.cpp
   MapVectorTest.cpp
   OptionalTest.cpp
   PackedVectorTest.cpp
diff --git a/llvm/unittests/ADT/MappedIteratorTest.cpp b/llvm/unittests/ADT/MappedIteratorTest.cpp
new file mode 100644
index 00000000000..d8c48a7d625
--- /dev/null
+++ b/llvm/unittests/ADT/MappedIteratorTest.cpp
@@ -0,0 +1,51 @@
+//===------ MappedIteratorTest.cpp - Unit tests for mapped_iterator -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(MappedIteratorTest, ApplyFunctionOnDereference) {
+  std::vector<int> V({0});
+
+  auto I = map_iterator(V.begin(), [](int X) { return X + 1; });
+
+  EXPECT_EQ(*I, 1) << "should have applied function in dereference";
+}
+
+TEST(MappedIteratorTest, ApplyFunctionOnArrow) {
+  struct S {
+    int Z = 0;
+  };
+
+  std::vector<int> V({0});
+  S Y;
+  S* P = &Y;
+
+  auto I = map_iterator(V.begin(), [&](int X) -> S& { return *(P + X); });
+
+  I->Z = 42;
+
+  EXPECT_EQ(Y.Z, 42) << "should have applied function during arrow";
+}
+
+TEST(MappedIteratorTest, FunctionPreservesReferences) {
+  std::vector<int> V({1});
+  std::map<int, int> M({ {1, 1} });
+
+  auto I = map_iterator(V.begin(), [&](int X) -> int& { return M[X]; });
+  *I = 42;
+
+  EXPECT_EQ(M[1], 42) << "assignment should have modified M";
+}
+
+} // anonymous namespace