1 files changed, 84 insertions, 19 deletions

diff --git a/tests/test_memory.cpp b/tests/test_memory.cpp
index bd80ca1..85d6e86 100644
--- a/tests/test_memory.cpp
+++ b/tests/test_memory.cpp
@@ -1,30 +1,38 @@
 #include "common.hpp"
 
 #include "writer_string.hpp"
+#include "allocator.hpp"
 
 #include <string>
+#include <vector>
 
 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 >= 16384;
+	}
 
 	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 +45,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_string(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 +76,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 +88,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 +98,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 +124,72 @@ 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(page_allocs == page_deallocs); // no live pages left
+	}
+
+	CHECK(page_allocs == page_deallocs); // everything is freed
+
+	// restore old functions
+	set_memory_management_functions(old_allocate, old_deallocate);
+}
+
+TEST(memory_page_management)
+{
+	page_allocs = page_deallocs = 0;
+
+	// remember old functions
+	allocation_function old_allocate = get_memory_allocation_function();
+	deallocation_function old_deallocate = get_memory_deallocation_function();
+
+	// replace functions
+	set_memory_management_functions(allocate, deallocate);
+
+	{
+		xml_document doc;
+
+		CHECK(page_allocs == 0 && page_deallocs == 0);
+
+		// initial fill
+		std::vector<xml_node> nodes;
+
+		for (size_t i = 0; i < 4000; ++i)
+		{
+			xml_node node = doc.append_child(STR("node"));
+			CHECK(node);
+
+			nodes.push_back(node);
+		}
+
+		CHECK(page_allocs > 0 && page_deallocs == 0);
+
+		// grow-prune loop
+		size_t offset = 0;
+		size_t prime = 15485863;
+
+		while (nodes.size() > 0)
+		{
+			offset = (offset + prime) % nodes.size();
+
+			doc.remove_child(nodes[offset]);
+
+			nodes[offset] = nodes.back();
+			nodes.pop_back();
+		}
+
+		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);