#include "ref.h"
#include "blob.h"
#include <library/cpp/yt/malloc/malloc.h>
#include <util/system/info.h>
#include <util/system/align.h>
namespace NYT {
////////////////////////////////////////////////////////////////////////////////
namespace NDetail {
// N.B. We would prefer these arrays to be zero sized
// but zero sized arrays are not supported in MSVC.
const char EmptyRefData[1] = {0};
char MutableEmptyRefData[1] = {0};
} // namespace NDetail
////////////////////////////////////////////////////////////////////////////////
class TBlobHolder
: public TSharedRangeHolder
{
public:
explicit TBlobHolder(TBlob&& blob)
: Blob_(std::move(blob))
{ }
// TSharedRangeHolder overrides.
std::optional<size_t> GetTotalByteSize() const override
{
return Blob_.Capacity();
}
private:
const TBlob Blob_;
};
////////////////////////////////////////////////////////////////////////////////
class TStringHolder
: public TSharedRangeHolder
{
public:
TStringHolder(TString&& string, TRefCountedTypeCookie cookie)
: String_(std::move(string))
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
, Cookie_(cookie)
#endif
{
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
TRefCountedTrackerFacade::AllocateTagInstance(Cookie_);
TRefCountedTrackerFacade::AllocateSpace(Cookie_, String_.length());
#endif
}
~TStringHolder()
{
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
TRefCountedTrackerFacade::FreeTagInstance(Cookie_);
TRefCountedTrackerFacade::FreeSpace(Cookie_, String_.length());
#endif
}
const TString& String() const
{
return String_;
}
// TSharedRangeHolder overrides.
std::optional<size_t> GetTotalByteSize() const override
{
return String_.capacity();
}
private:
const TString String_;
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
const TRefCountedTypeCookie Cookie_;
#endif
};
////////////////////////////////////////////////////////////////////////////////
template <class TDerived>
class TAllocationHolderBase
: public TSharedRangeHolder
{
public:
~TAllocationHolderBase()
{
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
TRefCountedTrackerFacade::FreeTagInstance(Cookie_);
TRefCountedTrackerFacade::FreeSpace(Cookie_, Size_);
#endif
}
TMutableRef GetRef()
{
return TMutableRef(static_cast<TDerived*>(this)->GetBegin(), Size_);
}
protected:
size_t Size_;
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
TRefCountedTypeCookie Cookie_;
#endif
void Initialize(
size_t size,
TSharedMutableRefAllocateOptions options,
TRefCountedTypeCookie cookie)
{
Size_ = size;
Cookie_ = cookie;
if (options.InitializeStorage) {
::memset(static_cast<TDerived*>(this)->GetBegin(), 0, Size_);
}
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
TRefCountedTrackerFacade::AllocateTagInstance(Cookie_);
TRefCountedTrackerFacade::AllocateSpace(Cookie_, Size_);
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
class TDefaultAllocationHolder
: public TAllocationHolderBase<TDefaultAllocationHolder>
, public TWithExtraSpace<TDefaultAllocationHolder>
{
public:
TDefaultAllocationHolder(
size_t size,
TSharedMutableRefAllocateOptions options,
TRefCountedTypeCookie cookie)
{
if (options.ExtendToUsableSize) {
if (auto usableSize = GetUsableSpaceSize(); usableSize != 0) {
size = usableSize;
}
}
Initialize(size, options, cookie);
}
char* GetBegin()
{
return static_cast<char*>(GetExtraSpacePtr());
}
// TSharedRangeHolder overrides.
std::optional<size_t> GetTotalByteSize() const override
{
return Size_;
}
};
////////////////////////////////////////////////////////////////////////////////
class TPageAlignedAllocationHolder
: public TAllocationHolderBase<TPageAlignedAllocationHolder>
{
public:
TPageAlignedAllocationHolder(
size_t size,
TSharedMutableRefAllocateOptions options,
TRefCountedTypeCookie cookie)
: Begin_(static_cast<char*>(::aligned_malloc(size, GetPageSize())))
{
Initialize(size, options, cookie);
}
~TPageAlignedAllocationHolder()
{
::free(Begin_);
}
char* GetBegin()
{
return Begin_;
}
// TSharedRangeHolder overrides.
std::optional<size_t> GetTotalByteSize() const override
{
return AlignUp(Size_, GetPageSize());
}
private:
char* const Begin_;
};
////////////////////////////////////////////////////////////////////////////////
TRef TRef::FromBlob(const TBlob& blob)
{
return TRef(blob.Begin(), blob.Size());
}
bool TRef::AreBitwiseEqual(TRef lhs, TRef rhs)
{
if (lhs.Size() != rhs.Size()) {
return false;
}
if (lhs.Size() == 0) {
return true;
}
return ::memcmp(lhs.Begin(), rhs.Begin(), lhs.Size()) == 0;
}
////////////////////////////////////////////////////////////////////////////////
TMutableRef TMutableRef::FromBlob(TBlob& blob)
{
return TMutableRef(blob.Begin(), blob.Size());
}
////////////////////////////////////////////////////////////////////////////////
TSharedRef TSharedRef::FromString(TString str, TRefCountedTypeCookie tagCookie)
{
auto holder = New<TStringHolder>(std::move(str), tagCookie);
auto ref = TRef::FromString(holder->String());
return TSharedRef(ref, std::move(holder));
}
TSharedRef TSharedRef::FromBlob(TBlob&& blob)
{
auto ref = TRef::FromBlob(blob);
auto holder = New<TBlobHolder>(std::move(blob));
return TSharedRef(ref, std::move(holder));
}
TSharedRef TSharedRef::MakeCopy(TRef ref, TRefCountedTypeCookie tagCookie)
{
if (!ref) {
return {};
}
if (ref.Empty()) {
return TSharedRef::MakeEmpty();
}
auto result = TSharedMutableRef::Allocate(ref.Size(), {.InitializeStorage = false}, tagCookie);
::memcpy(result.Begin(), ref.Begin(), ref.Size());
return result;
}
std::vector<TSharedRef> TSharedRef::Split(size_t partSize) const
{
YT_VERIFY(partSize > 0);
std::vector<TSharedRef> result;
result.reserve(Size() / partSize + 1);
auto sliceBegin = Begin();
while (sliceBegin < End()) {
auto sliceEnd = sliceBegin + partSize;
if (sliceEnd < sliceBegin || sliceEnd > End()) {
sliceEnd = End();
}
result.push_back(Slice(sliceBegin, sliceEnd));
sliceBegin = sliceEnd;
}
return result;
}
////////////////////////////////////////////////////////////////////////////////
TSharedMutableRef TSharedMutableRef::Allocate(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie)
{
auto holder = NewWithExtraSpace<TDefaultAllocationHolder>(size, size, options, tagCookie);
auto ref = holder->GetRef();
return TSharedMutableRef(ref, std::move(holder));
}
TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie)
{
auto holder = New<TPageAlignedAllocationHolder>(size, options, tagCookie);
auto ref = holder->GetRef();
return TSharedMutableRef(ref, std::move(holder));
}
TSharedMutableRef TSharedMutableRef::FromBlob(TBlob&& blob)
{
auto ref = TMutableRef::FromBlob(blob);
auto holder = New<TBlobHolder>(std::move(blob));
return TSharedMutableRef(ref, std::move(holder));
}
TSharedMutableRef TSharedMutableRef::MakeCopy(TRef ref, TRefCountedTypeCookie tagCookie)
{
if (!ref) {
return {};
}
if (ref.Empty()) {
return TSharedMutableRef::MakeEmpty();
}
auto result = Allocate(ref.Size(), {.InitializeStorage = false}, tagCookie);
::memcpy(result.Begin(), ref.Begin(), ref.Size());
return result;
}
////////////////////////////////////////////////////////////////////////////////
TString ToString(TRef ref)
{
return TString(ref.Begin(), ref.End());
}
TString ToString(const TMutableRef& ref)
{
return ToString(TRef(ref));
}
TString ToString(const TSharedRef& ref)
{
return ToString(TRef(ref));
}
TString ToString(const TSharedMutableRef& ref)
{
return ToString(TRef(ref));
}
size_t GetPageSize()
{
static const size_t PageSize = NSystemInfo::GetPageSize();
return PageSize;
}
size_t RoundUpToPage(size_t bytes)
{
return AlignUp<size_t>(bytes, GetPageSize());
}
size_t GetByteSize(const TSharedRefArray& array)
{
size_t size = 0;
if (array) {
for (const auto& part : array) {
size += part.Size();
}
}
return size;
}
////////////////////////////////////////////////////////////////////////////////
i64 TSharedRefArray::ByteSize() const
{
i64 result = 0;
if (*this) {
for (const auto& part : *this) {
result += part.Size();
}
}
return result;
}
std::vector<TSharedRef> TSharedRefArray::ToVector() const
{
if (!Impl_) {
return {};
}
return std::vector<TSharedRef>(Begin(), End());
}
TString TSharedRefArray::ToString() const
{
if (!Impl_) {
return {};
}
TString result;
size_t size = 0;
for (const auto& part : *this) {
size += part.size();
}
result.ReserveAndResize(size);
char* ptr = result.begin();
for (const auto& part : *this) {
size += part.size();
::memcpy(ptr, part.begin(), part.size());
ptr += part.size();
}
return result;
}
TSharedRefArray TSharedRefArray::MakeCopy(
const TSharedRefArray& array,
TRefCountedTypeCookie tagCookie)
{
TSharedRefArrayBuilder builder(
array.Size(),
array.ByteSize(),
tagCookie);
for (const auto& part : array) {
auto partCopy = builder.AllocateAndAdd(part.Size());
::memcpy(partCopy.Begin(), part.Begin(), part.Size());
}
return builder.Finish();
}
////////////////////////////////////////////////////////////////////////////////
TSharedRefArrayBuilder::TSharedRefArrayBuilder(
size_t size,
size_t poolCapacity,
TRefCountedTypeCookie tagCookie)
: AllocationCapacity_(poolCapacity)
, Impl_(TSharedRefArray::NewImpl(
size,
poolCapacity,
tagCookie,
size))
, CurrentAllocationPtr_(Impl_->GetBeginAllocationPtr())
{ }
void TSharedRefArrayBuilder::Add(TSharedRef part)
{
YT_ASSERT(CurrentPartIndex_ < Impl_->Size());
Impl_->MutableBegin()[CurrentPartIndex_++] = std::move(part);
}
TMutableRef TSharedRefArrayBuilder::AllocateAndAdd(size_t size)
{
YT_ASSERT(CurrentPartIndex_ < Impl_->Size());
YT_ASSERT(CurrentAllocationPtr_ + size <= Impl_->GetBeginAllocationPtr() + AllocationCapacity_);
TMutableRef ref(CurrentAllocationPtr_, size);
CurrentAllocationPtr_ += size;
TSharedRangeHolderPtr holder(Impl_.Get(), false);
TSharedRef sharedRef(ref, std::move(holder));
Add(std::move(sharedRef));
return ref;
}
TSharedRefArray TSharedRefArrayBuilder::Finish()
{
return TSharedRefArray(std::move(Impl_));
}
////////////////////////////////////////////////////////////////////////////////
} // namespace NYT