From d9eda8eb1873db01883b725324d6678c9fe69229 Mon Sep 17 00:00:00 2001
From: "Michael D. Lowis" <mike@mdlowis.com>
Date: Sat, 4 Apr 2015 13:32:46 -0400
Subject: [PATCH] Finished segment find and mark routine

---
 source/runtime/segment.c | 37 ++++++++++++++++++++++++++----------
 tests/test_segment.c     | 41 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/source/runtime/segment.c b/source/runtime/segment.c
index ee9226b..fa69bc6 100644
--- a/source/runtime/segment.c
+++ b/source/runtime/segment.c
@@ -39,6 +39,21 @@ bool segment_full(segment_t* seg) {
     return (0u == seg->blockmap[0]);
 }
 
+static void* alloc_block(segment_t* seg, uintptr_t root_idx, uintptr_t submap_idx) {
+    void* object = NULL;
+    uintptr_t submap_mask = (1u << submap_idx);
+    if (0u != (seg->blockmap[root_idx + 1] & submap_mask)) {
+        /* Get the object's pointer */
+        uintptr_t offset = seg->blocksize * ((root_idx * 16u) + submap_idx);
+        object = seg->start + offset;
+        /* Mark it as allocated */
+        seg->blockmap[root_idx + 1] &= ~(submap_mask);
+        if (0u == seg->blockmap[root_idx + 1])
+            seg->blockmap[0] &= ~(1u << root_idx);
+    }
+    return object;
+}
+
 void* segment_alloc(segment_t* seg) {
     void* obj = NULL;
     /* Check if the segment has any free blocks */
@@ -46,13 +61,8 @@ void* segment_alloc(segment_t* seg) {
         /* Find the free block */
         uint8_t root_idx   = get_free_index(seg->blockmap[0]);
         uint8_t submap_idx = get_free_index(seg->blockmap[root_idx + 1]);
-        /* Calculate it's address */
-        uintptr_t offset = seg->blocksize * ((root_idx * 16u) + submap_idx);
-        obj = seg->start + offset;
-        /* Mark it allocated */
-        seg->blockmap[root_idx + 1] &= ~(1u << submap_idx);
-        if (0u == seg->blockmap[root_idx + 1])
-            seg->blockmap[0] &= ~(1u << root_idx);
+        /* Allocate and return its address */
+        obj = alloc_block(seg, root_idx, submap_idx);
     }
     return obj;
 }
@@ -63,8 +73,15 @@ void segment_clear_map(segment_t* seg)
 }
 
 void* segment_find_and_mark(segment_t* seg, uintptr_t addr) {
-    (void)seg;
-    (void)addr;
-    return NULL;
+    void* obj = NULL;
+    uintptr_t start = (uintptr_t)&(seg->start[0]);
+    uintptr_t end   = (uintptr_t)seg->end;
+    if ((start <= addr) && (addr < end)) {
+        uintptr_t block_num  = (addr - start) / (seg->blocksize * sizeof(uintptr_t));
+        uintptr_t root_idx   = block_num / 16u;
+        uintptr_t submap_idx = block_num % 16u;
+        obj = alloc_block(seg, root_idx, submap_idx);
+    }
+    return obj;
 }
 
diff --git a/tests/test_segment.c b/tests/test_segment.c
index 016e455..4bbd7ea 100644
--- a/tests/test_segment.c
+++ b/tests/test_segment.c
@@ -1,6 +1,8 @@
 #include "atf.h"
 #include "segment.h"
 
+#include <stdio.h>
+
 TEST_SUITE(Segment) {
     /* Verify: segment_create
      *************************************************************************/
@@ -74,4 +76,43 @@ TEST_SUITE(Segment) {
         CHECK(obj == NULL);
         segment_destroy(seg);
     }
+
+    /* Verify: segment_find_and_mark
+     *************************************************************************/
+    TEST(Verify_find_and_mark_should_find_and_mark_a_block) {
+        segment_t* seg = segment_create(2u, NULL);
+        CHECK(&seg->start[2] == segment_find_and_mark(seg, (uintptr_t)&seg->start[2]));
+        segment_destroy(seg);
+    }
+
+    TEST(Verify_find_and_mark_should_find_and_mark_first_block) {
+        segment_t* seg = segment_create(2u, NULL);
+        CHECK(&seg->start[0] == segment_find_and_mark(seg, (uintptr_t)&seg->start[0]));
+        segment_destroy(seg);
+    }
+
+    TEST(Verify_find_and_mark_should_find_and_mark_last_block) {
+        segment_t* seg = segment_create(2u, NULL);
+        CHECK(&seg->start[510] == segment_find_and_mark(seg, ((uintptr_t)&seg->start[511]) + (sizeof(uintptr_t) - 1) ));
+        segment_destroy(seg);
+    }
+
+    TEST(Verify_find_and_mark_should_return_null_when_block_is_already_in_use) {
+        segment_t* seg = segment_create(2u, NULL);
+        seg->blockmap[1] &= ~(2);
+        CHECK(NULL == segment_find_and_mark(seg, (uintptr_t)&seg->start[2]));
+        segment_destroy(seg);
+    }
+
+    TEST(Verify_find_and_mark_should_return_null_when_address_is_lower_than_segment_range) {
+        segment_t* seg = segment_create(2u, NULL);
+        CHECK(NULL == segment_find_and_mark(seg, (uintptr_t)&seg->start[-1]));
+        segment_destroy(seg);
+    }
+
+    TEST(Verify_find_and_mark_should_return_null_when_address_is_greater_than_segment_range) {
+        segment_t* seg = segment_create(2u, NULL);
+        CHECK(NULL == segment_find_and_mark(seg, (uintptr_t)&seg->start[512]));
+        segment_destroy(seg);
+    }
 }
-- 
2.54.0