diff --git a/kernel/include/thread/Task.h b/kernel/include/thread/Task.h
index e445b298..1c3d1395 100644
--- a/kernel/include/thread/Task.h
+++ b/kernel/include/thread/Task.h
@@ -56,4 +56,8 @@ struct Task
 
     void save_floating();
     void restore_floating();
+
+    void switch_to_address_space();
+
+    bool has_died();
 };
\ No newline at end of file
diff --git a/kernel/src/sys/exec.cpp b/kernel/src/sys/exec.cpp
index 74cf52a4..a0922cc2 100644
--- a/kernel/src/sys/exec.cpp
+++ b/kernel/src/sys/exec.cpp
@@ -99,7 +99,7 @@ void sys_exec(Context* context, const char* pathname)
 
     Scheduler::reset_task(task, image);
 
-    set_context_from_task(*task, context);
+    task->restore_context(context);
 
     kfree(kpathname);
 
diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
index 0896e5ce..c9d66fdb 100644
--- a/kernel/src/thread/Scheduler.cpp
+++ b/kernel/src/thread/Scheduler.cpp
@@ -320,7 +320,7 @@ void Scheduler::task_yield(Context* context)
 {
     ASSERT(Interrupts::is_in_handler());
     Interrupts::disable();
-    get_context_to_task(*sched_current_task, context);
+    sched_current_task->save_context(context);
     bool was_idle = false;
     if (sched_current_task->state == sched_current_task->Idle)
     {
@@ -334,15 +334,14 @@ void Scheduler::task_yield(Context* context)
         {
             if (sched_current_task->id != original_task->id || was_idle)
             {
-                if (!was_idle && original_task->is_user_task() && original_task->state != original_task->Exited)
+                if (!was_idle && original_task->is_user_task() && !original_task->has_died())
                 {
-                    task_save_floating(*original_task);
+                    original_task->save_floating();
                 }
                 if (sched_current_task->is_user_task())
                 {
-                    VMM::switch_to_user_address_space(sched_current_task->address_space);
-                    VMM::apply_address_space();
-                    task_restore_floating(*sched_current_task);
+                    sched_current_task->switch_to_address_space();
+                    sched_current_task->restore_floating();
                 }
                 else if (!was_idle && original_task->is_user_task() && !sched_current_task->is_user_task())
                 {
@@ -351,17 +350,17 @@ void Scheduler::task_yield(Context* context)
                 }
             }
             sched_current_task->task_time = 20;
-            set_context_from_task(*sched_current_task, context);
+            sched_current_task->restore_context(context);
             return;
         }
     } while (sched_current_task != original_task);
     if (!was_idle && original_task->is_user_task() && original_task->state != original_task->Exited)
     {
-        task_save_floating(*original_task);
+        original_task->save_floating();
     }
     sched_current_task = &idle_task;
     sched_current_task->task_time = frequency;
-    if (!was_idle) { set_context_from_task(*sched_current_task, context); }
+    if (!was_idle) { sched_current_task->restore_context(context); }
     return;
 }
 
diff --git a/kernel/src/thread/Task.cpp b/kernel/src/thread/Task.cpp
index c470d216..0010e7d8 100644
--- a/kernel/src/thread/Task.cpp
+++ b/kernel/src/thread/Task.cpp
@@ -1,4 +1,5 @@
 #include "thread/Task.h"
+#include "memory/VMM.h"
 #include "std/string.h"
 
 void Task::restore_context(Context* context)
@@ -53,4 +54,15 @@ int Task::alloc_fd_greater_than_or_equal(int base_fd)
     if (fd == TASK_MAX_FDS) { return -1; }
 
     return fd;
+}
+
+void Task::switch_to_address_space()
+{
+    VMM::switch_to_user_address_space(address_space);
+    VMM::apply_address_space();
+}
+
+bool Task::has_died()
+{
+    return state == Exited;
 }
\ No newline at end of file