rewrite pmem

12 files changed, 464 insertions, 446 deletions

diff --git a/arch/riscv64/kernel/vmem.c b/arch/riscv64/kernel/vmem.c
index 772251e..8954ee3 100644
--- a/arch/riscv64/kernel/vmem.c
+++ b/arch/riscv64/kernel/vmem.c
@@ -38,23 +38,23 @@
  * @param f Flags to use.
  * @return Corresponding page table entry.
  */
-#define to_pte(p, f) ((pm_to_pnum(p) << 10) | (f))
+#define to_pte(p, f) ((((p) >> page_shift()) << 10) | (f))
 
 /**
- * Get virtual address in page table entry.
+ * Get physical address in page table entry.
  *
  * @param pte Page table entry.
- * @return Corresponding virtual address.
+ * @return Corresponding physical address.
  */
-#define pte_addr(pte) __va(pnum_to_pm(pte_ppn(pte)))
+#define pte_paddr(pte) (pte_ppn(pte) << page_shift())
 
 /**
- * Get physical address in page table entry.
+ * Get virtual address in page table entry.
  *
  * @param pte Page table entry.
- * @return Corresponding physical address.
+ * @return Corresponding virtual address.
  */
-#define pte_paddr(pte) (pnum_to_pm(pte_ppn(pte)))
+#define pte_addr(pte) __va(pte_paddr(pte))
 
 /**
  * Virtual memory address to page order index.
@@ -166,7 +166,7 @@ stat_t stat_vpage(struct vmem *branch, vm_t vaddr, pm_t *paddr,
  */
 static struct vmem *__create_leaf()
 {
-	pm_t new_leaf = alloc_page(MM_KPAGE, 0);
+	pm_t new_leaf = alloc_page(MM_KPAGE);
 	memset((void *)new_leaf, 0, sizeof(struct vmem));
 	return (struct vmem *)to_pte((pm_t)__pa(new_leaf), VM_V);
 }
@@ -244,17 +244,18 @@ void flush_tlb_all()
 static void __use_vmem(struct vmem *branch, enum mm_mode m)
 {
 	branch = (struct vmem *)__pa(branch);
+	pm_t pn = (pm_t)(branch) >> page_shift();
+
+	pm_t mode = DEFAULT_Sv_MODE;
 
 	if (m == Sv32)
-		csr_write(CSR_SATP,
-		          SATP_MODE_Sv32 | pm_to_pnum((pm_t)(branch)));
+		mode = SATP_MODE_Sv32;
 	else if (m == Sv39)
-		csr_write(CSR_SATP,
-		          SATP_MODE_Sv39 | pm_to_pnum((pm_t)(branch)));
-	else
-		csr_write(CSR_SATP,
-		          SATP_MODE_Sv48 | pm_to_pnum((pm_t)(branch)));
+		mode = SATP_MODE_Sv39;
+	else if (m == Sv48)
+		mode = SATP_MODE_Sv48;
 
+	csr_write(CSR_SATP, mode | pn);
 	flush_tlb_all();
 	/* Sv57 && Sv64 in the future? */
 	/** @todo ASID table for maybe faster context switches? */
@@ -271,7 +272,7 @@ struct vmem *init_vmem(void *fdt)
 
 struct vmem *create_vmem()
 {
-	struct vmem *b = (struct vmem *)alloc_page(MM_KPAGE, 0);
+	struct vmem *b = (struct vmem *)alloc_page(MM_KPAGE);
 	memset(b, 0, MM_KPAGE_SIZE);
 
 	populate_kvmem(b);
diff --git a/common/bits.c b/common/bits.c
index 2efbe2f..5691617 100644
--- a/common/bits.c
+++ b/common/bits.c
@@ -36,3 +36,23 @@ __weak uint64_t __bswap64(const uint64_t u)
 	        (u & 0x000000000000ff00ULL) << 40 |
 	        (u & 0x00000000000000ffULL) << 56;
 }
+
+#undef ffs
+__weak int ffs(int v)
+{
+	/* http://graphics.stanford.edu/~seander/bithacks.html#ZerosOnRightParallel */
+	if (v == 0)
+		return 0;
+
+	int c = 32;
+	v &= -v;
+
+	if (v) c--;
+	if (v & 0x0000FFFF) c -= 16;
+	if (v & 0x00FF00FF) c -= 8;
+	if (v & 0x0F0F0F0F) c -= 4;
+	if (v & 0x33333333) c -= 2;
+	if (v & 0x55555555) c -= 1;
+
+	return c + 1;
+}
diff --git a/common/mem.c b/common/mem.c
index 1f893d4..6d032b4 100644
--- a/common/mem.c
+++ b/common/mem.c
@@ -15,20 +15,20 @@ size_t __mm_shifts[10];
 size_t __mm_widths[10];
 size_t __mm_sizes[10];
 size_t __mm_page_shift;
-size_t __mm_max_order;
+enum mm_order __mm_max_order;
 
 void init_mem(size_t max_order, size_t bits[10], size_t page_shift)
 {
 	__mm_max_order = max_order;
 	__mm_page_shift = page_shift;
 
-	__mm_shifts[0] = 0;
+	__mm_shifts[0] = page_shift;
 	__mm_widths[0] = 1 << bits[0];
 	__mm_sizes[0] = 1 << __mm_page_shift;
 
 	for (size_t i = 1; i <= __mm_max_order; ++i) {
 		__mm_widths[i] = 1 << bits[i];
 		__mm_shifts[i] = __mm_shifts[i - 1] + bits[i - 1];
-		__mm_sizes[i] = 1UL << __mm_shifts[i] << __mm_page_shift;
+		__mm_sizes[i] = 1UL << __mm_shifts[i];
 	}
 }
diff --git a/common/mem_regions.c b/common/mem_regions.c
index 0101ac2..8d066e7 100644
--- a/common/mem_regions.c
+++ b/common/mem_regions.c
@@ -55,7 +55,7 @@ static struct mem_region *__insert_free_region(struct mem_region_root *r,
                                                struct mem_region *m)
 {
 	struct sp_node *n = sp_root(&r->free_regions), *p = NULL;
-	size_t start = m->start;
+	vm_t start = m->start;
 	size_t size = m->end - m->start;
 	enum sp_dir d = LEFT;
 
diff --git a/common/nodes.c b/common/nodes.c
index 9cfd2a2..60dec79 100644
--- a/common/nodes.c
+++ b/common/nodes.c
@@ -57,7 +57,7 @@
  */
 static struct node_region *__create_region()
 {
-	struct node_region *r = (struct node_region *)alloc_page(BASE_PAGE, 0);
+	struct node_region *r = (struct node_region *)alloc_page(BASE_PAGE);
 	memset(r, FREE, BASE_PAGE_SIZE);
 	return r;
 }
diff --git a/common/pmem.c b/common/pmem.c
index 748a1c8..b8a453a 100644
--- a/common/pmem.c
+++ b/common/pmem.c
@@ -41,106 +41,79 @@
  */
 
 /**
- * Loop through all page usage bits in current bitmap.
+ * Beauty macro for looping over all page indexes.
+ * The current page index is stored in \c page.
  *
- * @param var Memory leaf or branch containing bitmap.
- * @param start Start looking from this index.
- * @param end Stop looking before this index.
- * @param attr Attribute name of bitmap.
- * @param neg Negate whether we're looking for full or empty pages.
- *
- * \note These are all for pnum_t, i.e. O0_SHIFT is from 0
+ * @param num Number of pages in branch.
  */
-#define __foreach_page(var, start, end, attr, neg)          \
-	for (pnum_t page = start; page < end; ++page)       \
-	if (neg (bitmap_is_set(var->attr, page))) continue; \
-	else                                                \
-
-/** Easier to read negation. */
-#define NEG !
+#define foreach_page(num) \
+	for (pm_t page = 0; page < num; ++page)
 
 /**
- * Loop through all full pages.
+ * Loop over orders, giving the iterator the name \p iter.
  *
- * @param var Memory leaf or branch containing bitmap.
- * @param start Start looking from this index.
+ * @param iter Name of iterator.
  */
-#define foreach_full_page(var, start) \
-	__foreach_page(var, start, var->entries, full, NEG)
+#define foreach_order(iter) \
+	for (enum mm_order iter = MM_O0; iter <= max_order(); ++iter)
 
 /**
- * Loop through all not full pages.
+ * Loop over orders, with already initialized start iterator \p iter.
  *
- * @param var Memory leaf or branch containing bitmap.
- * @param start Start looking from this index.
+ * @param iter Name of iterator.
  */
-#define foreach_not_full_page(var, start) \
-	__foreach_page(var, start, var->entries, full, )
+#define foreach_order_init(iter) \
+	for (; iter <= max_order(); ++iter)
 
 /**
- * Loop through all used pages.
+ * Loop over orders in reverse, starting with highest, giving the iterator the
+ * name \p iter.
  *
- * @param var Memory leaf or branch containing bitmap.
- * @param start Start looking from this index.
+ * @param iter Name of iterator.
  */
-#define foreach_used_page(var, start) \
-	__foreach_page(var, start, var->entries, used, NEG)
+#define reverse_foreach_order(iter) \
+	for (enum mm_order iter = max_order(); iter != MM_MIN; --iter)
 
 /**
- * Loop through all not used pages.
+ * Loop over orders in reverse, starting with highest, with already initialized
+ * start iterator \p iter.
  *
- * @param var Memory leaf or branch containing bitmap.
- * @param start Start looking from this index.
+ * @param iter Name of iterator.
  */
-#define foreach_not_used_page(var, start) \
-	__foreach_page(var, start, var->entries, used, )
-
-/* curiously, all my optimisation efforts were in vain, and eight bits is the
- * best alternative. */
-
-/** Memory bitmap base size. */
-typedef uint8_t mm_info_t;
-
-/** Beauty typedef for void *, used for bitmaps in this file. */
-typedef void mm_node_t;
-
-/** Memory page leaf. */
-struct mm_leaf {
-	/** Number of entries in leaf. */
-	pnum_t entries;
+#define reverse_foreach_order_init(iter) \
+	for (; iter != MM_MIN; --iter)
 
-	/** Bitmap of used pages. */
-	mm_info_t *used;
-};
+/** Beauty typedef for uint8_t *, used for bitmaps in this file. */
+typedef uint8_t mm_bitmap_t[];
 
 /** Memory page branch. */
 struct mm_branch {
-	/** Number of entries in branch. */
-	pnum_t entries;
+	/** Number of entries in leaf. */
+	size_t num;
 
-	/** Bitmap of full nodes. */
-	mm_info_t *full;
+	/** Size of one whole span of one sub branch. */
+	size_t size;
 
-	/** Pointer to array of next order indexes. */
-	mm_node_t **next;
+	/** Bitmap of used pages. */
+	mm_bitmap_t used;
 };
 
 /** Order map. */
-struct mm_omap {
+struct mm_bucket {
 	/** Base address of map. */
 	pm_t base;
 
-	/** Pointer to array of nodes. */
-	mm_node_t **orders;
-
 	/** Order of map. */
 	enum mm_order order;
+
+	/** Pointer to array of nodes. */
+	struct mm_branch *tree[MM_NUM];
 };
 
 /** Physical map. */
 struct mm_pmap {
-	/** Order map, one per order up to maximum order. */
-	struct mm_omap *omap[NUM_ORDERS];
+	/** Buckets, one per order up to maximum order. */
+	struct mm_bucket *bucket[NUM_ORDERS];
 };
 
 /** Static physical map address. \note If I support NUMA, this should not be
@@ -148,359 +121,390 @@ struct mm_pmap {
 static struct mm_pmap *pmap = 0;
 
 /**
- * Helper function for marking a page used.
+ * Calculate size of branch structure plus bitmap for branch.
+ *
+ * @param num Number of elements in branch.
+ * @return Size in bytes of a branch.
+ */
+static size_t sizeof_branch(size_t num)
+{
+	return align_up(sizeof(struct mm_branch) + (num + 8) / 8,
+	                sizeof(struct mm_branch));
+}
+
+/**
+ * Get top of current branch, that is, the start of a following branch.
+ *
+ * @param branch Branch whose top to calculate.
+ * @return Top of \p branch.
+ */
+static struct mm_branch *branch_top(struct mm_branch *branch)
+{
+	return (struct mm_branch *)(sizeof_branch(branch->num) + (pm_t)branch);
+}
+
+/**
+ * Get the branch under \p branch at \p index.
+ *
+ * @param branch Branch whose sub branches to access.
+ * @param index Index of sub branch to access.
+ * @return Pointer to sub branch.
+ */
+static struct mm_branch *sub_branch(struct mm_branch *branch, size_t index)
+{
+	struct mm_branch *sub_start = branch_top(branch);
+	return (struct mm_branch *)(index * branch->size + (pm_t)sub_start);
+}
+
+/**
+ * Mark a page free in tree.
  *
- * @param op Order node pointer.
- * @param pnum Physical page number to mark free.
- * @param tgt Target order.
- * @param src Source order.
- * @param dst Destination order.
+ * @param branch Branch wherein some part of \p page lies.
+ * @param page Page address relative to start of RAM.
+ * @param req_order Order of page to be marked free.
+ * @param cur_order Current page order.
+ * @param tree_order Order context we're in.
  */
-static void __mark_free(mm_node_t *op, pnum_t pnum, enum mm_order tgt,
-                        enum mm_order src, enum mm_order dst)
+static void __mark_free(struct mm_branch *branch, pm_t page,
+                        enum mm_order req_order,
+                        enum mm_order cur_order,
+                        enum mm_order tree_order)
 {
-	size_t idx = pnum_to_index(pnum, src);
+	size_t idx = pm_to_index(page, cur_order);
 
-	if (src == dst) {
-		struct mm_leaf *o = (struct mm_leaf *)op;
-		bitmap_clear(o->used, idx);
+	if (cur_order == req_order) {
+		bitmap_clear(branch->used, idx);
 		return;
 	}
 
-	struct mm_branch *o = (struct mm_branch *)op;
-	if (src != tgt)
-		__mark_free(o->next[idx], pnum, tgt, src - 1, dst);
+	if(cur_order != tree_order)
+		__mark_free(sub_branch(branch, idx), page,
+		            req_order, cur_order - 1, tree_order);
+
+	/* freeing a page results in always clearing a full bit */
+	bitmap_clear(branch->used, idx);
+}
+
+/**
+ * Mark page in bucket free in all trees.
+ *
+ * @param bucket Bucket page lies in.
+ * @param order Order of page to free.
+ * @param addr Physical address of page.
+ */
+static void __mark_bucket_page_free(struct mm_bucket *bucket,
+                                    enum mm_order order, pm_t addr)
+{
+	pm_t fixup_addr = addr - bucket->base;
+	enum mm_order iter = bucket->order;
 
-	/* freeing a page results in always clearing a full bit? */
-	bitmap_clear(o->full, idx);
+	reverse_foreach_order_init(iter) {
+		__mark_free(bucket->tree[iter], fixup_addr,
+		            order, bucket->order, iter);
+	}
 }
 
-void free_page(enum mm_order order, pm_t paddr)
+void free_page(enum mm_order order, pm_t addr)
 {
-	/** \todo This could probably use an int for status, but eh */
-	for (size_t i = MM_O0; i <= __mm_max_order; ++i) {
-		if (!pmap->omap[i])
+	foreach_order(iter) {
+		struct mm_bucket *bucket = pmap->bucket[iter];
+		if (!bucket)
 			continue;
 
-		struct mm_omap *omap = pmap->omap[i];
-		if (paddr < omap->base)
+		if (addr < bucket->base)
 			continue;
 
-		for (size_t j = 0; j < omap->order; ++j)
-			__mark_free(omap->orders[j],
-			            pm_to_pnum(paddr - omap->base), order,
-			            omap->order, j);
-
+		__mark_bucket_page_free(bucket, order, addr);
 		return;
 	}
 }
 
 /**
- * Helper function for marking a page used.
+ * Mark page used in tree.
  *
- * @param op Order node pointer.
- * @param pnum Page number to mark used.
- * @param tgt Target order.
- * @param src Source order.
- * @param dst Destination order.
- * @return \ref true if order is filled, \ref false otherwise.
+ * @param branch Current branch.
+ * @param page Page address relative to start of RAM.
+ * @param req_order Page order.
+ * @param cur_order Current order.
+ * @param tree_order Order of context we're in.
+ * @return Whether the branch below got filled up.
  */
-static bool __mark_used(mm_node_t *op, pnum_t pnum, enum mm_order tgt,
-                        enum mm_order src, enum mm_order dst)
+static bool __mark_used(struct mm_branch *branch, pm_t page,
+                        enum mm_order req_order,
+                        enum mm_order cur_order,
+                        enum mm_order tree_order)
 {
-	size_t idx = pnum_to_index(pnum, src);
+	size_t idx = pm_to_index(page, cur_order);
 
-	if (src == dst) {
-		struct mm_leaf *o = (struct mm_leaf *)op;
-		bitmap_set(o->used, idx);
+	if (cur_order == req_order || cur_order == tree_order) {
+		bitmap_set(branch->used, idx);
 
-		if (idx == max_index(src))
+		if (idx == max_index(cur_order))
 			return true;
 
 		return false;
 	}
 
-	struct mm_branch *o = (struct mm_branch *)op;
-	if (src == tgt) {
-		bitmap_set(o->full, idx);
+	bool r = __mark_used(sub_branch(branch, idx), page,
+	                     req_order,
+	                     cur_order - 1,
+	                     tree_order);
 
-		if (idx == max_index(src))
-			return true;
-
-		return false;
-	}
-
-	if (__mark_used(o->next[idx], pnum, tgt, src - 1, dst)) {
-		bitmap_set(o->full, idx);
+	if (r) {
+		bitmap_set(branch->used, idx);
 
-		if (idx == max_index(src))
+		if (idx == max_index(cur_order))
 			return true;
 	}
 
 	return false;
 }
 
-void mark_used(enum mm_order order, pm_t paddr)
+/**
+ * Mark page in bucket used.
+ *
+ * @param bucket Bucket in which \p page lies.
+ * @param order Order of \p page.
+ * @param addr Physical address of \p page.
+ */
+static void __mark_bucket_page_used(struct mm_bucket *bucket,
+                                    enum mm_order order, pm_t addr)
 {
-	for (size_t i = MM_O0; i <= __mm_max_order; ++i) {
-		if (!pmap->omap[i])
-			continue;
+	pm_t fixed_addr = addr - bucket->base;
+	reverse_foreach_order(iter) {
+		__mark_used(bucket->tree[iter], fixed_addr,
+		            order, bucket->order, iter);
+	}
+}
 
-		struct mm_omap *omap = pmap->omap[i];
-		if (paddr < omap->base)
+void mark_used(enum mm_order order, pm_t addr)
+{
+	enum mm_order iter = order;
+	foreach_order_init(iter) {
+		struct mm_bucket *bucket = pmap->bucket[iter];
+		if (!bucket)
 			continue;
 
-		for (size_t j = 0; j <= omap->order; ++j)
-			__mark_used(omap->orders[j],
-			            pm_to_pnum(paddr - omap->base), order,
-			            omap->order, j);
+		if (addr < bucket->base)
+			continue;
 
+		__mark_bucket_page_used(bucket, iter, addr);
 		return;
 	}
 }
 
 /**
- * Look for next free page.
+ * Find first unused page on branch.
+ * Helper for converting between bitmap and pmem error conditions.
  *
- * @param op Operand node pointer.
- * @param offset Offset to where to start looking for available pages from.
- * @param src Source order.
- * @param dst Destination order.
- * @return Page number of found index.
+ * @param branch Branch to look for unused pages on.
+ * @return \c -1 if there are no free pages, otherwise the index of first unused
+ * page.
  */
-static pnum_t __enum_order(mm_node_t *op, pnum_t offset, enum mm_order src,
-                           enum mm_order dst)
+static pm_t __branch_find_first_unset(struct mm_branch *branch)
 {
-	size_t idx = pnum_to_index(offset, src);
+	size_t r = bitmap_find_first_unset(branch->used, branch->num);
+	if (r > branch->num)
+		return -1;
 
-	if (src == dst) {
-		struct mm_leaf *o = (struct mm_leaf *)op;
-		foreach_not_used_page(o, idx)
-		{
-			return page << order_offset(src);
-		}
+	return r;
+}
 
+/**
+ * Search for unused pages in tree.
+ *
+ * @param branch Current branch.
+ * @param cur_order Current order of branch.
+ * @param req_order Requested page order.
+ * @return \c -1 if there are no free pages, otherwise the address of the lower
+ * order page found.
+ */
+static pm_t __search_tree(struct mm_branch *branch,
+                          enum mm_order cur_order,
+                          enum mm_order req_order)
+{
+	pm_t page = __branch_find_first_unset(branch);
+	if (page == (pm_t)(-1))
 		return -1;
-	}
 
-	struct mm_branch *o = (struct mm_branch *)op;
-	foreach_not_full_page(o, idx)
-	{
-		/* if the suggested search index is full, the following level
-		 * would get an incorrect offset if trying to follow the original
-		 * suggestion. */
-		if (page != (pnum_t)idx)
-			offset = 0;
+	if (cur_order == req_order)
+		return page << order_shift(cur_order);
 
-		pnum_t ret = __enum_order(o->next[page], offset, src - 1, dst);
+	pm_t r = __search_tree(sub_branch(branch, page),
+	                       cur_order - 1, req_order);
 
-		if (!(ret < 0))
-			return (page << order_offset(src)) + ret;
-	}
+	if (r == (pm_t)(-1))
+		return -1;
 
-	return -1;
+	return (page << order_shift(cur_order)) + r;
 }
 
-pm_t alloc_page(enum mm_order order, pm_t offset)
+pm_t alloc_page(enum mm_order order)
 {
-	if (order > __mm_max_order)
+	if (order > max_order())
 		return 0;
 
-	pnum_t pnum = -1;
-	pm_t base = 0;
-	struct mm_omap *omap;
-	for (size_t i = order; i <= __mm_max_order; ++i) {
-		if (!pmap->omap[i])
+	pm_t p = -1;
+	struct mm_bucket *bucket;
+	enum mm_order iter = order;
+	foreach_order_init(iter) {
+		bucket = pmap->bucket[iter];
+		if (!bucket)
 			continue;
 
-		omap = pmap->omap[i];
-		if (offset != 0)
-			base = offset - omap->base;
-
-		pnum = __enum_order(omap->orders[order], pm_to_pnum(base),
-		                    omap->order, order);
+		p = __search_tree(bucket->tree[order], bucket->order, order);
 
-		if (!(pnum < 0))
+		if (p != (pm_t)(-1))
 			break;
 	}
 
-	if (pnum < 0)
+	if (p == (pm_t)(-1))
 		return 0;
 
-	pm_t paddr = pnum_to_pm(pnum) + omap->base;
-	mark_used(order, paddr);
-	return paddr;
+	p = p + bucket->base;
+	__mark_bucket_page_used(bucket, order, p);
+	return p;
 }
 
 /**
- * Populate order node map.
+ * Populate tree.
  *
- * @param op Address to where to write order node pointer.
- * @param cont Physical address where to continue writing data to.
- * @param src Source order.
- * @param dst Destination order.
- * @param num Number of nodes in this order to populate.
- * @return Physical address to continue from.
+ * @param cont Address at which to continue placing data.
+ * @param num Number of elements in branch.
+ * @param cur_order Current branch order.
+ * @param req_order Requested tree order.
+ * @return Top of tree.
  */
-static pm_t __populate_order(mm_node_t **op, pm_t cont, enum mm_order src,
-                             enum mm_order dst, size_t num)
+static pm_t __populate_tree(pm_t cont, size_t num,
+                            enum mm_order cur_order, enum mm_order req_order)
 {
-	/* unfortunate that populating the mm info is so complicated */
-	if (src == dst) {
-		struct mm_leaf *o = (struct mm_leaf *)move_forward(
-			cont, sizeof(struct mm_leaf));
-
-		o->entries = num;
-		o->used = (mm_info_t *)move_forward(cont, state_elems(num));
-		memset(o->used, 0, state_elems(num));
+	size_t s = sizeof_branch(num);
+	struct mm_branch *branch = (struct mm_branch *)cont;
+	memset(branch, 0, s);
+	cont += s;
 
-		*op = (mm_node_t *)o;
+	branch->num = num;
+	if (cur_order == req_order)
 		return cont;
-	}
-
-	struct mm_branch *o = (struct mm_branch *)move_forward(
-		cont, sizeof(struct mm_branch));
-
-	o->entries = num;
-	o->full = (mm_info_t *)move_forward(cont, state_elems(num));
-	cont = align_up(cont, sizeof(void *));
-	o->next = (mm_node_t **)move_forward(cont, next_elems(num));
-	memset(o->full, 0, state_elems(num));
-	memset(o->next, 0, next_elems(num));
 
-	for (size_t i = 0; i < num; ++i) {
-		cont = __populate_order(&o->next[i], cont, src - 1, dst,
-		                        order_width(src - 1));
+	pm_t prev = cont;
+	foreach_page(num) {
+		prev = cont;
+		cont = __populate_tree(cont, order_width(cur_order - 1),
+		                       cur_order - 1, req_order);
 	}
 
-	*op = (mm_node_t *)o;
+	branch->size = cont - prev;
 	return cont;
 }
 
 /**
- * Probe order node map.
+ * Populate bucket.
  *
- * @param cont Number of bytes written so far.
- * @param src Source order.
- * @param dst Destination order.
- * @param num Number of nodes in this order to calculate.
- * @return Number of bytes written so far.
+ * @param cont Address at which to continue placing data.
+ * @param base Base of bucket.
+ * @param num Number of elements in top level trees.
+ * @param order Bucket order.
+ * @return Top of bucket.
  */
-static pm_t __probe_order(pm_t cont, enum mm_order src, enum mm_order dst,
-                          size_t num)
+static pm_t __populate_bucket(pm_t cont, pm_t base, size_t num,
+                              enum mm_order order)
 {
-	if (src == dst) {
-		cont += sizeof(struct mm_leaf);
-		cont += state_elems(num);
-		return cont;
-	}
+	struct mm_bucket *bucket = (struct mm_bucket *)cont;
+	memset(bucket, 0, sizeof(*bucket));
+	cont += sizeof(*bucket);
 
-	cont += sizeof(struct mm_branch);
-	cont += state_elems(num);
-	cont = align_up(cont, sizeof(void *));
-	cont += next_elems(num);
+	bucket->order = order;
+	bucket->base = base;
 
-	for (size_t i = 0; i < num; ++i)
-		cont = __probe_order(cont, src - 1, dst, order_width(src - 1));
+	enum mm_order iter = order;
+	reverse_foreach_order_init(iter) {
+		bucket->tree[iter] = (struct mm_branch *)cont;
+		cont = __populate_tree(cont, num, order, iter);
+	}
 
 	return cont;
 }
 
-/** Populate order map.
- *
- * @param omap Address where to write order map pointer.
- * @param cont Address where to continue writing map data.
- * @param base Base of order map.
- * @param entries Number of order node entries in order map.
- * @param order Order of this order map.
- * @return Address to continue writing data to.
- */
-static pm_t __populate_omap(struct mm_omap **omap, pm_t cont, pm_t base,
-                            size_t entries, enum mm_order order)
+pm_t populate_pmap(pm_t ram_base, size_t ram_size, pm_t cont)
 {
-	struct mm_omap *lomap = (struct mm_omap *)move_forward(
-		cont, sizeof(struct mm_omap));
-	memset(lomap, 0, sizeof(struct mm_omap));
+	pm_t start = cont;
 
-	lomap->orders = (mm_node_t **)move_forward(
-		cont, (order + 1) * sizeof(mm_node_t **));
-	memset(lomap->orders, 0, (order + 1) * sizeof(mm_node_t **));
+	pmap = (struct mm_pmap *)cont;
+	memset(pmap, 0, sizeof(*pmap));
+	cont += sizeof(*pmap);
 
-	lomap->order = order;
-	lomap->base = base;
+	reverse_foreach_order(iter) {
+		size_t num = ram_size / order_size(iter);
+		if (num == 0)
+			continue;
 
-	for (size_t i = 0; i <= order; ++i)
-		cont = __populate_order(&lomap->orders[i], cont, order, i,
-		                        entries);
+		pmap->bucket[iter] = (struct mm_bucket *)cont;
+		cont = __populate_bucket(cont, ram_base, num, iter);
 
-	*omap = lomap;
-	return cont;
+		ram_size -= order_size(iter) * num;
+		ram_base += order_size(iter) * num;
+	}
+
+	return cont - start;
 }
 
 /**
- * Probe order map.
+ * Probe tree size.
  *
- * @param cont Number of bytes written so far.
- * @param entries Number of order node entries in this order map.
- * @param order Order of this order map.
- * @return Number of bytes written so far.
+ * @param cont Size to continue adding to.
+ * @param num Number of elements in tree.
+ * @param cur_order Current branch order.
+ * @param req_order Requested tree order.
+ * @return Size of tree added to \p cont.
  */
-static pm_t __probe_omap(pm_t cont, size_t entries, enum mm_order order)
+static pm_t __probe_tree(pm_t cont, size_t num, enum mm_order cur_order,
+                         enum mm_order req_order)
 {
-	cont += sizeof(struct mm_omap);
-	cont += (order + 1) * sizeof(mm_node_t **);
+	cont += sizeof_branch(num);
 
-	for (size_t i = 0; i <= order; ++i)
-		cont = __probe_order(cont, order, i, entries);
+	if (cur_order == req_order)
+		return cont;
+
+	foreach_page(num) {
+		cont = __probe_tree(cont, order_width(cur_order - 1),
+		                    cur_order - 1, req_order);
+	}
 
 	return cont;
 }
 
-/* only call from init */
-pm_t populate_pmap(pm_t ram_base, size_t ram_size, pm_t cont)
+/**
+ * Probe bucket size.
+ *
+ * @param cont Size to continue adding to.
+ * @param num Number of elements in bucket.
+ * @param order Bucket order.
+ * @return Size of bucket added to \p cont.
+ */
+static pm_t __probe_bucket(pm_t cont, size_t num, enum mm_order order)
 {
-	pm_t start = cont;
-	pmap = (struct mm_pmap *)move_forward(cont, sizeof(struct mm_pmap));
-	memset(pmap, 0, sizeof(struct mm_pmap));
+	cont += sizeof(struct mm_bucket);
 
-	pm_t ram_region = ram_base;
-	size_t ram_left = ram_size;
-	for (ssize_t i = __mm_max_order; i >= MM_O0; --i) {
-		size_t entries = ram_left / __mm_sizes[i];
-		if (entries == 0)
-			continue;
-
-		cont = __populate_omap(&pmap->omap[i], cont, ram_region,
-		                       entries, i);
-
-		ram_left -= __mm_sizes[i] * entries;
-		ram_region += (__mm_sizes[i] * entries);
+	reverse_foreach_order(iter) {
+		cont = __probe_tree(cont, num, order, iter);
 	}
 
-	return cont - start;
+	return cont;
 }
 
-/* not a huge fan of having a separate probe_pmap function as that seems like an
- * easy way to cause weird bugs. Should always at least check that probe_pmap
- * returns the same value as populate_pmap, or possibly even add in some method
- * to combine the two? */
-pm_t probe_pmap(pm_t ram_base, size_t ram_size)
+pm_t probe_pmap(size_t ram_size)
 {
-	pm_t cont = 0;
-
-	cont += sizeof(struct mm_pmap);
+	pm_t cont = sizeof(struct mm_pmap);
 
-	pm_t ram_region = ram_base;
-	size_t ram_left = ram_size;
-	for (ssize_t i = __mm_max_order; i >= MM_O0; --i) {
-		size_t entries = ram_left / __mm_sizes[i];
-		if (entries == 0)
+	reverse_foreach_order(iter) {
+		size_t num = ram_size / order_size(iter);
+		if (num == 0)
 			continue;
 
-		cont = __probe_omap(cont, entries, i);
+		cont = __probe_bucket(cont, num, iter);
 
-		ram_left -= __mm_sizes[i] * entries;
-		ram_region += (__mm_sizes[i] * entries);
+		ram_size -= order_size(iter) * num;
 	}
 
 	return cont;
@@ -615,7 +619,7 @@ void init_pmem(void *fdt)
 	 * ram map */
 	pm_t pmap_base = align_up(MAX(initrd_top, fdt_top), sizeof(int));
 
-	size_t probe_size = probe_pmap(ram_base, ram_size);
+	size_t probe_size = probe_pmap(ram_size);
 	size_t actual_size = populate_pmap(ram_base, ram_size, pmap_base);
 
 	if (probe_size != actual_size)
diff --git a/common/tcb.c b/common/tcb.c
index 1c1a623..9000941 100644
--- a/common/tcb.c
+++ b/common/tcb.c
@@ -43,7 +43,7 @@ void init_tcbs()
 	/* MM_O1 is 2MiB on riscv64, so 262144 different possible thread ids.
 	 * Should be enough, if we're really strapped for memory I might try
 	 * something smaller but this is fine for now. */
-	tcbs = (struct tcb **)alloc_page(MM_O1, 0);
+	tcbs = (struct tcb **)alloc_page(MM_O1);
 	num_tids = order_size(MM_O1) / sizeof(struct tcb *);
 	memset(tcbs, 0, order_size(MM_O1));
 }
@@ -94,7 +94,7 @@ static vm_t __setup_rpc_stack(struct tcb *t, size_t bytes)
 	size_t pages = __pages(bytes);
 	vmflags_t flags = VM_V | VM_R | VM_W | VM_U;
 	for (size_t i = 1; i <= pages; ++i) {
-		offset = alloc_page(BASE_PAGE, offset);
+		offset = alloc_page(BASE_PAGE);
 		map_vpage(t->proc.vmem, offset,
 		          RPC_STACK_TOP - BASE_PAGE_SIZE * i,
 		          flags, BASE_PAGE);
@@ -139,7 +139,7 @@ struct tcb *create_thread(struct tcb *p)
 {
 	hard_assert(tcbs, 0);
 
-	vm_t bottom = alloc_page(KERNEL_STACK_PAGE_ORDER, 0);
+	vm_t bottom = alloc_page(KERNEL_STACK_PAGE_ORDER);
 	/* move tcb to top of kernel stack, keeping alignment in check
 	 * (hopefully) */
 	/** \todo check alignment */
diff --git a/common/uapi/ipc.c b/common/uapi/ipc.c
index 01d65ed..34f707a 100644
--- a/common/uapi/ipc.c
+++ b/common/uapi/ipc.c
@@ -47,7 +47,7 @@ static struct sys_ret do_ipc(sys_arg_t pid,
 	r = get_rproc(r);
 
 	if (!r->callback)
-		return SYS_RET1(ERR_INIT);
+		return SYS_RET1(ERR_NOINIT);
 
 	/** \todo place data on rpc stack and clone into virtual memory */
 	set_return(t, r->callback);
diff --git a/common/vmem.c b/common/vmem.c
index fe0178e..73bc672 100644
--- a/common/vmem.c
+++ b/common/vmem.c
@@ -162,7 +162,7 @@ stat_t free_uvmem(struct tcb *r, vm_t va)
 stat_t alloc_uvmem_wrapper(struct vmem *b, pm_t *offset, vm_t vaddr,
                            vmflags_t flags, enum mm_order order, void *data)
 {
-	*offset = alloc_page(order, *offset);
+	*offset = alloc_page(order);
 	if (!*offset)
 		return INFO_TRGN; /* try again */
 
@@ -180,7 +180,7 @@ stat_t alloc_shared_wrapper(struct vmem *b, pm_t *offset, vm_t vaddr,
 	if (order != MM_O0)
 		return INFO_TRGN;
 
-	*offset = alloc_page(MM_O0, *offset);
+	*offset = alloc_page(MM_O0);
 
 	stat_t *status = (stat_t *)data, ret;
 	ret = map_vpage(b, *offset, vaddr, flags, order);
diff --git a/include/apos/bits.h b/include/apos/bits.h
index ce3aba1..adb8e85 100644
--- a/include/apos/bits.h
+++ b/include/apos/bits.h
@@ -14,6 +14,17 @@
 
 /** @name Arithmetic integer bit manipulation. */
 /** @{ */
+/**
+ * Find first set bit in \c int.
+ *
+ * @param v Integer to find first set bit in.
+ * @return Index of least significant bit + 1 or 0 if \p v is 0.
+ */
+#if __has_builtin(__builtin_ffs)
+#define ffs(v) __builtin_ffs(v)
+#else
+int ffs(int v);
+#endif
 
 /**
  * Check if bits are set.
@@ -134,6 +145,67 @@ static inline void bitmap_clear(void *bmap, size_t n)
 	clear_nbit(bitmap[i], r);
 }
 
+/**
+ * Find first bit, either set or unset, in bitmap.
+ *
+ * @param bmap Bitmap.
+ * @param n Size of bitmap in bits.
+ * @param set Wether to seek for set or unset bits.
+ * @return Index of found bit + 1 or \p n + 1 if no bit was found.
+ */
+static inline size_t bitmap_find_first(void *bmap, size_t n, bool set)
+{
+	size_t i = n / (sizeof(int) * 8);
+
+	size_t c = 0;
+	int *imap = (int *)bmap;
+
+	int target = set ? 0 : -1;
+	for (; c < i; ++c)
+		if (imap[c] != target)
+			break;
+
+	size_t b = c * sizeof(int) * 8;
+	int check = set ? imap[c] : ~imap[c];
+	if (c != i)
+		return b + ffs(check) - 1;
+
+	size_t r = n - (i * sizeof(int) * 8);
+	if (!r)
+		return n + 1;
+
+	bool comp = set ? true : false;
+	for (size_t a = 0; a < r; ++a)
+		if (bitmap_is_set(bmap, b + a) == comp)
+			return b + a;
+
+	return n + 1;
+}
+
+/**
+ * Convenience wrapper around bitmap_find_first().
+ *
+ * @param bmap \see bitmap_find_first().
+ * @param n \see bitmap_find_first().
+ * @return \see bitmap_find_first().
+ */
+static inline size_t bitmap_find_first_unset(void *bmap, size_t n)
+{
+	return bitmap_find_first(bmap, n, false);
+}
+
+/**
+ * Convenience wrapper around bitmap_find_first().
+ *
+ * @param bmap \see bitmap_find_first().
+ * @param n \see bitmap_find_first().
+ * @return \see bitmap_find_first().
+ */
+static inline size_t bitmap_find_first_set(void *bmap, size_t n)
+{
+	return bitmap_find_first(bmap, n, true);
+}
+
 /** @} */
 
 /**
diff --git a/include/apos/mem.h b/include/apos/mem.h
index 0eed025..da88edb 100644
--- a/include/apos/mem.h
+++ b/include/apos/mem.h
@@ -12,99 +12,15 @@
 #include <apos/utils.h>
 #include <apos/types.h>
 
-/** Helper macro for getting bit width of \ref mm_info_t. */
-#define MM_OINFO_WIDTH (sizeof(mm_info_t) * 8)
-
-/**
- * Extract index of page order \c order from base page index \c pnum.
- *
- * @param pnum Base page order index.
- * @param order Order page index to convert to.
- * @return Index of page order \c order.
- */
-#define pnum_to_index(pnum, order) \
-	(((pnum) >> order_offset(order)) & (order_width(order) - 1))
-
 /**
  * Convert physical memory address \c paddr to index of page order \c order.
  *
- * @param paddr Physical memory address.
+ * @param p Physical memory address.
  * @param order Order page index to convert to.
  * @return Index of page order \c order.
  */
-#define pm_to_index(paddr, order) \
-	(pnum_to_index(pm_to_pnum(paddr), (order)))
-
-/**
- * Convert physical memory address \c paddr to corresponding page number.
- *
- * @param paddr Physical memory address.
- * @return Corresponding page number.
- */
-#define pm_to_pnum(paddr) ((paddr) >> __mm_page_shift)
-
-/**
- * Convert page number to physical address.
- * Note that since a page number is the base page an address lies in,
- * @code pnum_to_pm(pm_to_pnum(p)) != p @endcode
- *
- * @param pnum Page number.
- * @return Corresponding physical address.
- */
-#define pnum_to_pm(pnum) ((pnum) << __mm_page_shift)
-
-/**
- * Add \c num to \c var and return value before addition.
- *
- * @param var Variable to add \c num to.
- * @param num Number to add to \c var.
- * @return Value of \c var before addition.
- */
-#define move_forward(var, num) (((var) += (num)) - (num))
-
-/**
- * Helper for calculating highest index of elements in order info map.
- * Since the number of entries is stored with the granularity of \c MM_OINFO_WIDTH,
- * the highest index element is \c num rounded up to the nearest index multiple
- * of \c MM_OINFO_WIDTH. This is due to some data access optimizations over in
- * common/pmem.c.
- *
- * @param num Number of entries in map.
- * @return Highest index of element in map.
- */
-#define num_elems(num) \
-	(((num) + MM_OINFO_WIDTH - 1) / MM_OINFO_WIDTH)
-/**
- * Helper for calculating starting index of element.
- *
- * @param num Number of starting entry.
- * @return Index of starting element in which entry resides.
- */
-#define num_indexes(num) ((num) / MM_OINFO_WIDTH)
-
-/**
- * Helper for calculating index of element from entry number.
- *
- * @param num Entry number.
- * @return Index of element in which entry resides.
- */
-#define index_elems(num) ((num) / MM_OINFO_WIDTH)
-
-/**
- * Helper for calculating size of state for storing elements in.
- *
- * @param num Number of entries.
- * @return Size of element state buffer.
- */
-#define state_elems(num) (sizeof(mm_info_t) * (num_elems(num)))
-
-/**
- * Helper for calculating size of pointer buffer.
- *
- * @param num Number of entries.
- * @return Size of pointer buffer.
- */
-#define next_elems(num) (sizeof(void *) * (num))
+#define pm_to_index(p, order) \
+	((p >> order_shift(order)) & (order_width(order) - 1))
 
 /**
  * Get highest possible index in an order.
@@ -120,7 +36,7 @@
  * @param order Order to query.
  * @return Starting offset of order bits.
  */
-#define order_offset(order) (__mm_shifts[order])
+#define order_shift(order) (__mm_shifts[order])
 
 /**
  * Get number of order bits in an address.
@@ -139,6 +55,20 @@
 #define order_size(order) (__mm_sizes[order])
 
 /**
+ * Get highest order supported by the current configuration.
+ *
+ * @return Max supported order.
+ */
+#define max_order() (__mm_max_order)
+
+/**
+ * Get base page shift.
+ *
+ * @return Page shift.
+ */
+#define page_shift() (__mm_page_shift)
+
+/**
  * Get number of elements needed to represent this order.
  *
  * @param order Order to query.
@@ -154,14 +84,6 @@
  */
 #define order_container(idx) ((idx) / MM_OINFO_WIDTH)
 
-/**
- * Entry index within the element that contains it.
- *
- * @param idx Index of entry.
- * @return Index of entry within its containing element.
- */
-#define order_bit(idx) ((idx) & (MM_OINFO_WIDTH - 1))
-
 /** \todo Get rid of slightly ugly __* syntax, as these aren't static. */
 
 /**
@@ -183,8 +105,6 @@
 /**
  * Get page number of physical address.
  *
- * \todo Isn't this the same as \ref pm_to_pnum()?
- *
  * @param x Physical address.
  * @return Corresponding page number.
  */
@@ -225,23 +145,11 @@
 /** Maximum number of page orders allowed. Likely massively overkill. */
 #define NUM_ORDERS 10
 
-/** Gives access to global page order shift information. \global */
-extern size_t __mm_shifts[NUM_ORDERS];
-
-/** Gives access to global page order width information. \global */
-extern size_t __mm_widths[NUM_ORDERS];
-
-/** Gives access to global page order size information. \global */
-extern size_t __mm_sizes[NUM_ORDERS];
-
-/** Gives access to global base page shift. \global */
-extern size_t __mm_page_shift;
-
-/** Gives access to global maximum order size. \global */
-extern size_t __mm_max_order;
-
 /** Give names to page orders. */
 enum mm_order {
+	/** NULL marker. */
+	MM_MIN = -1,
+
 	/** Base order. */
 	MM_O0 = 0,
 
@@ -271,10 +179,25 @@ enum mm_order {
 
 	/** Order 9. */
 	MM_O9 = 9,
+
+	/** Number of orders */
+	MM_NUM,
 };
 
-/** Page number. */
-typedef ssize_t pnum_t;
+/** Gives access to global page order shift information. \global */
+extern size_t __mm_shifts[NUM_ORDERS];
+
+/** Gives access to global page order width information. \global */
+extern size_t __mm_widths[NUM_ORDERS];
+
+/** Gives access to global page order size information. \global */
+extern size_t __mm_sizes[NUM_ORDERS];
+
+/** Gives access to global base page shift. \global */
+extern size_t __mm_page_shift;
+
+/** Gives access to global maximum order size. \global */
+extern enum mm_order __mm_max_order;
 
 /**
  * Initialize memory subsystem data. Populates __mm_* with data given.
diff --git a/include/apos/pmem.h b/include/apos/pmem.h
index fe03242..6278b57 100644
--- a/include/apos/pmem.h
+++ b/include/apos/pmem.h
@@ -17,17 +17,17 @@
  * Free physical page.
  *
  * @param order Order of page to free.
- * @param paddr Physical address of page.
+ * @param addr Physical address of page.
  */
-void free_page(enum mm_order order, pm_t paddr);
+void free_page(enum mm_order order, pm_t addr);
 
 /**
  * Mark page used.
  *
  * @param order Order of page to mark.
- * @param paddr Physical address of page.
+ * @param addr Physical address of page.
  */
-void mark_used(enum mm_order order, pm_t paddr);
+void mark_used(enum mm_order order, pm_t addr);
 
 /**
  * Allocate physical page.
@@ -36,10 +36,9 @@ void mark_used(enum mm_order order, pm_t paddr);
  * allows us to skip already checked pages when allocating a second page.
  *
  * @param order Order of page to allocate.
- * @param offset Hint as to which address to start looking from.
  * @return pm_t Physical address of page when succesful, else \c NULL.
  */
-pm_t alloc_page(enum mm_order order, pm_t offset);
+pm_t alloc_page(enum mm_order order);
 
 /**
  * Populate physical RAM usage map.
@@ -56,12 +55,11 @@ pm_t populate_pmap(pm_t ram_base, size_t ram_size, pm_t cont);
 /**
  * Probe size of RAM usage map.
  *
- * @param ram_base Base physical address of RAM.
  * @param ram_size Size of physical RAM.
  * @return Size of physical map. Check that it matches with \ref
  * populate_pmap().
  */
-pm_t probe_pmap(pm_t ram_base, size_t ram_size);
+pm_t probe_pmap(size_t ram_size);
 
 /**
  * Initialize physical memory subsystem.