4 files changed, 45 insertions, 9 deletions

diff --git a/Source/Kernel/Melon.ke b/Source/Kernel/Melon.ke
index c8b784a..71955a6 100755
--- a/Source/Kernel/Melon.ke
+++ b/Source/Kernel/Melon.ke
diff --git a/Source/Kernel/MemoryManager/Heap.class.cpp b/Source/Kernel/MemoryManager/Heap.class.cpp
index f5bbe0a..7331cd0 100644
--- a/Source/Kernel/MemoryManager/Heap.class.cpp
+++ b/Source/Kernel/MemoryManager/Heap.class.cpp
@@ -1,7 +1,7 @@
 #include "Heap.class.h"
 #include <MemoryManager/PageDirectory.class.h>
 
-Heap::Heap() {
+Heap::Heap() : m_mutex(MUTEX_FALSE) {
 	m_usable = false;
 	m_index.data = 0;
 	m_index.size = 0;
@@ -12,6 +12,8 @@ Heap::~Heap() {
 }
 
 void Heap::create(u32int start, u32int size, u32int idxsize, PageDirectory* pagedir, bool user, bool rw) {
+	if (m_usable) return;
+
 	if (start & 0x0FFF) start = (start & 0xFFFFF000) + 0x1000;
 	if (size & 0x0FFF) size = (size & 0xFFFFF000) + 0x1000;
 	m_start = start + idxsize;	//m_start is start of real data, start is start of index.
@@ -42,6 +44,8 @@ void Heap::create(u32int start, u32int size, u32int idxsize, PageDirectory* page
 
 	m_usable = true;
 	m_free = (m_end - m_start);
+
+	m_mutex.unlock();
 }
 
 void Heap::expand(u32int quantity) {
@@ -113,6 +117,8 @@ void Heap::contract() {	//Automatically work out how much we can contract
 }
 
 void *Heap::alloc(u32int sz, bool no_expand) {
+	m_mutex.waitLock();
+
 	u32int newsize = sz + sizeof(heap_header_t) + sizeof(heap_footer_t);
 	u32int iterator = 0;
 	while (iterator < m_index.size) {
@@ -120,8 +126,12 @@ void *Heap::alloc(u32int sz, bool no_expand) {
 		iterator++;
 	}
 	if (iterator == m_index.size) { //No hole is big enough
-		if (no_expand) return 0;
+		if (no_expand) {
+			m_mutex.unlock();
+			return 0;
+		}
 		expand((sz & 0xFFFFF000) + 0x1000);
+		m_mutex.unlock();
 		return alloc(sz, true);	//Recurse call
 	}
 
@@ -152,6 +162,8 @@ void *Heap::alloc(u32int sz, bool no_expand) {
 	
 	m_free -= loc->size;
 
+	m_mutex.unlock();
+
 	return (void*)((u32int)loc + sizeof(heap_header_t));
 }
 
@@ -162,6 +174,8 @@ void Heap::free(void *ptr) {
 	heap_footer_t *footer = (heap_footer_t*)((u32int)header + header->size - sizeof(heap_footer_t));
 	if (header->magic != HEAP_MAGIC or footer->magic != HEAP_MAGIC) return;
 
+	m_mutex.waitLock();
+
 	m_free += header->size;
 
 	//Unify left
@@ -192,4 +206,6 @@ void Heap::free(void *ptr) {
 			header->size >= 0x2000 and (m_end - m_start > HEAP_MIN_SIZE)) {
 		contract();
 	}
+
+	m_mutex.unlock();
 }
diff --git a/Source/Kernel/MemoryManager/Heap.class.h b/Source/Kernel/MemoryManager/Heap.class.h
index a1feaa1..930a589 100644
--- a/Source/Kernel/MemoryManager/Heap.class.h
+++ b/Source/Kernel/MemoryManager/Heap.class.h
@@ -2,6 +2,7 @@
 #define DEF_HEAP_CLASS_H
 
 #include <Core/common.wtf.h>
+#include <TaskManager/Mutex.class.h>
 
 //Heap minimum size : 2M
 #define HEAP_MIN_SIZE 0x00200000
@@ -33,6 +34,8 @@ class Heap {
 	heap_index_t m_index;
 	PageDirectory* m_pagedir;
 
+	Mutex m_mutex;
+
 	void insertIntoIndex(heap_header_t *e);
 	u32int findIndexEntry(heap_header_t *e);
 	void removeFromIndex(u32int idx);
@@ -41,20 +44,35 @@ class Heap {
 	void expand(u32int quantity);
 	void contract();	//Quantity is automatically calculated
 	
-
 	public:
 	Heap();
 	~Heap();
 
-	bool usable() { return m_usable; }
-
 	void create(u32int start, u32int size, u32int idxsize, PageDirectory* pagedir, bool user, bool rw);
 
 	void* alloc(u32int sz, bool no_expand = false);
 	void free(void* ptr);
 
-	u32int size() { return m_end - m_start; }
-	u32int free() { return m_free; }
+	bool usable() { 
+		m_mutex.waitLock();
+		bool ret = m_usable;
+		m_mutex.unlock();
+		return ret;
+	}
+
+	u32int size() {
+		m_mutex.waitLock();
+		u32int ret = m_end - m_start;
+		m_mutex.unlock();
+		return ret;
+	}
+
+	u32int free() {
+		m_mutex.waitLock();
+		u32int ret = m_free;
+		m_mutex.unlock();
+	   	return ret;
+	}
 };
 
 
diff --git a/Source/Kernel/TaskManager/Mutex.class.cpp b/Source/Kernel/TaskManager/Mutex.class.cpp
index a6c3ba8..8ba274f 100644
--- a/Source/Kernel/TaskManager/Mutex.class.cpp
+++ b/Source/Kernel/TaskManager/Mutex.class.cpp
@@ -13,8 +13,10 @@ bool Mutex::lock() {
 }
 
 void Mutex::waitLock() {
-	while (atomic_exchange(&m_locked, MUTEX_TRUE) == MUTEX_TRUE)
-		Task::currThread()->sleep(10);	//Wait 10ms
+	while (atomic_exchange(&m_locked, MUTEX_TRUE) == MUTEX_TRUE) {
+		if (Task::currThread() != 0) Task::currThread()->sleep(10);	//Wait 10ms
+		else return;
+	}
 }
 
 void Mutex::unlock() {