import cpu
import t, c


# FPU上下文结构体
@c.Attribute(t.attr.aligned(16))
class fpu_context:
    data: t.CUInt8T[512] # FXSAVE区域大小


@t.Object
class _FPUObject:
    CPUObject: cpu._CPUObject | t.CPtr
    # 初始化FPU
    def __init__(self, CPUObject: cpu._CPUObject | t.CPtr):
        self.enable()
        self.CPUObject = CPUObject
        c.Asm("fninit", [t.ASM_DESCR.CLOBBER_MEMORY])
    # 启用FPU
    def enable(self):
        cr0: t.CUInt64T = self.CPUObject.readCR0()
        # 清除EM位（禁用仿真），设置MP位（监控协处理器）
        cr0 &= ~(1 << 2); # 清除EM位
        cr0 |= (1 << 1);  # 设置MP位
        # 写入CR0寄存器
        self.CPUObject.writeCR0(cr0)
        # 读取CR4寄存器
        cr4: t.CUInt64T = self.CPUObject.readCR4()
        # 设置OSFXSR位和OSXMMEXCPT位，启用SSE
        cr4 |= (1 << 9) | (1 << 10)
        # 写入CR4寄存器
        self.CPUObject.writeCR4(cr4)
        # 初始化FPU
        c.Asm("fninit", [t.ASM_DESCR.CLOBBER_MEMORY])
    # 禁用FPU
    def disable(self):
        cr0: t.CUInt64T = self.CPUObject.readCR0()
        # 设置EM位（启用仿真），清除MP位（禁用监控协处理器）
        cr0 |= (1 << 2)  # 设置EM位
        cr0 &= ~(1 << 1) # 清除MP位
        # 写入CR0寄存器
        self.CPUObject.writeCR0(cr0)
    # 保存FPU上下文
    def save_context(self, context: fpu_context | t.CPtr):
        # 保存FPU状态到内存
        c.Asm(f"fxsave [{c.AsmInp(context, t.ASM_DESCR.REG_ANY)}]", [t.ASM_DESCR.CLOBBER_MEMORY])
    # 恢复FPU上下文
    def restore_context(self, context: t.CConst | fpu_context | t.CPtr):
        # 从内存恢复FPU状态
        c.Asm(f"fxrstor [{c.AsmInp(context, t.ASM_DESCR.REG_ANY)}]", [t.ASM_DESCR.CLOBBER_MEMORY])