From d0091a4df780ae58e3d202ecfc11dc8c8cf1abc3 Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Wed, 28 Dec 2022 23:03:30 +0000
Subject: [PATCH] Implemented macros for iterating over memblock regions,
 including free regions

---
 sandbox/memblock/main.c     | 22 ++++++++++----
 sandbox/memblock/memblock.c | 60 +++++++++++++++++++++++++++++++++++++
 sandbox/memblock/memblock.h | 17 +++++++++++
 3 files changed, 93 insertions(+), 6 deletions(-)

diff --git a/sandbox/memblock/main.c b/sandbox/memblock/main.c
index 601be27..7b0d8c1 100644
--- a/sandbox/memblock/main.c
+++ b/sandbox/memblock/main.c
@@ -29,19 +29,29 @@ int main(int argc, const char **argv)
 	memblock_reserve(0x10000, 0x40000);
 	memblock_reserve(0x60000, 0x20000);
 	memblock_reserve(0x30000, 0x40000);
+	memblock_reserve(0x100000, 0x10000);
 
 	printf("memory regions:\n");
-	for (unsigned int i = 0; i < memblock.memory.count; i++) {
+
+	memblock_iter_t it;
+	for_each_mem_range(&it, 0, 0x100000) {
 		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
-				memblock.memory.regions[i].base,
-				memblock.memory.regions[i].limit);
+				it.base,
+				it.limit);
 	}
 
 	printf("reserved regions:\n");
-	for (unsigned int i = 0; i < memblock.reserved.count; i++) {
+	for_each_reserved_mem_range(&it, 0, 0x100000) {
 		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
-				memblock.reserved.regions[i].base,
-				memblock.reserved.regions[i].limit);
+				it.base,
+				it.limit);
+	}
+
+	printf("free regions:\n");
+	for_each_free_mem_range(&it, 0, 0x100000) {
+		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
+				it.base,
+				it.limit);
 	}
 
 	munmap(system_memory, MEMORY_SIZE);
diff --git a/sandbox/memblock/memblock.c b/sandbox/memblock/memblock.c
index 2b823be..8d6ca21 100644
--- a/sandbox/memblock/memblock.c
+++ b/sandbox/memblock/memblock.c
@@ -1,7 +1,18 @@
 #include <stdbool.h>
+#include <limits.h>
 #include <string.h>
 #include "memblock.h"
 
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+
+#define ITER(a, b) ((uint64_t)(a) | ((uint64_t)(b) << 32))
+#define ITER_END ULLONG_MAX
+#define IDX_A(idx) ((idx) & 0xFFFFFFFF)
+#define IDX_B(idx) (((idx) >> 32) & 0xFFFFFFFF)
+
+#define ADDR_MAX 0xFFFFFFFF
+
 static memblock_region_t init_memory_regions[MEMBLOCK_INIT_MEMORY_REGION_COUNT];
 static memblock_region_t init_reserved_regions[MEMBLOCK_INIT_RESERVED_REGION_COUNT];
 
@@ -142,3 +153,52 @@ uintptr_t memblock_alloc(size_t size)
 {
 	return 0;
 }
+
+void __next_memory_region(memblock_iter_t *it, memblock_type_t *type_a, memblock_type_t *type_b, uintptr_t start, uintptr_t end)
+{
+	unsigned int idx_a = IDX_A(it->idx);
+	unsigned int idx_b = IDX_B(it->idx);
+
+	for (; idx_a < type_a->count; idx_a++) {
+		memblock_region_t *m = &type_a->regions[idx_a];
+		if (!type_b) {
+			it->base = m->base;
+			it->limit = m->limit;
+
+			it->idx = ITER(idx_a + 1, idx_b);
+			return;
+		}
+
+		uintptr_t m_start = m->base;
+		uintptr_t m_end = m->limit;
+
+		for (; idx_b < type_b->count + 1; idx_b++) {
+			memblock_region_t *r = &type_b->regions[idx_b];
+
+			/* r_start and r_end delimit the region of memory between the current and previous reserved regions */
+			uintptr_t r_start = idx_b > 0 ? r[-1].limit + 1 : 0;
+			uintptr_t r_end = idx_b < type_b->count ? r->base - 1 : ADDR_MAX;
+
+			if (r_start >= m_end) {
+				/* we've gone past the end of the current memory region, and need to go to the next one */
+				break;
+			}
+
+			if (m_start < r_end) {
+				it->base = MAX(m_start, r_start);
+				it->limit = MIN(m_end, r_end);
+
+				if (m_end <= r_end) {
+					idx_a++;
+				} else {
+					idx_b++;
+				}
+
+				it->idx = ITER(idx_a, idx_b);
+				return;
+			}
+		}
+	}
+
+	it->idx = ITER_END;
+}
diff --git a/sandbox/memblock/memblock.h b/sandbox/memblock/memblock.h
index 2148a6b..ad96b7d 100644
--- a/sandbox/memblock/memblock.h
+++ b/sandbox/memblock/memblock.h
@@ -2,12 +2,25 @@
 #define MEMBLOCK_H_
 
 #include <stddef.h>
+#include <limits.h>
 #include <stdint.h>
 
 #define MEMBLOCK_INIT_MEMORY_REGION_COUNT 128
 #define MEMBLOCK_INIT_RESERVED_REGION_COUNT 128
 
 #define __for_each_mem_range(i, type_a, type_b, p_start, p_end) \
+	for ((i)->idx = 0, __next_memory_region(i, type_a, type_b, p_start, p_end); \
+			(i)->idx != ULLONG_MAX; \
+			__next_memory_region(i, type_a, type_b, p_start, p_end))
+
+#define for_each_mem_range(i, p_start, p_end) \
+	__for_each_mem_range(i, &memblock.memory, NULL, p_start, p_end)
+
+#define for_each_reserved_mem_range(i, p_start, p_end) \
+	__for_each_mem_range(i, &memblock.reserved, NULL, p_start, p_end)
+
+#define for_each_free_mem_range(i, p_start, p_end) \
+	__for_each_mem_range(i, &memblock.memory, &memblock.reserved, p_start, p_end)
 
 typedef uint64_t memblock_index_t;
 
@@ -45,4 +58,8 @@ extern int memblock_reserve(uintptr_t base, size_t size);
 
 extern uintptr_t memblock_alloc(size_t size);
 
+extern void __next_memory_region(memblock_iter_t *it, \
+		memblock_type_t *type_a, memblock_type_t *type_b,
+		uintptr_t start, uintptr_t end);
+
 #endif