/* SPDX-License-Identifier: GPL-2.0 OR MIT */
/*
 * Copyright (C) 2016-2018 René van Dorst <opensource@vdorst.com>. All Rights Reserved.
 * Copyright (C) 2015-2019 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 */

#define MASK_U32		0x3c
#define CHACHA20_BLOCK_SIZE	64
#define STACK_SIZE		32

#define X0	$t0
#define X1	$t1
#define X2	$t2
#define X3	$t3
#define X4	$t4
#define X5	$t5
#define X6	$t6
#define X7	$t7
#define X8	$t8
#define X9	$t9
#define X10	$v1
#define X11	$s6
#define X12	$s5
#define X13	$s4
#define X14	$s3
#define X15	$s2
/* Use regs which are overwritten on exit for Tx so we don't leak clear data. */
#define T0	$s1
#define T1	$s0
#define T(n)	T ## n
#define X(n)	X ## n

/* Input arguments */
#define STATE		$a0
#define OUT		$a1
#define IN		$a2
#define BYTES		$a3

/* Output argument */
/* NONCE[0] is kept in a register and not in memory.
 * We don't want to touch original value in memory.
 * Must be incremented every loop iteration.
 */
#define NONCE_0		$v0

/* SAVED_X and SAVED_CA are set in the jump table.
 * Use regs which are overwritten on exit else we don't leak clear data.
 * They are used to handling the last bytes which are not multiple of 4.
 */
#define SAVED_X		X15
#define SAVED_CA	$s7

#define IS_UNALIGNED	$s7

#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
#define MSB 0
#define LSB 3
#define	CPU_TO_LE32(n) \
	wsbh	n, n; \
	rotr	n, 16;
#else
#define MSB 3
#define LSB 0
#define CPU_TO_LE32(n)
#endif

#define FOR_EACH_WORD(x) \
	x( 0); \
	x( 1); \
	x( 2); \
	x( 3); \
	x( 4); \
	x( 5); \
	x( 6); \
	x( 7); \
	x( 8); \
	x( 9); \
	x(10); \
	x(11); \
	x(12); \
	x(13); \
	x(14); \
	x(15);

#define FOR_EACH_WORD_REV(x) \
	x(15); \
	x(14); \
	x(13); \
	x(12); \
	x(11); \
	x(10); \
	x( 9); \
	x( 8); \
	x( 7); \
	x( 6); \
	x( 5); \
	x( 4); \
	x( 3); \
	x( 2); \
	x( 1); \
	x( 0);

#define PLUS_ONE_0	 1
#define PLUS_ONE_1	 2
#define PLUS_ONE_2	 3
#define PLUS_ONE_3	 4
#define PLUS_ONE_4	 5
#define PLUS_ONE_5	 6
#define PLUS_ONE_6	 7
#define PLUS_ONE_7	 8
#define PLUS_ONE_8	 9
#define PLUS_ONE_9	10
#define PLUS_ONE_10	11
#define PLUS_ONE_11	12
#define PLUS_ONE_12	13
#define PLUS_ONE_13	14
#define PLUS_ONE_14	15
#define PLUS_ONE_15	16
#define PLUS_ONE(x)	PLUS_ONE_ ## x
#define _CONCAT3(a,b,c)	a ## b ## c
#define CONCAT3(a,b,c)	_CONCAT3(a,b,c)

#define STORE_UNALIGNED(x) \
CONCAT3(.Lchacha_mips_xor_unaligned_, PLUS_ONE(x), _b: ;) \
	.if (x != 12); \
		lw	T0, (x*4)(STATE); \
	.endif; \
	lwl	T1, (x*4)+MSB ## (IN); \
	lwr	T1, (x*4)+LSB ## (IN); \
	.if (x == 12); \
		addu	X ## x, NONCE_0; \
	.else; \
		addu	X ## x, T0; \
	.endif; \
	CPU_TO_LE32(X ## x); \
	xor	X ## x, T1; \
	swl	X ## x, (x*4)+MSB ## (OUT); \
	swr	X ## x, (x*4)+LSB ## (OUT);

#define STORE_ALIGNED(x) \
CONCAT3(.Lchacha_mips_xor_aligned_, PLUS_ONE(x), _b: ;) \
	.if (x != 12); \
		lw	T0, (x*4)(STATE); \
	.endif; \
	lw	T1, (x*4) ## (IN); \
	.if (x == 12); \
		addu	X ## x, NONCE_0; \
	.else; \
		addu	X ## x, T0; \
	.endif; \
	CPU_TO_LE32(X ## x); \
	xor	X ## x, T1; \
	sw	X ## x, (x*4) ## (OUT);

/* Jump table macro.
 * Used for setup and handling the last bytes, which are not multiple of 4.
 * X15 is free to store Xn
 * Every jumptable entry must be equal in size.
 */
#define JMPTBL_ALIGNED(x) \
.Lchacha_mips_jmptbl_aligned_ ## x: ; \
	.set	noreorder; \
	b	.Lchacha_mips_xor_aligned_ ## x ## _b; \
	.if (x == 12); \
		addu	SAVED_X, X ## x, NONCE_0; \
	.else; \
		addu	SAVED_X, X ## x, SAVED_CA; \
	.endif; \
	.set	reorder

#define JMPTBL_UNALIGNED(x) \
.Lchacha_mips_jmptbl_unaligned_ ## x: ; \
	.set	noreorder; \
	b	.Lchacha_mips_xor_unaligned_ ## x ## _b; \
	.if (x == 12); \
		addu	SAVED_X, X ## x, NONCE_0; \
	.else; \
		addu	SAVED_X, X ## x, SAVED_CA; \
	.endif; \
	.set	reorder

#define AXR(A, B, C, D,  K, L, M, N,  V, W, Y, Z,  S) \
	addu	X(A), X(K); \
	addu	X(B), X(L); \
	addu	X(C), X(M); \
	addu	X(D), X(N); \
	xor	X(V), X(A); \
	xor	X(W), X(B); \
	xor	X(Y), X(C); \
	xor	X(Z), X(D); \
	rotr	X(V), 32 - S; \
	rotr	X(W), 32 - S; \
	rotr	X(Y), 32 - S; \
	rotr	X(Z), 32 - S;

.text
.set	reorder
.set	noat
.globl	chacha_crypt_arch
.ent	chacha_crypt_arch
chacha_crypt_arch:
	.frame	$sp, STACK_SIZE, $ra

	/* Load number of rounds */
	lw	$at, 16($sp)

	addiu	$sp, -STACK_SIZE

	/* Return bytes = 0. */
	beqz	BYTES, .Lchacha_mips_end

	lw	NONCE_0, 48(STATE)

	/* Save s0-s7 */
	sw	$s0,  0($sp)
	sw	$s1,  4($sp)
	sw	$s2,  8($sp)
	sw	$s3, 12($sp)
	sw	$s4, 16($sp)
	sw	$s5, 20($sp)
	sw	$s6, 24($sp)
	sw	$s7, 28($sp)

	/* Test IN or OUT is unaligned.
	 * IS_UNALIGNED = ( IN | OUT ) & 0x00000003
	 */
	or	IS_UNALIGNED, IN, OUT
	andi	IS_UNALIGNED, 0x3

	b	.Lchacha_rounds_start

.align 4
.Loop_chacha_rounds:
	addiu	IN,  CHACHA20_BLOCK_SIZE
	addiu	OUT, CHACHA20_BLOCK_SIZE
	addiu	NONCE_0, 1

.Lchacha_rounds_start:
	lw	X0,  0(STATE)
	lw	X1,  4(STATE)
	lw	X2,  8(STATE)
	lw	X3,  12(STATE)

	lw	X4,  16(STATE)
	lw	X5,  20(STATE)
	lw	X6,  24(STATE)
	lw	X7,  28(STATE)
	lw	X8,  32(STATE)
	lw	X9,  36(STATE)
	lw	X10, 40(STATE)
	lw	X11, 44(STATE)

	move	X12, NONCE_0
	lw	X13, 52(STATE)
	lw	X14, 56(STATE)
	lw	X15, 60(STATE)

.Loop_chacha_xor_rounds:
	addiu	$at, -2
	AXR( 0, 1, 2, 3,  4, 5, 6, 7, 12,13,14,15, 16);
	AXR( 8, 9,10,11, 12,13,14,15,  4, 5, 6, 7, 12);
	AXR( 0, 1, 2, 3,  4, 5, 6, 7, 12,13,14,15,  8);
	AXR( 8, 9,10,11, 12,13,14,15,  4, 5, 6, 7,  7);
	AXR( 0, 1, 2, 3,  5, 6, 7, 4, 15,12,13,14, 16);
	AXR(10,11, 8, 9, 15,12,13,14,  5, 6, 7, 4, 12);
	AXR( 0, 1, 2, 3,  5, 6, 7, 4, 15,12,13,14,  8);
	AXR(10,11, 8, 9, 15,12,13,14,  5, 6, 7, 4,  7);
	bnez	$at, .Loop_chacha_xor_rounds

	addiu	BYTES, -(CHACHA20_BLOCK_SIZE)

	/* Is data src/dst unaligned? Jump */
	bnez	IS_UNALIGNED, .Loop_chacha_unaligned

	/* Set number rounds here to fill delayslot. */
	lw	$at, (STACK_SIZE+16)($sp)

	/* BYTES < 0, it has no full block. */
	bltz	BYTES, .Lchacha_mips_no_full_block_aligned

	FOR_EACH_WORD_REV(STORE_ALIGNED)

	/* BYTES > 0? Loop again. */
	bgtz	BYTES, .Loop_chacha_rounds

	/* Place this here to fill delay slot */
	addiu	NONCE_0, 1

	/* BYTES < 0? Handle last bytes */
	bltz	BYTES, .Lchacha_mips_xor_bytes

.Lchacha_mips_xor_done:
	/* Restore used registers */
	lw	$s0,  0($sp)
	lw	$s1,  4($sp)
	lw	$s2,  8($sp)
	lw	$s3, 12($sp)
	lw	$s4, 16($sp)
	lw	$s5, 20($sp)
	lw	$s6, 24($sp)
	lw	$s7, 28($sp)

	/* Write NONCE_0 back to right location in state */
	sw	NONCE_0, 48(STATE)

.Lchacha_mips_end:
	addiu	$sp, STACK_SIZE
	jr	$ra

.Lchacha_mips_no_full_block_aligned:
	/* Restore the offset on BYTES */
	addiu	BYTES, CHACHA20_BLOCK_SIZE

	/* Get number of full WORDS */
	andi	$at, BYTES, MASK_U32

	/* Load upper half of jump table addr */
	lui	T0, %hi(.Lchacha_mips_jmptbl_aligned_0)

	/* Calculate lower half jump table offset */
	ins	T0, $at, 1, 6

	/* Add offset to STATE */
	addu	T1, STATE, $at

	/* Add lower half jump table addr */
	addiu	T0, %lo(.Lchacha_mips_jmptbl_aligned_0)

	/* Read value from STATE */
	lw	SAVED_CA, 0(T1)

	/* Store remaining bytecounter as negative value */
	subu	BYTES, $at, BYTES

	jr	T0

	/* Jump table */
	FOR_EACH_WORD(JMPTBL_ALIGNED)


.Loop_chacha_unaligned:
	/* Set number rounds here to fill delayslot. */
	lw	$at, (STACK_SIZE+16)($sp)

	/* BYTES > 0, it has no full block. */
	bltz	BYTES, .Lchacha_mips_no_full_block_unaligned

	FOR_EACH_WORD_REV(STORE_UNALIGNED)

	/* BYTES > 0? Loop again. */
	bgtz	BYTES, .Loop_chacha_rounds

	/* Write NONCE_0 back to right location in state */
	sw	NONCE_0, 48(STATE)

	.set noreorder
	/* Fall through to byte handling */
	bgez	BYTES, .Lchacha_mips_xor_done
.Lchacha_mips_xor_unaligned_0_b:
.Lchacha_mips_xor_aligned_0_b:
	/* Place this here to fill delay slot */
	addiu	NONCE_0, 1
	.set reorder

.Lchacha_mips_xor_bytes:
	addu	IN, $at
	addu	OUT, $at
	/* First byte */
	lbu	T1, 0(IN)
	addiu	$at, BYTES, 1
	xor	T1, SAVED_X
	sb	T1, 0(OUT)
	beqz	$at, .Lchacha_mips_xor_done
	/* Second byte */
	lbu	T1, 1(IN)
	addiu	$at, BYTES, 2
	rotr	SAVED_X, 8
	xor	T1, SAVED_X
	sb	T1, 1(OUT)
	beqz	$at, .Lchacha_mips_xor_done
	/* Third byte */
	lbu	T1, 2(IN)
	rotr	SAVED_X, 8
	xor	T1, SAVED_X
	sb	T1, 2(OUT)
	b	.Lchacha_mips_xor_done

.Lchacha_mips_no_full_block_unaligned:
	/* Restore the offset on BYTES */
	addiu	BYTES, CHACHA20_BLOCK_SIZE

	/* Get number of full WORDS */
	andi	$at, BYTES, MASK_U32

	/* Load upper half of jump table addr */
	lui	T0, %hi(.Lchacha_mips_jmptbl_unaligned_0)

	/* Calculate lower half jump table offset */
	ins	T0, $at, 1, 6

	/* Add offset to STATE */
	addu	T1, STATE, $at

	/* Add lower half jump table addr */
	addiu	T0, %lo(.Lchacha_mips_jmptbl_unaligned_0)

	/* Read value from STATE */
	lw	SAVED_CA, 0(T1)

	/* Store remaining bytecounter as negative value */
	subu	BYTES, $at, BYTES

	jr	T0

	/* Jump table */
	FOR_EACH_WORD(JMPTBL_UNALIGNED)
.end chacha_crypt_arch
.set at

/* Input arguments
 * STATE	$a0
 * OUT		$a1
 * NROUND	$a2
 */

#undef X12
#undef X13
#undef X14
#undef X15

#define X12	$a3
#define X13	$at
#define X14	$v0
#define X15	STATE

.set noat
.globl	hchacha_block_arch
.ent	hchacha_block_arch
hchacha_block_arch:
	.frame	$sp, STACK_SIZE, $ra

	addiu	$sp, -STACK_SIZE

	/* Save X11(s6) */
	sw	X11, 0($sp)

	lw	X0,  0(STATE)
	lw	X1,  4(STATE)
	lw	X2,  8(STATE)
	lw	X3,  12(STATE)
	lw	X4,  16(STATE)
	lw	X5,  20(STATE)
	lw	X6,  24(STATE)
	lw	X7,  28(STATE)
	lw	X8,  32(STATE)
	lw	X9,  36(STATE)
	lw	X10, 40(STATE)
	lw	X11, 44(STATE)
	lw	X12, 48(STATE)
	lw	X13, 52(STATE)
	lw	X14, 56(STATE)
	lw	X15, 60(STATE)

.Loop_hchacha_xor_rounds:
	addiu	$a2, -2
	AXR( 0, 1, 2, 3,  4, 5, 6, 7, 12,13,14,15, 16);
	AXR( 8, 9,10,11, 12,13,14,15,  4, 5, 6, 7, 12);
	AXR( 0, 1, 2, 3,  4, 5, 6, 7, 12,13,14,15,  8);
	AXR( 8, 9,10,11, 12,13,14,15,  4, 5, 6, 7,  7);
	AXR( 0, 1, 2, 3,  5, 6, 7, 4, 15,12,13,14, 16);
	AXR(10,11, 8, 9, 15,12,13,14,  5, 6, 7, 4, 12);
	AXR( 0, 1, 2, 3,  5, 6, 7, 4, 15,12,13,14,  8);
	AXR(10,11, 8, 9, 15,12,13,14,  5, 6, 7, 4,  7);
	bnez	$a2, .Loop_hchacha_xor_rounds

	/* Restore used register */
	lw	X11, 0($sp)

	sw	X0,  0(OUT)
	sw	X1,  4(OUT)
	sw	X2,  8(OUT)
	sw	X3,  12(OUT)
	sw	X12, 16(OUT)
	sw	X13, 20(OUT)
	sw	X14, 24(OUT)
	sw	X15, 28(OUT)

	addiu	$sp, STACK_SIZE
	jr	$ra
.end hchacha_block_arch
.set at