/*
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2021 Warner Losh
 * Copyright (c) 2023 Stormshield
 * Copyright (c) 2023 Klara, Inc.
 */

#include <sys/syscall.h>

#define STDOUT_FILENO   1

#define MUTEX_LOCKED    0x01
#define MUTEX_UNLOCKED  0x00

#define STACK_SIZE      4096
#define TLS_SIZE        4096

        .text
        .file "swp_test.S"
        .syntax unified
        .globl main
        .p2align 2
        .type main,%function
        .code 32

main:
        /*
         * Stack slots:
         * 0 - Sync word
         * 1 - Thread id
         * 2 - Shared word
         */
        sub sp, sp, #12

        /* Print a message */
        movw r0, :lower16:.L.mainmsg
        movt r0, :upper16:.L.mainmsg
        ldr r1, =(.L.mainmsgEnd - .L.mainmsg - 1)
        bl print

        /* Create two secondary threads */
        mov r0, #1
        str r0, [sp, #4]        /* Thread ID */
        movw r0, :lower16:secondary_thread
        movt r0, :upper16:secondary_thread
        mov r1, sp
        movw r2, :lower16:stack1
        movt r2, :upper16:stack1
        movw r3, :lower16:tls1
        movt r3, :upper16:tls1
        bl create_thr

1:
        /*
         * Wait for the first new thread to ack its existence by
         * incrementing the thread id.
         */
        ldr r0, [sp, #4]
        cmp r0, #1
        bne 2f
        ldr r7, =SYS_sched_yield
        swi 0
        b 1b

2:
        /* Create thread #2 */
        movw r0, :lower16:secondary_thread
        movt r0, :upper16:secondary_thread
        mov r1, sp
        movw r2, :lower16:stack2
        movt r2, :upper16:stack2
        movw r3, :lower16:tls2
        movt r3, :upper16:tls2
        bl create_thr

3:
        /*
         * Wait for the first new thread to ack its existence by
         * incrementing the thread id.
         */
        ldr r0, [sp, #4]
        cmp r0, #2
        bne 4f
        ldr r7, =SYS_sched_yield
        swi 0
        b 3b

        /* Loop */
4:
        mov r0, sp
        mov r1, #0      /* Thread loop */
        add r2, sp, #8
        bl thread_loop
        b 4b

        /* UNREACHABLE */
        mov r0, #0
        ldr r7, =SYS_exit
        swi 0

        .p2align 2
        .type secondary_thread,%function
        .code 32
secondary_thread:
        /*
         * On entry, r0 is where we stashed our sync word and
         * ack word (thread ID).
         *
         * Stash the sync word in r4, thread ID in r5.
         */
        mov r4, r0
        ldr r5, [r0, #4]

        /* Print a message */
        movw r0, :lower16:.L.secondarymsg
        movt r0, :upper16:.L.secondarymsg
        ldr r1, =(.L.secondarymsgEnd - .L.secondarymsg - 1)
        bl print

        /* Acknowledge that we started */
        add r0, r5, #1
        str r0, [r4, #4]

1:
        mov r0, r4
        mov r1, r5
        add r2, r4, #8
        bl thread_loop
        b 1b

        .p2align 2
        .type thread_loop,%function
        .code 32
thread_loop:
        push {r4, r5, r6, r7, r8, lr}

        /*
         * r0 == sync word
         * r1 == thread ID
         * r2 == shared word
         */
        mov r4, r0
        mov r5, r1
        mov r6, r2
        bl lock_mutex_swp
        str r5, [r6] /* Write the thread ID */
        bl random_cycles

        # Save off the now cycle count */
        mov r8, r0

        /* Print the thread ID and cycle count */
        mov r0, r5
        mov r1, #0
        bl printnum

        /* Separator */
        movw r0, :lower16:.L.idsep
        movt r0, :upper16:.L.idsep
        ldr r1, =(.L.idsepEnd - .L.idsep - 1)
        bl print

        /* Cycle count */
        mov r0, r8
        mov r1, #1
        bl printnum

1:
        ldr r0, [r6]
        cmp r0, r5      /* Check against the thread ID */
        bne 2f
        str r5, [r6]

        /*
         * Check if the count hit 0, otherwise go again.
         */
        cmp r8, #0
        beq 3f
        sub r8, r8, #1
        b 1b

2:
        /* exit(1) */
        mov r0, #1
        ldr r7, =SYS_exit
        swi 0

3:
        mov r0, r4
        bl unlock_mutex_swp

        /*
         * Yield to lower the chance that we end up re-acquiring, the other two
         * threads are still actively trying to acquire the lock.
         */
        ldr r7, =SYS_sched_yield
        swi 0

        pop {r4, r5, r6, r7, r8, lr}
        bx lr

        .p2align 2
        .type random_cycles,%function
        .code 32
random_cycles:
        /* Return a random number < 4k */
        sub sp, sp, #4
        mov r0, sp
        mov r1, #4
        mov r2, #0
        ldr r7, =SYS_getrandom
        swi 0

        /*
         * Just truncate the result of getrandom(2)
         * to put us within range.  Naive, but functional.
         */
        ldr r0, [sp]
        mov r1, #0xfff
        and r0, r0, r1
        add sp, sp, #4
        bx lr

        /*
         * lock_mutex_swp and unlock_mutex_swp lifted from
         * ARM documentation on SWP/SWPB.
         */
        .p2align 2
        .type lock_mutex_swp,%function
        .code 32
lock_mutex_swp:
        mov r2, #MUTEX_LOCKED
        swp r1, r2, [r0]        /* Swap in lock value. */
        cmp r1, r2              /* Check if we were locked already. */
        beq lock_mutex_swp      /* Retry if so */
        bx lr                   /* Return locked */

        .p2align 2
        .type unlock_mutex_swp,%function
        .code 32
unlock_mutex_swp:
        mov r1, #MUTEX_UNLOCKED
        str r1, [r0]            /* Move in unlocked */
        bx lr

        .p2align 2
        .type create_thr,%function
        .code 32
create_thr:
        /*
         * r0 == start_func
         * r1 == arg
         * r2 == stack_base
         * r3 == tls_base
         */
        sub sp, sp, #56
        str r0, [sp, #4]        /* start_func */
        str r1, [sp, #8]        /* arg */
        str r2, [sp, #12]       /* stack_base */
        mov r0, #STACK_SIZE
        str r0, [sp, #16]       /* stack_size */
        str r3, [sp, #20]       /* tls_base */
        mov r0, #TLS_SIZE
        str r0, [sp, #24]       /* tls_size */
        mov r0, #0
        str r0, [sp, #28]
        str r0, [sp, #32]
        str r0, [sp, #36]
        str r0, [sp, #40]

        add r0, sp, #4  /* &thrp */
        mov r1, #52     /* sizeof(thrp) */
        ldr r7, =SYS_thr_new
        swi 0

        add sp, sp, #56
        bx lr

        .p2align 2
        .type printnum,%function
        .code 32
printnum:
        push {r4, r5, r6, r7, r8, r10, lr}
        sub sp, #4

        /* 1000000000 */
        movw r6, #0xca00
        movt r6, #0x3b9a

        udiv r5, r0, r6
        cmp r5, #9
        bhi abort

        /* r4 is our accumulator */
        mov r4, r0
        /* r5 to be used as our "significant bit" */
        mov r5, #0
        /* r10 is "output_newline" */
        mov r10, r1

1:
        cmp r6, #0
        beq 4f

        /* Divide by current place */
        udiv r0, r4, r6
        /* Significant already? print anyways */
        cmp r5, #0
        bne 2f

        /*
         * Not significant, maybe print.  If we made it all the way to 1, we
         * need to just print the 0 anyways.
         */
        cmp r6, #1
        beq 2f

        cmp r0, #0
        bne 2f
        b 3f    /* Proceed */

        /* Print */
2:
        mov r5, #1
        mov r8, r0
        add r0, r0, #0x30
        str r0, [sp]
        mov r0, sp
        mov r1, #1
        bl print

        /* Multiply back into place and subtract from accumulator */
        mul r0, r8, r6
        sub r4, r4, r0

3:
        mov r3, #10
        udiv r6, r6, r3
        b 1b

4:
        cmp r10, #0
        beq 5f

        /* newline */
        mov r0, #0x0a
        str r0, [sp]
        mov r0, sp
        mov r1, #1
        bl print

5:
        add sp, sp, #4
        pop {r4, r5, r6, r7, r8, r10, lr}
        bx lr

abort:
        movw r0, :lower16:.L.badnum
        movt r0, :upper16:.L.badnum
        ldr r1, =(.L.badnumEnd - .L.badnum - 1)
        bl print

        mov r0, #1
        ldr r7, =SYS_exit
        swi 0

        .p2align 2
        .type print,%function
        .code 32
print:
        /* r0 - string, r1 = size */
        mov r2, r1
        mov r1, r0
        ldr r0, =STDOUT_FILENO
        ldr r7, =SYS_write
        swi 0

        bx lr

.L.mainmsg:
        .asciz "Main thread\n"
.L.mainmsgEnd:
        .size .L.mainmsg, .L.mainmsgEnd - .L.mainmsg
.L.secondarymsg:
        .asciz "Secondary thread\n"
.L.secondarymsgEnd:
        .size .L.secondarymsg, .L.secondarymsgEnd - .L.secondarymsg
.L.badnum:
        .asciz "Bad number\n"
.L.badnumEnd:
        .size .L.badnum, .L.badnumEnd - .L.badnum
.L.idsep:
        .asciz " - cycles "
.L.idsepEnd:
        .size .L.idsep, .L.idsepEnd - .L.idsep

        .type stack1,%object
        .local stack1
        .comm stack1,STACK_SIZE,1
        .type tls1,%object
        .local tls1
        .comm tls1,TLS_SIZE,1

        .type stack2,%object
        .local stack2
        .comm stack2,STACK_SIZE,1
        .type tls2,%object
        .local tls2
        .comm tls2,TLS_SIZE,1