diff --git a/runtime/include/arch/aarch64/context.h b/runtime/include/arch/aarch64/context.h
index 839483a..b665962 100644
--- a/runtime/include/arch/aarch64/context.h
+++ b/runtime/include/arch/aarch64/context.h
@@ -43,8 +43,8 @@ arch_context_restore(mcontext_t *active_context, struct arch_context *sandbox_co
 	/* Assumption: Base Context is only ever used by arch_context_switch */
 	assert(sandbox_context != &worker_thread_base_context);
 
-	assert(sandbox_context->sp);
-	assert(sandbox_context->pc);
+	assert(sandbox_context->regs[ureg_rsp]);
+	assert(sandbox_context->regs[ureg_rip]);
 
 	/* Transitioning from Fast -> Running */
 	assert(sandbox_context->variant == arch_context_variant_fast);
@@ -97,23 +97,28 @@ arch_context_switch(struct arch_context *a, struct arch_context *b)
 	             "adr x1, reset%=\n\t"
 	             "str x1, [%[a], 8]\n\t"
 	             "str x0, [%[a]]\n\t"
-	             "ldr x2, [%[b]]\n\t"
-	             "cbz x2, slow%=\n\t"
-	             "ldr x3, [%[b], 8]\n\t"
-	             "mov sp, x2\n\t"
-	             "br x3\n\t"
+		     "mov x0, #1\n\t"
+		     "str x0, [%[av]]\n\t"
+	             "ldr x1, [%[bv]]\n\t"
+		     "sub x1, x1, #2\n\t"
+	             "cbz x1, slow%=\n\t"
+		     "ldr x0, [%[b]]\n\t"
+	             "ldr x1, [%[b], 8]\n\t"
+	             "mov sp, x0\n\t"
+	             "br x1\n\t"
 	             "slow%=:\n\t"
 	             "br %[slowpath]\n\t"
 	             ".align 8\n\t"
 	             "reset%=:\n\t"
-	             "mov x1, #0\n\t"
-	             "str x1, [%[b]]\n\t"
+	             "mov x1, #3\n\t"
+	             "str x1, [%[bv]]\n\t"
 	             ".align 8\n\t"
 	             "exit%=:\n\t"
 	             :
-	             : [ a ] "r"(a_registers), [ b ] "r"(b_registers), [ slowpath ] "r"(&arch_context_restore_preempted)
+	             : [ a ] "r"(a_registers), [ b ] "r"(b_registers), [ av ] "r"(&a->variant), [ bv ] "r"(&b->variant), 
+		       [ slowpath ] "r"(&arch_context_restore_preempted)
 	             : "memory", "cc", "x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7", "x8", "x9", "x10", "x11", "x12",
-	               "x13", "x14", "x15", "x16", "x17", "x18", "x19", "x20", "x21", "x22", "x23", "x24", "x25", "x26",
+	               "x13", "x14", "x15", "x16", "x17", "x18", "x19", "x20", "x21", "x22", "x23", "x24",
 	               "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15");
 
 	return 0;