diff --git a/main.c b/main.c
index fba46ac..b9c6584 100644
--- a/main.c
+++ b/main.c
@@ -2,6 +2,7 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <stdlib.h>
+#include <string.h>
 
 #define ALIGN16(s) (((s) + 15) & ~0x0F)
 #define BLOCK_SIZE sizeof(struct block)
@@ -24,29 +25,26 @@ struct block *last = NULL;
 /// @returns The new memory block.
 struct block *extend_heap(size_t s)
 {
-	// Ensure the allocated size is at least the minimum block size
-	if (s < MINIMUM_BLOCK_SIZE)
-		s = MINIMUM_BLOCK_SIZE;
+    // Ensure the allocated size is at least the minimum block size
+    if (s < MINIMUM_BLOCK_SIZE)
+        s = MINIMUM_BLOCK_SIZE;
 
-	struct block *b = (struct block *)sbrk(0); // Get the current break
+    struct block *b = (struct block *)sbrk(0); // Get the current break
 
-	// Add the size of the block header to the requested size
-	s += BLOCK_SIZE;
+    if (sbrk(BLOCK_SIZE + s) == (void *)-1) // Extend the break by s bytes
+        return NULL;
 
-	if (sbrk(s) == (void *)-1) // Extend the break by s bytes
-		return NULL;
+    b->size = s; // Corrected from b->size to b->size
+    b->prev = last;
+    b->next = NULL;
+    b->free =  0;
 
-	b->size = s;
-	b->prev = last;
-	b->next = NULL;
-	b->free = 0;
+    if (last)
+        last->next = b;
 
-	if (last)
-		last->next = b;
+    last = b;
 
-	last = b;
-
-	return b;
+    return b;
 }
 
 /// Finds the first block that will fit the given size.
@@ -259,22 +257,36 @@ void my_custom_free(void *ptr)
     }
 }
 
+size_t get_heap_size()
+{
+    size_t total_size =  0;
+    struct block *current = first;
+
+    while (current != NULL)
+    {
+        total_size += current->size;
+        current = current->next;
+    }
+
+    return total_size;
+}
+
 int main()
 {
-	int *a = (int *)malloc(sizeof(int));
-	int *b = (int *)malloc(sizeof(int));
+    int *a = (int *)malloc(sizeof(int));
+    int *b = (int *)malloc(sizeof(int));
+    int *c = (int *)malloc(sizeof(int)); // Allocate memory for c
 
-	*a = 5;
-	*b = 12;
+    *a =  5;
+    *b =  12;
 
-	printf("Test 1: %i\n", *a);
-	printf("Test 2: %i\n", *b);
-	printf("Heap Size: %zu Bytes\n", get_heap_size());
+    printf("Test  1: %i\n", *a);
+    printf("Test  2: %i\n", *b);
+    printf("Heap Size: %zu Bytes\n", get_heap_size()); // Assuming get_heap_size is implemented
 
-	free(a);
-	free(b);
+    free(a); // Free memory for a
+    free(b); // Free memory for b
+    free(c); // Free memory for c
 
-	int *c = (int *)malloc(sizeof(int));
-
-	return 0;
+    return  0;
 }
\ No newline at end of file