zone::push_finalizer reverts memory allocation on exception

3 files changed, 74 insertions, 40 deletions

diff --git a/c/zone.c b/c/zone.c
index 877f4ed..1aaad9f 100644
--- a/c/zone.c
+++ b/c/zone.c
@@ -34,7 +34,7 @@ static inline bool init_chunk_array(msgpack_zone_chunk_array* ca, size_t chunk_s
 	const size_t sz = chunk_size;
 
 	char* ptr = (char*)malloc(sz);
-	if(!ptr) {
+	if(ptr == NULL) {
 		free(array);
 		return NULL;
 	}
@@ -88,7 +88,7 @@ void* msgpack_zone_malloc_expand(msgpack_zone* zone, size_t size)
 
 		chunk = (msgpack_zone_chunk*)realloc(ca->array,
 				sizeof(msgpack_zone_chunk) * nnext);
-		if(!chunk) {
+		if(chunk == NULL) {
 			return NULL;
 		}
 
@@ -104,7 +104,7 @@ void* msgpack_zone_malloc_expand(msgpack_zone* zone, size_t size)
 	}
 
 	char* ptr = (char*)malloc(sz);
-	if(!ptr) {
+	if(ptr == NULL) {
 		return NULL;
 	}
 
@@ -144,44 +144,38 @@ static inline void clear_finalizer_array(msgpack_zone_finalizer_array* fa)
 	fa->tail = fa->array;
 }
 
-bool msgpack_zone_push_finalizer(msgpack_zone* zone,
+bool msgpack_zone_push_finalizer_expand(msgpack_zone* zone,
 		void (*func)(void* data), void* data)
 {
 	msgpack_zone_finalizer_array* const fa = &zone->finalizer_array;
 
-	msgpack_zone_finalizer* fin = fa->tail;
-
-	if(fin == fa->end) {
-		// fa->arrayに空きがない
-		// fa->arrayを拡張する
-
-		size_t nnext;
-		const size_t nused = fa->end - fa->array;
+	const size_t nused = fa->end - fa->array;
 
-		if(nused == 0) {
-			// 初回の呼び出し：fa->tail == fa->end == fa->array == NULL
+	size_t nnext;
+	if(nused == 0) {
+		// 初回の呼び出し：fa->tail == fa->end == fa->array == NULL
 
-			// glibcは72バイト以下のmallocが高速
-			nnext = (sizeof(msgpack_zone_finalizer) < 72/2) ?
-					72 / sizeof(msgpack_zone_finalizer) : 8;
+		// glibcは72バイト以下のmallocが高速
+		nnext = (sizeof(msgpack_zone_finalizer) < 72/2) ?
+				72 / sizeof(msgpack_zone_finalizer) : 8;
 
-		} else {
-			nnext = (fa->end - fa->array) * 2;
-		}
+	} else {
+		nnext = nused * 2;
+	}
 
-		fin = (msgpack_zone_finalizer*)realloc(fa->array,
+	msgpack_zone_finalizer* tmp =
+		(msgpack_zone_finalizer*)realloc(fa->array,
 				sizeof(msgpack_zone_finalizer) * nnext);
-		if(!fin) {
-			return false;
-		}
-
-		fa->array      = fin;
-		fa->end        = fin + nnext;
-		fin = fa->tail = fin + nused;
+	if(tmp == NULL) {
+		return false;
 	}
 
-	fin->func = func;
-	fin->data = data;
+	fa->array  = tmp;
+	fa->end    = tmp + nnext;
+	fa->tail   = tmp + nused;
+
+	fa->tail->func = func;
+	fa->tail->data = data;
 
 	++fa->tail;
 
diff --git a/c/zone.h b/c/zone.h
index 93a4642..a3dfe41 100644
--- a/c/zone.h
+++ b/c/zone.h
@@ -67,7 +67,7 @@ void msgpack_zone_free(msgpack_zone* zone);
 
 static inline void* msgpack_zone_malloc(msgpack_zone* zone, size_t size);
 
-bool msgpack_zone_push_finalizer(msgpack_zone* zone,
+static inline bool msgpack_zone_push_finalizer(msgpack_zone* zone,
 		void (*func)(void* data), void* data);
 
 bool msgpack_zone_is_empty(msgpack_zone* zone);
@@ -88,14 +88,37 @@ void* msgpack_zone_malloc(msgpack_zone* zone, size_t size)
 
 	msgpack_zone_chunk* chunk = zone->chunk_array.tail;
 
-	if(chunk->free > size) {
-		char* ptr = chunk->ptr;
-		chunk->ptr  += size;
-		chunk->free -= size;
-		return ptr;
+	if(chunk->free < size) {
+		return msgpack_zone_malloc_expand(zone, size);
+	}
+
+	char* ptr = chunk->ptr;
+	chunk->ptr  += size;
+	chunk->free -= size;
+
+	return ptr;
+}
+
+
+bool msgpack_zone_push_finalizer_expand(msgpack_zone* zone,
+		void (*func)(void* data), void* data);
+
+bool msgpack_zone_push_finalizer(msgpack_zone* zone,
+		void (*func)(void* data), void* data)
+{
+	msgpack_zone_finalizer_array* const fa = &zone->finalizer_array;
+	msgpack_zone_finalizer* fin = fa->tail;
+
+	if(fin == fa->end) {
+		return msgpack_zone_push_finalizer_expand(zone, func, data);
 	}
 
-	return msgpack_zone_malloc_expand(zone, size);
+	fin->func = func;
+	fin->data = data;
+
+	++fa->tail;
+
+	return true;
 }
 
 
diff --git a/cpp/zone.hpp.erb b/cpp/zone.hpp.erb
index 930c8e8..874c900 100644
--- a/cpp/zone.hpp.erb
+++ b/cpp/zone.hpp.erb
@@ -45,6 +45,8 @@ public:
 	<%}%>
 
 private:
+	void undo_malloc(size_t s);
+
 	template <typename T>
 	static void object_destructor(void* obj);
 
@@ -91,14 +93,29 @@ void zone::object_destructor(void* obj)
 	reinterpret_cast<T*>(obj)->~T();
 }
 
+inline void zone::undo_malloc(size_t s)
+{
+	msgpack_zone_chunk* chunk = chunk_array.tail;
+	chunk->ptr  -= s;
+	chunk->free += s;
+}
+
 <%0.upto(GENERATION_LIMIT) {|i|%>
 template <typename T<%1.upto(i) {|j|%>, typename A<%=j%><%}%>>
 T* zone::allocate(<%=(1..i).map{|j|"A#{j} a#{j}"}.join(', ')%>)
 {
 	void* x = malloc(sizeof(T));
-	push_finalizer(&zone::object_destructor<T>, x);
-	try { return new (x) T(<%=(1..i).map{|j|"a#{j}"}.join(', ')%>); }
-	catch (...) { --finalizer_array.tail; throw; }
+	if(!msgpack_zone_push_finalizer(this, &zone::object_destructor<T>, x)) {
+		undo_malloc(sizeof(T));
+		throw std::bad_alloc();
+	}
+	try {
+		return new (x) T(<%=(1..i).map{|j|"a#{j}"}.join(', ')%>);
+	} catch (...) {
+		--finalizer_array.tail;
+		undo_malloc(sizeof(T));
+		throw;
+	}
 }
 <%}%>