// 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)