memory cleanup and track bucket transience

We want to make sure that if a user calls PkeNew, the determined
bucket is not a transient one.

1 files changed, 19 insertions, 11 deletions

diff --git a/src/memory.cpp b/src/memory.cpp
index 716c6a6..775b30a 100644
--- a/src/memory.cpp
+++ b/src/memory.cpp
@@ -6,21 +6,33 @@ struct pke::MemBucket {
 	int64_t head;
 	int64_t free;
 	int64_t allocs;
+	bool transient;
 	char * ptr;
 };
 
-pke::MemBucket buckets[128]; // 32 gigs
+pke::MemBucket buckets[128];
 int64_t bucketHead = 0;
 
-int64_t InitNewBucket(int64_t sz) {
+int64_t InitNewBucket(int64_t sz, bool transient = false) {
 	buckets[bucketHead].size = sz;
-	buckets[bucketHead].ptr = reinterpret_cast<char *>(std::malloc(sz));
-	buckets[bucketHead].free = 0;
 	buckets[bucketHead].head = 0;
+	buckets[bucketHead].free = 0;
 	buckets[bucketHead].allocs = 0;
+	buckets[bucketHead].transient = transient;
+	buckets[bucketHead].ptr = reinterpret_cast<char *>(std::malloc(sz));
 	return bucketHead++;
 }
 
+void DestroyBucket(pke::MemBucket *bkt) {
+	std::free(bkt->ptr);
+	bkt->size = 0;
+	bkt->head = 0;
+	bkt->free = 0;
+	bkt->allocs = 0;
+	bkt->transient = false;
+	bkt->ptr = reinterpret_cast<char *>(0xCAFEBABE);
+}
+
 void *pke::PkeNew(std::size_t sz, pke::MemBucket *bkt) {
 	void *ptr = bkt->ptr + bkt->head;
 	bkt->head += sz;
@@ -31,7 +43,7 @@ void *pke::PkeNew(std::size_t sz, pke::MemBucket *bkt) {
 void *pke::PkeNew(std::size_t sz) {
 	MemBucket *bkt = nullptr;
 	for (long i = 0; i < bucketHead; ++i) {
-		if (buckets[i].size - buckets[i].head > sz) {
+		if (buckets[i].transient == false && buckets[i].size - buckets[i].head > sz) {
 			bkt = &buckets[i];
 		}
 	}
@@ -66,7 +78,7 @@ void pke::PkeDelete(const void *ptr, std::size_t sz) {
 }
 
 pke::MemBucket *pke::BeginTransientBucket(int64_t sz) {
-	return &buckets[InitNewBucket(sz)];
+	return &buckets[InitNewBucket(sz, true)];
 }
 
 void pke::EndTransientBucket(pke::MemBucket *bkt) {
@@ -74,11 +86,7 @@ void pke::EndTransientBucket(pke::MemBucket *bkt) {
 	for (int64_t i = 0; i < bucketHead; ++i) {
 		if (&buckets[i] == bkt) {
 			foundIndex = i;
-			std::free(bkt->ptr);
-			bkt->size = 0;
-			bkt->head = 0;
-			bkt->free = 0;
-			bkt->ptr = reinterpret_cast<char *>(0xCAFEBABE);
+			DestroyBucket(&buckets[i]);
 			break;
 		}
 	}