chore: import upstream snapshot with attribution

include/dgl/array_iterator.h

@@ -0,0 +1,173 @@
+/**
+ *  Copyright (c) 2020 by Contributors
+ * @file dgl/array_iterator.h
+ * @brief Various iterators.
+ */
+#ifndef DGL_ARRAY_ITERATOR_H_
+#define DGL_ARRAY_ITERATOR_H_
+
+#ifdef __CUDA_ARCH__
+#define CUB_INLINE __host__ __device__ __forceinline__
+#else
+#define CUB_INLINE inline
+#endif  // __CUDA_ARCH__
+
+#include <algorithm>
+#include <iterator>
+#include <utility>
+
+namespace dgl {
+namespace aten {
+
+using std::swap;
+
+// Make std::pair work on both host and device
+template <typename DType>
+struct Pair {
+  Pair() = default;
+  Pair(const Pair& other) = default;
+  Pair(Pair&& other) = default;
+  CUB_INLINE Pair(DType a, DType b) : first(a), second(b) {}
+  CUB_INLINE Pair& operator=(const Pair& other) {
+    first = other.first;
+    second = other.second;
+    return *this;
+  }
+  CUB_INLINE operator std::pair<DType, DType>() const {
+    return std::make_pair(first, second);
+  }
+  CUB_INLINE bool operator==(const Pair& other) const {
+    return (first == other.first) && (second == other.second);
+  }
+  CUB_INLINE void swap(const Pair& other) const {
+    std::swap(first, other.first);
+    std::swap(second, other.second);
+  }
+  DType first, second;
+};
+
+template <typename DType>
+CUB_INLINE void swap(const Pair<DType>& r1, const Pair<DType>& r2) {
+  r1.swap(r2);
+}
+
+// PairRef and PairIterator that serves as an iterator over a pair of arrays in
+// a zipped fashion like zip(a, b).
+template <typename DType>
+struct PairRef {
+  PairRef() = delete;
+  PairRef(const PairRef& other) = default;
+  PairRef(PairRef&& other) = default;
+  CUB_INLINE PairRef(DType* const r, DType* const c) : a(r), b(c) {}
+  CUB_INLINE PairRef& operator=(const PairRef& other) {
+    *a = *other.a;
+    *b = *other.b;
+    return *this;
+  }
+  CUB_INLINE PairRef& operator=(const Pair<DType>& val) {
+    *a = val.first;
+    *b = val.second;
+    return *this;
+  }
+  CUB_INLINE operator Pair<DType>() const { return Pair<DType>(*a, *b); }
+  CUB_INLINE operator std::pair<DType, DType>() const {
+    return std::make_pair(*a, *b);
+  }
+  CUB_INLINE bool operator==(const PairRef& other) const {
+    return (*a == *(other.a)) && (*b == *(other.b));
+  }
+  CUB_INLINE void swap(const PairRef& other) const {
+    std::swap(*a, *other.a);
+    std::swap(*b, *other.b);
+  }
+  DType *a, *b;
+};
+
+template <typename DType>
+CUB_INLINE void swap(const PairRef<DType>& r1, const PairRef<DType>& r2) {
+  r1.swap(r2);
+}
+
+template <typename DType>
+struct PairIterator : public std::iterator<
+                          std::random_access_iterator_tag, Pair<DType>,
+                          std::ptrdiff_t, Pair<DType*>, PairRef<DType>> {
+  PairIterator() = default;
+  PairIterator(const PairIterator& other) = default;
+  PairIterator(PairIterator&& other) = default;
+  CUB_INLINE PairIterator(DType* x, DType* y) : a(x), b(y) {}
+  PairIterator& operator=(const PairIterator& other) = default;
+  PairIterator& operator=(PairIterator&& other) = default;
+  ~PairIterator() = default;
+  CUB_INLINE bool operator==(const PairIterator& other) const {
+    return a == other.a;
+  }
+  CUB_INLINE bool operator!=(const PairIterator& other) const {
+    return a != other.a;
+  }
+  CUB_INLINE bool operator<(const PairIterator& other) const {
+    return a < other.a;
+  }
+  CUB_INLINE bool operator>(const PairIterator& other) const {
+    return a > other.a;
+  }
+  CUB_INLINE bool operator<=(const PairIterator& other) const {
+    return a <= other.a;
+  }
+  CUB_INLINE bool operator>=(const PairIterator& other) const {
+    return a >= other.a;
+  }
+  CUB_INLINE PairIterator& operator+=(const std::ptrdiff_t& movement) {
+    a += movement;
+    b += movement;
+    return *this;
+  }
+  CUB_INLINE PairIterator& operator-=(const std::ptrdiff_t& movement) {
+    a -= movement;
+    b -= movement;
+    return *this;
+  }
+  CUB_INLINE PairIterator& operator++() {
+    ++a;
+    ++b;
+    return *this;
+  }
+  CUB_INLINE PairIterator& operator--() {
+    --a;
+    --b;
+    return *this;
+  }
+  CUB_INLINE PairIterator operator++(int) {
+    PairIterator ret(*this);
+    operator++();
+    return ret;
+  }
+  CUB_INLINE PairIterator operator--(int) {
+    PairIterator ret(*this);
+    operator--();
+    return ret;
+  }
+  CUB_INLINE PairIterator operator+(const std::ptrdiff_t& movement) const {
+    return PairIterator(a + movement, b + movement);
+  }
+  CUB_INLINE PairIterator operator-(const std::ptrdiff_t& movement) const {
+    return PairIterator(a - movement, b - movement);
+  }
+  CUB_INLINE std::ptrdiff_t operator-(const PairIterator& other) const {
+    return a - other.a;
+  }
+  CUB_INLINE PairRef<DType> operator*() const { return PairRef<DType>(a, b); }
+  CUB_INLINE PairRef<DType> operator*() { return PairRef<DType>(a, b); }
+  CUB_INLINE PairRef<DType> operator[](size_t offset) const {
+    return PairRef<DType>(a + offset, b + offset);
+  }
+  CUB_INLINE PairRef<DType> operator[](size_t offset) {
+    return PairRef<DType>(a + offset, b + offset);
+  }
+  DType *a, *b;
+};
+
+};  // namespace aten
+};  // namespace dgl
+
+#endif  // DGL_ARRAY_ITERATOR_H_