Allow vm_space to merge ranges

src/libraries/kutil/vm_space.cpp

@@ -1,16 +1,18 @@
 #include <algorithm>
 #include "kutil/logger.h"
+#include "kutil/vector.h"
 #include "kutil/vm_space.h"
 
 namespace kutil {
 
 using node_type = kutil::avl_node<vm_range>;
-
+using node_vec = kutil::vector<node_type*>;
 
 vm_space::vm_space(uintptr_t start, size_t size, allocator &alloc) :
+	m_slab(alloc),
 	m_alloc(alloc)
 {
-	node_type *node = m_alloc.pop();
+	node_type *node = m_slab.pop();
 	node->address = start;
 	node->size = size;
 	node->state = vm_state::none;
@@ -23,8 +25,8 @@ vm_space::vm_space(uintptr_t start, size_t size, allocator &alloc) :
 inline static bool
 overlaps(node_type *node, uintptr_t start, size_t size)
 {
-	return start <= node->end() &&
-		   (start + size) >= node->address;
+	return start < node->end() &&
+		   (start + size) > node->address;
 }
 
 static node_type *
@@ -46,7 +48,7 @@ node_type *
 vm_space::split_out(node_type *node, uintptr_t start, size_t size, vm_state state)
 {
 	// No cross-boundary splits allowed for now
-	bool contained = 
+	bool contained =
 		start >= node->address &&
 		start+size <= node->end();
 
@@ -59,23 +61,32 @@ vm_space::split_out(node_type *node, uintptr_t start, size_t size, vm_state stat
 	if (state == old_state)
 		return node;
 
+	node->state = state;
+
 	log::debug(logs::vmem, "Splitting out region %016llx-%016llx[%d] from %016llx-%016llx[%d]",
 			start, start+size, state, node->address, node->end(), old_state);
 
+	bool do_consolidate = false;
 	if (node->address < start) {
 		// Split off rest into new node
 		size_t leading = start - node->address;
 
-		node_type *next = m_alloc.pop();
+		node_type *next = m_slab.pop();
 		next->state = state;
 		next->address = start;
 		next->size = node->size - leading;
+
 		node->size = leading;
+		node->state = old_state;
+
+		log::debug(logs::vmem,
+			"   leading region %016llx-%016llx[%d]",
+			node->address, node->address + node->size, node->state);
 
 		m_ranges.insert(next);
 		node = next;
 	} else {
-		// TODO: check if this can merge with prev node
+		do_consolidate = true;
 	}
 
 	if (node->end() > start + size) {
@@ -83,27 +94,79 @@ vm_space::split_out(node_type *node, uintptr_t start, size_t size, vm_state stat
 		size_t trailing =  node->size - size;
 		node->size -= trailing;
 
-		node_type *next = m_alloc.pop();
+		node_type *next = m_slab.pop();
 		next->state = old_state;
 		next->address = node->end();
 		next->size = trailing;
 
+		log::debug(logs::vmem,
+			"   tailing region %016llx-%016llx[%d]",
+			next->address, next->address + next->size, next->state);
+
 		m_ranges.insert(next);
 	} else {
-		// TODO: check if this can merge with next node
+		do_consolidate = true;
 	}
 
+	if (do_consolidate)
+		node = consolidate(node);
+
 	return node;
 }
 
+inline void gather(node_type *node, node_vec &vec)
+{
+	if (node) {
+		gather(node->left(), vec);
+		vec.append(node);
+		gather(node->right(), vec);
+	}
+}
+
+node_type *
+vm_space::consolidate(node_type *needle)
+{
+	node_vec nodes(m_ranges.count(), m_alloc);
+	gather(m_ranges.root(), nodes);
+
+	node_type *prev = nullptr;
+	for (auto *node : nodes) {
+		log::debug(logs::vmem,
+			"* Existing region %016llx-%016llx[%d]",
+			node->address, node->address + node->size, node->state);
+
+		if (prev && node->address == prev->end() && node->state == prev->state) {
+			log::debug(logs::vmem,
+				"Joining regions %016llx-%016llx[%d] %016llx-%016llx[%d]",
+				prev->address, prev->address + prev->size, prev->state,
+				node->address, node->address + node->size, node->state);
+
+			prev->size += node->size;
+			if (needle == node)
+				needle = prev;
+			m_ranges.remove(node, m_slab);
+		} else {
+			prev = node;
+		}
+	}
+
+	return needle;
+}
+
 uintptr_t
 vm_space::reserve(uintptr_t start, size_t size)
 {
 	log::debug(logs::vmem, "Reserving region %016llx-%016llx", start, start+size);
 
 	node_type *node = find_overlapping(m_ranges.root(), start, size);
-	if (!node || node->state != vm_state::none)
+	if (!node) {
+		log::debug(logs::vmem, "  found no match");
 		return 0;
+	} else if (node->state != vm_state::none) {
+		log::debug(logs::vmem, "  found wrong state %016llx-%016llx[%d]",
+			node->address, node->address + node->size, node->state);
+		return 0;
+	}
 
 	node = split_out(node, start, size, vm_state::reserved);
 	return node ? start : 0;