3 files changed, 105 insertions, 16 deletions

diff --git a/kernel/memory/paging.c b/kernel/memory/paging.c
index 58c5091..24a9ea7 100644
--- a/kernel/memory/paging.c
+++ b/kernel/memory/paging.c
@@ -796,16 +796,36 @@ void *mem_alloc_pages(mem_ctx_t ctx, size_t count, unsigned int flags)
 void *mem_alloc_pages_at(mem_ctx_t ctx, size_t count, void *virt,
 						 unsigned int flags)
 {
-	void *phys = alloc_phys_pages(count);
-	if (phys == NULL)
-		return NULL;
+	size_t pages_needed = count;
+	uint8_t *virtual_address = virt;
 
-	if (map_pages((volatile struct pml4 *)ctx->pml4, virt, phys, flags,
-				  count)) {
-		if (phys)
-			free_phys_pages(phys, count);
-		return NULL;
+	void *phys_start = NULL;
+
+	while (pages_needed > 0) {
+		struct phys_page_slice phys_pages =
+			alloc_phys_page_withextra(pages_needed);
+		if (phys_pages.pagestart == NULL) {
+			free_phys_pages(phys_start ? phys_start : phys_pages.pagestart,
+							count - pages_needed);
+			return NULL;
+		}
+
+		if (!phys_start)
+			phys_start = phys_pages.pagestart;
+
+		assert(pages_needed >= phys_pages.num_pages, "overflow");
+		pages_needed -= phys_pages.num_pages;
+		virtual_address += phys_pages.num_pages * PAGE_SIZE;
+
+		if (map_pages((volatile struct pml4 *)ctx->pml4,
+					  (void *)virtual_address, phys_pages.pagestart, flags,
+					  phys_pages.num_pages)) {
+            assert(phys_start, "expected something allocated");
+			free_phys_pages(phys_start, count - pages_needed);
+			return NULL;
+		}
 	}
+
 	return virt;
 }
 
diff --git a/kernel/memory/physalloc.c b/kernel/memory/physalloc.c
index 60e7017..856a627 100644
--- a/kernel/memory/physalloc.c
+++ b/kernel/memory/physalloc.c
@@ -58,7 +58,7 @@ static long page_idx(void *page)
 	return -1;
 }
 
-static inline bool bitmap_get(int i)
+static inline bool bitmap_get(size_t i)
 {
 	return (bitmap[i / 64] >> i % 64) & 1;
 }
@@ -76,17 +76,17 @@ static inline void bitmap_set(int i, bool v)
 
 void *alloc_phys_page(void)
 {
-	return alloc_phys_pages(1);
+	return alloc_phys_pages_exact(1);
 }
 
-void *alloc_phys_pages(size_t pages)
+void *alloc_phys_pages_exact(size_t pages)
 {
 	if (pages < 1)
 		return NULL;
 
 	size_t n_contiguous = 0;
-	int free_region_start = 0;
-	for (uint64_t i = 0; i < page_count; i++) {
+	size_t free_region_start = 0;
+	for (size_t i = 0; i < page_count; i++) {
 		bool free = !bitmap_get(i);
 
 		if (free) {
@@ -105,11 +105,53 @@ void *alloc_phys_pages(size_t pages)
 	return NULL;
 }
 
+struct phys_page_slice alloc_phys_page_withextra(size_t max_pages)
+{
+	if (max_pages == 0)
+		return PHYS_PAGE_SLICE_NULL;
+
+	for (size_t i = 0; i < page_count; i++) {
+		const bool free = !bitmap_get(i);
+		if (!free)
+			continue;
+
+		// now allocated
+		bitmap_set(i, true);
+
+		// found at least one page, guaranteed to return valid slice at this
+		// point
+		struct phys_page_slice out = {
+			.pagestart = page_at(i),
+			.num_pages = 1,
+		};
+
+		// add some extra pages if possible
+		for (; out.num_pages < MIN(page_count - i, max_pages);
+			 ++out.num_pages) {
+			// early return if max_pages isn't available
+			if (bitmap_get(i + out.num_pages)) {
+				return out;
+			}
+			bitmap_set(i + out.num_pages, true);
+		}
+
+		return out;
+	}
+
+	// only reachable if there is not a single free page in the bitmap
+	return PHYS_PAGE_SLICE_NULL;
+}
+
 void free_phys_page(void *ptr)
 {
 	free_phys_pages(ptr, 1);
 }
 
+void free_phys_pages_slice(struct phys_page_slice slice)
+{
+	free_phys_pages(slice.pagestart, slice.num_pages);
+}
+
 void free_phys_pages(void *ptr, size_t pages)
 {
 	if (ptr == NULL)
diff --git a/kernel/memory/physalloc.h b/kernel/memory/physalloc.h
index d91c57a..e279409 100644
--- a/kernel/memory/physalloc.h
+++ b/kernel/memory/physalloc.h
@@ -11,11 +11,31 @@
 
 #include <comus/memory.h>
 
+/// Represents some contiguous physical pages
+struct phys_page_slice {
+	void *pagestart;
+	size_t num_pages;
+};
+
+#define PHYS_PAGE_SLICE_NULL \
+	((struct phys_page_slice){ .pagestart = NULL, .num_pages = 0 })
+
 /**
  * Initalize the physical page allocator
  */
 void physalloc_init(struct memory_map *map);
 
+/*
+ * Allocates the first page(s) it finds. Returns a pointer to that page
+ * and, if there are (up to max_pages) extra pages free after it, it allocates
+ * them as well.
+ *
+ * @param max_pages - the maximum number of pages to mark as allocated
+ * @returns a slice of all of the allocated pages, num_pages will be
+ * <= max_pages
+ */
+struct phys_page_slice alloc_phys_page_withextra(size_t max_pages);
+
 /**
  * Allocates a single physical page in memory
  * @preturns the physical address of the page
@@ -23,10 +43,11 @@ void physalloc_init(struct memory_map *map);
 void *alloc_phys_page(void);
 
 /**
- * Allocates count physical pages in memory
- * @returns the physical address of the first page
+ * Allocates count contiguous physical pages in memory
+ * @returns the physical address of the first page, or NULL if no
+ * contiguous pages exist.
  */
-void *alloc_phys_pages(size_t count);
+void *alloc_phys_pages_exact(size_t count);
 
 /**
 * Frees a single physical page in memory
@@ -41,4 +62,10 @@ void free_phys_page(void *ptr);
  */
 void free_phys_pages(void *ptr, size_t count);
 
+/**
+ * Frees a slice of physical pages in memory
+ * @param slice - the pages to free
+ */
+void free_phys_pages_slice(struct phys_page_slice slice);
+
 #endif /* physalloc.h */