Fix malloc-free logic for correct memory management

include/malloc_from_scratch/memory_internal.h

@@ -42,6 +42,7 @@ void* getPayloadAddress(MemoryBlock* block);
 MemoryBlock* getMetadata(void* payload_address);
 MemoryBlock* getMemoryBlockFromAddress(void* address);
 size_t getSizeOfAllocatedMemoryBlock(MemoryBlock* block);
+bool isPointerInHeap(void* ptr);
 void* getErrorCodeInVoidPtr(size_t error_code);
 
 // Test helper functions to inspect allocator state

src/free.cpp

@@ -14,18 +14,22 @@ void free(void* ptr)
         return;
     }
 
-    void* current_program_break = sbrk(0);
-
-    if (ptr < internal::heap_start || ptr >= current_program_break)
+    if (!internal::isPointerInHeap(ptr))
     {
         return;
     }
 
     internal::MemoryBlock* block_to_free = internal::getMetadata(ptr);
     if (block_to_free == nullptr || block_to_free->allocated_ == false)
     {
+        if (block_to_free)
+        {
+            exit(-1);
+        }
         return;
     }
+
+    internal::total_memory_allocated -= block_to_free->size_;
     block_to_free->allocated_ = false;
 
     internal::mergeFreeMemoryBlocks();
@@ -62,15 +66,22 @@ void decreaseHeap(MemoryBlock* block_heap_end)
         return;
     }
 
-    size_t release_size = block_heap_end->size_ + BLOCK_METADATA_SIZE;
-    void* program_break = sbrk(-static_cast<intptr_t>(release_size));
+    void* current_break = sbrk(0);
+    void* block_start = reinterpret_cast<void*>(block_heap_end);
+    intptr_t release_size =
+        reinterpret_cast<char*>(current_break) - reinterpret_cast<char*>(block_start);
+
+    if (release_size <= 0)
+    {
+        return;
+    }
+
+    void* program_break = sbrk(-release_size);
 
     if (program_break == getErrorCodeInVoidPtr(-1) || program_break == getErrorCodeInVoidPtr(0))
     {
         exit(-1);
     }
-
-    total_memory_allocated -= release_size;
 }
 
 void mergeFreeMemoryBlocks()

src/malloc.cpp

@@ -9,13 +9,16 @@ namespace mem
 
 void* malloc(size_t size)
 {
+    if (size == 0)
+    {
+        return nullptr;
+    }
+
     if (internal::heap_start == nullptr)
     {
         internal::heap_start = sbrk(0);
     }
 
-    internal::total_memory_allocated += size;
-
     internal::MemoryBlock* found_block =
         internal::findLargeEnoughFreeMemoryBlock(&internal::block_list_head, size);
 
@@ -60,6 +63,7 @@ void* malloc(size_t size)
         new_block->allocated_ = true;
         new_block->next_ = nullptr;
         internal::insertMemoryBlockAtEnd(&internal::block_list_head, new_block);
+        internal::total_memory_allocated += size;
     }
 
     return internal::getPayloadAddress(new_block);
@@ -73,7 +77,7 @@ void* increaseHeap(size_t size)
     void* result = sbrk(static_cast<intptr_t>(size));
     if (result == getErrorCodeInVoidPtr(-1) || result == getErrorCodeInVoidPtr(0))
     {
-        exit(-1);
+        return nullptr;
     }
 
     return result;
@@ -116,15 +120,17 @@ void* splitFreeMemoryBlockIfPossible(MemoryBlock* new_block, size_t size)
     }
     else
     {
+        new_block->size_ = new_block->size_;
         new_block->allocated_ = true;
     }
 
+    total_memory_allocated += new_block->size_;
     return getPayloadAddress(new_block);
 }
 
 void* getMemoryBlockSplitAddress(MemoryBlock* new_block, size_t size)
 {
-    return (reinterpret_cast<char*>(getPayloadAddress(new_block)) + size);
+    return (reinterpret_cast<char*>(new_block) + BLOCK_METADATA_SIZE + size);
 }
 
 void insertMemoryBlockAtEnd(MemoryBlock** block_list_head, MemoryBlock* new_block)
@@ -171,6 +177,17 @@ size_t getSizeOfAllocatedMemoryBlock(MemoryBlock* block)
     return reinterpret_cast<size_t>(getPayloadAddress(block));
 }
 
+bool isPointerInHeap(void* ptr)
+{
+    void* current_program_break = sbrk(0);
+    if (ptr < heap_start || ptr >= current_program_break)
+    {
+        return false;
+    }
+
+    return true;
+}
+
 void* getErrorCodeInVoidPtr(size_t error_code) { return reinterpret_cast<void*>(error_code); }
 
 } // namespace internal