; Returns the index of the first element matching the provided predicate function (or -1 if not found)
; r1 - The array pointer
; r2 - The array length (number of items)
; r3 - The stride (size of one item) - either 1, 2 or 4 (bytes)
; r4 - The predicate
; r5 - Predicate context
; The predicate function should follow the stdlib calling convention
;  and receive one argument and return either a zero when the value is not the one we search for
; , or any other result if it is the searched-for item.
pub find_index:
    push r12 ; We will store the predicate pointer here
    push r11 ; We will store the current pointer here
    push r10 ; We will store the stride here
    push r9 ; We will store the final address here
    push r8 ; We will store the mask here
    push r1 ; We need the array pointer to calculate the item index

    mov r12, r4
    mov r11, r1
    mov r10, r3
    mov r9, r2
    lsr r6, r3, 1 ; We turn the stride into a byte shift
    lsl r9, r9, r6 ; We calculate bytes left
    add r9, r9, r1 ; We add the start address to get the final address

    push r13 ; We save up the return address because we will provide our own to the predicate
    counter r13
    add r13, r13, 52 ; Point to just after the predicate call - we can set this up now so we don't waste loop cycles

    nand r8, zr, zr ; We create a mask of 0xFFFFFFFF
    mov r6, 4
    sub r6, r6, r3 ; We create a "negative stride", e.g. 4 -> 0, 2 -> 2, 1 -> 3
    lsl r6, r6, 3
    lsr r8, r8, r6 ; We shift the mask by the negative stride to obtain the proper mask for a value
    ; e.g. stride 4 -> mask is 0xFFFFFFFF
    ; stride 2 -> mask is 0x0000FFFF
    ; stride 1 -> mask is 0x000000FF

    push r5 ; We save the predicate context on the stack
    find_index_loop:
        load_32 r1, [r11] ; We load the element
        and r1, r1, r8 ; We mask it to handle stride 2 and 1 cases
        load_32 r2, [sp] ; We load the predicate context into r2
        jmp r12 ; We call the predicate
        cmp r1, zr
        jne find_index_found_item ; If we found the item, we jump out
        ; If we didn't, move to next item
        add r11, r11, r10 ; We add the stride to the pointer
        cmp r11, r9 ; We compare with the final address
        je find_index_not_found ; If we reached the end we're done
        jmp find_index_loop
    
    find_index_not_found:
        add sp, sp, 4 ; The predicate context is not useful
        pop r13 ; We get our return address
        nand r1, zr, zr ; We put -1 in r1
        add sp, sp, 4 ; The old array pointer are not useful
        jmp find_index_postamble
    
    find_index_found_item:
        add sp, sp, 4 ; The predicate context is not useful
        pop r13 ; We get our return address
        pop r1 ; We get the array pointer
        sub r1, r11, r1 ; We calculate the bytes from the start
        lsr r10, r10, 1 ; We shift the stride to get the amount to shift the bytes for
        lsr r1, r1, r10 ; We shift to get the index of the item

    find_index_postamble:
        pop r8
        pop r9
        pop r10
        pop r11
        pop r12
        jmp r13 ; Return