IListRefIterator Class — pytorch Architecture

Source Code

aten/src/ATen/core/IListRef.h lines 364–498

class IListRefIterator {
 private:
#define DEFINE_FRIEND_CLASS(TAG, ...)                        \
  friend class detail::IListRefTagImpl<IListRefTag::TAG, T>; \
  friend class detail::IListRefTagImplBase<                  \
      IListRefTag::TAG,                                      \
      T,                                                     \
      typename detail::IListRefTagImpl<IListRefTag::TAG, T>::elem_type>;
  TORCH_ILISTREF_FORALL_TAGS(DEFINE_FRIEND_CLASS)
#undef DEFINE_FRIEND_CLASS

 public:
  // C++17 friendly std::iterator implementation
  using iterator_category = std::bidirectional_iterator_tag;
  using value_type = T;
  using difference_type = std::ptrdiff_t;
  using pointer = T*;
  using reference = T&;

  using unboxed_iterator_type = typename detail::
      IListRefTagImpl<IListRefTag::Unboxed, T>::list_type::const_iterator;
  using boxed_iterator_type = typename detail::
      IListRefTagImpl<IListRefTag::Boxed, T>::list_type::const_iterator;
  using materialized_iterator_type =
      typename detail::MaterializedIListRef<T>::const_iterator;

  IListRefIterator() : tag_(IListRefTag::None) {}

#if defined(_MSC_VER) && _ITERATOR_DEBUG_LEVEL != 0
  // See [Note: MSVC Iterator Debug]
  IListRefIterator(const IListRefIterator& iterator)
      : tag_(iterator.tag_) {
    switch (tag_) {
      case IListRefTag::Boxed:
        payload_.boxed_iterator = iterator.payload_.boxed_iterator;
        break;
      case IListRefTag::Unboxed:
        payload_.unboxed_iterator = iterator.payload_.unboxed_iterator;
        break;
      case IListRefTag::Materialized:
        payload_.materialized_iterator = iterator.payload_.materialized_iterator;
        break;
      default:
        TORCH_INTERNAL_ASSERT(false, "invalid IListRef tag.");
    }
  }
#endif

#if defined(_MSC_VER) && _ITERATOR_DEBUG_LEVEL == 2
  // See [Note: MSVC Iterator Debug]
  ~IListRefIterator() noexcept(false) {
    switch (tag_) {
      case IListRefTag::Boxed:
        payload_.boxed_iterator.~boxed_iterator_type();
        break;
      case IListRefTag::Unboxed:
        payload_.unboxed_iterator.~unboxed_iterator_type();
        break;
      case IListRefTag::Materialized:
        payload_.materialized_iterator.~materialized_iterator_type();
        break;
      default:
        TORCH_INTERNAL_ASSERT(false, "invalid IListRef tag.");
    }
  }
#endif

  IListRefIterator(boxed_iterator_type boxed) : tag_(IListRefTag::Boxed) {
    payload_.boxed_iterator = boxed;
  }

  IListRefIterator(unboxed_iterator_type unboxed) : tag_(IListRefTag::Unboxed) {
    payload_.unboxed_iterator = unboxed;
  }

  IListRefIterator(materialized_iterator_type materialized) : tag_(IListRefTag::Materialized) {
    payload_.materialized_iterator = materialized;
  }

  detail::IListRefConstRef<T> operator*() const {
    TORCH_ILISTREF_UNWRAP(tag_, { return ImplT::iterator_get(this_); });
  }

  IListRefIterator& operator++() {
    TORCH_ILISTREF_UNWRAP(tag_, { ++this_; });
    return *this;
  }

  IListRefIterator operator++(int) {
    auto old = *this;
    TORCH_ILISTREF_UNWRAP(tag_, { ++this_; });
    return old;
  }

  IListRefIterator& operator--() {
    TORCH_ILISTREF_UNWRAP(tag_, { --this_; });
    return *this;
  }

  IListRefIterator operator--(int) {
    auto old = *this;
    TORCH_ILISTREF_UNWRAP(tag_, { --this_; });
    return old;
  }

  bool operator==(const IListRefIterator& rhs) const {
    if (tag_ != rhs.tag_) {
      return false;
    }
    TORCH_ILISTREF_UNWRAP(tag_, {
      auto& rhs_it = ImplT::unwrap(rhs);
      return this_ == rhs_it;
    });
  }

  bool operator!=(const IListRefIterator& rhs) const {
    return !(*this == rhs);
  }

 private:
  union Payload {
    boxed_iterator_type boxed_iterator;
    unboxed_iterator_type unboxed_iterator;
    materialized_iterator_type materialized_iterator;
    void* _init_ptr;
    Payload() : _init_ptr(nullptr) {}
#if defined(_MSC_VER)
    // See [Note: MSVC Iterator Debug]
    ~Payload() {}
#endif
  };

  Payload payload_;
  IListRefTag tag_;
};