[REFACTPR] Update the MapObj to be generic and allow subclass

Author: tqchen

include/tvm/ffi/container/container_details.h

@@ -36,134 +36,6 @@
 namespace tvm {
 namespace ffi {
 namespace details {
-/*!
- * \brief Base template for classes with array like memory layout.
- *
- *        It provides general methods to access the memory. The memory
- *        layout is ArrayType + [ElemType]. The alignment of ArrayType
- *        and ElemType is handled by the memory allocator.
- *
- * \tparam ArrayType The array header type, contains object specific metadata.
- * \tparam ElemType The type of objects stored in the array right after
- * ArrayType.
- *
- * \code{.cpp}
- * // Example usage of the template to define a simple array wrapper
- * class ArrayObj : public tvm::ffi::details::InplaceArrayBase<ArrayObj, Elem> {
- * public:
- *  // Wrap EmplaceInit to initialize the elements
- *  template <typename Iterator>
- *  void Init(Iterator begin, Iterator end) {
- *   size_t num_elems = std::distance(begin, end);
- *   auto it = begin;
- *   this->size = 0;
- *   for (size_t i = 0; i < num_elems; ++i) {
- *     InplaceArrayBase::EmplaceInit(i, *it++);
- *     this->size++;
- *   }
- *  }
- * }
- *
- * void test_function() {
- *   vector<Elem> fields;
- *   auto ptr = make_inplace_array_object<ArrayObj, Elem>(fields.size());
- *   ptr->Init(fields.begin(), fields.end());
- *
- *   // Access the 0th element in the array.
- *   assert(ptr->operator[](0) == fields[0]);
- * }
- * \endcode
- */
-template <typename ArrayType, typename ElemType>
-class InplaceArrayBase {
- public:
-  /*!
-   * \brief Access element at index
-   * \param idx The index of the element.
-   * \return Const reference to ElemType at the index.
-   */
-  const ElemType& operator[](size_t idx) const {
-    size_t size = Self()->GetSize();
-    if (idx > size) {
-      TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds " << size;
-    }
-    return *(reinterpret_cast<ElemType*>(AddressOf(idx)));
-  }
-
-  /*!
-   * \brief Access element at index
-   * \param idx The index of the element.
-   * \return Reference to ElemType at the index.
-   */
-  ElemType& operator[](size_t idx) {
-    size_t size = Self()->GetSize();
-    if (idx > size) {
-      TVM_FFI_THROW(IndexError) << "Index " << idx << " out of bounds " << size;
-    }
-    return *(reinterpret_cast<ElemType*>(AddressOf(idx)));
-  }
-
-  /*!
-   * \brief Destroy the Inplace Array Base object
-   */
-  ~InplaceArrayBase() {
-    if constexpr (!(std::is_standard_layout_v<ElemType> && std::is_trivial_v<ElemType>)) {
-      size_t size = Self()->GetSize();
-      for (size_t i = 0; i < size; ++i) {
-        ElemType* fp = reinterpret_cast<ElemType*>(AddressOf(i));
-        fp->ElemType::~ElemType();
-      }
-    }
-  }
-
- private:
-  InplaceArrayBase() = default;
-  friend ArrayType;
-
- protected:
-  /*!
-   * \brief Construct a value in place with the arguments.
-   *
-   * \tparam Args Type parameters of the arguments.
-   * \param idx Index of the element.
-   * \param args Arguments to construct the new value.
-   *
-   * \note Please make sure ArrayType::GetSize returns 0 before first call of
-   * EmplaceInit, and increment GetSize by 1 each time EmplaceInit succeeds.
-   */
-  template <typename... Args>
-  void EmplaceInit(size_t idx, Args&&... args) {
-    void* field_ptr = AddressOf(idx);
-    new (field_ptr) ElemType(std::forward<Args>(args)...);
-  }
-
-  /*!
-   * \brief Return the self object for the array.
-   *
-   * \return Pointer to ArrayType.
-   */
-  inline ArrayType* Self() const {
-    return static_cast<ArrayType*>(const_cast<InplaceArrayBase*>(this));
-  }
-
-  /*!
-   * \brief Return the raw pointer to the element at idx.
-   *
-   * \param idx The index of the element.
-   * \return Raw pointer to the element.
-   */
-  void* AddressOf(size_t idx) const {
-    static_assert(
-        alignof(ArrayType) % alignof(ElemType) == 0 && sizeof(ArrayType) % alignof(ElemType) == 0,
-        "The size and alignment of ArrayType should respect "
-        "ElemType's alignment.");
-
-    size_t kDataStart = sizeof(ArrayType);
-    ArrayType* self = Self();
-    char* data_start = reinterpret_cast<char*>(self) + kDataStart;
-    return data_start + idx * sizeof(ElemType);
-  }
-};
 
 /*!
  * \brief iterator adapter that adapts TIter to return another type.

include/tvm/ffi/container/map.h

@@ -31,12 +31,25 @@
 #include <tvm/ffi/object.h>
 #include <tvm/ffi/optional.h>
 
-#include <type_traits>
 #include <unordered_map>
 
 namespace tvm {
 namespace ffi {
 
+/*! \brief Map object */
+class MapObj : public MapBaseObj {
+ public:
+  /// \cond Doxygen_Suppress
+  static constexpr const int32_t _type_index = TypeIndex::kTVMFFIMap;
+  static const constexpr bool _type_final = true;
+  TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIMap, MapObj, Object);
+  /// \endcond
+
+ protected:
+  template <typename, typename, typename>
+  friend class Map;
+};
+
 /*!
  * \brief Map container of NodeRef->NodeRef in DSL graph.
  *  Map implements copy on write semantics, which means map is mutable
@@ -64,7 +77,7 @@ class Map : public ObjectRef {
   /*!
    * \brief default constructor
    */
-  Map() { data_ = MapObj::Empty(); }
+  Map() { data_ = MapObj::Empty<MapObj>(); }
   /*!
    * \brief move constructor
    * \param other source
@@ -159,22 +172,22 @@ class Map : public ObjectRef {
    */
   template <typename IterType>
   Map(IterType begin, IterType end) {
-    data_ = MapObj::CreateFromRange(begin, end);
+    data_ = MapObj::CreateFromRange<MapObj>(begin, end);
   }
   /*!
    * \brief constructor from initializer list
    * \param init The initalizer list
    */
   Map(std::initializer_list<std::pair<K, V>> init) {
-    data_ = MapObj::CreateFromRange(init.begin(), init.end());
+    data_ = MapObj::CreateFromRange<MapObj>(init.begin(), init.end());
   }
   /*!
    * \brief constructor from unordered_map
    * \param init The unordered_map
    */
   template <typename Hash, typename Equal>
   Map(const std::unordered_map<K, V, Hash, Equal>& init) {  // NOLINT(*)
-    data_ = MapObj::CreateFromRange(init.begin(), init.end());
+    data_ = MapObj::CreateFromRange<MapObj>(init.begin(), init.end());
   }
   /*!
    * \brief Read element from map.
@@ -206,7 +219,7 @@ class Map : public ObjectRef {
   void clear() {
     MapObj* n = GetMapObj();
     if (n != nullptr) {
-      data_ = MapObj::Empty();
+      data_ = MapObj::Empty<MapObj>();
     }
   }
   /*!
@@ -216,7 +229,7 @@ class Map : public ObjectRef {
    */
   void Set(const K& key, const V& value) {
     CopyOnWrite();
-    MapObj::InsertMaybeReHash(MapObj::KVType(key, value), &data_);
+    MapObj::InsertMaybeReHash<MapObj>(MapObj::KVType(key, value), &data_);
   }
   /*! \return begin iterator */
   iterator begin() const { return iterator(GetMapObj()->begin()); }
@@ -249,9 +262,9 @@ class Map : public ObjectRef {
    */
   MapObj* CopyOnWrite() {
     if (data_.get() == nullptr) {
-      data_ = MapObj::Empty();
+      data_ = MapObj::Empty<MapObj>();
     } else if (!data_.unique()) {
-      data_ = MapObj::CopyFrom(GetMapObj());
+      data_ = MapObj::CopyFrom<MapObj>(GetMapObj());
     }
     return GetMapObj();
   }