--- a/cos/kernel/malloc.c	Mon Jan 16 21:38:55 2012 +0100
+++ b/cos/kernel/malloc.c	Tue Jan 24 17:54:16 2012 +0100
@@ -5,104 +5,153 @@
 
 uint64_t placement_address = (uint64_t)&kernel_end;
 
-void* kmalloc_int(uint64_t size)
-{
-   uint64_t tmp = placement_address;
-   placement_address += size;
-   return (void*)tmp;
-}
-
-void* kmalloc_aligned_int(uint64_t size)
-{
-   if ( (placement_address | 0xFFF) != 0 )
-   {
-      placement_address &= ~(0xFFF);
-      placement_address += 0x1000;
-   }
-   uint64_t tmp = placement_address;
-   placement_address += size;
-   return (void*)tmp;
-}
-
-
 // ================= Other malloc
 // TODO: move this to user space?
 
 #define HEAP_MAGIC 0xc0ffee
-#define HEAP_START 0x400000
-#define HEAP_SIZE  0x200000
 #define HEAP_INUSE 1
 #define HEAP_FREE 0
 
-typedef struct {
-   uint64_t magic;
-   uint64_t state;
-   uint64_t size;
-} heap_t;
-
 /*
   malloc and free divide the chunks of memory present at the heap
   of the kernel into smaller parts.
   The heap is located at: 0x
 */
 
-static heap_t* kernel_heap = (heap_t*) 0x400000; // 4 MB - 6 MB is heap
+static heap_t* kernel_heap = 0; 
+
+void expand_heap(heap_t *heap, uint64_t newsize);
+
 /* Allocates 'size' bytes and returns the pointer if succesfull.
+   Before the kernel is initialized, placement malloc is used.
    Kernelpanic in case of failure..
 */
-
-void* kmalloc(uint64_t size) 
+static void* kmalloc_int(uint64_t size, uint64_t aligned) 
 {
-   return kmalloc_int(size);
-
-   // Start at the beginning of our heap and search a free block:
-   heap_t *current = kernel_heap;
-   while (current->magic == HEAP_MAGIC)
+   // Check if there exists a kernel heap. If not, use placement malloc for now.
+   if (kernel_heap == 0)
    {
-      if ((current->state == HEAP_FREE) && (current->size >= size))
+      // Before that the heap is initialized, use placement malloc
+      if (aligned == 1)
       {
-         // Mark block as used:
-         current->state = HEAP_INUSE;
+         // Align placement address
+         if ( (placement_address | 0xFFF) != 0 )
+         {
+            placement_address &= ~(0xFFF);
+            placement_address += 0x1000;
+         }
+      }
+      uint64_t tmp = placement_address;
+      placement_address += size;
+      return (void*)tmp;
+   }
+   else
+   {
+      // We have a kernel heap, search the heap for a block of suitable size.
 
-         // Insert a heap header if required:
-         if (current->size > size + sizeof(heap_t) + 1)
+      // Start at the beginning of our heap and search a free block:
+      heap_header_t *current = kernel_heap->first_block;
+      while (current->magic == HEAP_MAGIC)
+      {
+         if ((current->state == HEAP_FREE) && (current->size >= size))
          {
-            // Calculate location of the inserted header:
-            heap_t *newheader = (heap_t*) (((char*)current)+size+sizeof(heap_t));
+            // We found a hole with suitable size. Determine if split is of any use.
+
+            // Insert a heap header if required:
+            if (current->size > size + sizeof(heap_t) + 1)
+            {
+               // Calculate location of the inserted header:
+               heap_header_t *newheader = (heap_header_t*) (((uint64_t)current)+size+sizeof(heap_header_t));
 
-            // Set the new header fields:
-            newheader->size = current->size - size - sizeof(heap_t);
-            newheader->state = HEAP_FREE;
-            newheader->magic = HEAP_MAGIC;
+               // Set the new header fields:
+               newheader->size = current->size - size - sizeof(heap_header_t);
+               newheader->state = HEAP_FREE;
+               newheader->magic = HEAP_MAGIC;
 
-            // Set the size of this block
-            current->size = size; 
+               // Set the size of this block
+               current->size = size; 
+            }
+            else
+            {
+               // We allocate this whole block
+               // Mark block as used:
+               current->state = HEAP_INUSE;
+            }
+            // Calculate the size of the block:
+            char *address = ((char*)current)+sizeof(heap_header_t);
+            return address;
+            
          }
-         else
-         {
-            // We allocate this whole block
-         }
-         // Calculate the size of the block:
-         char *address = ((char*)current)+sizeof(heap_t);
-         return address;
-         
+         // Goto next heap block:
+         current = (heap_header_t*)(((uint64_t) current) + current->size + sizeof(heap_header_t));
       }
-      // Goto next heap block:
-      current = (heap_t*)(((char*) current) + current->size + sizeof(heap_t));
+
+      // We did not find a block large enough. Expand the heap and try again.
+      expand_heap(kernel_heap, 0x10000000);
+      return kmalloc_int(size, aligned); 
    }
-   return 0x0;
+}
+
+// Convenient wrappers:
+void* kmalloc(uint64_t size)
+{
+   return kmalloc_int(size, 0);
+}
+
+void* kmalloc_a(uint64_t size)
+{
+   return kmalloc_int(size, 1);
 }
 
 void kfree(void* ptr) 
 {
-  printf("Free address %x\n", ptr);
+   printf("Free address %x\n", ptr);
+   if (kernel_heap == 0)
+   {
+      return;
+   }
+   else
+   {
+      // TODO
+   }
+   // TODO: free blocks:
 }
 
-void init_heap(void)
+heap_t* create_heap(uint64_t location, uint64_t size)
 {
-   // Initialize the kernel heap:
-   kernel_heap->magic = HEAP_MAGIC;
-   kernel_heap->state = HEAP_FREE;
-   kernel_heap->size = HEAP_SIZE - sizeof(heap_t);
+   // Make sure that the new heap location is mapped into the address space:
+   uint64_t i = 0;
+   while (i < size)
+   {
+      alloc_frame( get_page(location + i, kernel_map));
+      i += 0x1000;
+   }
+
+   // create new heap structure:
+   heap_t* h = (heap_t*)kmalloc(sizeof(heap_t));
+   
+   // Create one big hole:
+   h->first_block = (heap_header_t*)location; // Place header at beginning of heap
+   h->first_block->magic = HEAP_MAGIC;
+   h->first_block->state = HEAP_FREE;
+   h->first_block->size = size - sizeof(heap_header_t);
+   return h;
 }
 
+// Make the total heap bigger:
+void expand_heap(heap_t *heap, uint64_t newsize)
+{
+   uint64_t oldsize = heap->end_address - heap->start_address;
+   if (newsize < oldsize)
+   {
+      panic("Cannot expand heap to a smaller size\n");
+   }
+
+   uint64_t i = oldsize;
+   while (i < newsize)
+   {
+      alloc_frame( get_page(heap->start_address + i, kernel_map));
+      i += 0x1000;
+   }
+}
+