/*

 * Copyright Supranational LLC

 * Licensed under the Apache License, Version 2.0, see LICENSE for details.

 * SPDX-License-Identifier: Apache-2.0

 */

#ifndef __BLS12_381_ASM_SHA256_H__

#define __BLS12_381_ASM_SHA256_H__

#include "vect.h"

#if (defined(__x86_64__) || defined(__x86_64) || defined(_M_X64)) && \

     defined(__SHA__) /* -msha */ && !defined(__BLST_PORTABLE__)

# define sha256_block_data_order blst_sha256_block_data_order_shaext

#elif defined(__aarch64__) && \

      defined(__ARM_FEATURE_CRYPTO) && !defined(__BLST_PORTABLE__)

# define sha256_block_data_order blst_sha256_block_armv8

#else

# define sha256_block_data_order blst_sha256_block_data_order

#endif

#define sha256_hcopy blst_sha256_hcopy

#define sha256_bcopy blst_sha256_bcopy

#define sha256_emit  blst_sha256_emit

void sha256_block_data_order(unsigned int *h, const void *inp, size_t blocks);

void sha256_hcopy(unsigned int dst[8], const unsigned int src[8]);

void sha256_bcopy(void *dst, const void *src, size_t len);

/*

 * If SHA256_CTX conflicts with something, just redefine it to alternative

 * custom name prior including this header.

 */

typedef struct {

    unsigned int h[8];

    unsigned long long N;

    unsigned char buf[64];

    size_t off;

} SHA256_CTX;

static void sha256_init_h(unsigned int h[8])

{

    h[0] = 0x6a09e667U;

    h[1] = 0xbb67ae85U;

    h[2] = 0x3c6ef372U;

    h[3] = 0xa54ff53aU;

    h[4] = 0x510e527fU;

    h[5] = 0x9b05688cU;

    h[6] = 0x1f83d9abU;

    h[7] = 0x5be0cd19U;

}

static void sha256_init(SHA256_CTX *ctx)

{

    sha256_init_h(ctx->h);

    ctx->N = 0;

    vec_zero(ctx->buf, sizeof(ctx->buf));

    ctx->off = 0;

}

static void sha256_update(SHA256_CTX *ctx, const void *_inp, size_t len)

{

    size_t n;

    const unsigned char *inp = _inp;

    ctx->N += len;

    if ((len != 0) & ((n = ctx->off) != 0)) {

        size_t rem = sizeof(ctx->buf) - n;

        if (rem > len) {

            sha256_bcopy(ctx->buf + n, inp, len);

            ctx->off += len;

            return;

        } else {

            sha256_bcopy(ctx->buf + n, inp, rem);

            inp += rem;

            len -= rem;

            sha256_block_data_order(ctx->h, ctx->buf, 1);

            vec_zero(ctx->buf, sizeof(ctx->buf));

            ctx->off = 0;

        }

    }

    n = len / sizeof(ctx->buf);

    if (n > 0) {

        sha256_block_data_order(ctx->h, inp, n);

        n *= sizeof(ctx->buf);

        inp += n;

        len -= n;

    }

    if (len)

        sha256_bcopy(ctx->buf, inp, ctx->off = len);

}

#define __TOBE32(ptr, val) ((ptr)[0] = (unsigned char)((val)>>24), \

                            (ptr)[1] = (unsigned char)((val)>>16), \

                            (ptr)[2] = (unsigned char)((val)>>8),  \

                            (ptr)[3] = (unsigned char)(val))

#if 1

void sha256_emit(unsigned char md[32], const unsigned int h[8]);

#else

static void sha256_emit(unsigned char md[32], const unsigned int h[8])

{

    unsigned int h_i;

    h_i = h[0]; __TOBE32(md + 0, h_i);

    h_i = h[1]; __TOBE32(md + 4, h_i);

    h_i = h[2]; __TOBE32(md + 8, h_i);

    h_i = h[3]; __TOBE32(md + 12, h_i);

    h_i = h[4]; __TOBE32(md + 16, h_i);

    h_i = h[5]; __TOBE32(md + 20, h_i);

    h_i = h[6]; __TOBE32(md + 24, h_i);

    h_i = h[7]; __TOBE32(md + 28, h_i);

}

#endif

static void sha256_final(unsigned char md[32], SHA256_CTX *ctx)

{

    unsigned long long bits = ctx->N * 8;

    size_t n = ctx->off;

    unsigned char *tail;

    ctx->buf[n++] = 0x80;

    if (n > (sizeof(ctx->buf) - 8)) {

        sha256_block_data_order(ctx->h, ctx->buf, 1);

        vec_zero(ctx->buf, sizeof(ctx->buf));

    }

    tail = ctx->buf + sizeof(ctx->buf) - 8;

    __TOBE32(tail, (unsigned int)(bits >> 32));

    __TOBE32(tail + 4, (unsigned int)bits);

    sha256_block_data_order(ctx->h, ctx->buf, 1);

    sha256_emit(md, ctx->h);

}

#undef __TOBE32

#endif