Nucleus

diff --git a/task/scheduler.c b/task/scheduler.c
index f3bf5b8..c927eb3 100644
--- a/task/scheduler.c
+++ b/task/scheduler.c
@@ -14,7 +14,7 @@
 
 static void scheduler_new(Object *);
 static void scheduler_delete(Object *);
-void context_switch(uintptr_t eip, page_dir_t pageDir, uintptr_t ebx,
+void context_switch(uintptr_t eip, uintptr_t ebx,
                     uintptr_t esi, uintptr_t edi,
                     uintptr_t ebp, uintptr_t esp);
 
@@ -54,8 +54,6 @@ scheduler_delete(Object *obj)
 static void
 switch_to_task(Task *task)
 {
-	page_dir_t pageDir = task->pageDir;
-
 	/* Save current task state */
 	if (__builtin_expect(!!current, 1)) {
 		lock(current);
@@ -66,15 +64,16 @@ switch_to_task(Task *task)
 		asm volatile("mov %%ebp, %0" : "=r" (current->ebp));
 		current->eip = (uintptr_t) &&end;
 		unlock(current);
-		if (pageDir == current->pageDir)
-			pageDir = 0;
 		put(current);
 	}
 
 	/* Switch to new context */
+	switch_to_mm(task->vm);
 	current = task; /* Given reference, so no get() */
-	context_switch(current->eip, pageDir, current->ebx,
-	               current->esi, current->edi, current->ebp, current->esp);
+	set_kernel_stack((uintptr_t) current + KERNEL_STACK_SIZE);
+	context_switch(current->eip, current->ebx,
+	               current->esi, current->edi,
+	               current->ebp, current->esp);
 end:
 	/* This prevents GCC from optimising the jump to be after the return */
 	asm volatile("":::"memory");
@@ -86,7 +85,7 @@ highest_priority_queue(Scheduler *s)
 {
 	enum Priority p;
 	for (p = PRIORITY_COUNT - 1; p > 0; p--) {
-		if (count(s->queue[p ]))
+		if (count(s->queue[p]))
 			return p;
 	}
 	return 0;