3 files changed, 76 insertions, 20 deletions

diff --git a/dev/kernel/src/FS/HeFS.cc b/dev/kernel/src/FS/HeFS.cc
index cfa8baae..7864bc6f 100644
--- a/dev/kernel/src/FS/HeFS.cc
+++ b/dev/kernel/src/FS/HeFS.cc
@@ -4,7 +4,6 @@
 
 ------------------------------------------- */
 
-#include "NewKit/Defines.h"
 #ifdef __FSKIT_INCLUDES_HEFS__
 
 #include <modules/AHCI/AHCI.h>
diff --git a/dev/kernel/src/MemoryMgr.cc b/dev/kernel/src/MemoryMgr.cc
index 8020ca1c..eff54435 100644
--- a/dev/kernel/src/MemoryMgr.cc
+++ b/dev/kernel/src/MemoryMgr.cc
@@ -108,7 +108,7 @@ namespace Kernel
 		if (!ptr_heap || new_sz < 1)
 			return nullptr;
 
-		kout << "This function is not implemented by the kernel itself.\r";
+		kout << "This function is not implemented by the kernel yet.\r";
 
 		ke_panic(RUNTIME_CHECK_INVALID);
 
@@ -129,8 +129,8 @@ namespace Kernel
 
 		sz_fix += sizeof(Detail::MM_INFORMATION_BLOCK);
 
-		PageMgr heap_mgr;
-		auto	wrapper = heap_mgr.Request(wr, user, No, sz_fix, pad_amount);
+		PageMgr page_mgr;
+		auto	wrapper = page_mgr.Request(wr, user, No, sz_fix, pad_amount);
 
 		Detail::MM_INFORMATION_BLOCK_PTR heap_info_ptr =
 			reinterpret_cast<Detail::MM_INFORMATION_BLOCK_PTR>(
@@ -244,8 +244,8 @@ namespace Kernel
 			PTEWrapper		page_wrapper(No, No, No, reinterpret_cast<UIntPtr>(heap_info_ptr) - sizeof(Detail::MM_INFORMATION_BLOCK));
 			Ref<PTEWrapper> pte_address{page_wrapper};
 
-			PageMgr heap_mgr;
-			heap_mgr.Free(pte_address);
+			PageMgr page_mgr;
+			page_mgr.Free(pte_address);
 
 			return kErrorSuccess;
 		}
diff --git a/dev/kernel/src/UserProcessScheduler.cc b/dev/kernel/src/UserProcessScheduler.cc
index cdf7d91f..660dd544 100644
--- a/dev/kernel/src/UserProcessScheduler.cc
+++ b/dev/kernel/src/UserProcessScheduler.cc
@@ -105,11 +105,33 @@ namespace Kernel
 	}
 
 	/***********************************************************************************/
-	/** @brief Allocate pointer to track list. */
+	/** @brief Allocate pointer to heap tree. */
+	/***********************************************************************************/
+
+	USER_PROCESS::USER_HEAP_TREE* sched_try_go_upper_heap_tree(USER_PROCESS::USER_HEAP_TREE* tree)
+	{
+		if (tree)
+		{
+			tree = tree->MemoryNext;
+
+			if (!tree)
+			{
+				return nullptr;
+			}
+		}
+
+		return tree;
+	}
+
+	/***********************************************************************************/
+	/** @brief Allocate pointer to heap tree. */
 	/***********************************************************************************/
 
 	ErrorOr<VoidPtr> USER_PROCESS::New(SizeT sz, SizeT pad_amount)
 	{
+		if (this->UsedMemory > kSchedMaxMemoryLimit)
+			return ErrorOr<VoidPtr>(-kErrorHeapOutOfMemory);
+
 #ifdef __NE_VIRTUAL_MEMORY_SUPPORT__
 		auto vm_register = hal_read_cr3();
 		hal_write_cr3(this->VMRegister);
@@ -125,33 +147,68 @@ namespace Kernel
 		{
 			this->HeapTree = new USER_HEAP_TREE();
 
-			this->HeapTree->MemoryEntryPad	 = pad_amount;
+			this->HeapTree->MemoryColor = USER_HEAP_TREE::kBlackMemory;
+
+			this->HeapTree->MemoryEntryPad	= pad_amount;
 			this->HeapTree->MemoryEntrySize = sz;
 
 			this->HeapTree->MemoryEntry = ptr;
 
-			this->HeapTree->MemoryPrev = nullptr;
-			this->HeapTree->MemoryNext = nullptr;
+			this->HeapTree->MemoryPrev	 = nullptr;
+			this->HeapTree->MemoryNext	 = nullptr;
+			this->HeapTree->MemoryParent = nullptr;
+			this->HeapTree->MemoryChild	 = nullptr;
 		}
 		else
 		{
-			USER_HEAP_TREE* entry = this->HeapTree;
+			USER_HEAP_TREE* entry	   = this->HeapTree;
+			USER_HEAP_TREE* prev_entry = entry;
+
+			BOOL is_parent = NO;
 
 			while (entry)
 			{
-				if (entry->MemoryEntry == nullptr)
-					break; // chose to break here, when we get an already allocated memory entry for our needs.
+				if (entry->MemoryEntrySize < 1)
+					break;
 
-				entry = entry->MemoryNext;
+				prev_entry = entry;
+
+				if (entry->MemoryNext)
+				{
+					entry = entry->MemoryNext;
+				}
+				else if (entry->MemoryChild)
+				{
+					entry	  = entry->MemoryChild;
+					is_parent = YES;
+				}
+				else
+				{
+					entry = entry->MemoryParent;
+					entry = sched_try_go_upper_heap_tree(entry);
+				}
 			}
 
-			entry->MemoryNext			   = new USER_HEAP_TREE();
-			entry->MemoryNext->MemoryEntry = ptr;
-			entry->MemoryEntrySize = sz;
-			entry->MemoryEntryPad = pad_amount;
+			if (!entry)
+				entry			   = new USER_HEAP_TREE();
 
-			entry->MemoryNext->MemoryPrev = entry;
-			entry->MemoryNext->MemoryNext = nullptr;
+			entry->MemoryEntry = ptr;
+			entry->MemoryEntrySize		   = sz;
+			entry->MemoryEntryPad		   = pad_amount;
+
+			if (is_parent)
+			{
+				entry->MemoryParent = prev_entry;
+				prev_entry->MemoryChild				= entry;
+
+				prev_entry->MemoryColor = USER_HEAP_TREE::kBlackMemory;
+				entry->MemoryColor = USER_HEAP_TREE::kRedMemory;
+			}
+			else
+			{
+				prev_entry->MemoryNext = entry;
+				entry->MemoryPrev = prev_entry;
+			}
 		}
 
 		this->UsedMemory += sz;