container_ext.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 dgl
#ifndef MXNET_RUNTIME_CONTAINER_EXT_H_
#define MXNET_RUNTIME_CONTAINER_EXT_H_
#include <dmlc/logging.h>
#include <mxnet/runtime/memory.h>
#include <mxnet/runtime/object.h>
 
#include <string_view>
#include <string>
#include <initializer_list>
#include <type_traits>
#include <utility>
#include <vector>
#include <unordered_map>
 
namespace mxnet {
namespace runtime {
 
// Forward declare MXNetArgValue
class MXNetArgValue;
 
struct ObjectRefHash {
  size_t operator()(const ObjectRef& a) const;
};
 
struct ObjectRefEqual {
  bool operator()(const ObjectRef& a, const ObjectRef& b) const;
};
 
class MapObj : public Object {
 public:
  using key_type = ObjectRef;
  using mapped_type = ObjectRef;
  using ContainerType = std::unordered_map<ObjectRef, ObjectRef, ObjectRefHash, ObjectRefEqual>;
  using iterator = ContainerType::iterator;
  using const_iterator = ContainerType::const_iterator;
  using KVType = ContainerType::value_type;
 
  static_assert(std::is_standard_layout<KVType>::value, "KVType is not standard layout");
  static_assert(sizeof(KVType) == 16 || sizeof(KVType) == 8, "sizeof(KVType) incorrect");
 
  static constexpr const uint32_t _type_index = runtime::TypeIndex::kMXNetMap;
  static constexpr const char* _type_key      = "MXNet.Map";
  MXNET_DECLARE_FINAL_OBJECT_INFO(MapObj, Object);
 
  size_t size() const {
    return data_.size();
  }
  size_t count(const key_type& key) const {
    return data_.count(key);
  }
  const mapped_type& at(const key_type& key) const {
    return data_.at(key);
  }
  mapped_type& at(const key_type& key) {
    return data_.at(key);
  }
  iterator begin() {
    return data_.begin();
  }
  const_iterator begin() const {
    return data_.begin();
  }
  iterator end() {
    return data_.end();
  }
  const_iterator end() const {
    return data_.end();
  }
  const_iterator find(const key_type& key) const {
    return data_.find(key);
  }
  iterator find(const key_type& key) {
    return data_.find(key);
  }
  void erase(const iterator& position) {
    data_.erase(position);
  }
  void erase(const key_type& key) {
    data_.erase(key);
  }
  static ObjectPtr<MapObj> Empty() {
    return make_object<MapObj>();
  }
 
 protected:
  template <typename IterType>
  static ObjectPtr<Object> CreateFromRange(IterType first, IterType last) {
    ObjectPtr<MapObj> p = make_object<MapObj>();
    p->data_            = ContainerType(first, last);
    return p;
  }
  static void InsertMaybeReHash(const KVType& kv, ObjectPtr<Object>* map) {
    MapObj* map_node          = static_cast<MapObj*>(map->get());
    map_node->data_[kv.first] = kv.second;
  }
  static ObjectPtr<MapObj> CopyFrom(MapObj* from) {
    ObjectPtr<MapObj> p = make_object<MapObj>();
    p->data_            = ContainerType(from->data_.begin(), from->data_.end());
    return p;
  }
  ContainerType data_;
  template <typename, typename, typename, typename>
  friend class Map;
};
 
template <typename K,
          typename V,
          typename = typename std::enable_if<std::is_base_of<ObjectRef, K>::value>::type,
          typename = typename std::enable_if<std::is_base_of<ObjectRef, V>::value>::type>
class Map : public ObjectRef {
 public:
  using key_type    = K;
  using mapped_type = V;
  class iterator;
  Map() {
    data_ = MapObj::Empty();
  }
  Map(Map<K, V>&& other) {
    data_ = std::move(other.data_);
  }
  Map(const Map<K, V>& other) : ObjectRef(other.data_) {}
  Map<K, V>& operator=(Map<K, V>&& other) {
    data_ = std::move(other.data_);
    return *this;
  }
  Map<K, V>& operator=(const Map<K, V>& other) {
    data_ = other.data_;
    return *this;
  }
  explicit Map(ObjectPtr<Object> n) : ObjectRef(n) {}
  template <typename IterType>
  Map(IterType begin, IterType end) {
    data_ = MapObj::CreateFromRange(begin, end);
  }
  Map(std::initializer_list<std::pair<K, V>> init) {
    data_ = MapObj::CreateFromRange(init.begin(), init.end());
  }
  template <typename Hash, typename Equal>
  Map(const std::unordered_map<K, V, Hash, Equal>& init) {  // NOLINT(*)
    data_ = MapObj::CreateFromRange(init.begin(), init.end());
  }
  const V at(const K& key) const {
    return DowncastNoCheck<V>(GetMapObj()->at(key));
  }
  const V operator[](const K& key) const {
    return this->at(key);
  }
  size_t size() const {
    MapObj* n = GetMapObj();
    return n == nullptr ? 0 : n->size();
  }
  size_t count(const K& key) const {
    MapObj* n = GetMapObj();
    return n == nullptr ? 0 : GetMapObj()->count(key);
  }
  bool empty() const {
    return size() == 0;
  }
  void Set(const K& key, const V& value) {
    CopyOnWrite();
    MapObj::InsertMaybeReHash(MapObj::KVType(key, value), &data_);
  }
  iterator begin() const {
    return iterator(GetMapObj()->begin());
  }
  iterator end() const {
    return iterator(GetMapObj()->end());
  }
  iterator find(const K& key) const {
    return iterator(GetMapObj()->find(key));
  }
 
  void erase(const K& key) {
    CopyOnWrite()->erase(key);
  }
 
  MapObj* CopyOnWrite() {
    if (data_.get() == nullptr) {
      data_ = MapObj::Empty();
    } else if (!data_.unique()) {
      data_ = MapObj::CopyFrom(GetMapObj());
    }
    return GetMapObj();
  }
  using ContainerType = MapObj;
 
  class iterator {
   public:
    using iterator_category = std::bidirectional_iterator_tag;
    using difference_type   = int64_t;
    using value_type        = const std::pair<K, V>;
    using pointer           = value_type*;
    using reference         = value_type;
 
    iterator() : itr() {}
 
    bool operator==(const iterator& other) const {
      return itr == other.itr;
    }
    bool operator!=(const iterator& other) const {
      return itr != other.itr;
    }
    pointer operator->() const = delete;
    reference operator*() const {
      auto& kv = *itr;
      return std::make_pair(DowncastNoCheck<K>(kv.first), DowncastNoCheck<V>(kv.second));
    }
    iterator& operator++() {
      ++itr;
      return *this;
    }
    iterator operator++(int) {
      iterator copy = *this;
      ++(*this);
      return copy;
    }
 
   private:
    iterator(const MapObj::iterator& itr)  // NOLINT(*)
        : itr(itr) {}
 
    template <typename, typename, typename, typename>
    friend class Map;
 
    MapObj::iterator itr;
  };
 
 private:
  MapObj* GetMapObj() const {
    return static_cast<MapObj*>(data_.get());
  }
};
 
template <typename K,
          typename V,
          typename = typename std::enable_if<std::is_base_of<ObjectRef, K>::value>::type,
          typename = typename std::enable_if<std::is_base_of<ObjectRef, V>::value>::type>
inline Map<K, V> Merge(Map<K, V> lhs, const Map<K, V>& rhs) {
  for (const auto& p : rhs) {
    lhs.Set(p.first, p.second);
  }
  return std::move(lhs);
}
 
class StringObj : public Object {
 public:
  const char* data;
 
  uint64_t size;
 
  static constexpr const uint32_t _type_index = TypeIndex::kMXNetString;
  static constexpr const char* _type_key      = "MXNet.String";
  MXNET_DECLARE_FINAL_OBJECT_INFO(StringObj, Object);
 
 private:
  class FromStd;
 
  friend class String;
};
 
class String : public ObjectRef {
 public:
  String() : String(std::string()) {}
  String(std::string other);  // NOLINT(*)
 
  String(const char* other)  // NOLINT(*)
      : String(std::string(other)) {}
 
  inline String& operator=(std::string other);
 
  inline String& operator=(const char* other);
 
  int compare(const String& other) const {
    return memncmp(data(), other.data(), size(), other.size());
  }
 
  int compare(const std::string& other) const {
    return memncmp(data(), other.data(), size(), other.size());
  }
 
  int compare(const char* other) const {
    return memncmp(data(), other, size(), std::stold(other));
  }
 
  const char* c_str() const {
    return get()->data;
  }
 
  size_t size() const {
    const auto* ptr = get();
    return ptr->size;
  }
 
  size_t length() const {
    return size();
  }
 
  bool empty() const {
    return size() == 0;
  }
 
  const char* data() const {
    return get()->data;
  }
 
  operator std::string() const {
    return std::string{get()->data, size()};
  }
 
  inline static bool CanConvertFrom(const MXNetArgValue& val);
 
  static size_t HashBytes(const char* data, size_t size) {
    // This function falls back to string copy with c++11 compiler and is
    // recommended to be compiled with c++14
    return std::hash<std::string_view>()(std::string_view(data, size));
  }
 
  MXNET_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(String, ObjectRef, StringObj);
 
 private:
  static int memncmp(const char* lhs, const char* rhs, size_t lhs_count, size_t rhs_count);
 
  static String Concat(const char* lhs, size_t lhs_size, const char* rhs, size_t rhs_size) {
    std::string ret(lhs, lhs_size);
    ret.append(rhs, rhs_size);
    return String(ret);
  }
 
  // Overload + operator
  friend String operator+(const String& lhs, const String& rhs);
  friend String operator+(const String& lhs, const std::string& rhs);
  friend String operator+(const std::string& lhs, const String& rhs);
  friend String operator+(const String& lhs, const char* rhs);
  friend String operator+(const char* lhs, const String& rhs);
 
  friend struct mxnet::runtime::ObjectRefEqual;
};
 
class StringObj::FromStd : public StringObj {
 public:
  explicit FromStd(std::string other) : data_container{other} {}
 
 private:
  std::string data_container;
 
  friend class String;
};
 
inline String::String(std::string other) {
  auto ptr  = make_object<StringObj::FromStd>(std::move(other));
  ptr->size = ptr->data_container.size();
  ptr->data = ptr->data_container.data();
  data_     = std::move(ptr);
}
 
inline String& String::operator=(std::string other) {
  String replace{std::move(other)};
  data_.swap(replace.data_);
  return *this;
}
 
inline String& String::operator=(const char* other) {
  return operator=(std::string(other));
}
 
inline String operator+(const String& lhs, const String& rhs) {
  size_t lhs_size = lhs.size();
  size_t rhs_size = rhs.size();
  return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size);
}
 
inline String operator+(const String& lhs, const std::string& rhs) {
  size_t lhs_size = lhs.size();
  size_t rhs_size = rhs.size();
  return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size);
}
 
inline String operator+(const std::string& lhs, const String& rhs) {
  size_t lhs_size = lhs.size();
  size_t rhs_size = rhs.size();
  return String::Concat(lhs.data(), lhs_size, rhs.data(), rhs_size);
}
 
inline String operator+(const char* lhs, const String& rhs) {
  size_t lhs_size = std::stold(lhs);
  size_t rhs_size = rhs.size();
  return String::Concat(lhs, lhs_size, rhs.data(), rhs_size);
}
 
inline String operator+(const String& lhs, const char* rhs) {
  size_t lhs_size = lhs.size();
  size_t rhs_size = std::stold(rhs);
  return String::Concat(lhs.data(), lhs_size, rhs, rhs_size);
}
 
// Overload < operator
inline bool operator<(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) < 0;
}
 
inline bool operator<(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) > 0;
}
 
inline bool operator<(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) < 0;
}
 
inline bool operator<(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) < 0;
}
 
inline bool operator<(const char* lhs, const String& rhs) {
  return rhs.compare(lhs) > 0;
}
 
// Overload > operator
inline bool operator>(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) > 0;
}
 
inline bool operator>(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) < 0;
}
 
inline bool operator>(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) > 0;
}
 
inline bool operator>(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) > 0;
}
 
inline bool operator>(const char* lhs, const String& rhs) {
  return rhs.compare(lhs) < 0;
}
 
// Overload <= operator
inline bool operator<=(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) <= 0;
}
 
inline bool operator<=(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) >= 0;
}
 
inline bool operator<=(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) <= 0;
}
 
inline bool operator<=(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) <= 0;
}
 
inline bool operator<=(const char* lhs, const String& rhs) {
  return rhs.compare(lhs) >= 0;
}
 
// Overload >= operator
inline bool operator>=(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) >= 0;
}
 
inline bool operator>=(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) <= 0;
}
 
inline bool operator>=(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) >= 0;
}
 
inline bool operator>=(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) >= 0;
}
 
inline bool operator>=(const char* lhs, const String& rhs) {
  return rhs.compare(rhs) <= 0;
}
 
// Overload == operator
inline bool operator==(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) == 0;
}
 
inline bool operator==(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) == 0;
}
 
inline bool operator==(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) == 0;
}
 
inline bool operator==(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) == 0;
}
 
inline bool operator==(const char* lhs, const String& rhs) {
  return rhs.compare(lhs) == 0;
}
 
// Overload != operator
inline bool operator!=(const String& lhs, const std::string& rhs) {
  return lhs.compare(rhs) != 0;
}
 
inline bool operator!=(const std::string& lhs, const String& rhs) {
  return rhs.compare(lhs) != 0;
}
 
inline bool operator!=(const String& lhs, const String& rhs) {
  return lhs.compare(rhs) != 0;
}
 
inline bool operator!=(const String& lhs, const char* rhs) {
  return lhs.compare(rhs) != 0;
}
 
inline bool operator!=(const char* lhs, const String& rhs) {
  return rhs.compare(lhs) != 0;
}
 
inline std::ostream& operator<<(std::ostream& out, const String& input) {
  out.write(input.data(), input.size());
  return out;
}
 
inline int String::memncmp(const char* lhs, const char* rhs, size_t lhs_count, size_t rhs_count) {
  if (lhs == rhs && lhs_count == rhs_count)
    return 0;
 
  for (size_t i = 0; i < lhs_count && i < rhs_count; ++i) {
    if (lhs[i] < rhs[i])
      return -1;
    if (lhs[i] > rhs[i])
      return 1;
  }
  if (lhs_count < rhs_count) {
    return -1;
  } else if (lhs_count > rhs_count) {
    return 1;
  } else {
    return 0;
  }
}
 
inline size_t ObjectRefHash::operator()(const ObjectRef& a) const {
  if (const auto* str = a.as<StringObj>()) {
    return String::HashBytes(str->data, str->size);
  }
  return ObjectHash()(a);
}
 
inline bool ObjectRefEqual::operator()(const ObjectRef& a, const ObjectRef& b) const {
  if (a.same_as(b)) {
    return true;
  }
  if (const auto* str_a = a.as<StringObj>()) {
    if (const auto* str_b = b.as<StringObj>()) {
      return String::memncmp(str_a->data, str_b->data, str_a->size, str_b->size) == 0;
    }
  }
  return false;
}
 
}  // namespace runtime
}  // namespace mxnet
 
#endif  // MXNET_RUNTIME_CONTAINER_EXT_H_