beautification

2 files changed, 12 insertions, 1 deletions

diff --git a/arch/riscv64/kernel/cpu.c b/arch/riscv64/kernel/cpu.c
index 4273d02..a730b87 100644
--- a/arch/riscv64/kernel/cpu.c
+++ b/arch/riscv64/kernel/cpu.c
@@ -7,10 +7,21 @@
  */
 
 #include <apos/tcb.h>
+#include <apos/atomic.h>
+
 #include <arch/cpu.h>
 
+static atomic_long cpus = 0;
+
 void cpu_assign(struct tcb *t)
 {
+	/* bounce cpu id forward, or if first time here, get a cpu id */
+	struct tcb *c = cur_tcb();
+	if (likely(c))
+		t->cpu_id = c->cpu_id;
+	else
+		t->cpu_id = cpus++;
+
 	__asm__ volatile ("mv tp, %0\n" : : "r" (t) :);
 }
 
diff --git a/arch/riscv64/kernel/proc.c b/arch/riscv64/kernel/proc.c
index 10f8bfa..b047e2c 100644
--- a/arch/riscv64/kernel/proc.c
+++ b/arch/riscv64/kernel/proc.c
@@ -40,7 +40,7 @@ void set_args(struct tcb *t, struct sys_ret a)
 struct sys_ret get_args(struct tcb *t)
 {
 	struct riscv_regs *r = (struct riscv_regs *)(--t);
-	return (struct sys_ret){r->a0, r->a1, r->a2, r->a3, r->a4, r->a5};
+	return SYS_RET6(r->a0, r->a1, r->a2, r->a3, r->a4, r->a5);
 }
 
 void set_thread(struct tcb *t)