graph_node_list.h

#pragma once

#include "torch/csrc/jit/assertions.h"

namespace torch { namespace jit {

// Intrusive doubly linked lists with sane reverse iterators.
// The header file is named generic_graph_node_list.h because it is ONLY
// used for Graph's Node lists, and if you want to use it for other
// things, you will have to do some refactoring.
//
// At the moment, the templated type T must support a few operations:
//
//  - It must have a field: T* next_in_graph[2] = { nullptr, nullptr };
//    which are used for the intrusive linked list pointers.
//
//  - It must have a method 'destroy()', which removes T from the
//    list and frees a T.
//
// In practice, we are only using it with Node and const Node.  'destroy()'
// needs to be renegotiated if you want to use this somewhere else.
//
// Besides the benefits of being intrusive, unlike std::list, these lists handle
// forward and backward iteration uniformly because we require a
// "before-first-element" sentinel.  This means that reverse iterators
// physically point to the element they logically point to, rather than
// the off-by-one behavior for all standard library reverse iterators.

static constexpr int kNextDirection = 0;
static constexpr int kPrevDirection = 1;

template <typename T>
struct generic_graph_node_list;

template <typename T>
struct generic_graph_node_list_iterator;

struct Node;
using graph_node_list = generic_graph_node_list<Node>;
using const_graph_node_list = generic_graph_node_list<const Node>;
using graph_node_list_iterator = generic_graph_node_list_iterator<Node>;
using const_graph_node_list_iterator = generic_graph_node_list_iterator<const Node>;

template <typename T>
struct generic_graph_node_list_iterator {
  generic_graph_node_list_iterator()
    : cur(nullptr), d(kNextDirection) {}
  generic_graph_node_list_iterator(T * cur, int d)
    : cur(cur), d(d) {}
  generic_graph_node_list_iterator(const generic_graph_node_list_iterator & rhs) = default;
  generic_graph_node_list_iterator(generic_graph_node_list_iterator && rhs) = default;
  generic_graph_node_list_iterator& operator=(const generic_graph_node_list_iterator & rhs) = default;
  generic_graph_node_list_iterator& operator=(generic_graph_node_list_iterator && rhs) = default;
  T * operator*() const { return cur; }
  T * operator->() const { return cur; }
  generic_graph_node_list_iterator & operator++() {
    JIT_ASSERT(cur);
    cur = cur->next_in_graph[d];
    return *this;
  }
  generic_graph_node_list_iterator operator++(int) {
    generic_graph_node_list_iterator old = *this;
    ++(*this);
    return old;
  }
  generic_graph_node_list_iterator & operator--() {
    JIT_ASSERT(cur);
    cur = cur->next_in_graph[reverseDir()];
    return *this;
  }
  generic_graph_node_list_iterator operator--(int) {
    generic_graph_node_list_iterator old = *this;
    --(*this);
    return old;
  }

  // erase cur without invalidating this iterator
  // named differently from destroy so that ->/. bugs do not
  // silently cause the wrong one to be called.
  // iterator will point to the previous entry after call
  void destroyCurrent() {
    T * n = cur;
    cur = cur->next_in_graph[reverseDir()];
    n->destroy();
  }
  generic_graph_node_list_iterator reverse() {
    return generic_graph_node_list_iterator(cur, reverseDir());
  }
private:
  int reverseDir() {
    return d == kNextDirection ? kPrevDirection : kNextDirection;
  }
  T * cur;
  int d; //direction 0 is forward 1 is reverse, see next_in_graph
};

template <typename T>
struct generic_graph_node_list {
  using iterator = generic_graph_node_list_iterator<T>;
  using const_iterator = generic_graph_node_list_iterator<const T>;
  generic_graph_node_list_iterator<T> begin() {
    return generic_graph_node_list_iterator<T>(head->next_in_graph[d], d);
  }
  generic_graph_node_list_iterator<const T> begin() const {
    return generic_graph_node_list_iterator<const T>(head->next_in_graph[d], d);
  }
  generic_graph_node_list_iterator<T> end() {
    return generic_graph_node_list_iterator<T>(head,d);
  }
  generic_graph_node_list_iterator<const T> end() const {
    return generic_graph_node_list_iterator<const T>(head,d);
  }
  generic_graph_node_list_iterator<T> rbegin() {
    return reverse().begin();
  }
  generic_graph_node_list_iterator<const T> rbegin() const {
    return reverse().begin();
  }
  generic_graph_node_list_iterator<T> rend() {
    return reverse().end();
  }
  generic_graph_node_list_iterator<const T> rend() const {
    return reverse().end();
  }
  generic_graph_node_list reverse() {
    return generic_graph_node_list(head, d == kNextDirection ? kPrevDirection : kNextDirection);
  }
  const generic_graph_node_list reverse() const {
    return generic_graph_node_list(head, d == kNextDirection ? kPrevDirection : kNextDirection);
  }
  T* front() { return head->next_in_graph[d]; }
  const T* front() const { return head->next_in_graph[d]; }
  T* back() { return head->next_in_graph[!d]; }
  const T* back() const { return head->next_in_graph[!d]; }
  generic_graph_node_list(T * head, int d)
    : head(head), d(d) {}
private:
  T * head;
  int d;
};

template <typename T>
static inline bool operator==(generic_graph_node_list_iterator<T> a, generic_graph_node_list_iterator<T> b) {
  return *a == *b;
}

template <typename T>
static inline bool operator!=(generic_graph_node_list_iterator<T> a, generic_graph_node_list_iterator<T> b) {
  return *a != *b;
}

}}

namespace std {

template<typename T>
struct iterator_traits<torch::jit::generic_graph_node_list_iterator<T>> {
  using difference_type = int64_t;
  using value_type = T*;
  using pointer = T**;
  using reference = T*&;
  using iterator_category = bidirectional_iterator_tag;
};

}