/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2024 Getz Mikalsen <getz@FreeBSD.org>
*/

#include <machine/asm.h>

        .weak   strcspn
        .set    strcspn, __strcspn
        .text

ENTRY(__strcspn)
        stp     x29, x30, [sp, #-16]!
        mov     x29, sp
        mov     x15, #1                 // preload register with 1 for stores

        /* check for special cases */
        ldrb    w4, [x1]                // first character in the set
        cbz     w4, .Lstrlen

        movi    v0.16b, #0

        ldrb    w5, [x1, #1]            // second character in the set
        cbz     w5, .Lstrchr

        sub     sp, sp, #256            // allocate 256 bytes on the stack

        /* no special case matches -- prepare lookup table */
        mov     w3, #20
        .p2align 4
0:      add     x9, sp, x3, lsl #3
        stp     xzr, xzr, [x9]
        stp     xzr, xzr, [x9, #16]
        subs    w3, w3, #4
        b.cs    0b

        /* utilize SIMD stores to speed up zeroing the table */
        stp     q0, q0, [sp, #6*32]
        stp     q0, q0, [sp, #7*32]

        add     x1, x1, #2
        strb    w15, [sp, x4]           // register first chars in the set
        strb    w15, [sp, x5]

        mov     x4, x0                  // stash a copy of src

        /* process remaining chars in set */
        .p2align 4
0:      ldrb    w5, [x1]
        strb    w15, [sp, x5]
        cbz     w5, 1f                  // end of set?

        ldrb    w5, [x1, #1]
        strb    w15, [sp, x5]
        cbz     w5, 1f

        add     x1, x1, #2
        b       0b

        /* find match */
        .p2align 4
1:      ldrb    w8, [x0]
        ldrb    w9, [sp, x8]
        cbnz    w9, 2f

        ldrb    w8, [x0, #1]
        ldrb    w9, [sp, x8]
        cbnz    w9, 3f

        ldrb    w8, [x0, #2]
        ldrb    w9, [sp, x8]
        cbnz    w9, 4f

        ldrb    w8, [x0, #3]
        ldrb    w9, [sp, x8]
        add     x0, x0, #4
        cbz     w9, 1b

        sub     x0, x0, #3              // fix up return value
4:      sub     x4, x4, #1
3:      add     x0, x0, #1
2:      sub     x0, x0, x4
        mov     sp, x29
        ldp     x29, x30, [sp], #16     // restore sp and lr
        ret

        /* set is empty, degrades to strlen */
        .p2align 4
.Lstrlen:
        mov     sp, x29
        ldp     x29, x30, [sp], #16     // restore sp and lr
        b       strlen

        /* just one character in set, degrades to strchrnul */
        .p2align 4
.Lstrchr:
        stp     x0, x1, [sp, #-16]!
        mov     x1, x4

        bl      strchrnul

        ldp     x18, x17, [sp], #16     // restore stashed src
        sub     x0, x0, x18

        ldp     x29, x30, [sp], #16     // Restore sp and lr
        ret

END(__strcspn)