From d2f6e8254dc359be96a0fbf44e47d4fe6e307960 Mon Sep 17 00:00:00 2001
From: Stanislav Paskalev <spaskalev@protonmail.com>
Date: Fri, 28 Jun 2024 00:29:15 +0300
Subject: [PATCH] Add buddy_get_embed_at function. Fixes #114. (#118)

---
 buddy_alloc.h | 24 ++++++++++++++++++++++++
 tests.c       | 19 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/buddy_alloc.h b/buddy_alloc.h
index c22ffa4..d0ef7dd 100644
--- a/buddy_alloc.h
+++ b/buddy_alloc.h
@@ -56,6 +56,12 @@ struct buddy *buddy_init_alignment(unsigned char *at, unsigned char *main, size_
  */
 struct buddy *buddy_embed(unsigned char *main, size_t memory_size);
 
+/*
+ * Returns the address of a previously-created buddy allocator at the arena.
+ * Use to get a new handle to the allocator when the arena is moved or copied.
+ */
+struct buddy *buddy_get_embed_at(unsigned char *main, size_t memory_size);
+
 /*
  * Initializes a binary buddy memory allocator embedded in the specified arena
  * using a non-default alignment.
@@ -63,6 +69,12 @@ struct buddy *buddy_embed(unsigned char *main, size_t memory_size);
  */
 struct buddy *buddy_embed_alignment(unsigned char *main, size_t memory_size, size_t alignment);
 
+/*
+ * Returns the address of a previously-created buddy allocator at the arena.
+ * Use to get a new handle to the allocator when the arena is moved or copied.
+ */
+struct buddy *buddy_get_embed_at_alignment(unsigned char *main, size_t memory_size, size_t alignment);
+
 /* Resizes the arena and metadata to a new size. */
 struct buddy *buddy_resize(struct buddy *buddy, size_t new_memory_size);
 
@@ -502,6 +514,10 @@ struct buddy *buddy_embed(unsigned char *main, size_t memory_size) {
     return buddy_embed_alignment(main, memory_size, BUDDY_ALLOC_ALIGN);
 }
 
+struct buddy *buddy_get_embed_at(unsigned char *main, size_t memory_size) {
+    return buddy_get_embed_at_alignment(main, memory_size, BUDDY_ALLOC_ALIGN);
+}
+
 struct buddy *buddy_embed_alignment(unsigned char *main, size_t memory_size, size_t alignment) {
     struct buddy_embed_check check_result;
     struct buddy *buddy;
@@ -527,6 +543,14 @@ struct buddy *buddy_embed_alignment(unsigned char *main, size_t memory_size, siz
     return buddy;
 }
 
+struct buddy *buddy_get_embed_at_alignment(unsigned char *main, size_t memory_size, size_t alignment) {
+    struct buddy_embed_check check_result = buddy_embed_offset(memory_size, alignment);
+    if (!check_result.can_fit) {
+        return NULL;
+    }
+    return (struct buddy *)(main + check_result.offset);
+}
+
 struct buddy *buddy_resize(struct buddy *buddy, size_t new_memory_size) {
     if (new_memory_size == buddy->memory_size) {
         return buddy;
diff --git a/tests.c b/tests.c
index 4edc346..9cc3d8d 100644
--- a/tests.c
+++ b/tests.c
@@ -1378,6 +1378,23 @@ void test_buddy_embedded_malloc_alignment(void) {
 	assert(buddy_arena_size(buddy) == 3584);
 }
 
+void test_buddy_embed_at(void) {
+	unsigned char buf1[4096] = {0};
+	unsigned char buf2[4096] = {0};
+	struct buddy* buddy;
+	start_test;
+	buddy = buddy_embed(buf1, 4096);
+	buddy_malloc(buddy, 2048);
+	memcpy(buf2, buf1, 4096);
+	buddy = buddy_get_embed_at(buf2, 0);
+	assert(buddy == NULL);
+	buddy = buddy_get_embed_at(buf2, 4096);
+	assert(buddy != NULL);
+	assert(buddy_malloc(buddy, 2048) == NULL);
+	buddy_free(buddy, buf2);
+	assert(buddy_malloc(buddy, 2048) == buf2);
+}
+
 void test_buddy_mixed_use_01(void) {
 	unsigned char *buddy_buf = malloc(buddy_sizeof(512));
 	unsigned char data_buf[512];
@@ -2426,6 +2443,8 @@ int main(void) {
 		test_buddy_embedded_malloc_01();
 		test_buddy_embedded_malloc_alignment();
 
+		test_buddy_embed_at();
+
 		test_buddy_mixed_use_01();
 		test_buddy_mixed_use_02();
 		test_buddy_mixed_use_03();