Simplify the entry address computation

We actually don't need to preserve the initially computed base address
for each binary, but we can just preserve the final entry address.

Signed-off-by: Miquel Sabaté Solà <mikisabate@gmail.com>

3 files changed, 25 insertions, 22 deletions

diff --git a/include/fbos/sched.h b/include/fbos/sched.h
index e05d0a1..3d35f9f 100644
--- a/include/fbos/sched.h
+++ b/include/fbos/sched.h
@@ -12,18 +12,30 @@ enum task_id {
 	TASK_FIZZBUZZ = 3,
 };
 
-// TODO: if we only care about the absolute address, it can be further
-// simplified.
+// All the information we need to grab for processes.
 struct task_struct {
+	// The stack allocated for the process. As you can see when initializing
+	// each process on `kernel/main.c`, we go over the top for its size. There
+	// is also the fact that we need to keep this as the first attribue to allow
+	// for simple `sp` values.
 	void *stack;
-	const void *addr;
-	uint64_t entry_offset;
+
+	// The address for the binary entry.
+	const void *entry_addr;
 };
 
 // Tasks available on this kernel.
 extern struct task_struct tasks[4];
 
-// Switch execution to the given task id.
-void switch_to(int task_id);
+// Switch execution to the given U-mode task. Note that this function will not
+// do the actual returning, but it prepares the relevant registers for an
+// eventual jump.
+__kernel __always_inline void switch_to(int task_id)
+{
+	asm volatile("csrc sstatus, %0\n\t"
+				 "mv ra, %1\n\t"
+				 "csrw sepc, ra" ::"r"(1 << 8),
+				 "r"(tasks[task_id].entry_addr));
+}
 
 #endif // __FBOS_SCHED_H_
diff --git a/kernel/initrd.c b/kernel/initrd.c
index ff1f52b..27d0ece 100644
--- a/kernel/initrd.c
+++ b/kernel/initrd.c
@@ -101,8 +101,9 @@ __kernel void extract_elf(int task_id, const unsigned char *const addr, size_t s
 	/* }; */
 
 #ifdef __KERNEL__
-	tasks[task_id].addr = (const void *)addr;
-	tasks[task_id].entry_offset = (uint64_t)addr[0x18];
+	uint64_t offset = (uint64_t)addr[0x18];
+
+	tasks[task_id].entry_addr = (const void *)(addr + offset);
 #endif
 }
 
diff --git a/kernel/main.c b/kernel/main.c
index ba4ada0..58d67f5 100644
--- a/kernel/main.c
+++ b/kernel/main.c
@@ -7,22 +7,12 @@
 unsigned long init_stack[4][THREAD_SIZE / sizeof(unsigned long)];
 
 struct task_struct tasks[4] = {
-	[TASK_INIT] = { .stack = init_stack[0], .addr = nullptr, .entry_offset = 0, },
-	[TASK_FIZZ] = { .stack = init_stack[1], .addr = nullptr, .entry_offset = 0, },
-	[TASK_BUZZ] = { .stack = init_stack[2], .addr = nullptr, .entry_offset = 0, },
-	[TASK_FIZZBUZZ] = { .stack = init_stack[3], .addr = nullptr, .entry_offset = 0, },
+	[TASK_INIT] = { .stack = init_stack[0], .entry_addr = nullptr, },
+	[TASK_FIZZ] = { .stack = init_stack[1], .entry_addr = nullptr, },
+	[TASK_BUZZ] = { .stack = init_stack[2], .entry_addr = nullptr, },
+	[TASK_FIZZBUZZ] = { .stack = init_stack[3], .entry_addr = nullptr, },
 };
 
-// TODO
-__kernel void switch_to(int task_id)
-{
-	const char *ddr = (const char *)tasks[task_id].addr + tasks[task_id].entry_offset;
-
-	asm volatile("csrc sstatus, %[mask]" : : [mask] "r"(1 << 8));
-	asm volatile("mv ra, %0" : : "r"(ddr));
-	asm volatile("csrw sepc, ra");
-}
-
 /*
  * This is the main entry point of the kernel after head.S is done. This
  * function can (and will) assume that everything has been reset and that we can