Made SharedBuffer/UniqueBuffer share guts with IoBuffer

This enables way more efficient marshaling of compact binary objects and attachments

8 files changed, 101 insertions, 125 deletions

diff --git a/zencore/compactbinary.cpp b/zencore/compactbinary.cpp
index 4ee9e9281..cae6c506d 100644
--- a/zencore/compactbinary.cpp
+++ b/zencore/compactbinary.cpp
@@ -1079,7 +1079,7 @@ LoadCompactBinary(BinaryReader& Ar, BufferAllocator Allocator)
 	// Allocate the buffer, copy the header, and read the remainder of the field.
 	UniqueBuffer Buffer = Allocator(FieldSize);
 	ZEN_ASSERT(Buffer.GetSize() == FieldSize);
-	MutableMemoryView View = Buffer.GetView();
+	MutableMemoryView View = Buffer.GetMutableView();
 	memcpy(View.GetData(), HeaderBytes.data(), HeaderBytes.size());
 	View.RightChopInline(HeaderBytes.size());
 	if (!View.IsEmpty())
diff --git a/zencore/compactbinarypackage.cpp b/zencore/compactbinarypackage.cpp
index 9eb67ac93..f84137ff6 100644
--- a/zencore/compactbinarypackage.cpp
+++ b/zencore/compactbinarypackage.cpp
@@ -145,7 +145,7 @@ CbAttachment::Load(BinaryReader& Reader, BufferAllocator Allocator)
 		std::vector<uint8_t> HashBuffer;
 		CbField				 HashField = LoadCompactBinary(Reader, [&HashBuffer](uint64_t Size) -> UniqueBuffer {
 			 HashBuffer.resize(Size);
-			 return UniqueBuffer::MakeView(HashBuffer.data(), Size);
+			 return UniqueBuffer::MakeMutableView(HashBuffer.data(), Size);
 		 });
 		Hash						   = HashField.AsAttachment();
 		ZEN_ASSERT(!HashField.HasError(), "Attachments must be a non-empty binary value with a content hash.");
@@ -351,7 +351,7 @@ CbPackage::Load(BinaryReader& Reader, BufferAllocator Allocator, AttachmentResol
 	const auto StackAllocator = [&Allocator, &StackBuffer](uint64_t Size) -> UniqueBuffer {
 		if (Size <= sizeof(StackBuffer))
 		{
-			return UniqueBuffer::MakeView(StackBuffer, Size);
+			return UniqueBuffer::MakeMutableView(StackBuffer, Size);
 		}
 
 		return Allocator(Size);
diff --git a/zencore/include/zencore/iobuffer.h b/zencore/include/zencore/iobuffer.h
index c93d27959..f38af3c27 100644
--- a/zencore/include/zencore/iobuffer.h
+++ b/zencore/include/zencore/iobuffer.h
@@ -65,13 +65,38 @@ public:
 			Materialize();
 	}
 
-	inline bool IsOwned() const { return !!(m_Flags & kIsOwned); }
+	inline bool IsOwnedByThis() const { return !!(m_Flags & kIsOwnedByThis); }
+
+	inline void SetIsOwnedByThis(bool NewState)
+	{
+		if (NewState)
+		{
+			m_Flags |= kIsOwnedByThis;
+		}
+		else
+		{
+			m_Flags &= ~kIsOwnedByThis;
+		}
+	}
+
+	inline bool IsOwned() const
+	{
+		bool ThisIsOwned = !!(m_Flags & kIsOwnedByThis);
+		if (ThisIsOwned)
+		{
+			return true;
+		}
+		return m_OuterCore && m_OuterCore->IsOwned();
+	}
+
 	inline bool IsImmutable() const { return !(m_Flags & kIsMutable); }
 	inline bool IsNull() const { return m_DataBytes == 0; }
 
 	inline IoBufferExtendedCore*	   ExtendedCore();
 	inline const IoBufferExtendedCore* ExtendedCore() const;
 
+	ZENCORE_API void* MutableDataPointer() const;
+
 	inline const void* DataPointer() const
 	{
 		EnsureDataValid();
@@ -86,18 +111,6 @@ public:
 		m_DataBytes = Sz;
 	}
 
-	inline void SetIsOwned(bool NewState)
-	{
-		if (NewState)
-		{
-			m_Flags |= kIsOwned;
-		}
-		else
-		{
-			m_Flags &= ~kIsOwned;
-		}
-	}
-
 	inline void SetIsImmutable(bool NewState)
 	{
 		if (!NewState)
@@ -121,7 +134,7 @@ protected:
 
 	enum Flags
 	{
-		kIsOwned		= 1 << 0,
+		kIsOwnedByThis	= 1 << 0,
 		kIsMutable		= 1 << 1,
 		kIsExtended		= 1 << 2,  // Is actually a SharedBufferExtendedCore
 		kIsMaterialized = 1 << 3,  // Data pointers are valid
@@ -240,7 +253,7 @@ public:
 
 	inline IoBuffer(EAssumeOwnershipTag, const void* DataPtr, size_t Sz) : m_Core(new IoBufferCore(DataPtr, Sz))
 	{
-		m_Core->SetIsOwned(true);
+		m_Core->SetIsOwnedByThis(true);
 	}
 
 	ZENCORE_API IoBuffer(EFileTag, void* FileHandle, uint64_t ChunkFileOffset, uint64_t ChunkSize);
@@ -255,6 +268,8 @@ public:
 
 private:
 	RefPtr<IoBufferCore> m_Core = new IoBufferCore;
+
+	friend class SharedBuffer;
 };
 
 class IoBufferBuilder
diff --git a/zencore/include/zencore/sharedbuffer.h b/zencore/include/zencore/sharedbuffer.h
index c6206f780..093f43231 100644
--- a/zencore/include/zencore/sharedbuffer.h
+++ b/zencore/include/zencore/sharedbuffer.h
@@ -4,6 +4,7 @@
 
 #include "zencore.h"
 
+#include <zencore/iobuffer.h>
 #include <zencore/memory.h>
 #include <zencore/refcount.h>
 
@@ -11,49 +12,10 @@
 
 namespace zen {
 
-class BufferOwner : public RefCounted
-{
-protected:
-	inline BufferOwner(void* DataPtr, uint64_t DataSize, bool IsOwned, BufferOwner* OuterBuffer = nullptr)
-	: m_IsOwned(IsOwned)
-	, m_Data(DataPtr)
-	, m_Size(DataSize)
-	, m_Outer(OuterBuffer)
-	{
-	}
-
-	virtual ~BufferOwner();
-
-	// Ownership is a transitive property, and m_IsOwned currently only flags that this instance is responsible
-	// for managing the allocated memory, so we need to make recursive calls. Could be optimized slightly by
-	// adding a dedicated flag
-	inline bool IsOwned() const
-	{
-		if (m_IsOwned)
-		{
-			return true;
-		}
-		else
-		{
-			return m_Outer && m_Outer->IsOwned();
-		}
-	}
-
-	BufferOwner(const BufferOwner&) = delete;
-	BufferOwner& operator=(const BufferOwner&) = delete;
-
-private:
-	bool				m_IsOwned;
-	void*				m_Data;
-	uint64_t			m_Size;
-	RefPtr<BufferOwner> m_Outer;
-
-	friend class UniqueBuffer;
-	friend class SharedBuffer;
-};
-
 /**
- * Reference to a memory buffer with a single owner (see std::unique_ptr)
+ * Reference to a memory buffer with a single owner
+ *
+ * Internally
  */
 class UniqueBuffer
 {
@@ -62,24 +24,25 @@ public:
 	UniqueBuffer& operator=(const UniqueBuffer&) = delete;
 
 	UniqueBuffer() = default;
-	ZENCORE_API explicit UniqueBuffer(BufferOwner* Owner);
+	ZENCORE_API explicit UniqueBuffer(IoBufferCore* Owner);
 
-	void*		GetData() { return m_buffer->m_Data; }
-	const void* GetData() const { return m_buffer->m_Data; }
-	size_t		GetSize() const { return m_buffer->m_Size; }
-				operator MutableMemoryView() { return GetView(); }
-				operator MemoryView() const { return MemoryView(m_buffer->m_Data, m_buffer->m_Size); }
+	void*		GetData() { return m_Buffer->MutableDataPointer(); }
+	const void* GetData() const { return m_Buffer->DataPointer(); }
+	size_t		GetSize() const { return m_Buffer->DataBytes(); }
+				operator MutableMemoryView() { return GetMutableView(); }
+				operator MemoryView() const { return GetView(); }
 
-	MutableMemoryView GetView() { return MutableMemoryView(m_buffer->m_Data, m_buffer->m_Size); }
+	inline MutableMemoryView GetMutableView() { return MutableMemoryView(m_Buffer->MutableDataPointer(), m_Buffer->DataBytes()); }
+	inline MemoryView		 GetView() const { return MemoryView(m_Buffer->MutableDataPointer(), m_Buffer->DataBytes()); }
 
 	/** Make an uninitialized owned buffer of the specified size. */
 	ZENCORE_API static UniqueBuffer Alloc(uint64_t Size);
 
 	/** Make a non-owned view of the input. */
-	ZENCORE_API static UniqueBuffer MakeView(void* DataPtr, uint64_t Size);
+	ZENCORE_API static UniqueBuffer MakeMutableView(void* DataPtr, uint64_t Size);
 
 private:
-	RefPtr<BufferOwner> m_buffer;
+	RefPtr<IoBufferCore> m_Buffer;
 
 	friend class SharedBuffer;
 };
@@ -92,46 +55,38 @@ class SharedBuffer
 public:
 	SharedBuffer() = default;
 	ZENCORE_API explicit SharedBuffer(UniqueBuffer&&);
-	inline explicit SharedBuffer(BufferOwner* Owner) : m_buffer(Owner) {}
-
-	void* GetData()
-	{
-		if (m_buffer)
-		{
-			return m_buffer->m_Data;
-		}
-		return nullptr;
-	}
+	inline explicit SharedBuffer(IoBufferCore* Owner) : m_Buffer(Owner) {}
+	ZENCORE_API explicit SharedBuffer(IoBuffer&& Buffer) : m_Buffer(std::move(Buffer.m_Core)) {}
 
 	const void* GetData() const
 	{
-		if (m_buffer)
+		if (m_Buffer)
 		{
-			return m_buffer->m_Data;
+			return m_Buffer->DataPointer();
 		}
 		return nullptr;
 	}
 
 	size_t GetSize() const
 	{
-		if (m_buffer)
+		if (m_Buffer)
 		{
-			return m_buffer->m_Size;
+			return m_Buffer->DataBytes();
 		}
 		return 0;
 	}
 
 	ZENCORE_API void MakeOwned();
-	bool			 IsOwned() const { return m_buffer && m_buffer->IsOwned(); }
-	inline explicit	 operator bool() const { return m_buffer; }
-	inline bool		 IsNull() const { return !m_buffer; }
-	inline void		 Reset() { m_buffer = nullptr; }
+	bool			 IsOwned() const { return m_Buffer && m_Buffer->IsOwned(); }
+	inline explicit	 operator bool() const { return m_Buffer; }
+	inline bool		 IsNull() const { return !m_Buffer; }
+	inline void		 Reset() { m_Buffer = nullptr; }
 
 	MemoryView GetView() const
 	{
-		if (m_buffer)
+		if (m_Buffer)
 		{
-			return MemoryView(m_buffer->m_Data, m_buffer->m_Size);
+			return MemoryView(m_Buffer->DataPointer(), m_Buffer->DataBytes());
 		}
 		else
 		{
@@ -143,7 +98,7 @@ public:
 
 	SharedBuffer& operator=(UniqueBuffer&& Rhs)
 	{
-		m_buffer = std::move(Rhs.m_buffer);
+		m_Buffer = std::move(Rhs.m_Buffer);
 		return *this;
 	}
 
@@ -161,7 +116,7 @@ public:
 	ZENCORE_API static SharedBuffer Clone(MemoryView View);
 
 private:
-	RefPtr<BufferOwner> m_buffer;
+	RefPtr<IoBufferCore> m_Buffer;
 };
 
 void sharedbuffer_forcelink();
diff --git a/zencore/iobuffer.cpp b/zencore/iobuffer.cpp
index a42dd83f4..ec5d599b4 100644
--- a/zencore/iobuffer.cpp
+++ b/zencore/iobuffer.cpp
@@ -52,7 +52,7 @@ IoBufferCore::IoBufferCore(size_t InSize)
 	m_DataPtr	= AllocateBuffer(InSize, sizeof(void*));
 	m_DataBytes = InSize;
 
-	SetIsOwned(true);
+	SetIsOwnedByThis(true);
 }
 
 IoBufferCore::IoBufferCore(size_t InSize, size_t Alignment)
@@ -60,12 +60,12 @@ IoBufferCore::IoBufferCore(size_t InSize, size_t Alignment)
 	m_DataPtr	= AllocateBuffer(InSize, Alignment);
 	m_DataBytes = InSize;
 
-	SetIsOwned(true);
+	SetIsOwnedByThis(true);
 }
 
 IoBufferCore::~IoBufferCore()
 {
-	if (IsOwned() && m_DataPtr)
+	if (IsOwnedByThis() && m_DataPtr)
 	{
 		FreeBuffer();
 		m_DataPtr = nullptr;
@@ -104,12 +104,20 @@ IoBufferCore::MakeOwned(bool Immutable)
 		memcpy(OwnedDataPtr, m_DataPtr, m_DataBytes);
 
 		m_DataPtr = OwnedDataPtr;
-		SetIsOwned(true);
+		SetIsOwnedByThis(true);
 	}
 
 	SetIsImmutable(Immutable);
 }
 
+void*
+IoBufferCore::MutableDataPointer() const
+{
+	EnsureDataValid();
+	ZEN_ASSERT(!IsImmutable());
+	return const_cast<void*>(m_DataPtr);
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 IoBufferExtendedCore::IoBufferExtendedCore(void* FileHandle, uint64_t Offset, uint64_t Size, bool TransferHandleOwnership)
@@ -117,7 +125,7 @@ IoBufferExtendedCore::IoBufferExtendedCore(void* FileHandle, uint64_t Offset, ui
 , m_FileHandle(FileHandle)
 , m_FileOffset(Offset)
 {
-	m_Flags |= kIsOwned | kIsExtended;
+	m_Flags |= kIsOwnedByThis | kIsExtended;
 
 	if (TransferHandleOwnership)
 	{
@@ -130,7 +138,7 @@ IoBufferExtendedCore::IoBufferExtendedCore(const IoBufferExtendedCore* Outer, ui
 , m_FileHandle(Outer->m_FileHandle)
 , m_FileOffset(Outer->m_FileOffset + Offset)
 {
-	m_Flags |= kIsOwned | kIsExtended;
+	m_Flags |= kIsOwnedByThis | kIsExtended;
 }
 
 IoBufferExtendedCore::~IoBufferExtendedCore()
diff --git a/zencore/sharedbuffer.cpp b/zencore/sharedbuffer.cpp
index bc991053d..db6deff38 100644
--- a/zencore/sharedbuffer.cpp
+++ b/zencore/sharedbuffer.cpp
@@ -10,68 +10,65 @@
 
 namespace zen {
 
-BufferOwner::~BufferOwner()
-{
-	if (m_IsOwned)
-	{
-		Memory::Free(m_Data);
-	}
-}
-
 //////////////////////////////////////////////////////////////////////////
 
 UniqueBuffer
 UniqueBuffer::Alloc(uint64_t Size)
 {
-	void*		 Buffer = Memory::Alloc(Size, 16);
-	BufferOwner* Owner	= new BufferOwner(Buffer, Size, /* owned */ true);
+	void*		  Buffer = Memory::Alloc(Size, 16);
+	IoBufferCore* Owner	 = new IoBufferCore(Buffer, Size);
+	Owner->SetIsOwnedByThis(true);
+	Owner->SetIsImmutable(false);
 
 	return UniqueBuffer(Owner);
 }
 
 UniqueBuffer
-UniqueBuffer::MakeView(void* DataPtr, uint64_t Size)
+UniqueBuffer::MakeMutableView(void* DataPtr, uint64_t Size)
 {
-	return UniqueBuffer(new BufferOwner(DataPtr, Size, /* owned */ false));
+	IoBufferCore* Owner = new IoBufferCore(DataPtr, Size);
+	Owner->SetIsImmutable(false);
+	return UniqueBuffer(Owner);
 }
 
-UniqueBuffer::UniqueBuffer(BufferOwner* Owner) : m_buffer(Owner)
+UniqueBuffer::UniqueBuffer(IoBufferCore* Owner) : m_Buffer(Owner)
 {
 }
 
 //////////////////////////////////////////////////////////////////////////
 
-SharedBuffer::SharedBuffer(UniqueBuffer&& InBuffer) : m_buffer(std::move(InBuffer.m_buffer))
+SharedBuffer::SharedBuffer(UniqueBuffer&& InBuffer) : m_Buffer(std::move(InBuffer.m_Buffer))
 {
 }
 
 void
 SharedBuffer::MakeOwned()
 {
-	if (IsOwned() || !m_buffer)
+	if (IsOwned() || !m_Buffer)
 		return;
 
-	const uint64_t Size		= m_buffer->m_Size;
+	const uint64_t Size		= m_Buffer->DataBytes();
 	void*		   Buffer	= Memory::Alloc(Size, 16);
-	auto		   NewOwner = new BufferOwner(Buffer, Size, /* owned */ true);
+	auto		   NewOwner = new IoBufferCore(Buffer, Size);
+	NewOwner->SetIsOwnedByThis(true);
 
-	memcpy(Buffer, m_buffer->m_Data, Size);
+	memcpy(Buffer, m_Buffer->DataPointer(), Size);
 
-	m_buffer = NewOwner;
+	m_Buffer = NewOwner;
 }
 
 SharedBuffer
-SharedBuffer::MakeView(MemoryView View, SharedBuffer Buffer)
+SharedBuffer::MakeView(MemoryView View, SharedBuffer OuterBuffer)
 {
 	// Todo: verify that view is within the shared buffer
 
-	return SharedBuffer(new BufferOwner(const_cast<void*>(View.GetData()), View.GetSize(), /* owned */ false, Buffer.m_buffer));
+	return SharedBuffer(new IoBufferCore(OuterBuffer.m_Buffer, const_cast<void*>(View.GetData()), View.GetSize()));
 }
 
 SharedBuffer
 SharedBuffer::MakeView(const void* Data, uint64_t Size)
 {
-	return SharedBuffer(new BufferOwner(const_cast<void*>(Data), Size, /* owned */ false));
+	return SharedBuffer(new IoBufferCore(const_cast<void*>(Data), Size));
 }
 
 SharedBuffer
@@ -79,8 +76,9 @@ SharedBuffer::Clone()
 {
 	const uint64_t Size		= GetSize();
 	void*		   Buffer	= Memory::Alloc(Size, 16);
-	auto		   NewOwner = new BufferOwner(Buffer, Size, /* owned */ true);
-	memcpy(Buffer, m_buffer->m_Data, Size);
+	auto		   NewOwner = new IoBufferCore(Buffer, Size);
+	NewOwner->SetIsOwnedByThis(true);
+	memcpy(Buffer, m_Buffer->DataPointer(), Size);
 
 	return SharedBuffer(NewOwner);
 }
@@ -90,7 +88,8 @@ SharedBuffer::Clone(MemoryView View)
 {
 	const uint64_t Size		= View.GetSize();
 	void*		   Buffer	= Memory::Alloc(Size, 16);
-	auto		   NewOwner = new BufferOwner(Buffer, Size, /* owned */ true);
+	auto		   NewOwner = new IoBufferCore(Buffer, Size);
+	NewOwner->SetIsOwnedByThis(true);
 	memcpy(Buffer, View.GetData(), Size);
 
 	return SharedBuffer(NewOwner);
diff --git a/zenserver/projectstore.cpp b/zenserver/projectstore.cpp
index cd85e6d22..cb6be69e8 100644
--- a/zenserver/projectstore.cpp
+++ b/zenserver/projectstore.cpp
@@ -1062,8 +1062,7 @@ HttpProjectService::HttpProjectService(CasStore& Store, ProjectStore* Projects)
 
 				if (IoBuffer Data = IoBufferBuilder::MakeFromFile(AttachmentPath.native().c_str()))
 				{
-					// TODO: this should use the IoBuffer directly to avoid mapping the chunk into memory
-					return SharedBuffer::Clone(MemoryView(Data.Data(), Data.Size()));
+					return SharedBuffer(std::move(Data));
 				}
 				else
 				{
diff --git a/zenserver/projectstore.h b/zenserver/projectstore.h
index 0b41b837b..38c53ea6e 100644
--- a/zenserver/projectstore.h
+++ b/zenserver/projectstore.h
@@ -87,7 +87,7 @@ public:
 
 		const std::string& OplogId() const { return m_OplogId; }
 
-		const std::wstring& TempDir() const { return m_TempPath.native(); }
+		const std::wstring&			 TempDir() const { return m_TempPath.native(); }
 		const std::filesystem::path& TempPath() const { return m_TempPath; }
 
 		spdlog::logger& Log() { return m_OuterProject->Log(); }