/*

 * Copyright (c) 2025 Fastly, Kazuho Oku

 *

 * Permission is hereby granted, free of charge, to any person obtaining a copy

 * of this software and associated documentation files (the "Software"), to

 * deal in the Software without restriction, including without limitation the

 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or

 * sell copies of the Software, and to permit persons to whom the Software is

 * furnished to do so, subject to the following conditions:

 *

 * The above copyright notice and this permission notice shall be included in

 * all copies or substantial portions of the Software.

 *

 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING

 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS

 * IN THE SOFTWARE.

 */

#ifndef picotls_quiclb_h

#define picotls_quiclb_h

#if defined(__x86_64__) || defined(_M_X64)

#include <emmintrin.h>

#define PICOTLS_QUICLB_HAVE_SSE2 1

#endif

union picotls_quiclb_block {

    uint8_t bytes[PTLS_AES_BLOCK_SIZE];

    uint64_t u64[PTLS_AES_BLOCK_SIZE / sizeof(uint64_t)];

#if PICOTLS_QUICLB_HAVE_SSE2

    __m128i m128;

#endif

};

/**

 * encrypts one block of AES, assuming the context is `ptls_cipher_context_t` backed by ptls_foo_aes128ecb

 */

static inline void picotls_quiclb_cipher_aes(void *aesecb, union picotls_quiclb_block *block)

{

    ptls_cipher_encrypt(aesecb, block->bytes, block->bytes, PTLS_AES_BLOCK_SIZE);

}

/**

 * calculates X ^ AES(mask_and_expand(Y))

 */

static inline void picotls_quiclb_one_round(void (*aesecb_func)(void *aesecb, union picotls_quiclb_block *), void *aesecb_ctx,

                                            union picotls_quiclb_block *dest, const union picotls_quiclb_block *x,

                                            const union picotls_quiclb_block *y, const union picotls_quiclb_block *mask,

                                            const union picotls_quiclb_block *len_pass)

{

#if PICOTLS_QUICLB_HAVE_SSE2

    dest->m128 = _mm_or_si128(_mm_and_si128(y->m128, mask->m128), len_pass->m128);

#else

    for (size_t i = 0; i < PTLS_ELEMENTSOF(dest->u64); ++i)

        dest->u64[i] = (y->u64[i] & mask->u64[i]) | len_pass->u64[i];

#endif

    aesecb_func(aesecb_ctx, dest);

#if PICOTLS_QUICLB_HAVE_SSE2

    dest->m128 = _mm_xor_si128(dest->m128, x->m128);

#else

    for (size_t i = 0; i < PTLS_ELEMENTSOF(dest->u64); ++i)

        dest->u64[i] ^= x->u64[i];

#endif

}

static inline void picotls_quiclb_split_input(union picotls_quiclb_block *l, union picotls_quiclb_block *r, const uint8_t *input,

                                              size_t len)

{

    size_t i;

    for (i = 0; i < (len + 1) / 2; ++i)

        l->bytes[i] = input[i];

    for (; i < PTLS_ELEMENTSOF(l->bytes); ++i)

        l->bytes[i] = 0;

    for (i = 0; i < (len + 1) / 2; ++i)

        r->bytes[i] = input[i + len / 2];

    for (; i < PTLS_ELEMENTSOF(r->bytes); ++i)

        r->bytes[i] = 0;

}

static inline void picotls_quiclb_merge_output(uint8_t *output, size_t len, const union picotls_quiclb_block *l,

                                               const union picotls_quiclb_block *r)

{

    uint8_t *outp = output;

    for (size_t i = 0; i < len / 2; ++i)

        *outp++ = l->bytes[i];

    if (len % 2 == 0) {

        for (size_t i = 0; i < len / 2; ++i)

            *outp++ = r->bytes[i];

    } else {

        *outp++ = (l->bytes[len / 2] & 0xf0) | (r->bytes[0] & 0x0f);

        for (size_t i = 0; i < len / 2; ++i)

            *outp++ = r->bytes[i + 1];

    }

}

static inline void picotls_quiclb_do_init(ptls_cipher_context_t *ctx, const void *iv)

{

    /* no-op */

}

static inline void picotls_quiclb_transform(void (*aesecb_func)(void *aesecb, union picotls_quiclb_block *), void *aesecb_ctx,

                                            void *output, const void *input, size_t len, int encrypt)

{

    static const struct quiclb_mask_t {

        union picotls_quiclb_block l, r;

    } masks[] = {

        {{{0xff, 0xff, 0xff, 0xf0}}, {{0x0f, 0xff, 0xff, 0xff}}},                                                 /* 7 (MIN_LEN) */

        {{{0xff, 0xff, 0xff, 0xff}}, {{0xff, 0xff, 0xff, 0xff}}},                                                 /* 8 */

        {{{0xff, 0xff, 0xff, 0xff, 0xf0}}, {{0x0f, 0xff, 0xff, 0xff, 0xff}}},                                     /* 9 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff}}, {{0xff, 0xff, 0xff, 0xff, 0xff}}},                                     /* 10 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xf0}}, {{0x0f, 0xff, 0xff, 0xff, 0xff, 0xff}}},                         /* 11 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}, {{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}},                         /* 12 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xf0}}, {{0x0f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}},             /* 13 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}, {{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}},             /* 14 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xf0}}, {{0x0f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}}, /* 15 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}, {{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}}, /* 16 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xf0}},

         {{0x0f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}}, /* 17 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}},

         {{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}}, /* 18 */

        {{{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xf0}},

         {{0x0f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}}} /* 19 */

    };

    assert(PTLS_QUICLB_MIN_BLOCK_SIZE <= len && len <= PTLS_QUICLB_MAX_BLOCK_SIZE);

    PTLS_BUILD_ASSERT(PTLS_QUICLB_MAX_BLOCK_SIZE == PTLS_QUICLB_MIN_BLOCK_SIZE + PTLS_ELEMENTSOF(masks) - 1);

    const struct quiclb_mask_t *mask = &masks[len - PTLS_QUICLB_MIN_BLOCK_SIZE];

    union picotls_quiclb_block l0, r0, r1, l1, r2, l2, len_pass = {{0}};

    len_pass.bytes[14] = (uint8_t)len;

#define ROUND(rnd, dest, x, y, mask_side)                                                                                          \

    do {                                                                                                                           \

        len_pass.bytes[15] = (rnd);                                                                                                \

        picotls_quiclb_one_round(aesecb_func, aesecb_ctx, &dest, &x, &y, &mask->mask_side, &len_pass);                             \

    } while (0)

    if (encrypt) {

        picotls_quiclb_split_input(&l0, &r0, input, len);

        ROUND(1, r1, r0, l0, l);

        ROUND(2, l1, l0, r1, r);

        ROUND(3, r2, r1, l1, l);

        ROUND(4, l2, l1, r2, r);

        picotls_quiclb_merge_output(output, len, &l2, &r2);

    } else {

        picotls_quiclb_split_input(&l2, &r2, input, len);

        ROUND(4, l1, l2, r2, r);

        ROUND(3, r1, r2, l1, l);

        ROUND(2, l0, l1, r1, r);

        ROUND(1, r0, r1, l0, l);

        picotls_quiclb_merge_output(output, len, &l0, &r0);

    }

#undef ROUND

}

#endif