From 84efc4471070711467536c1a6106625f8b2a34fe Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Thu, 29 Dec 2022 10:19:57 +0000
Subject: [PATCH] memblock region iteration can now be bounded to a certain
 memory range

---
 sandbox/memblock/main.c     |  2 +-
 sandbox/memblock/memblock.c | 59 ++++++++++++++++++++++++++++---------
 sandbox/memblock/memblock.h | 11 +++++++
 3 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/sandbox/memblock/main.c b/sandbox/memblock/main.c
index 7b0d8c1..866e910 100644
--- a/sandbox/memblock/main.c
+++ b/sandbox/memblock/main.c
@@ -48,7 +48,7 @@ int main(int argc, const char **argv)
 	}
 
 	printf("free regions:\n");
-	for_each_free_mem_range(&it, 0, 0x100000) {
+	for_each_free_mem_range(&it, 0, ULLONG_MAX) {
 		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
 				it.base,
 				it.limit);
diff --git a/sandbox/memblock/memblock.c b/sandbox/memblock/memblock.c
index 8d6ca21..bde37bc 100644
--- a/sandbox/memblock/memblock.c
+++ b/sandbox/memblock/memblock.c
@@ -11,7 +11,10 @@
 #define IDX_A(idx) ((idx) & 0xFFFFFFFF)
 #define IDX_B(idx) (((idx) >> 32) & 0xFFFFFFFF)
 
-#define ADDR_MAX 0xFFFFFFFF
+/* the maximum possible value for a pointer type.
+ * Note that any pointers returned by the memblock API will still
+ * be bounded by the defined memory regions, and not by this constant. */
+#define ADDR_MAX (~(uintptr_t)0)
 
 static memblock_region_t init_memory_regions[MEMBLOCK_INIT_MEMORY_REGION_COUNT];
 static memblock_region_t init_reserved_regions[MEMBLOCK_INIT_RESERVED_REGION_COUNT];
@@ -172,10 +175,22 @@ void __next_memory_region(memblock_iter_t *it, memblock_type_t *type_a, memblock
 		uintptr_t m_start = m->base;
 		uintptr_t m_end = m->limit;
 
+		if (m_end < start) {
+			/* we haven't reached the requested memory range yet */
+			continue;
+		}
+
+		if (m_start > end) {
+			/* we have gone past the requested memory range and can now stop */
+			break;
+		}
+
 		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 */
+			/* r_start and r_end delimit the region of memory between the current and previous reserved regions.
+			 * if we have gone past the last reserved region, these variables delimit the range between the end
+			 * of the last reserved region and the end of memory. */
 			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;
 
@@ -184,21 +199,37 @@ void __next_memory_region(memblock_iter_t *it, memblock_type_t *type_a, memblock
 				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;
+			/* we've already gone past this free memory region. move to the next one */
+			if (m_start >= r_end) {
+				continue;
 			}
+
+			/* we want the area that is overlapped by both
+			 *	region M (m_start - m_end) : The region defined as system memory.
+			 *	region R (r_start - r_end) : The region defined as free / outside of any reserved regions.
+			 */
+			it->base = MAX(m_start, r_start);
+			it->limit = MIN(m_end, r_end);
+
+			/* further limit the region to the intersection between the region itself and the
+			 * specified iteration bounds */
+			it->base = MAX(it->base, start);
+			it->limit = MIN(it->limit, end);
+
+			/* whichever region is smaller, increment the pointer for that type, so we can
+			 * compare the larger region with the next region of the incremented type. */
+			if (m_end <= r_end) {
+				idx_a++;
+			} else {
+				idx_b++;
+			}
+
+			/* store the position for the next iteration */
+			it->idx = ITER(idx_a, idx_b);
+			return;
 		}
 	}
 
+	/* ULLONG_MAX signals the end of the iteration */
 	it->idx = ITER_END;
 }
diff --git a/sandbox/memblock/memblock.h b/sandbox/memblock/memblock.h
index ad96b7d..c0806c1 100644
--- a/sandbox/memblock/memblock.h
+++ b/sandbox/memblock/memblock.h
@@ -24,6 +24,17 @@
 
 typedef uint64_t memblock_index_t;
 
+typedef enum memblock_region_status {
+	/* Used in memblock.memory regions, indicates that the memory region exists */
+	MEMBLOCK_MEMORY = 0,
+	/* Used in memblock.reserved regions, indicates that the memory region was reserved
+	 * by a call to memblock_alloc() */
+	MEMBLOCK_ALLOC,
+	/* Used in memblock.reserved regions, indicates that the memory region was reserved
+	 * by a call to memblock_reserve() */
+	MEMBLOCK_RESERVED,
+} memblock_region_status_t;
+
 typedef struct memblock_region {
 	/* the address of the first byte that makes up the region */
 	uintptr_t base;