/*      $NetBSD: chacha_neon.c,v 1.9 2023/08/07 01:07:36 rin Exp $      */

/*-
 * Copyright (c) 2020 The NetBSD Foundation, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/types.h>
#include <sys/endian.h>

#include <crypto/arch/arm/arm_neon.h>
#include <crypto/arch/arm/arm_neon_imm.h>
#include "chacha_neon.h"

/*
 * Tempting to use VSHL/VSRI instead of VSHL/VSHR/VORR, but in practice
 * it hurts performance at least on Cortex-A8.
 */
#if 1
#define vrolq_n_u32(x, n)       (vshlq_n_u32(x, n) | vshrq_n_u32(x, 32 - (n)))
#else
#define vrolq_n_u32(x, n)       vsriq_n_u32(vshlq_n_u32(x, n), x, 32 - (n))
#endif

static inline uint32x4_t
rol16(uint32x4_t x)
{
        uint16x8_t y16, x16 = vreinterpretq_u16_u32(x);

        y16 = vrev32q_u16(x16);

        return vreinterpretq_u32_u16(y16);
}

static inline uint32x4_t
rol12(uint32x4_t x)
{

        return vrolq_n_u32(x, 12);
}

static inline uint32x4_t
rol8(uint32x4_t x)
{
#if defined(__aarch64__)
        static const uint8x16_t rol8_tab = VQ_N_U8(
                  3, 0, 1, 2,  7, 4, 5, 6,
                 11, 8, 9,10, 15,12,13,14
        );
        uint8x16_t y8, x8 = vreinterpretq_u8_u32(x);

        y8 = vqtbl1q_u8(x8, rol8_tab);

        return vreinterpretq_u32_u8(y8);
#elif 0
        /*
         * GCC does a lousy job with this, spilling two 64-bit vector
         * registers to the stack every time.  There should be plenty
         * of vector registers free, requiring no spills at all, and
         * GCC should be able to hoist the load of rol8_tab out of any
         * loops, but it doesn't and so attempting to use VTBL hurts
         * more than it helps.
         */
        static const uint8x8_t rol8_tab = V_N_U8(
                 3, 0, 1, 2,  7, 4, 5, 6
        );

        uint64x2_t y64, x64 = vreinterpretq_u64_u32(x);

        y64 = (uint64x2_t) {
                (uint64_t)vtbl1_u8((uint8x8_t)x64[0], rol8_tab),
                (uint64_t)vtbl1_u8((uint8x8_t)x64[1], rol8_tab),
        };

        return vreinterpretq_u32_u64(y64);
#else
        return vrolq_n_u32(x, 8);
#endif
}

static inline uint32x4_t
rol7(uint32x4_t x)
{

        return vrolq_n_u32(x, 7);
}

static inline void
chacha_permute(uint32x4_t *p0, uint32x4_t *p1, uint32x4_t *p2, uint32x4_t *p3,
    unsigned nr)
{
        uint32x4_t r0, r1, r2, r3;
        uint32x4_t c0, c1, c2, c3;

        r0 = *p0;
        r1 = *p1;
        r2 = *p2;
        r3 = *p3;

        for (; nr > 0; nr -= 2) {
                r0 = vaddq_u32(r0, r1); r3 ^= r0; r3 = rol16(r3);
                r2 = vaddq_u32(r2, r3); r1 ^= r2; r1 = rol12(r1);
                r0 = vaddq_u32(r0, r1); r3 ^= r0; r3 = rol8(r3);
                r2 = vaddq_u32(r2, r3); r1 ^= r2; r1 = rol7(r1);

                c0 = r0;
                c1 = vextq_u32(r1, r1, 1);
                c2 = vextq_u32(r2, r2, 2);
                c3 = vextq_u32(r3, r3, 3);

                c0 = vaddq_u32(c0, c1); c3 ^= c0; c3 = rol16(c3);
                c2 = vaddq_u32(c2, c3); c1 ^= c2; c1 = rol12(c1);
                c0 = vaddq_u32(c0, c1); c3 ^= c0; c3 = rol8(c3);
                c2 = vaddq_u32(c2, c3); c1 ^= c2; c1 = rol7(c1);

                r0 = c0;
                r1 = vextq_u32(c1, c1, 3);
                r2 = vextq_u32(c2, c2, 2);
                r3 = vextq_u32(c3, c3, 1);
        }

        *p0 = r0;
        *p1 = r1;
        *p2 = r2;
        *p3 = r3;
}

void
chacha_core_neon(uint8_t out[restrict static 64],
    const uint8_t in[static 16],
    const uint8_t k[static 32],
    const uint8_t c[static 16],
    unsigned nr)
{
        uint32x4_t in0, in1, in2, in3;
        uint32x4_t r0, r1, r2, r3;

        r0 = in0 = vreinterpretq_u32_u8(vld1q_u8(c));
        r1 = in1 = vreinterpretq_u32_u8(vld1q_u8(k + 0));
        r2 = in2 = vreinterpretq_u32_u8(vld1q_u8(k + 16));
        r3 = in3 = vreinterpretq_u32_u8(vld1q_u8(in));

        chacha_permute(&r0, &r1, &r2, &r3, nr);

        vst1q_u8(out + 0, vreinterpretq_u8_u32(vaddq_u32(r0, in0)));
        vst1q_u8(out + 16, vreinterpretq_u8_u32(vaddq_u32(r1, in1)));
        vst1q_u8(out + 32, vreinterpretq_u8_u32(vaddq_u32(r2, in2)));
        vst1q_u8(out + 48, vreinterpretq_u8_u32(vaddq_u32(r3, in3)));
}

void
hchacha_neon(uint8_t out[restrict static 32],
    const uint8_t in[static 16],
    const uint8_t k[static 32],
    const uint8_t c[static 16],
    unsigned nr)
{
        uint32x4_t r0, r1, r2, r3;

        r0 = vreinterpretq_u32_u8(vld1q_u8(c));
        r1 = vreinterpretq_u32_u8(vld1q_u8(k + 0));
        r2 = vreinterpretq_u32_u8(vld1q_u8(k + 16));
        r3 = vreinterpretq_u32_u8(vld1q_u8(in));

        chacha_permute(&r0, &r1, &r2, &r3, nr);

        vst1q_u8(out + 0, vreinterpretq_u8_u32(r0));
        vst1q_u8(out + 16, vreinterpretq_u8_u32(r3));
}

void
chacha_stream_neon(uint8_t *restrict s, size_t n,
    uint32_t blkno,
    const uint8_t nonce[static 12],
    const uint8_t k[static 32],
    unsigned nr)
{

        for (; n >= 256; s += 256, n -= 256, blkno += 4)
                chacha_stream256_neon(s, blkno, nonce, k, chacha_const32, nr);

        if (n) {
                const uint32x4_t blkno_inc = /* (1,0,0,0) */
                    vsetq_lane_u32(1, vdupq_n_u32(0), 0);
                uint32x4_t in0, in1, in2, in3;
                uint32x4_t r0, r1, r2, r3;

                in0 = vreinterpretq_u32_u8(vld1q_u8(chacha_const32));
                in1 = vreinterpretq_u32_u8(vld1q_u8(k + 0));
                in2 = vreinterpretq_u32_u8(vld1q_u8(k + 16));
                in3 = (uint32x4_t) VQ_N_U32(
                        blkno,
                        le32dec(nonce),
                        le32dec(nonce + 4),
                        le32dec(nonce + 8)
                );

                for (; n; s += 64, n -= 64) {
                        r0 = in0;
                        r1 = in1;
                        r2 = in2;
                        r3 = in3;
                        chacha_permute(&r0, &r1, &r2, &r3, nr);
                        r0 = vaddq_u32(r0, in0);
                        r1 = vaddq_u32(r1, in1);
                        r2 = vaddq_u32(r2, in2);
                        r3 = vaddq_u32(r3, in3);

                        if (n < 64) {
                                uint8_t buf[64] __aligned(16);

                                vst1q_u8(buf + 0, vreinterpretq_u8_u32(r0));
                                vst1q_u8(buf + 16, vreinterpretq_u8_u32(r1));
                                vst1q_u8(buf + 32, vreinterpretq_u8_u32(r2));
                                vst1q_u8(buf + 48, vreinterpretq_u8_u32(r3));
                                memcpy(s, buf, n);

                                break;
                        }

                        vst1q_u8(s + 0, vreinterpretq_u8_u32(r0));
                        vst1q_u8(s + 16, vreinterpretq_u8_u32(r1));
                        vst1q_u8(s + 32, vreinterpretq_u8_u32(r2));
                        vst1q_u8(s + 48, vreinterpretq_u8_u32(r3));
                        in3 = vaddq_u32(in3, blkno_inc);
                }
        }
}

void
chacha_stream_xor_neon(uint8_t *s, const uint8_t *p, size_t n,
    uint32_t blkno,
    const uint8_t nonce[static 12],
    const uint8_t k[static 32],
    unsigned nr)
{

        for (; n >= 256; s += 256, p += 256, n -= 256, blkno += 4)
                chacha_stream_xor256_neon(s, p, blkno, nonce, k,
                    chacha_const32, nr);

        if (n) {
                const uint32x4_t blkno_inc = /* (1,0,0,0) */
                    vsetq_lane_u32(1, vdupq_n_u32(0), 0);
                uint32x4_t in0, in1, in2, in3;
                uint32x4_t r0, r1, r2, r3;

                in0 = vreinterpretq_u32_u8(vld1q_u8(chacha_const32));
                in1 = vreinterpretq_u32_u8(vld1q_u8(k + 0));
                in2 = vreinterpretq_u32_u8(vld1q_u8(k + 16));
                in3 = (uint32x4_t) VQ_N_U32(
                        blkno,
                        le32dec(nonce),
                        le32dec(nonce + 4),
                        le32dec(nonce + 8)
                );

                for (; n; s += 64, p += 64, n -= 64) {
                        r0 = in0;
                        r1 = in1;
                        r2 = in2;
                        r3 = in3;
                        chacha_permute(&r0, &r1, &r2, &r3, nr);
                        r0 = vaddq_u32(r0, in0);
                        r1 = vaddq_u32(r1, in1);
                        r2 = vaddq_u32(r2, in2);
                        r3 = vaddq_u32(r3, in3);

                        if (n < 64) {
                                uint8_t buf[64] __aligned(16);
                                unsigned i;

                                vst1q_u8(buf + 0, vreinterpretq_u8_u32(r0));
                                vst1q_u8(buf + 16, vreinterpretq_u8_u32(r1));
                                vst1q_u8(buf + 32, vreinterpretq_u8_u32(r2));
                                vst1q_u8(buf + 48, vreinterpretq_u8_u32(r3));

                                for (i = 0; i < n - n%4; i += 4)
                                        le32enc(s + i,
                                            le32dec(p + i) ^ le32dec(buf + i));
                                for (; i < n; i++)
                                        s[i] = p[i] ^ buf[i];

                                break;
                        }

                        r0 ^= vreinterpretq_u32_u8(vld1q_u8(p + 0));
                        r1 ^= vreinterpretq_u32_u8(vld1q_u8(p + 16));
                        r2 ^= vreinterpretq_u32_u8(vld1q_u8(p + 32));
                        r3 ^= vreinterpretq_u32_u8(vld1q_u8(p + 48));
                        vst1q_u8(s + 0, vreinterpretq_u8_u32(r0));
                        vst1q_u8(s + 16, vreinterpretq_u8_u32(r1));
                        vst1q_u8(s + 32, vreinterpretq_u8_u32(r2));
                        vst1q_u8(s + 48, vreinterpretq_u8_u32(r3));
                        in3 = vaddq_u32(in3, blkno_inc);
                }
        }
}

void
xchacha_stream_neon(uint8_t *restrict s, size_t nbytes,
    uint32_t blkno,
    const uint8_t nonce[static 24],
    const uint8_t k[static 32],
    unsigned nr)
{
        uint8_t subkey[32];
        uint8_t subnonce[12];

        hchacha_neon(subkey, nonce/*[0:16)*/, k, chacha_const32, nr);
        memset(subnonce, 0, 4);
        memcpy(subnonce + 4, nonce + 16, 8);
        chacha_stream_neon(s, nbytes, blkno, subnonce, subkey, nr);
}

void
xchacha_stream_xor_neon(uint8_t *restrict c, const uint8_t *p, size_t nbytes,
    uint32_t blkno,
    const uint8_t nonce[static 24],
    const uint8_t k[static 32],
    unsigned nr)
{
        uint8_t subkey[32];
        uint8_t subnonce[12];

        hchacha_neon(subkey, nonce/*[0:16)*/, k, chacha_const32, nr);
        memset(subnonce, 0, 4);
        memcpy(subnonce + 4, nonce + 16, 8);
        chacha_stream_xor_neon(c, p, nbytes, blkno, subnonce, subkey, nr);
}