/*
 * Copyright (c) The mldsa-native project authors
 * Copyright (c) The mlkem-native project authors
 * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
 */

/*yaml
  Name: rej_uniform_eta4_aarch64_asm
  Description: AArch64 rejection sampling of eta=4 secret coefficients
  Signature: uint64_t mld_rej_uniform_eta4_aarch64_asm(int32_t *r, const uint8_t *buf, unsigned buflen, const uint8_t *table)
  ABI:
    Architecture: aarch64
    CallingConvention: AAPCS64
    x0:
      type: buffer
      size_bytes: 1024
      permissions: write-only
      c_parameter: int32_t *r
      description: Output buffer (256 x int32_t)
    x1:
      type: buffer
      size_bytes: x2
      permissions: read-only
      c_parameter: const uint8_t *buf
      description: Input buffer
    x2:
      type: scalar
      c_parameter: unsigned buflen
      description: Length of input buffer (must be a multiple of 8)
      test_with: 272  # MLD_AARCH64_REJ_UNIFORM_ETA4_BUFLEN
    x3:
      type: buffer
      size_bytes: 4096
      permissions: read-only
      c_parameter: const uint8_t *table
      description: Lookup table (4096 x uint8_t)
*/


/*
 * WARNING: This file is auto-derived from the mldsa-native source file
 *   dev/aarch64_opt/src/rej_uniform_eta4_aarch64_asm.S using scripts/simpasm. Do not modify it directly.
 */

.text
.balign 4
#ifdef __APPLE__
.global _mld_rej_uniform_eta4_aarch64_asm
_mld_rej_uniform_eta4_aarch64_asm:
#else
.global mld_rej_uniform_eta4_aarch64_asm
mld_rej_uniform_eta4_aarch64_asm:
#endif

        .cfi_startproc
        sub sp, sp, #0x240
        .cfi_adjust_cfa_offset 0x240
        mov x7, #0x1                // =1
        movk x7, #0x2, lsl #16
        movk x7, #0x4, lsl #32
        movk x7, #0x8, lsl #48
        mov v31.d[0], x7
        mov x7, #0x10               // =16
        movk x7, #0x20, lsl #16
        movk x7, #0x40, lsl #32
        movk x7, #0x80, lsl #48
        mov v31.d[1], x7
        movi v30.8h, #0x9
        movi v7.8h, #0x4
        mov x8, sp
        mov x7, x8
        mov x11, #0x0               // =0
        eor v16.16b, v16.16b, v16.16b

Lrej_uniform_eta4_initial_zero:
        str q16, [x7], #0x40
        stur q16, [x7, #-0x30]
        stur q16, [x7, #-0x20]
        stur q16, [x7, #-0x10]
        add x11, x11, #0x20
        cmp x11, #0x100
        b.lt Lrej_uniform_eta4_initial_zero
        mov x7, x8
        mov x9, #0x0                // =0
        mov x4, #0x100              // =256

Lrej_uniform_eta4_loop8:
        cmp x9, x4
        b.hs Lrej_uniform_eta4_memory_copy
        sub x2, x2, #0x8
        ld1 { v0.8b }, [x1], #8
        movi v26.8b, #0xf
        and v27.8b, v0.8b, v26.8b
        ushr v28.8b, v0.8b, #0x4
        zip1 v26.8b, v27.8b, v28.8b
        zip2 v29.8b, v27.8b, v28.8b
        ushll v16.8h, v26.8b, #0x0
        ushll v17.8h, v29.8b, #0x0
        cmhi v4.8h, v30.8h, v16.8h
        cmhi v5.8h, v30.8h, v17.8h
        and v4.16b, v4.16b, v31.16b
        and v5.16b, v5.16b, v31.16b
        uaddlv s20, v4.8h
        uaddlv s21, v5.8h
        fmov w12, s20
        fmov w13, s21
        ldr q24, [x3, x12, lsl #4]
        ldr q25, [x3, x13, lsl #4]
        cnt v4.16b, v4.16b
        cnt v5.16b, v5.16b
        uaddlv s20, v4.8h
        uaddlv s21, v5.8h
        fmov w12, s20
        fmov w13, s21
        tbl v16.16b, { v16.16b }, v24.16b
        tbl v17.16b, { v17.16b }, v25.16b
        st1 { v16.8h }, [x7]
        add x7, x7, x12, lsl #1
        st1 { v17.8h }, [x7]
        add x7, x7, x13, lsl #1
        add x12, x12, x13
        add x9, x9, x12
        cmp x2, #0x8
        b.hs Lrej_uniform_eta4_loop8

Lrej_uniform_eta4_memory_copy:
        cmp x9, x4
        csel x9, x9, x4, lo
        mov x11, #0x0               // =0
        mov x7, x8

Lrej_uniform_eta4_final_copy:
        ldr q16, [x7], #0x20
        ldur q18, [x7, #-0x10]
        sub v16.8h, v7.8h, v16.8h
        sub v18.8h, v7.8h, v18.8h
        sshll2 v17.4s, v16.8h, #0x0
        sshll v16.4s, v16.4h, #0x0
        sshll2 v19.4s, v18.8h, #0x0
        sshll v18.4s, v18.4h, #0x0
        str q16, [x0], #0x40
        stur q17, [x0, #-0x30]
        stur q18, [x0, #-0x20]
        stur q19, [x0, #-0x10]
        add x11, x11, #0x10
        cmp x11, #0x100
        b.lt Lrej_uniform_eta4_final_copy
        mov x0, x9
        add sp, sp, #0x240
        .cfi_adjust_cfa_offset -0x240
        ret
        .cfi_endproc

#if defined(__ELF__)
.section .note.GNU-stack,"",%progbits
#endif
