ASM Sieve of Eratosthenes

.global main
.extern printf
.extern fprintf
.extern stderr
.extern atoi
.extern malloc
.extern free

.text
fmt: .asciz "%d\n"
usage: .asciz "Usage: %s N\n"
errmsg: .asciz "N must be equal to or greater than 2\n"


// Print number in rdx
// This saves registers
print_number:
    push    %rbp
    mov     %rsp, %rbp

    push    %rax
    push    %rsi
    push    %rdi
    push    %rcx
    push    %rdx

    mov     %rdx, %rsi
    mov     $fmt, %rdi
    xor     %rax, %rax
    call    printf

    pop     %rdx
    pop     %rcx
    pop     %rdi
    pop     %rsi
    pop     %rax

    leave
    ret


// Initialize array
// rdi = N
// rsi = array address
init_array:
    push    %rbp
    mov     %rsp, %rbp
    push    %rdi
    push    %rsi

    mov     (%rsp), %rdi
    xor     %rax, %rax

_init_loop:
    cmp     %rax, 8(%rsp)
    je      _init_loop_exit

    mov     %eax, (%rsi, %rax, 4)

    inc     %rax
    jmp     _init_loop

_init_loop_exit:
    leave
    ret


// Sieve
// rdi = N
// rsi = array address
// rdx = skip
sieve:
    push    %rbp
    mov     %rsp, %rbp
    push    %rdx
    push    %rax

    mov     %rdx, %rax

_sieve_loop:
    cmp     %rdi, %rax
    jge     _sieve_exit

    movl    $0, (%rsi, %rax, 4)

    add     %rdx, %rax
    jmp     _sieve_loop

_sieve_exit:
    pop     %rax
    pop     %rdx
    leave
    ret


// Eratosthenes
// rdi = N
// rsi = array address
eratosthenes:
    push    %rbp
    mov     %rsp, %rbp

    mov     $2, %rcx

_loop:
    cmp     %rcx, %rdi
    je      _leave

    mov     (%rsi, %rcx, 4), %edx
    cmp     $0, %edx
    je      _next

    call    print_number
    call    sieve

_next:
    inc     %rcx
    jmp     _loop

_leave:
    leave
    ret


main:
    push    %rbp
    mov     %rsp, %rbp
    mov     %rsi, %r15

    cmp     $2, %rdi
    jne     give_usage

    // Get argument from argv[1]
    mov     8(%r15), %rdi
    call    atoi
    cmp     $2, %rax
    jl      print_error
    inc     %rax
    push    %rax            # save N

    // Call malloc
    mov     (%rsp), %rdi    # N
    shl     $2, %rdi        # N * sizeof(int)
    call    malloc
    push    %rax            # save pointer

    // Initialize array
    mov     8(%rsp), %rdi   # N
    mov     (%rsp), %rsi    # array
    call    init_array

    // Sieve of Eratosthenes
    mov     8(%rsp), %rdi
    mov     (%rsp), %rsi
    call    eratosthenes

    // Free memory
    mov     (%rsp), %rdi    # move array to arg0
    call    free

    leave
    xor     %rax, %rax
    ret

print_error:
    mov     stderr, %rdi
    mov     $errmsg, %rsi
    xor     %rax, %rax
    call    fprintf

give_usage:
    mov     (%r15), %rdx
    mov     $usage, %rsi
    mov     stderr, %rdi
    xor     %rax, %rax
    call    fprintf
    mov     $1, %rax

    leave
    ret