Implement MxMemoryPool (#671)

* Implement MxMemoryPool

* Naming fix

* Annotations and size asserts

* hex padding

LEGO1/omni/include/mxbitset.h

@@ -0,0 +1,95 @@
+#ifndef MXBITSET_H
+#define MXBITSET_H
+
+#pragma warning(disable : 4237)
+
+#include "mxtypes.h"
+
+#include <assert.h>
+#include <limits.h> // CHAR_BIT
+
+template <size_t N>
+class MxBitset {
+public:
+	MxBitset() { Tidy(); }
+
+	// SIZE 0x08
+	class Reference {
+		friend class MxBitset<N>;
+
+	public:
+		Reference& Flip()
+		{
+			m_bitset->Flip(m_offset);
+			return (*this);
+		}
+		bool operator~() const { return (!m_bitset->Test(m_offset)); }
+		operator bool() const { return (m_bitset->Test(m_offset)); }
+
+	private:
+		Reference(MxBitset<N>& p_bitset, size_t p_offset) : m_bitset(&p_bitset), m_offset(p_offset) {}
+		MxBitset<N>* m_bitset; // 0x00
+		size_t m_offset;       // 0x04
+	};
+
+	Reference operator[](size_t p_bit) { return (Reference(*this, p_bit)); }
+
+	MxBitset<N>& Flip(size_t p_bit)
+	{
+		if (N <= p_bit) {
+			Xran();
+		}
+		m_blocks[p_bit / e_bitsPerBlock] ^= 1 << p_bit % e_bitsPerBlock;
+		return (*this);
+	}
+
+	size_t Count()
+	{
+		// debug only, intentionally unimplemented
+		return 0;
+	}
+
+	bool Test(MxU32 p_bit)
+	{
+		if (p_bit >= N) {
+			Xran();
+		}
+
+		return (m_blocks[p_bit / e_bitsPerBlock] & (1 << p_bit % e_bitsPerBlock)) != 0;
+	}
+
+	MxU32 Size() const { return N; }
+
+private:
+	void Tidy(MxU32 p_value = 0)
+	{
+		for (MxS32 i = e_blocksRequired; i >= 0; --i) {
+			m_blocks[i] = p_value;
+		}
+
+		// No need to trim if all bits were zeroed out
+		if (p_value != 0) {
+			Trim();
+		}
+	}
+
+	// Apply bit mask to most significant block
+	void Trim()
+	{
+		if (N % e_bitsPerBlock != 0) {
+			m_blocks[e_blocksRequired] &= ((1 << (N % e_bitsPerBlock)) - 1);
+		}
+	}
+
+	void Xran() { assert("invalid MxBitset<N> position" == NULL); }
+
+	// Not a real enum. This is how STL BITSET defines these constants.
+	enum {
+		e_bitsPerBlock = CHAR_BIT * sizeof(MxU32),
+		e_blocksRequired = N == 0 ? 0 : (N - 1) / e_bitsPerBlock
+	};
+
+	MxU32 m_blocks[e_blocksRequired + 1]; // 0x00
+};
+
+#endif // MXBITSET_H

LEGO1/omni/include/mxdsbuffer.h

@@ -66,7 +66,7 @@ class MxDSBuffer : public MxCore {
 	inline MxU8* GetBuffer() { return m_pBuffer; }
 	inline MxU8** GetBufferRef() { return &m_pBuffer; }
 	inline undefined4 GetUnknown14() { return m_unk0x14; }
-	inline MxU16 GetRefCount() { return m_refcount; }
+	inline MxU16 GetRefCount() { return m_referenceCount; }
 	inline Type GetMode() { return m_mode; }
 	inline MxU32 GetWriteOffset() { return m_writeOffset; }
 	inline MxU32 GetBytesRemaining() { return m_bytesRemaining; }
@@ -85,7 +85,7 @@ class MxDSBuffer : public MxCore {
 	undefined4 m_unk0x14;           // 0x14
 	undefined4 m_unk0x18;           // 0x18
 	undefined4 m_unk0x1c;           // 0x1c
-	MxU16 m_refcount;               // 0x20
+	MxU16 m_referenceCount;         // 0x20
 	Type m_mode;                    // 0x24
 	MxU32 m_writeOffset;            // 0x28
 	MxU32 m_bytesRemaining;         // 0x2c

LEGO1/omni/include/mxmemorypool.h

@@ -0,0 +1,86 @@
+#ifndef MXMEMORYPOOL_H
+#define MXMEMORYPOOL_H
+
+#include "decomp.h"
+#include "mxbitset.h"
+#include "mxtypes.h"
+
+#include <assert.h>
+
+template <size_t BS, size_t NB>
+class MxMemoryPool {
+public:
+	MxMemoryPool() : m_pool(NULL), m_blockSize(BS) {}
+	~MxMemoryPool() { delete[] m_pool; }
+
+	MxResult Allocate();
+	MxU8* Get();
+	void Release(MxU8*);
+
+	MxU32 GetPoolSize() const { return m_blockRef.Size(); }
+
+private:
+	MxU8* m_pool;            // 0x00
+	MxU32 m_blockSize;       // 0x04
+	MxBitset<NB> m_blockRef; // 0x08
+};
+
+template <size_t BS, size_t NB>
+MxResult MxMemoryPool<BS, NB>::Allocate()
+{
+	assert(m_pool == NULL);
+	assert(m_blockSize);
+	assert(m_blockRef.Size());
+
+	m_pool = new MxU8[GetPoolSize() * m_blockSize * 1024];
+	assert(m_pool);
+
+	return m_pool ? SUCCESS : FAILURE;
+}
+
+template <size_t BS, size_t NB>
+MxU8* MxMemoryPool<BS, NB>::Get()
+{
+	assert(m_pool != NULL);
+	assert(m_blockSize);
+	assert(m_blockRef.Size());
+
+	for (MxU32 i = 0; i < GetPoolSize(); i++) {
+		if (!m_blockRef[i]) {
+			m_blockRef[i].Flip();
+
+#ifdef _DEBUG
+			// TODO: This is actually some debug print function, but
+			// we just need any func with variatic args to eliminate diff noise.
+			printf("Get> %d pool: busy %d blocks\n", m_blockSize, m_blockRef.Count());
+#endif
+
+			return &m_pool[i * m_blockSize * 1024];
+		}
+	}
+
+	return NULL;
+}
+
+template <size_t BS, size_t NB>
+void MxMemoryPool<BS, NB>::Release(MxU8* p_buf)
+{
+	assert(m_pool != NULL);
+	assert(m_blockSize);
+	assert(m_blockRef.Size());
+
+	MxU32 i = (MxU32) (p_buf - m_pool) / (m_blockSize * 1024);
+
+	assert(i >= 0 && i < GetPoolSize());
+	assert(m_blockRef[i]);
+
+	if (m_blockRef[i]) {
+		m_blockRef[i].Flip();
+	}
+
+#ifdef _DEBUG
+	printf("Release> %d pool: busy %d blocks\n", m_blockSize, m_blockRef.Count());
+#endif
+}
+
+#endif // MXMEMORYPOOL_H

LEGO1/omni/include/mxstreamer.h

@@ -4,50 +4,16 @@
 #include "decomp.h"
 #include "mxcore.h"
 #include "mxdsobject.h"
+#include "mxmemorypool.h"
 #include "mxnotificationparam.h"
 #include "mxstreamcontroller.h"
 #include "mxtypes.h"
 
+#include <assert.h>
 #include <list>
 
-// NOTE: This feels like some kind of templated class, maybe something from the
+typedef MxMemoryPool<64, 22> MxMemoryPool64;
-//       STL. But I haven't figured out what yet (it's definitely not a vector).
+typedef MxMemoryPool<128, 2> MxMemoryPool128;
-class MxStreamerSubClass1 {
-public:
-	inline MxStreamerSubClass1(undefined4 p_size)
-	{
-		m_buffer = NULL;
-		m_size = p_size;
-		undefined4* ptr = &m_unk0x08;
-		for (int i = 0; i >= 0; i--) {
-			ptr[i] = 0;
-		}
-	}
-
-	// FUNCTION: LEGO1 0x100b9110
-	~MxStreamerSubClass1() { delete[] m_buffer; }
-
-	undefined4 GetSize() const { return m_size; }
-
-	void SetBuffer(undefined* p_buf) { m_buffer = p_buf; }
-	inline undefined* GetBuffer() const { return m_buffer; }
-	inline undefined* GetUnk08Ref() const { return (undefined*) &m_unk0x08; }
-
-private:
-	undefined* m_buffer;
-	undefined4 m_size;
-	undefined4 m_unk0x08;
-};
-
-class MxStreamerSubClass2 : public MxStreamerSubClass1 {
-public:
-	inline MxStreamerSubClass2() : MxStreamerSubClass1(0x40) {}
-};
-
-class MxStreamerSubClass3 : public MxStreamerSubClass1 {
-public:
-	inline MxStreamerSubClass3() : MxStreamerSubClass1(0x80) {}
-};
 
 // VTABLE: LEGO1 0x100dc760
 class MxStreamerNotification : public MxNotificationParam {
@@ -105,13 +71,44 @@ class MxStreamer : public MxCore {
 	MxResult FUN_100b99b0(MxDSAction* p_action);
 	MxResult DeleteObject(MxDSAction* p_dsAction);
 
-	inline const MxStreamerSubClass2& GetSubclass1() { return m_subclass1; }
+	MxU8* GetMemoryBlock(MxU32 p_blockSize)
-	inline const MxStreamerSubClass3& GetSubclass2() { return m_subclass2; }
+	{
+		switch (p_blockSize) {
+		case 0x40:
+			return m_pool64.Get();
+
+		case 0x80:
+			return m_pool128.Get();
+
+		default:
+			assert("Invalid block size for memory pool" == NULL);
+			break;
+		}
+
+		return NULL;
+	}
+
+	void ReleaseMemoryBlock(MxU8* p_block, MxU32 p_blockSize)
+	{
+		switch (p_blockSize) {
+		case 0x40:
+			m_pool64.Release(p_block);
+			break;
+
+		case 0x80:
+			m_pool128.Release(p_block);
+			break;
+
+		default:
+			assert("Invalid block size for memory pool" == NULL);
+			break;
+		}
+	}
 
 private:
 	list<MxStreamController*> m_openStreams; // 0x08
-	MxStreamerSubClass2 m_subclass1;         // 0x14
+	MxMemoryPool64 m_pool64;                 // 0x14
-	MxStreamerSubClass3 m_subclass2;         // 0x20
+	MxMemoryPool128 m_pool128;               // 0x20
 };
 
 // clang-format off
@@ -119,6 +116,12 @@ class MxStreamer : public MxCore {
 // list<MxStreamController *,allocator<MxStreamController *> >::~list<MxStreamController *,allocator<MxStreamController *> >
 // clang-format on
 
+// TEMPLATE: LEGO1 0x100b9100
+// MxMemoryPool<64,22>::~MxMemoryPool<64,22>
+
+// TEMPLATE: LEGO1 0x100b9110
+// MxMemoryPool<128,2>::~MxMemoryPool<128,2>
+
 // SYNTHETIC: LEGO1 0x100b9120
 // MxStreamer::`scalar deleting destructor'
 

LEGO1/omni/src/stream/mxdsbuffer.cpp

@@ -14,7 +14,7 @@ DECOMP_SIZE_ASSERT(MxDSBuffer, 0x34);
 // FUNCTION: LEGO1 0x100c6470
 MxDSBuffer::MxDSBuffer()
 {
-	m_refcount = 0;
+	m_referenceCount = 0;
 	m_pBuffer = NULL;
 	m_pIntoBuffer = NULL;
 	m_pIntoBuffer2 = NULL;
@@ -30,6 +30,8 @@ MxDSBuffer::MxDSBuffer()
 // FUNCTION: LEGO1 0x100c6530
 MxDSBuffer::~MxDSBuffer()
 {
+	assert(m_referenceCount == 0);
+
 	if (m_pBuffer != NULL) {
 		switch (m_mode) {
 		case e_allocate:
@@ -37,41 +39,14 @@ MxDSBuffer::~MxDSBuffer()
 			delete[] m_pBuffer;
 			break;
 
-		case e_chunk: {
+		case e_chunk:
-			MxU32 offset = m_writeOffset / 1024;
+			Streamer()->ReleaseMemoryBlock(m_pBuffer, m_writeOffset / 1024);
-			MxStreamer* streamer = Streamer();
-
-			switch (offset) {
-			case 0x40: {
-				MxU32 a =
-					(m_pBuffer - streamer->GetSubclass1().GetBuffer()) / (streamer->GetSubclass1().GetSize() << 10);
-
-				MxU32 bit = 1 << ((MxU8) a & 0x1f);
-				MxU32 index = (a & ~0x18u) >> 3;
-
-				if ((*(MxU32*) (&streamer->GetSubclass1().GetUnk08Ref()[index])) & bit) {
-					MxU32* ptr = (MxU32*) (&streamer->GetSubclass1().GetUnk08Ref()[index]);
-					*ptr = *ptr ^ bit;
-				}
 			break;
-			}
-			case 0x80: {
-				MxU32 a =
-					(m_pBuffer - streamer->GetSubclass2().GetBuffer()) / (streamer->GetSubclass2().GetSize() << 10);
 
-				MxU32 bit = 1 << ((MxU8) a & 0x1f);
+		case e_preallocated:
-				MxU32 index = (a & ~0x18u) >> 3;
-
-				if ((*(MxU32*) (&streamer->GetSubclass2().GetUnk08Ref()[index])) & bit) {
-					MxU32* ptr = (MxU32*) (&streamer->GetSubclass2().GetUnk08Ref()[index]);
-					*ptr = *ptr ^ bit;
-				}
 			break;
 		}
 	}
-		}
-		}
-	}
 
 	m_unk0x14 = 0;
 	m_unk0x1c = 0;
@@ -85,58 +60,21 @@ MxResult MxDSBuffer::AllocateBuffer(MxU32 p_bufferSize, Type p_mode)
 	switch (p_mode) {
 	case e_allocate:
 		m_pBuffer = new MxU8[p_bufferSize];
+		assert(m_pBuffer); // m_firstRiffChunk?
 		break;
 
-	case e_chunk: {
+	case e_chunk:
-		MxStreamer* streamer = Streamer();
+		m_pBuffer = Streamer()->GetMemoryBlock(p_bufferSize / 1024);
-
-		switch (p_bufferSize / 1024) {
-		case 0x40: {
-			for (MxU32 i = 0; i < 22; i++) {
-				if (((1 << (i & 0x1f)) & (*(MxU32*) &streamer->GetSubclass1().GetUnk08Ref()[(i & ~0x18u) >> 3])) == 0) {
-					MxU32* ptr = (MxU32*) &streamer->GetSubclass1().GetUnk08Ref()[(i & 0xffffffe7) >> 3];
-
-					*ptr = *ptr ^ 1 << (i & 0x1f);
-
-					m_pBuffer =
-						(MxU8*) (streamer->GetSubclass1().GetSize() * i * 0x400 + streamer->GetSubclass1().GetBuffer());
-					goto done;
-				}
-			}
-
-			m_pBuffer = NULL;
 		break;
 	}
-		case 0x80: {
-			for (MxU32 i = 0; i < 2; i++) {
-				if (((1 << (i & 0x1f)) & (*(MxU32*) &streamer->GetSubclass2().GetUnk08Ref()[(i & ~0x18u) >> 3])) == 0) {
-					MxU32* ptr = (MxU32*) &streamer->GetSubclass2().GetUnk08Ref()[(i & 0xffffffe7) >> 3];
 
-					*ptr = *ptr ^ 1 << (i & 0x1f);
-
-					m_pBuffer =
-						(MxU8*) (streamer->GetSubclass2().GetSize() * i * 0x400 + streamer->GetSubclass2().GetBuffer());
-					goto done;
-				}
-			}
-
-			m_pBuffer = NULL;
-			break;
-		}
-		default:
-			m_pBuffer = NULL;
-		}
-	}
-	}
-
-done:
 	m_pIntoBuffer = m_pBuffer;
 	m_pIntoBuffer2 = m_pBuffer;
 
 	if (m_pBuffer != NULL) {
-		m_mode = p_mode;
 		m_bytesRemaining = p_bufferSize;
-		m_writeOffset = p_bufferSize;
+		m_writeOffset = m_bytesRemaining;
+		m_mode = p_mode;
 		result = SUCCESS;
 	}
 
@@ -444,8 +382,8 @@ MxU8* MxDSBuffer::SkipToData()
 // FUNCTION: LEGO1 0x100c6ec0
 MxU8 MxDSBuffer::ReleaseRef(MxDSChunk*)
 {
-	if (m_refcount != 0) {
+	if (m_referenceCount != 0) {
-		m_refcount--;
+		m_referenceCount--;
 	}
 	return 0;
 }
@@ -454,7 +392,7 @@ MxU8 MxDSBuffer::ReleaseRef(MxDSChunk*)
 void MxDSBuffer::AddRef(MxDSChunk* p_chunk)
 {
 	if (p_chunk) {
-		m_refcount++;
+		m_referenceCount++;
 	}
 }
 

LEGO1/omni/src/stream/mxstreamer.cpp

@@ -8,6 +8,10 @@
 #include <algorithm>
 
 DECOMP_SIZE_ASSERT(MxStreamer, 0x2c);
+DECOMP_SIZE_ASSERT(MxMemoryPool64, 0x0c);
+DECOMP_SIZE_ASSERT(MxMemoryPool128, 0x0c);
+DECOMP_SIZE_ASSERT(MxBitset<22>, 0x04);
+DECOMP_SIZE_ASSERT(MxBitset<2>, 0x04);
 
 // FUNCTION: LEGO1 0x100b8f00
 MxStreamer::MxStreamer()
@@ -18,17 +22,11 @@ MxStreamer::MxStreamer()
 // FUNCTION: LEGO1 0x100b9190
 MxResult MxStreamer::Create()
 {
-	undefined* b = new undefined[m_subclass1.GetSize() * 0x5800];
+	if (m_pool64.Allocate() || m_pool128.Allocate()) {
-	m_subclass1.SetBuffer(b);
+		return FAILURE;
-	if (b) {
-		b = new undefined[m_subclass2.GetSize() * 0x800];
-		m_subclass2.SetBuffer(b);
-		if (b) {
-			return SUCCESS;
-		}
 	}
 
-	return FAILURE;
+	return SUCCESS;
 }
 
 // FUNCTION: LEGO1 0x100b91d0