container.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
// Acknowledgement: This file originates from incubator-tvm
#ifndef MXNET_NODE_CONTAINER_H_
#define MXNET_NODE_CONTAINER_H_

#include <mxnet/node/node.h>

#include <type_traits>
#include <vector>
#include <initializer_list>
#include <unordered_map>
#include <utility>
#include <string>

namespace mxnet {

class ArrayNode : public Object {
 public:
  std::vector<ObjectRef> data;

  static constexpr const char* _type_key = "Array";
  MXNET_DECLARE_FINAL_OBJECT_INFO(ArrayNode, Object);
};

template<typename Converter,
         typename TIter>
class IterAdapter {
 public:
  using difference_type = typename std::iterator_traits<TIter>::difference_type;
  using value_type = typename Converter::ResultType;
  using pointer = typename Converter::ResultType*;
  using reference = typename Converter::ResultType&;   // NOLINT(*)
  using iterator_category = typename std::iterator_traits<TIter>::iterator_category;

  explicit IterAdapter(TIter iter) : iter_(iter) {}
  inline IterAdapter& operator++() {
    ++iter_;
    return *this;
  }
  inline IterAdapter operator+(difference_type offset) const {
    return IterAdapter(iter_ + offset);
  }

  template<typename T = IterAdapter>
  typename std::enable_if<std::is_same<iterator_category, std::random_access_iterator_tag>::value,
                          typename T::difference_type>::type
  inline operator-(const IterAdapter& rhs) const {
    return iter_ - rhs.iter_;
  }

  inline bool operator==(IterAdapter other) const {
    return iter_ == other.iter_;
  }
  inline bool operator!=(IterAdapter other) const {
    return !(*this == other);
  }
  inline const value_type operator*() const {
    return Converter::convert(*iter_);
  }

 private:
  TIter iter_;
};

template<typename T,
         typename = typename std::enable_if<std::is_base_of<ObjectRef, T>::value>::type >
class Array : public ObjectRef {
 public:
  Array() {
    data_ = make_object<ArrayNode>();
  }
  Array(Array<T> && other) {  // NOLINT(*)
    data_ = std::move(other.data_);
  }
  Array(const Array<T> &other) { // NOLINT(*)
    data_ = std::move(other.data_);
  }
  explicit Array(runtime::ObjectPtr<Object> n) : ObjectRef(n) {}
  template<typename IterType>
  Array(IterType begin, IterType end) {
    assign(begin, end);
  }
  Array(std::initializer_list<T> init) { // NOLINT(*)
    assign(init.begin(), init.end());
  }
  Array(const std::vector<T>& init) { // NOLINT(*)
    assign(init.begin(), init.end());
  }
  explicit Array(size_t n, const T& val) {
    auto tmp_node = make_object<ArrayNode>();
    for (size_t i = 0; i < n; ++i) {
      tmp_node->data.push_back(val);
    }
    data_ = std::move(tmp_node);
  }
  Array<T>& operator=(Array<T> && other) {
    data_ = std::move(other.data_);
    return *this;
  }
  Array<T>& operator=(const Array<T> & other) {
    data_ = other.data_;
    return *this;
  }
  template<typename IterType>
  void assign(IterType begin, IterType end) {
    auto n = make_object<ArrayNode>();
    for (IterType it = begin; it != end; ++it) {
      n->data.push_back(T(*it));
    }
    data_ = std::move(n);
  }
  inline const T operator[](size_t i) const {
    return DowncastNoCheck<T>(
        static_cast<const ArrayNode*>(data_.get())->data[i]);
  }
  inline size_t size() const {
    if (data_.get() == nullptr) return 0;
    return static_cast<const ArrayNode*>(data_.get())->data.size();
  }
  inline ArrayNode* CopyOnWrite() {
    if (data_.get() == nullptr || !data_.unique())  {
      runtime::ObjectPtr<ArrayNode> n = make_object<ArrayNode>();
      n->data = static_cast<ArrayNode*>(data_.get())->data;
      runtime::ObjectPtr<Object>(std::move(n)).swap(data_);
    }
    return static_cast<ArrayNode*>(data_.get());
  }
  inline void push_back(const T& item) {
    ArrayNode* n = this->CopyOnWrite();
    n->data.push_back(item);
  }
  inline void resize(size_t size) {
    ArrayNode* n = this->CopyOnWrite();
    n->data.resize(size);
  }
  inline void Set(size_t i, const T& value) {
    ArrayNode* n = this->CopyOnWrite();
    n->data[i] = value;
  }
  inline bool empty() const {
    return size() == 0;
  }
  template<typename F>
  inline void MutateByApply(F fmutate) {
    ArrayNode* ptr = static_cast<ArrayNode*>(data_.get());
    if (ptr == nullptr) return;
    if (data_.unique()) {
      // Copy on write optimization.
      // Perform inplace update because this is an unique copy.
      for (size_t i = 0; i < ptr->data.size(); ++i) {
        // It is important to use move here
        // to make prevent the element's ref count from increasing
        // so fmutate itself can perform copy-on-write optimization
        T old_elem = DowncastNoCheck<T>(std::move(ptr->data[i]));
        T new_elem = fmutate(std::move(old_elem));
        ptr->data[i] = std::move(new_elem);
      }
    } else {
      // lazily trigger copy if there is element change.
      runtime::ObjectPtr<ArrayNode> copy;
      for (size_t i = 0; i < ptr->data.size(); ++i) {
        T old_elem = DowncastNoCheck<T>(ptr->data[i]);
        T new_elem = fmutate(old_elem);
        if (!new_elem.same_as(ptr->data[i])) {
          // copy the old array
          if (copy == nullptr) {
            copy = runtime::make_object<ArrayNode>(*ptr);
          }
          copy->data[i] = std::move(new_elem);
        }
      }
      // replace the data with the new copy.
      if (copy != nullptr) {
        data_ = std::move(copy);
      }
    }
  }

  using ContainerType = ArrayNode;

  struct ValueConverter {
    using ResultType = T;
    static inline T convert(const ObjectRef& n) {
      return DowncastNoCheck<T>(n);
    }
  };
  using iterator = IterAdapter<ValueConverter,
                               std::vector<ObjectRef>::const_iterator>;

  using reverse_iterator = IterAdapter<
    ValueConverter,
    std::vector<ObjectRef>::const_reverse_iterator>;

  inline iterator begin() const {
    return iterator(static_cast<const ArrayNode*>(data_.get())->data.begin());
  }
  inline iterator end() const {
    return iterator(static_cast<const ArrayNode*>(data_.get())->data.end());
  }
  inline reverse_iterator rbegin() const {
    return reverse_iterator(static_cast<const ArrayNode*>(data_.get())->data.rbegin());
  }
  inline reverse_iterator rend() const {
    return reverse_iterator(static_cast<const ArrayNode*>(data_.get())->data.rend());
  }
};

}  // namespace mxnet
#endif  // MXNET_NODE_CONTAINER_H_