Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 56 additions & 7 deletions ffi/include/tvm/ffi/container/array.h
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,18 @@ namespace tvm {
namespace ffi {

/*! \brief array node content in array */
class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, Any> {
class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, TVMFFIAny> {
public:
~ArrayObj() {
Any* begin = MutableBegin();
for (int64_t i = 0; i < size_; ++i) {
(begin + i)->Any::~Any();
}
if (data_deleter_ != nullptr) {
data_deleter_(data_);
}
}

/*! \return The size of the array */
size_t size() const { return this->size_; }

Expand All @@ -52,10 +62,22 @@ class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, Any>
* \param i The index
* \return the i-th element.
*/
const Any at(int64_t i) const { return this->operator[](i); }
const Any& at(int64_t i) const { return this->operator[](i); }

/*!
* \brief Read i-th element from array.
* \param i The index
* \return the i-th element.
*/
const Any& operator[](int64_t i) const {
if (i >= size_) {
TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << size_;
}
return static_cast<Any*>(data_)[i];
}

/*! \return begin constant iterator */
const Any* begin() const { return static_cast<Any*>(InplaceArrayBase::AddressOf(0)); }
const Any* begin() const { return static_cast<Any*>(data_); }

/*! \return end constant iterator */
const Any* end() const { return begin() + size_; }
Expand All @@ -68,7 +90,12 @@ class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, Any>
* \param i The index
* \param item The value to be set
*/
void SetItem(int64_t i, Any item) { this->operator[](i) = std::move(item); }
void SetItem(int64_t i, Any item) {
if (i >= size_) {
TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << size_;
}
static_cast<Any*>(data_)[i] = std::move(item);
}

/*!
* \brief Constructs a container and copy from another
Expand Down Expand Up @@ -138,21 +165,31 @@ class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, Any>
size_t GetSize() const { return this->size_; }

/*! \return begin mutable iterator */
Any* MutableBegin() const { return static_cast<Any*>(InplaceArrayBase::AddressOf(0)); }
Any* MutableBegin() const { return static_cast<Any*>(this->data_); }

/*! \return end mutable iterator */
Any* MutableEnd() const { return MutableBegin() + size_; }

/*!
* \brief Emplace a new element at the back of the array
* \param args The arguments to construct the new element
*/
template <typename... Args>
void EmplaceInit(size_t idx, Args&&... args) {
Any* itr = MutableBegin() + idx;
new (itr) Any(std::forward<Args>(args)...);
}

/*!
* \brief Create an ArrayObj with the given capacity.
* \param n Required capacity
* \return Ref-counted ArrayObj requested
*/
static ObjectPtr<ArrayObj> Empty(int64_t n = kInitSize) {
TVM_FFI_ICHECK_GE(n, 0);
ObjectPtr<ArrayObj> p = make_inplace_array_object<ArrayObj, Any>(n);
p->capacity_ = n;
p->size_ = 0;
p->data_ = p->AddressOf(0);
return p;
}

Expand Down Expand Up @@ -235,11 +272,17 @@ class ArrayObj : public Object, public details::InplaceArrayBase<ArrayObj, Any>
return this;
}

/*! \brief Data pointer to the first element of the array */
void* data_;
/*! \brief Number of elements used */
int64_t size_;

/*! \brief Number of elements allocated */
int64_t capacity_;
/*!
* \brief Optional data deleter when data is allocated separately
* and its deletion is not managed by ArrayObj::deleter_.
*/
void (*data_deleter_)(void*) = nullptr;

/*! \brief Initial size of ArrayObj */
static constexpr int64_t kInitSize = 4;
Expand Down Expand Up @@ -472,6 +515,12 @@ class Array : public ObjectRef {
p->EmplaceInit(p->size_++, item);
}

template <typename... Args>
void emplace_back(Args&&... args) {
ArrayObj* p = CopyOnWrite(1);
p->EmplaceInit(p->size_++, std::forward<Args>(args)...);
}

/*!
* \brief Insert an element into the given position
* \param position An iterator pointing to the insertion point
Expand Down
Loading
Loading