1 files changed, 18 insertions, 23 deletions

diff --git a/src/memory.cpp b/src/memory.cpp
index 83f8483..abaea00 100644
--- a/src/memory.cpp
+++ b/src/memory.cpp
@@ -21,6 +21,13 @@ int64_t InitNewBucket(int64_t sz) {
 	return bucketHead++;
 }
 
+void *pke::PkeNew(std::size_t sz, pke::MemBucket *bkt) {
+	void *ptr = bkt->ptr + bkt->head;
+	bkt->head += sz;
+	bkt->allocs++;
+	return ptr;
+}
+
 void *pke::PkeNew(std::size_t sz) {
 	MemBucket *bkt = nullptr;
 	for (long i = 0; i < bucketHead; ++i) {
@@ -31,20 +38,20 @@ void *pke::PkeNew(std::size_t sz) {
 	if (bkt == nullptr) {
 		bkt = &buckets[InitNewBucket(DEFAULT_BUCKET_SIZE)];
 	}
-	void *ptr = bkt->ptr + bkt->head;
-	bkt->head += sz;
-	bkt->allocs++;
-	return ptr;
+	return PkeNew(sz, bkt);
 }
 
-void *pke::PkeNew(std::size_t sz, pke::MemBucket *bkt) {
-	void *ptr = bkt->ptr + bkt->head;
-	bkt->head += sz;
-	bkt->allocs++;
-	return ptr;
+void pke::PkeDelete(const void *ptr, std::size_t sz, pke::MemBucket *bkt) {
+	bkt->allocs--;
+	bkt->free += sz;
+	if (ptr == bkt->ptr + bkt->head - sz) bkt->head -= sz;
+	if (bkt->allocs == 0) {
+		bkt->head = 0;
+		bkt->free = 0;
+	}
 }
 
-void pke::PkeDelete(void *ptr) {
+void pke::PkeDelete(const void *ptr, std::size_t sz) {
 	pke::MemBucket *bkt = nullptr;
 	for (long i = 0; i < bucketHead; ++i) {
 		bkt = &buckets[i];
@@ -54,19 +61,7 @@ void pke::PkeDelete(void *ptr) {
 		std::printf("bad pointer %llu", reinterpret_cast<unsigned long long>(ptr));
 		throw "bad pointer";
 	}
-	bkt->allocs--;
-	if (bkt->allocs == 0) {
-		bkt->head = 0;
-		bkt->free = 0;
-	}
-}
-
-void pke::PkeDelete(void *ptr, pke::MemBucket *bkt) {
-	bkt->allocs--;
-	if (bkt->allocs == 0) {
-		bkt->head = 0;
-		bkt->free = 0;
-	}
+	PkeDelete(ptr, sz, bkt);
 }
 
 pke::MemBucket *pke::BeginTransientBucket(int64_t sz) {