; ============================================================================
; memory.asm - 对齐内存分配器 (x86-64 MASM)
; 支持 16/32/64 字节对齐 (SIMD 需要)
; ============================================================================

include asm_macros.inc

.DATA
ALIGN 16
_alloc_header_size  EQU 32                 ; 分配头大小 (包含对齐偏移和原始指针)

.CODE

; ============================================================================
; 对齐内存分配
; ============================================================================

; void* mem_aligned_alloc(uint64_t size, uint64_t alignment)
; RCX = size, RDX = alignment
; 返回: RAX = 对齐指针
mem_aligned_alloc PROC
    push    rbx
    push    rdi
    push    rsi
    sub     rsp, 40                         ; 保留栈空间

    mov     rbx, rcx                        ; rbx = size
    mov     rdi, rdx                        ; rdi = alignment

    ; 验证 alignment 是 2 的幂
    mov     rax, rdi
    dec     rax
    test    rdi, rax
    jnz     alloc_invalid_alignment

    ; 确保 alignment >= 16
    cmp     rdi, 16
    jae     alloc_alignment_ok
    mov     rdi, 16

alloc_alignment_ok:
    ; 需要的总大小 = size + alignment + header
    mov     rax, rbx
    add     rax, rdi                         ; + alignment
    add     rax, _alloc_header_size         ; + header
    mov     rsi, rax                        ; rsi = total size

    ; 调用 malloc
    mov     rcx, rsi
    call    malloc
    test    rax, rax
    jz      alloc_failed

    ; rax = 原始指针
    ; 计算对齐后的地址
    mov     rcx, rax
    add     rcx, _alloc_header_size         ; 跳过 header
    ; 对齐
    mov     rdx, rdi
    dec     rdx
    not     rdx
    add     rcx, rdx                        ; 加上 alignment-1
    and     rcx, rdx                        ; 对齐到 alignment

    ; rcx = 对齐后的指针
    ; 在对齐指针前存储原始指针
    mov     [rcx - 8], rax                  ; 存储原始 malloc 指针
    mov     [rcx - 16], rbx                 ; 存储请求的 size
    mov     [rcx - 24], rdi                 ; 存储 alignment
    mov     [rcx - 32], rsi                 ; 存储 total size

    mov     rax, rcx                        ; 返回对齐指针
    jmp     alloc_done

alloc_invalid_alignment:
    xor     rax, rax                        ; 返回 NULL
    jmp     alloc_done

alloc_failed:
    xor     rax, rax

alloc_done:
    add     rsp, 40
    pop     rsi
    pop     rdi
    pop     rbx
    ret
mem_aligned_alloc ENDP

; void mem_aligned_free(void* ptr)
; RCX = ptr (对齐指针)
mem_aligned_free PROC
    test    rcx, rcx
    jz      free_done

    ; 获取原始 malloc 指针
    mov     rax, [rcx - 8]
    test    rax, rax
    jz      free_done

    ; 调用 free
    mov     rcx, rax
    call    free

free_done:
    ret
mem_aligned_free ENDP

; void* mem_aligned_realloc(void* ptr, uint64_t size, uint64_t alignment)
; RCX = ptr, RDX = size, R8 = alignment
; 返回: RAX = 新对齐指针
mem_aligned_realloc PROC
    push    rbx
    push    rdi
    push    rsi
    push    r12
    sub     rsp, 56                         ; shadow + 3 saved qwords

    mov     rbx, rdx                        ; rbx = new size
    mov     rdi, r8                         ; rdi = alignment

    ; 如果 ptr 为 NULL, 直接分配
    test    rcx, rcx
    jz      realloc_alloc

    ; 保存旧信息
    mov     [rsp + 32], rcx                 ; 旧对齐指针（即用户数据起始）
    mov     rax, [rcx - 8]
    mov     [rsp + 24], rax                 ; 旧原始 malloc 指针
    mov     rax, [rcx - 16]
    mov     [rsp + 16], rax                 ; 旧 size

    ; 分配新内存
    mov     rcx, rbx
    mov     rdx, rdi
    call    mem_aligned_alloc
    test    rax, rax
    jz      realloc_failed

    mov     rsi, rax                        ; rsi = 新对齐指针

    ; 复制旧数据 (取较小值)
    mov     rdx, [rsp + 16]                 ; 旧 size
    cmp     rdx, rbx
    jbe     realloc_copy_ok
    mov     rdx, rbx

realloc_copy_ok:
    ; fast_memcpy(dst, src, count): rcx=dst, rdx=src, r8=count
    mov     rcx, rsi
    mov     r8, rdx
    mov     rdx, [rsp + 32]                 ; 旧对齐指针
    call    fast_memcpy

    ; 释放旧内存（使用原始 malloc 指针）
    mov     rcx, [rsp + 24]
    call    mem_aligned_free

    mov     rax, rsi                        ; 返回新指针
    jmp     realloc_done

realloc_alloc:
    mov     rcx, rbx
    mov     rdx, rdi
    call    mem_aligned_alloc
    jmp     realloc_done

realloc_failed:
    xor     rax, rax

realloc_done:
    add     rsp, 56
    pop     r12
    pop     rsi
    pop     rdi
    pop     rbx
    ret
mem_aligned_realloc ENDP

END