203 lines
6.6 KiB
Python
203 lines
6.6 KiB
Python
import mm
|
|
import drivers.serial.uart.serial as serial
|
|
import asm
|
|
import t, c
|
|
|
|
CORO_READY: t.CDefine = 0
|
|
CORO_RUNNING: t.CDefine = 1
|
|
CORO_YIELDED: t.CDefine = 2
|
|
CORO_DONE: t.CDefine = 3
|
|
MAX_COROUTINES: t.CDefine = 32
|
|
CORO_STACK_SIZE: t.CDefine = 8192
|
|
|
|
|
|
@t.Object
|
|
class Coroutine:
|
|
rsp: t.CUInt64T
|
|
stack: t.CVoid | t.CPtr
|
|
stack_size: t.CUInt64T
|
|
func: t.CVoid | t.CPtr
|
|
arg: t.CVoid | t.CPtr
|
|
state: t.CInt
|
|
cid: t.CInt
|
|
caller_rsp: t.CUInt64T
|
|
ret_val: t.CInt
|
|
|
|
def resume(self) -> t.CInt:
|
|
global _coro_old_rsp, _coro_new_rsp
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
if self.state == CORO_DONE: return self.ret_val
|
|
if self.state != CORO_READY and self.state != CORO_YIELDED: return -1
|
|
prev_cid: t.CInt = m.current_cid
|
|
m.current_cid = self.cid
|
|
self.state = CORO_RUNNING
|
|
_coro_old_rsp = c.Addr(m.coroutines[prev_cid].rsp)
|
|
_coro_new_rsp = c.Addr(self.rsp)
|
|
self.caller_rsp = t.CUInt64T(_coro_old_rsp)
|
|
_coro_switch()
|
|
if self.state == CORO_DONE:
|
|
return self.ret_val
|
|
return 0
|
|
|
|
@staticmethod
|
|
def _yield():
|
|
global _coro_old_rsp, _coro_new_rsp
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
cid: t.CInt = m.current_cid
|
|
if cid < 0: return
|
|
m.coroutines[cid].state = CORO_YIELDED
|
|
prev_cid: t.CInt = -1
|
|
i: t.CInt
|
|
for i in range(m.count):
|
|
if m.coroutines[i].state == CORO_RUNNING and m.coroutines[i].caller_rsp != 0:
|
|
prev_cid = i
|
|
break
|
|
if prev_cid < 0: return
|
|
m.current_cid = prev_cid
|
|
_coro_old_rsp = c.Addr(m.coroutines[cid].rsp)
|
|
_coro_new_rsp = c.Addr(m.coroutines[prev_cid].rsp)
|
|
_coro_switch()
|
|
|
|
@property
|
|
def is_done(self) -> t.CInt:
|
|
if self.state == CORO_DONE:
|
|
return 1
|
|
return 0
|
|
|
|
@property
|
|
def is_yielded(self) -> t.CInt:
|
|
if self.state == CORO_YIELDED:
|
|
return 1
|
|
return 0
|
|
|
|
@property
|
|
def retval(self) -> t.CInt:
|
|
return self.ret_val
|
|
|
|
|
|
@t.Object
|
|
class CoroutineManager:
|
|
coroutines: t.CArray[Coroutine, MAX_COROUTINES]
|
|
current_cid: t.CInt
|
|
count: t.CInt
|
|
|
|
def __init__():
|
|
global _coro_mgr_ptr
|
|
_coro_mgr_ptr = c.Addr(_coro_mgr)
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
m.current_cid = 0
|
|
m.count = 1
|
|
m.coroutines[0].state = CORO_RUNNING
|
|
m.coroutines[0].cid = 0
|
|
m.coroutines[0].caller_rsp = 0
|
|
m.coroutines[0].ret_val = 0
|
|
m.coroutines[0].rsp = 0
|
|
m.coroutines[0].stack = None
|
|
m.coroutines[0].stack_size = 0
|
|
m.coroutines[0].func = None
|
|
m.coroutines[0].arg = None
|
|
main_rsp_ptr: t.CVoid | t.CPtr = c.Addr(m.coroutines[0].rsp)
|
|
c.Asm(f"""mov qword ptr [{c.AsmInp(main_rsp_ptr, t.ASM_DESCR.REG_ANY)}], rsp""",
|
|
op=[t.ASM_DESCR.CLOBBER_MEMORY])
|
|
|
|
@staticmethod
|
|
def current() -> Coroutine | t.CPtr:
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
return c.Addr(m.coroutines[m.current_cid])
|
|
|
|
@staticmethod
|
|
def get_coroutine(cid: t.CInt) -> Coroutine | t.CPtr:
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
if cid < 0 or cid >= m.count: return None
|
|
return c.Addr(m.coroutines[cid])
|
|
|
|
@staticmethod
|
|
def create(func: t.CVoid | t.CPtr, arg: t.CVoid | t.CPtr) -> Coroutine | t.CPtr:
|
|
m: CoroutineManager | t.CPtr = _coro_mgr_ptr
|
|
if m.count >= MAX_COROUTINES: return None
|
|
cid: t.CInt = m.count
|
|
m.count += 1
|
|
co: Coroutine | t.CPtr = c.Addr(m.coroutines[cid])
|
|
co.state = CORO_READY
|
|
co.cid = cid
|
|
co.func = func
|
|
co.arg = arg
|
|
co.caller_rsp = 0
|
|
co.ret_val = 0
|
|
stack_ptr: t.CVoid | t.CPtr = mm.malloc(CORO_STACK_SIZE)
|
|
if not stack_ptr: return None
|
|
co.stack = stack_ptr
|
|
co.stack_size = CORO_STACK_SIZE
|
|
stack_top: t.CUInt64T = t.CUInt64T(stack_ptr) + CORO_STACK_SIZE - 16
|
|
stack_top = stack_top & ~t.CUInt64T(0xF)
|
|
entry_addr: t.CUInt64T = t.CUInt64T(c.Addr(_coro_entry))
|
|
sp: t.CPtr = t.CPtr(stack_top)
|
|
c.Asm(f"""mov rax, {c.AsmInp(sp, t.ASM_DESCR.REG_ANY)}
|
|
mov qword ptr [rax - 56], 0
|
|
mov qword ptr [rax - 48], 0
|
|
mov qword ptr [rax - 40], 0
|
|
mov qword ptr [rax - 32], 0
|
|
mov qword ptr [rax - 24], 0
|
|
mov qword ptr [rax - 16], 0
|
|
mov qword ptr [rax - 8], 0
|
|
mov qword ptr [rax], {c.AsmInp(entry_addr, t.ASM_DESCR.REG_ANY)}""",
|
|
op=[t.ASM_DESCR.CLOBBER_MEMORY, t.ASM_DESCR.CLOBBER_RAX])
|
|
stack_top = stack_top - 56
|
|
co.rsp = stack_top
|
|
return co
|
|
|
|
|
|
_coro_mgr: CoroutineManager = CoroutineManager()
|
|
_coro_mgr_ptr: CoroutineManager | t.CPtr = c.Addr(_coro_mgr)
|
|
_coro_old_rsp: t.CVoid | t.CPtr
|
|
_coro_new_rsp: t.CVoid | t.CPtr
|
|
|
|
@c.Attribute(t.attr.naked)
|
|
def _coro_switch():
|
|
c.Asm(f"""push rax
|
|
push rbx
|
|
push rbp
|
|
push r12
|
|
push r13
|
|
push r14
|
|
push r15
|
|
mov rax, qword ptr [rip + _coro_old_rsp]
|
|
mov qword ptr [rax], rsp
|
|
mov rax, qword ptr [rip + _coro_new_rsp]
|
|
mov rsp, qword ptr [rax]
|
|
pop r15
|
|
pop r14
|
|
pop r13
|
|
pop r12
|
|
pop rbp
|
|
pop rbx
|
|
pop rax
|
|
ret""",
|
|
op=[t.ASM_DESCR.CLOBBER_MEMORY])
|
|
|
|
@c.Attribute(t.attr.naked)
|
|
def _coro_entry():
|
|
c.Asm(f"""sub rsp, 8
|
|
mov rax, qword ptr [rip + _coro_mgr_ptr]
|
|
mov ecx, dword ptr [rax + 1056]
|
|
movsxd rcx, ecx
|
|
shl rcx, 5
|
|
mov rdi, qword ptr [rax + rcx + 24]
|
|
mov rax, qword ptr [rax + rcx + 16]
|
|
call rax
|
|
mov rax, qword ptr [rip + _coro_mgr_ptr]
|
|
mov ecx, dword ptr [rax + 1056]
|
|
movsxd rcx, ecx
|
|
shl rcx, 5
|
|
mov dword ptr [rax + rcx + 32], 3
|
|
mov dword ptr [rax + rcx + 48], eax
|
|
add rsp, 8""",
|
|
op=[t.ASM_DESCR.CLOBBER_MEMORY, t.ASM_DESCR.CLOBBER_RAX,
|
|
t.ASM_DESCR.CLOBBER_RCX, t.ASM_DESCR.CLOBBER_RDX,
|
|
t.ASM_DESCR.CLOBBER_R8, t.ASM_DESCR.CLOBBER_R9,
|
|
t.ASM_DESCR.CLOBBER_R10, t.ASM_DESCR.CLOBBER_R11])
|
|
Coroutine._yield()
|
|
while True:
|
|
asm.sti()
|
|
asm.hlt()
|