Nucleus

diff --git a/task/scheduler.c b/task/scheduler.c
index d0fb83f..f3bf5b8 100644
--- a/task/scheduler.c
+++ b/task/scheduler.c
@@ -14,8 +14,8 @@
 
 static void scheduler_new(Object *);
 static void scheduler_delete(Object *);
-void context_switch(uintptr_t eip, page_dir_t pagedir,
-                    uintptr_t esi, uintptr_t edi, uintptr_t ebx,
+void context_switch(uintptr_t eip, page_dir_t pageDir, uintptr_t ebx,
+                    uintptr_t esi, uintptr_t edi,
                     uintptr_t ebp, uintptr_t esp);
 
 static _Atomic size_t tasks = 0;
@@ -54,6 +54,8 @@ 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);
@@ -64,28 +66,30 @@ 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 */
 	current = task; /* Given reference, so no get() */
-	context_switch(current->eip, current->pageDir, current->esi,
-	               current->edi, current->ebx, current->ebp, current->esp);
+	context_switch(current->eip, pageDir, 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");
 }
 
 /* Find the next schedulable ready queue */
-static ObjectList *
+static enum Priority
 highest_priority_queue(Scheduler *s)
 {
 	enum Priority p;
 	for (p = PRIORITY_COUNT - 1; p > 0; p--) {
-		if (count(s->queue[p]))
-			return s->queue[p];
+		if (count(s->queue[p ]))
+			return p;
 	}
-	return NULL;
+	return 0;
 }
 
 /* Schedule the next task */
@@ -95,21 +99,22 @@ schedule(void)
 	enter_critical_section();
 	Task *task = current;
 	Scheduler *s = cpu->scheduler;
-	ObjectList *queue = highest_priority_queue(s);
+	enum Priority highest = highest_priority_queue(s);
 
-	/* No schedulable tasks */
-	if (!queue) {
-		if (current && current->state == RUNNING)
-			goto end;
+	/* Continue current task */
+	if (current && current->state == RUNNING
+	 && (!highest || current->priority > highest))
+		goto end;
 
-		/* Idle */
+	/* Idle */
+	if (!highest) {
 		current = NULL;
 		if (task) {
 			tasks--;
 			s->tasks--;
 		}
 		asm volatile("sti");
-		while (!(queue = highest_priority_queue(s)))
+		while (!(highest = highest_priority_queue(s)))
 			asm volatile("hlt");
 		asm volatile("cli");
 		if (task) {
@@ -120,7 +125,7 @@ schedule(void)
 	}
 
 	/* Schedule next task */
-	task = pop_from_start(queue);
+	task = pop_from_start(s->queue[highest]);
 	task->state = RUNNING;
 	if (task == current)
 		goto end;
@@ -187,7 +192,7 @@ balance_scheduler(void)
 	enter_critical_section();
 	while (s->tasks > max) {
 		s->tasks--;
-		t = pop_from_start(highest_priority_queue(s));
+		t = pop_from_start(s->queue[highest_priority_queue(s)]);
 		t->scheduler = NULL;
 		enqueue_task(t);
 	}
diff --git a/task/switch.S b/task/switch.S
index 54769af..07f1e56 100644
--- a/task/switch.S
+++ b/task/switch.S
@@ -14,11 +14,14 @@ context_switch:
 	cli
 	mov 4(%esp), %ecx
 	mov 8(%esp), %eax
-	mov 12(%esp), %esi
-	mov 16(%esp), %edi
-	mov 20(%esp), %ebx
+	mov 12(%esp), %ebx
+	mov 16(%esp), %esi
+	mov 20(%esp), %edi
 	mov 24(%esp), %ebp
 	mov 28(%esp), %esp
+	test %eax, %eax
+	je 1f
 	mov %eax, %cr3
+1:
 	sti
 	jmp *%ecx