1 files changed, 26 insertions, 19 deletions

diff --git a/tests/test_memory.cpp b/tests/test_memory.cpp
index 32d395b..ae5eb8c 100644
--- a/tests/test_memory.cpp
+++ b/tests/test_memory.cpp
@@ -1,30 +1,37 @@
 #include "common.hpp"
 
 #include "writer_string.hpp"
+#include "allocator.hpp"
 
 #include <string>
 
 namespace
 {
-	int allocate_count = 0;
-	int deallocate_count = 0;
+	int page_allocs = 0;
+	int page_deallocs = 0;
+
+	bool is_page(size_t size)
+	{
+		return size >= 8192;
+	}
 
 	void* allocate(size_t size)
 	{
-		++allocate_count;
-		return new char[size];
+		void* ptr = memory_allocate(size);
+		page_allocs += is_page(memory_size(ptr));
+		return ptr;
 	}
 
 	void deallocate(void* ptr)
 	{
-		++deallocate_count;
-		delete[] reinterpret_cast<char*>(ptr);
+		page_deallocs += is_page(memory_size(ptr));
+		memory_deallocate(ptr);
 	}
 }
 
 TEST(memory_custom_memory_management)
 {
-	allocate_count = deallocate_count = 0;
+	page_allocs = page_deallocs = 0;
 
 	// remember old functions
 	allocation_function old_allocate = get_memory_allocation_function();
@@ -37,30 +44,30 @@ TEST(memory_custom_memory_management)
 		// parse document
 		xml_document doc;
 
-		CHECK(allocate_count == 0 && deallocate_count == 0);
+		CHECK(page_allocs == 0 && page_deallocs == 0);
 
 		CHECK(doc.load(STR("<node />")));
 	
-		CHECK(allocate_count == 2 && deallocate_count == 0);
+		CHECK(page_allocs == 1 && page_deallocs == 0);
 
 		// modify document (no new page)
 		CHECK(doc.first_child().set_name(STR("foobars")));
-		CHECK(allocate_count == 2 && deallocate_count == 0);
+		CHECK(page_allocs == 1 && page_deallocs == 0);
 
 		// modify document (new page)
 		std::basic_string<pugi::char_t> s(65536, 'x');
 
 		CHECK(doc.first_child().set_name(s.c_str()));
-		CHECK(allocate_count == 3 && deallocate_count == 0);
+		CHECK(page_allocs == 2 && page_deallocs == 0);
 
 		// modify document (new page, old one should die)
 		s += s;
 
 		CHECK(doc.first_child().set_name(s.c_str()));
-		CHECK(allocate_count == 4 && deallocate_count == 1);
+		CHECK(page_allocs == 3 && page_deallocs == 1);
 	}
 
-	CHECK(allocate_count == 4 && deallocate_count == 4);
+	CHECK(page_allocs == 3 && page_deallocs == 3);
 
 	// restore old functions
 	set_memory_management_functions(old_allocate, old_deallocate);
@@ -68,7 +75,7 @@ TEST(memory_custom_memory_management)
 
 TEST(memory_large_allocations)
 {
-	allocate_count = deallocate_count = 0;
+	page_allocs = page_deallocs = 0;
 
 	// remember old functions
 	allocation_function old_allocate = get_memory_allocation_function();
@@ -80,7 +87,7 @@ TEST(memory_large_allocations)
 	{
 		xml_document doc;
 
-		CHECK(allocate_count == 0 && deallocate_count == 0);
+		CHECK(page_allocs == 0 && page_deallocs == 0);
 
 		// initial fill
 		for (size_t i = 0; i < 128; ++i)
@@ -90,7 +97,7 @@ TEST(memory_large_allocations)
 			CHECK(doc.append_child(node_pcdata).set_value(s.c_str()));
 		}
 
-		CHECK(allocate_count > 0 && deallocate_count == 0);
+		CHECK(page_allocs > 0 && page_deallocs == 0);
 
 		// grow-prune loop
 		while (doc.first_child())
@@ -116,15 +123,15 @@ TEST(memory_large_allocations)
 			}
 		}
 
-		CHECK(allocate_count == deallocate_count + 1); // only one live page left (it waits for new allocations)
+		CHECK(page_allocs == page_deallocs + 1); // only one live page left (it waits for new allocations)
 
 		char buffer;
 		CHECK(doc.load_buffer_inplace(&buffer, 0, parse_fragment, get_native_encoding()));
 
-		CHECK(allocate_count == deallocate_count); // no live pages left
+		CHECK(page_allocs == page_deallocs); // no live pages left
 	}
 
-	CHECK(allocate_count == deallocate_count); // everything is freed
+	CHECK(page_allocs == page_deallocs); // everything is freed
 
 	// restore old functions
 	set_memory_management_functions(old_allocate, old_deallocate);