/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2024 Strahinja Stanisic <strajabot@FreeBSD.org>
 */

#include <machine/asm.h>

/*
 * a0 - const void *b
 * a1 - int c
 * a2 - size_t len
 */
ENTRY(memchr)
	/*
	 * a0 - const char *ptr
	 * a1 - char cccccccc[8]
	 * a2 - char iter[8]
	 * a3 - uint8_t *end
	 * a4 - uint64_t *end_align
	 * a5 - uint64_t *end_unroll
	 */

	beqz a2, .Lno_match

	/* c = (uint8_t) c */
	andi a1, a1, 0xFF

	/*
	 * t0 = 0x0101010101010101
	 * t1 = 0x8080808080808080
	 * t2 = b << 3
	 * cccccccc = (uint8_t)c * t0
	 * end = b + len;
	 * ptr = b & ~0b111
	 */
	add a3, a0, a2
	li t0, 0x01010101
	sltu t2, a0, a3
	slli t1, t0, 32
	neg t2, t2
	or t0, t0, t1
	and a3, a3, t2
	slli t1, t0, 7
	slli t2, a0, 3
	and a0, a0, ~0b111
	mul a1, t0, a1

	ld a2, (a0)

	/*
	 * mask_start = REP8_0x01 ^ (REP8_0x01 << t2)
	 * iter = iter ^ cccccccc
	 * iter = iter | mask_start
	 */
	sll t2, t0, t2
	xor a2, a2, a1
	xor t2, t2, t0
	or a2, a2, t2

	/* has_zero(iter)
	 * end_align = (end + 7) & ~0b111;
	 */
	addi a4, a3, 7
	not t2, a2
	sub a2, a2, t0
	and t2, t2, t1
	andi a4, a4, ~0b111
	and a2, a2, t2

	/* ptr = ptr + 8 */
	addi a0, a0, 8

	bnez a2, .Lfind_zero

	/* if(ptr == end_align) */
	beq a0, a4, .Lno_match

	/* end_unroll = end_align & ~0b1111 */
	andi a5, a4, ~0b1111

	/*
	 * Instead of branching to check if `ptr` is 16-byte aligned:
	 *   - Probe the next 8 bytes for `c`
	 *   - Align `ptr` down to the nearest 16-byte boundary
	 *
	 * If `ptr` was already 16-byte aligned, those 8 bytes will be
	 * checked again inside the unrolled loop.
	 *
	 * This removes an unpredictable branch and improves performance.
	 */

	ld a2, (a0)
	xor a2, a2, a1

	not t2, a2
	sub a2, a2, t0
	and t2, t2, t1
	and a2, a2, t2

	addi a0, a0, 8

	bnez a2, .Lfind_zero

	andi a0, a0, ~0b1111

	/* while(ptr != end_unroll) */
	beq a0, a5, .Lskip_loop
.Lloop:
	ld a2, (a0)
	ld t3, 8(a0)

	xor a2, a2, a1
	xor t3, t3, a1

	not t2, a2
	not t4, t3
	sub a2, a2, t0
	sub t3, t3, t0
	and t2, t2, t1
	and t4, t4, t1
	and a2, a2, t2
	and t3, t3, t4

	addi a0, a0, 8

	bnez a2, .Lfind_zero

	/* move into iter for find_zero */
	mv a2, t3

	addi a0, a0, 8

	bnez a2, .Lfind_zero

	bne a0, a5, .Lloop
.Lskip_loop:

	/* there might be one 8byte left */
	beq a0, a4, .Lno_match

	ld a2, (a0)
	xor a2, a2, a1

	not t2, a2
	sub a2, a2, t0
	and t2, t2, t1
	and a2, a2, t2

	addi a0, a0, 8

	beqz a2, .Lno_match

.Lfind_zero:
	/*
	 * ptr = ptr - 8
	 * t1 = 0x0001020304050607
	 * iter = iter & (-iter)
	 * iter = iter >> 7
	 * iter = iter * t1
	 * iter = iter >> 56
	 */
	li t1, 0x10203000
	neg t0, a2
	slli t1, t1, 4
	and a2, a2, t0
	addi t1, t1, 0x405
	srli a2, a2, 7
	slli t1, t1, 16
	addi a0, a0, -8
	addi t1, t1, 0x607
	mul a2, a2, t1
	srli a2, a2, 56

	/* left = end - ptr */
	sub t0, a3, a0

	/* return iter < left ? ptr + iter : NULL */
	sltu t1, a2, t0
	neg t1, t1
	add a0, a0, a2
	and a0, a0, t1
	ret

.Lno_match:
	li a0, 0
	ret
END(memchr)