// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2022-2024 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 */

#include <linux/linkage.h>
#include <asm/frame.h>

.section	.rodata, "a"
.align 16
CONSTANTS:	.octa 0x6b20657479622d323320646e61707865
.text

/*
 * Very basic SSE2 implementation of ChaCha20. Produces a given positive number
 * of blocks of output with a nonce of 0, taking an input key and 8-byte
 * counter. Importantly does not spill to the stack. Its arguments are:
 *
 *	rdi: output bytes
 *	rsi: 32-byte key input
 *	rdx: 8-byte counter input/output
 *	rcx: number of 64-byte blocks to write to output
 */
SYM_FUNC_START(__arch_chacha20_blocks_nostack)

.set	output,		%rdi
.set	key,		%rsi
.set	counter,	%rdx
.set	nblocks,	%rcx
.set	i,		%al
/* xmm registers are *not* callee-save. */
.set	temp,		%xmm0
.set	state0,		%xmm1
.set	state1,		%xmm2
.set	state2,		%xmm3
.set	state3,		%xmm4
.set	copy0,		%xmm5
.set	copy1,		%xmm6
.set	copy2,		%xmm7
.set	copy3,		%xmm8
.set	one,		%xmm9

	/* copy0 = "expand 32-byte k" */
	movaps		CONSTANTS(%rip),copy0
	/* copy1,copy2 = key */
	movups		0x00(key),copy1
	movups		0x10(key),copy2
	/* copy3 = counter || zero nonce */
	movq		0x00(counter),copy3
	/* one = 1 || 0 */
	movq		$1,%rax
	movq		%rax,one

.Lblock:
	/* state0,state1,state2,state3 = copy0,copy1,copy2,copy3 */
	movdqa		copy0,state0
	movdqa		copy1,state1
	movdqa		copy2,state2
	movdqa		copy3,state3

	movb		$10,i
.Lpermute:
	/* state0 += state1, state3 = rotl32(state3 ^ state0, 16) */
	paddd		state1,state0
	pxor		state0,state3
	movdqa		state3,temp
	pslld		$16,temp
	psrld		$16,state3
	por		temp,state3

	/* state2 += state3, state1 = rotl32(state1 ^ state2, 12) */
	paddd		state3,state2
	pxor		state2,state1
	movdqa		state1,temp
	pslld		$12,temp
	psrld		$20,state1
	por		temp,state1

	/* state0 += state1, state3 = rotl32(state3 ^ state0, 8) */
	paddd		state1,state0
	pxor		state0,state3
	movdqa		state3,temp
	pslld		$8,temp
	psrld		$24,state3
	por		temp,state3

	/* state2 += state3, state1 = rotl32(state1 ^ state2, 7) */
	paddd		state3,state2
	pxor		state2,state1
	movdqa		state1,temp
	pslld		$7,temp
	psrld		$25,state1
	por		temp,state1

	/* state1[0,1,2,3] = state1[1,2,3,0] */
	pshufd		$0x39,state1,state1
	/* state2[0,1,2,3] = state2[2,3,0,1] */
	pshufd		$0x4e,state2,state2
	/* state3[0,1,2,3] = state3[3,0,1,2] */
	pshufd		$0x93,state3,state3

	/* state0 += state1, state3 = rotl32(state3 ^ state0, 16) */
	paddd		state1,state0
	pxor		state0,state3
	movdqa		state3,temp
	pslld		$16,temp
	psrld		$16,state3
	por		temp,state3

	/* state2 += state3, state1 = rotl32(state1 ^ state2, 12) */
	paddd		state3,state2
	pxor		state2,state1
	movdqa		state1,temp
	pslld		$12,temp
	psrld		$20,state1
	por		temp,state1

	/* state0 += state1, state3 = rotl32(state3 ^ state0, 8) */
	paddd		state1,state0
	pxor		state0,state3
	movdqa		state3,temp
	pslld		$8,temp
	psrld		$24,state3
	por		temp,state3

	/* state2 += state3, state1 = rotl32(state1 ^ state2, 7) */
	paddd		state3,state2
	pxor		state2,state1
	movdqa		state1,temp
	pslld		$7,temp
	psrld		$25,state1
	por		temp,state1

	/* state1[0,1,2,3] = state1[3,0,1,2] */
	pshufd		$0x93,state1,state1
	/* state2[0,1,2,3] = state2[2,3,0,1] */
	pshufd		$0x4e,state2,state2
	/* state3[0,1,2,3] = state3[1,2,3,0] */
	pshufd		$0x39,state3,state3

	decb		i
	jnz		.Lpermute

	/* output0 = state0 + copy0 */
	paddd		copy0,state0
	movups		state0,0x00(output)
	/* output1 = state1 + copy1 */
	paddd		copy1,state1
	movups		state1,0x10(output)
	/* output2 = state2 + copy2 */
	paddd		copy2,state2
	movups		state2,0x20(output)
	/* output3 = state3 + copy3 */
	paddd		copy3,state3
	movups		state3,0x30(output)

	/* ++copy3.counter */
	paddq		one,copy3

	/* output += 64, --nblocks */
	addq		$64,output
	decq		nblocks
	jnz		.Lblock

	/* counter = copy3.counter */
	movq		copy3,0x00(counter)

	/* Zero out the potentially sensitive regs, in case nothing uses these again. */
	pxor		state0,state0
	pxor		state1,state1
	pxor		state2,state2
	pxor		state3,state3
	pxor		copy1,copy1
	pxor		copy2,copy2
	pxor		temp,temp

	ret
SYM_FUNC_END(__arch_chacha20_blocks_nostack)