// Copyright (c) 2017 Instituto de Pesquisas Eldorado. All rights reserved.

// DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

//

// This code is free software; you can redistribute it and/or modify it

// under the terms of the GNU General Public License version 2 only, as

// published by the Free Software Foundation.

//

// This code is distributed in the hope that it will be useful, but WITHOUT

// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

// FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

// version 2 for more details (a copy is included in the LICENSE file that

// accompanied this code).

//

// You should have received a copy of the GNU General Public License version

// 2 along with this work; if not, write to the Free Software Foundation,

// Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

//

// Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

// or visit www.oracle.com if you need additional information or have any

// questions.

// Implemented according to "Descriptions of SHA-256, SHA-384, and SHA-512"

// (http://www.iwar.org.uk/comsec/resources/cipher/sha256-384-512.pdf).

#include "asm/macroAssembler.inline.hpp"

#include "runtime/stubRoutines.hpp"

/**********************************************************************

 * SHA 256

 *********************************************************************/

void MacroAssembler::sha256_deque(const VectorRegister src,

                                  const VectorRegister dst1,

                                  const VectorRegister dst2,

                                  const VectorRegister dst3) {

  vsldoi (dst1, src, src, 12);

  vsldoi (dst2, src, src, 8);

  vsldoi (dst3, src, src, 4);

}

void MacroAssembler::sha256_round(const VectorRegister* hs,

                                  const int total_hs,

                                  int& h_cnt,

                                  const VectorRegister kpw) {

  // convenience registers: cycle from 0-7 downwards

  const VectorRegister a = hs[(total_hs + 0 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister b = hs[(total_hs + 1 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister c = hs[(total_hs + 2 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister d = hs[(total_hs + 3 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister e = hs[(total_hs + 4 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister f = hs[(total_hs + 5 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister g = hs[(total_hs + 6 - (h_cnt % total_hs)) % total_hs];

  const VectorRegister h = hs[(total_hs + 7 - (h_cnt % total_hs)) % total_hs];

  // temporaries

  VectorRegister ch  = VR0;

  VectorRegister maj = VR1;

  VectorRegister bsa = VR2;

  VectorRegister bse = VR3;

  VectorRegister vt0 = VR4;

  VectorRegister vt1 = VR5;

  VectorRegister vt2 = VR6;

  VectorRegister vt3 = VR7;

  vsel       (ch,  g,   f, e);

  vxor       (maj, a,   b);

  vshasigmaw (bse, e,   1, 0xf);

  vadduwm    (vt2, ch,  kpw);

  vadduwm    (vt1, h,   bse);

  vsel       (maj, b,   c, maj);

  vadduwm    (vt3, vt1, vt2);

  vshasigmaw (bsa, a,   1, 0);

  vadduwm    (vt0, bsa, maj);

  vadduwm    (d,   d,   vt3);

  vadduwm    (h,   vt3, vt0);

  // advance vector pointer to the next iteration

  h_cnt++;

}

void MacroAssembler::sha256_load_h_vec(const VectorRegister a,

                                       const VectorRegister e,

                                       const Register hptr) {

  // temporaries

  Register tmp = R8;

  VectorRegister vt0 = VR0;

  VectorRegister vRb = VR6;

  // labels

  Label sha256_aligned;

  andi_  (tmp,  hptr, 0xf);

  lvx    (a,    hptr);

  addi   (tmp,  hptr, 16);

  lvx    (e,    tmp);

  beq    (CCR0, sha256_aligned);

  // handle unaligned accesses

  load_perm(vRb, hptr);

  addi   (tmp, hptr, 32);

  vec_perm(a,   e,    vRb);

  lvx    (vt0,  tmp);

  vec_perm(e,   vt0,  vRb);

  // aligned accesses

  bind(sha256_aligned);

}

void MacroAssembler::sha256_load_w_plus_k_vec(const Register buf_in,

                                              const VectorRegister* ws,

                                              const int total_ws,

                                              const Register k,

                                              const VectorRegister* kpws,

                                              const int total_kpws) {

  Label w_aligned, after_w_load;

  Register tmp       = R8;

  VectorRegister vt0 = VR0;

  VectorRegister vt1 = VR1;

  VectorRegister vRb = VR6;

  andi_ (tmp, buf_in, 0xF);

  beq   (CCR0, w_aligned); // address ends with 0x0, not 0x8

  // deal with unaligned addresses

  lvx    (ws[0], buf_in);

  load_perm(vRb, buf_in);

  for (int n = 1; n < total_ws; n++) {

    VectorRegister w_cur = ws[n];

    VectorRegister w_prev = ws[n-1];

    addi (tmp, buf_in, n * 16);

    lvx  (w_cur, tmp);

    vec_perm(w_prev, w_cur, vRb);

  }

  addi   (tmp, buf_in, total_ws * 16);

  lvx    (vt0, tmp);

  vec_perm(ws[total_ws-1], vt0, vRb);

  b      (after_w_load);

  bind(w_aligned);

  // deal with aligned addresses

  lvx(ws[0], buf_in);

  for (int n = 1; n < total_ws; n++) {

    VectorRegister w = ws[n];

    addi (tmp, buf_in, n * 16);

    lvx  (w, tmp);

  }

  bind(after_w_load);

#if defined(VM_LITTLE_ENDIAN)

  // Byte swapping within int values

  li       (tmp, 8);

  lvsl     (vt0, tmp);

  vspltisb (vt1, 0xb);

  vxor     (vt1, vt0, vt1);

  for (int n = 0; n < total_ws; n++) {

    VectorRegister w = ws[n];

    vec_perm(w, w, vt1);

  }

#endif

  // Loading k, which is always aligned to 16-bytes

  lvx    (kpws[0], k);

  for (int n = 1; n < total_kpws; n++) {

    VectorRegister kpw = kpws[n];

    addi (tmp, k, 16 * n);

    lvx  (kpw, tmp);

  }

  // Add w to K

  assert(total_ws == total_kpws, "Redesign the loop below");

  for (int n = 0; n < total_kpws; n++) {

    VectorRegister kpw = kpws[n];

    VectorRegister w   = ws[n];

    vadduwm  (kpw, kpw, w);

  }

}

void MacroAssembler::sha256_calc_4w(const VectorRegister w0,

                                    const VectorRegister w1,

                                    const VectorRegister w2,

                                    const VectorRegister w3,

                                    const VectorRegister kpw0,

                                    const VectorRegister kpw1,

                                    const VectorRegister kpw2,

                                    const VectorRegister kpw3,

                                    const Register j,

                                    const Register k) {

  // Temporaries

  const VectorRegister  vt0  = VR0;

  const VectorRegister  vt1  = VR1;

  const VectorSRegister vsrt1 = vt1->to_vsr();

  const VectorRegister  vt2  = VR2;

  const VectorRegister  vt3  = VR3;

  const VectorSRegister vst3 = vt3->to_vsr();

  const VectorRegister  vt4  = VR4;

  // load to k[j]

  lvx        (vt0, j,   k);

  // advance j

  addi       (j,   j,   16); // 16 bytes were read

#if defined(VM_LITTLE_ENDIAN)

  // b = w[j-15], w[j-14], w[j-13], w[j-12]

  vsldoi     (vt1, w1,  w0, 12);

  // c = w[j-7], w[j-6], w[j-5], w[j-4]

  vsldoi     (vt2, w3,  w2, 12);

#else

  // b = w[j-15], w[j-14], w[j-13], w[j-12]

  vsldoi     (vt1, w0,  w1, 4);

  // c = w[j-7], w[j-6], w[j-5], w[j-4]

  vsldoi     (vt2, w2,  w3, 4);

#endif

  // d = w[j-2], w[j-1], w[j-4], w[j-3]

  vsldoi     (vt3, w3,  w3, 8);

  // b = s0(w[j-15]) , s0(w[j-14]) , s0(w[j-13]) , s0(w[j-12])

  vshasigmaw (vt1, vt1, 0,  0);

  // d = s1(w[j-2]) , s1(w[j-1]) , s1(w[j-4]) , s1(w[j-3])

  vshasigmaw (vt3, vt3, 0,  0xf);

  // c = s0(w[j-15]) + w[j-7],

  //     s0(w[j-14]) + w[j-6],

  //     s0(w[j-13]) + w[j-5],

  //     s0(w[j-12]) + w[j-4]

  vadduwm    (vt2, vt1, vt2);

  // c = s0(w[j-15]) + w[j-7] + w[j-16],

  //     s0(w[j-14]) + w[j-6] + w[j-15],

  //     s0(w[j-13]) + w[j-5] + w[j-14],

  //     s0(w[j-12]) + w[j-4] + w[j-13]

  vadduwm    (vt2, vt2, w0);

  // e = s0(w[j-15]) + w[j-7] + w[j-16] + s1(w[j-2]), // w[j]

  //     s0(w[j-14]) + w[j-6] + w[j-15] + s1(w[j-1]), // w[j+1]

  //     s0(w[j-13]) + w[j-5] + w[j-14] + s1(w[j-4]), // UNDEFINED

  //     s0(w[j-12]) + w[j-4] + w[j-13] + s1(w[j-3])  // UNDEFINED

  vadduwm    (vt4, vt2, vt3);

  // At this point, e[0] and e[1] are the correct values to be stored at w[j]

  // and w[j+1].

  // e[2] and e[3] are not considered.

  // b = s1(w[j]) , s1(s(w[j+1]) , UNDEFINED , UNDEFINED

  vshasigmaw (vt1, vt4, 0,  0xf);

  // v5 = s1(w[j-2]) , s1(w[j-1]) , s1(w[j]) , s1(w[j+1])

#if defined(VM_LITTLE_ENDIAN)

  xxmrgld    (vst3, vsrt1, vst3);

#else

  xxmrghd    (vst3, vst3, vsrt1);

#endif

  // c = s0(w[j-15]) + w[j-7] + w[j-16] + s1(w[j-2]), // w[j]

  //     s0(w[j-14]) + w[j-6] + w[j-15] + s1(w[j-1]), // w[j+1]

  //     s0(w[j-13]) + w[j-5] + w[j-14] + s1(w[j]),   // w[j+2]

  //     s0(w[j-12]) + w[j-4] + w[j-13] + s1(w[j+1])  // w[j+4]

  vadduwm    (vt2, vt2, vt3);

  // Updating w0 to w3 to hold the new previous 16 values from w.

  vmr        (w0,  w1);

  vmr        (w1,  w2);

  vmr        (w2,  w3);

  vmr        (w3,  vt2);

  // store k + w to v9 (4 values at once)

#if defined(VM_LITTLE_ENDIAN)

  vadduwm    (kpw0, vt2, vt0);

  vsldoi     (kpw1, kpw0, kpw0, 12);

  vsldoi     (kpw2, kpw0, kpw0, 8);

  vsldoi     (kpw3, kpw0, kpw0, 4);

#else

  vadduwm    (kpw3, vt2, vt0);

  vsldoi     (kpw2, kpw3, kpw3, 12);

  vsldoi     (kpw1, kpw3, kpw3, 8);

  vsldoi     (kpw0, kpw3, kpw3, 4);

#endif

}

void MacroAssembler::sha256_update_sha_state(const VectorRegister a,

                                             const VectorRegister b_,

                                             const VectorRegister c,

                                             const VectorRegister d,

                                             const VectorRegister e,

                                             const VectorRegister f,

                                             const VectorRegister g,

                                             const VectorRegister h,

                                             const Register hptr) {

  // temporaries

  VectorRegister vt0  = VR0;

  VectorRegister vt1  = VR1;

  VectorRegister vt2  = VR2;

  VectorRegister vt3  = VR3;

  VectorRegister vt4  = VR4;

  VectorRegister vt5  = VR5;

  VectorRegister vaux = VR6;

  VectorRegister vRb  = VR6;

  Register tmp        = R8;

  Register of16       = R8;

  Register of32       = R9;

  Label state_load_aligned;

  // Load hptr

  andi_   (tmp, hptr, 0xf);

  li      (of16, 16);

  lvx     (vt0, hptr);

  lvx     (vt5, of16, hptr);

  beq     (CCR0, state_load_aligned);

  // handle unaligned accesses

  li      (of32, 32);

  load_perm(vRb, hptr);

  vec_perm(vt0, vt5,  vRb);        // vt0 = hptr[0]..hptr[3]

  lvx     (vt1, hptr, of32);

  vec_perm(vt5, vt1,  vRb);        // vt5 = hptr[4]..hptr[7]

  // aligned accesses

  bind(state_load_aligned);

#if defined(VM_LITTLE_ENDIAN)

  vmrglw  (vt1, b_, a);            // vt1 = {a, b, ?, ?}

  vmrglw  (vt2, d, c);             // vt2 = {c, d, ?, ?}

  vmrglw  (vt3, f, e);             // vt3 = {e, f, ?, ?}

  vmrglw  (vt4, h, g);             // vt4 = {g, h, ?, ?}

  xxmrgld (vt1->to_vsr(), vt2->to_vsr(), vt1->to_vsr()); // vt1 = {a, b, c, d}

  xxmrgld (vt3->to_vsr(), vt4->to_vsr(), vt3->to_vsr()); // vt3 = {e, f, g, h}

  vadduwm (a,   vt0, vt1);         // a = {a+hptr[0], b+hptr[1], c+hptr[2], d+hptr[3]}

  vadduwm (e,   vt5, vt3);         // e = {e+hptr[4], f+hptr[5], g+hptr[6], h+hptr[7]}

  // Save hptr back, works for any alignment

  xxswapd (vt0->to_vsr(), a->to_vsr());

  stxvd2x (vt0->to_vsr(), hptr);

  xxswapd (vt5->to_vsr(), e->to_vsr());

  stxvd2x (vt5->to_vsr(), of16, hptr);

#else

  vmrglw  (vt1, a, b_);            // vt1 = {a, b, ?, ?}

  vmrglw  (vt2, c, d);             // vt2 = {c, d, ?, ?}

  vmrglw  (vt3, e, f);             // vt3 = {e, f, ?, ?}

  vmrglw  (vt4, g, h);             // vt4 = {g, h, ?, ?}

  xxmrgld (vt1->to_vsr(), vt1->to_vsr(), vt2->to_vsr()); // vt1 = {a, b, c, d}

  xxmrgld (vt3->to_vsr(), vt3->to_vsr(), vt4->to_vsr()); // vt3 = {e, f, g, h}

  vadduwm (d,   vt0, vt1);         // d = {a+hptr[0], b+hptr[1], c+hptr[2], d+hptr[3]}

  vadduwm (h,   vt5, vt3);         // h = {e+hptr[4], f+hptr[5], g+hptr[6], h+hptr[7]}

  // Save hptr back, works for any alignment

  stxvd2x (d->to_vsr(), hptr);

  stxvd2x (h->to_vsr(), of16, hptr);

#endif

}

static const uint32_t sha256_round_table[64] __attribute((aligned(16))) = {

  0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,

  0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,

  0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,

  0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,

  0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,

  0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,

  0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,

  0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,

  0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,

  0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,

  0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,

  0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,

  0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,

  0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,

  0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,

  0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,

};

static const uint32_t *sha256_round_consts = sha256_round_table;

//   R3_ARG1   - byte[]  Input string with padding but in Big Endian

//   R4_ARG2   - int[]   SHA.state (at first, the root of primes)

//   R5_ARG3   - int     offset

//   R6_ARG4   - int     limit

//

//   Internal Register usage:

//   R7        - k

//   R8        - tmp | j | of16

//   R9        - of32

//   VR0-VR8   - ch, maj, bsa, bse, vt0-vt3 | vt0-vt5, vaux/vRb

//   VR9-VR16  - a-h

//   VR17-VR20 - w0-w3

//   VR21-VR23 - vRb | vaux0-vaux2

//   VR24-VR27 - kpw0-kpw3

void MacroAssembler::sha256(bool multi_block) {

  static const ssize_t buf_size = 64;

  static const uint8_t w_size = sizeof(sha256_round_table)/sizeof(uint32_t);

#ifdef AIX

  // malloc provides 16 byte alignment

  if (((uintptr_t)sha256_round_consts & 0xF) != 0) {

    uint32_t *new_round_consts = (uint32_t*)malloc(sizeof(sha256_round_table));

    guarantee(new_round_consts, "oom");

    memcpy(new_round_consts, sha256_round_consts, sizeof(sha256_round_table));

    sha256_round_consts = (const uint32_t*)new_round_consts;

  }

#endif

  Register buf_in = R3_ARG1;

  Register state  = R4_ARG2;

  Register ofs    = R5_ARG3;

  Register limit  = R6_ARG4;

  Label sha_loop, core_loop;

  // Save non-volatile vector registers in the red zone

  static const VectorRegister nv[] = {

    VR20, VR21, VR22, VR23, VR24, VR25, VR26, VR27/*, VR28, VR29, VR30, VR31*/

  };

  static const uint8_t nv_size = sizeof(nv) / sizeof (VectorRegister);

  for (int c = 0; c < nv_size; c++) {

    Register tmp = R8;

    li  (tmp, (c - (nv_size)) * 16);

    stvx(nv[c], tmp, R1);

  }

  // Load hash state to registers

  VectorRegister a = VR9;

  VectorRegister b = VR10;

  VectorRegister c = VR11;

  VectorRegister d = VR12;

  VectorRegister e = VR13;

  VectorRegister f = VR14;

  VectorRegister g = VR15;

  VectorRegister h = VR16;

  static const VectorRegister hs[] = {a, b, c, d, e, f, g, h};

  static const int total_hs = sizeof(hs)/sizeof(VectorRegister);

  // counter for cycling through hs vector to avoid register moves between iterations

  int h_cnt = 0;

  // Load a-h registers from the memory pointed by state

#if defined(VM_LITTLE_ENDIAN)

  sha256_load_h_vec(a, e, state);

#else

  sha256_load_h_vec(d, h, state);

#endif

  // keep k loaded also during MultiBlock loops

  Register k = R7;

  assert(((uintptr_t)sha256_round_consts & 0xF) == 0, "k alignment");

  load_const_optimized(k, (address)sha256_round_consts, R0);

  // Avoiding redundant loads

  if (multi_block) {

    align(OptoLoopAlignment);

  }

  bind(sha_loop);

#if defined(VM_LITTLE_ENDIAN)

  sha256_deque(a, b, c, d);

  sha256_deque(e, f, g, h);

#else

  sha256_deque(d, c, b, a);

  sha256_deque(h, g, f, e);

#endif

  // Load 16 elements from w out of the loop.

  // Order of the int values is Endianess specific.

  VectorRegister w0 = VR17;

  VectorRegister w1 = VR18;

  VectorRegister w2 = VR19;

  VectorRegister w3 = VR20;

  static const VectorRegister ws[] = {w0, w1, w2, w3};

  static const int total_ws = sizeof(ws)/sizeof(VectorRegister);

  VectorRegister kpw0 = VR24;

  VectorRegister kpw1 = VR25;

  VectorRegister kpw2 = VR26;

  VectorRegister kpw3 = VR27;

  static const VectorRegister kpws[] = {kpw0, kpw1, kpw2, kpw3};

  static const int total_kpws = sizeof(kpws)/sizeof(VectorRegister);

  sha256_load_w_plus_k_vec(buf_in, ws, total_ws, k, kpws, total_kpws);

  // Cycle through the first 16 elements

  assert(total_ws == total_kpws, "Redesign the loop below");

  for (int n = 0; n < total_ws; n++) {

    VectorRegister vaux0 = VR21;

    VectorRegister vaux1 = VR22;

    VectorRegister vaux2 = VR23;

    sha256_deque(kpws[n], vaux0, vaux1, vaux2);

#if defined(VM_LITTLE_ENDIAN)

    sha256_round(hs, total_hs, h_cnt, kpws[n]);

    sha256_round(hs, total_hs, h_cnt, vaux0);

    sha256_round(hs, total_hs, h_cnt, vaux1);

    sha256_round(hs, total_hs, h_cnt, vaux2);

#else

    sha256_round(hs, total_hs, h_cnt, vaux2);

    sha256_round(hs, total_hs, h_cnt, vaux1);

    sha256_round(hs, total_hs, h_cnt, vaux0);

    sha256_round(hs, total_hs, h_cnt, kpws[n]);

#endif

  }

  Register tmp = R8;

  // loop the 16th to the 64th iteration by 8 steps

  li   (tmp, (w_size - 16) / total_hs);

  mtctr(tmp);

  // j will be aligned to 4 for loading words.

  // Whenever read, advance the pointer (e.g: when j is used in a function)

  Register j = R8;

  li   (j, 16*4);

  align(OptoLoopAlignment);

  bind(core_loop);

  // due to VectorRegister rotate, always iterate in multiples of total_hs

  for (int n = 0; n < total_hs/4; n++) {

    sha256_calc_4w(w0, w1, w2, w3, kpw0, kpw1, kpw2, kpw3, j, k);

    sha256_round(hs, total_hs, h_cnt, kpw0);

    sha256_round(hs, total_hs, h_cnt, kpw1);

    sha256_round(hs, total_hs, h_cnt, kpw2);

    sha256_round(hs, total_hs, h_cnt, kpw3);