#include "tx_port.h"

    .section .text
    .align  2

    .global _tx_thread_smp_protect
    .type   _tx_thread_smp_protect, @function
_tx_thread_smp_protect:

    /* Disable interrupts so we don't get preempted.  */
    csrr    a0, mstatus                               // Pickup current interrupt posture
    csrci   mstatus, 0x08                             // Lockout interrupts

    /* Pickup the hart ID.  */
    csrr    t2, mhartid                               // Pickup the current hart ID

    /* Build address to protection structure.  */
    la      t1, _tx_thread_smp_protection

    /* If this hart already owns protection, just nest the count.  */
    LWU     t3, 4(t1)                                 // Pickup owning hart
    beq     t3, t2, _owned                            // Already owned by this hart

    /* Try to get the protection.  */
    LWU     t4, 0(t1)                                 // Pickup protection flag
    bnez    t4, _protection_busy                      // If set, protection is busy
    li      t4, 1                                     // Build lock value
    amoswap.w.aq t5, t4, (t1)                         // Attempt to get protection
    bnez    t5, _protection_busy                      // If old value != 0, retry

    fence   rw, rw                                    // Ensure lock acquisition is visible
    sw      t2, 4(t1)                                 // Save owning hart

_owned:
    LWU     t5, 8(t1)                                 // Pickup ownership count
    addi    t5, t5, 1                                 // Increment ownership count
    sw      t5, 8(t1)                                 // Store ownership count
    fence   rw, rw                                    // Publish owner/count before return
    ret

_protection_busy:
    csrw    mstatus, a0                               // Restore interrupts
    j       _tx_thread_smp_protect                    // Restart the protection attempt