#ifndef ROUND_H_
#define ROUND_H_

#include "ascon.h"
#include "printstate.h"
#include "random.h"

__forceinline void KINIT(word_t* K0, word_t* K1, word_t* K2) {
  *K0 = TOSHARES(0);
  *K1 = TOSHARES(0);
  *K2 = TOSHARES(0);
}

__forceinline void PINIT(state_t* s) {
  randinit();
  s->x0 = TOSHARES(0);
  s->x1 = TOSHARES(0);
  s->x2 = TOSHARES(0);
  s->x3 = TOSHARES(0);
  s->x4 = TOSHARES(0);
  s->rx = TOSHARES(0);
}

#define TOFFOLI(a0, a1, a2, b0, b1, b2, c0, c1, c2) \
  do {                                              \
    (a0) ^= (~(b0)) & (c0);                         \
    (a0) ^= (b0) & (c2);                            \
    (a0) ^= (b2) & (c0);                            \
    (a1) ^= (~(b1)) & (c1);                         \
    (a1) ^= (b1) & (c0);                            \
    (a1) ^= (b0) & (c1);                            \
    (a2) ^= (~(b2)) & (c2);                         \
    (a2) ^= (b2) & (c1);                            \
    (a2) ^= (b1) & (c2);                            \
  } while (0)

__forceinline void ROUND(state_t* s, uint64_t C) {
  /* refresh randomness */
  /* s->rx = TOSHARES(0); */
  /* addition of round constant */
  s->x2.s0 ^= C;
  /* substitution layer */
  s->x0.s0 ^= s->x4.s0;
  s->x4.s0 ^= s->x3.s0;
  s->x2.s0 ^= s->x1.s0;
  s->x0.s1 ^= s->x4.s1;
  s->x4.s1 ^= s->x3.s1;
  s->x2.s1 ^= s->x1.s1;
  s->x0.s2 ^= s->x4.s2;
  s->x4.s2 ^= s->x3.s2;
  s->x2.s2 ^= s->x1.s2;
  /* start of shared keccak s-box from https://eprint.iacr.org/2019/536 */
  s->rx.s2 = s->rx.s0;
  s->rx.s0 ^= s->rx.s1;
  TOFFOLI(s->rx.s0, s->rx.s1, s->rx.s2, s->x4.s0, s->x4.s1, s->x4.s2, s->x0.s0,
          s->x0.s1, s->x0.s2);
  TOFFOLI(s->x0.s0, s->x0.s1, s->x0.s2, s->x1.s0, s->x1.s1, s->x1.s2, s->x2.s0,
          s->x2.s1, s->x2.s2);
  TOFFOLI(s->x2.s0, s->x2.s1, s->x2.s2, s->x3.s0, s->x3.s1, s->x3.s2, s->x4.s0,
          s->x4.s1, s->x4.s2);
  TOFFOLI(s->x4.s0, s->x4.s1, s->x4.s2, s->x0.s0, s->x0.s1, s->x0.s2, s->x1.s0,
          s->x1.s1, s->x1.s2);
  TOFFOLI(s->x1.s0, s->x1.s1, s->x1.s2, s->x2.s0, s->x2.s1, s->x2.s2, s->x3.s0,
          s->x3.s1, s->x3.s2);
  s->x3.s2 ^= s->rx.s2;
  s->x3.s1 ^= s->rx.s1;
  s->x3.s0 ^= s->rx.s0;
  /* end of shared keccak s-box */
  s->x1.s0 ^= s->x0.s0;
  s->x0.s0 ^= s->x4.s0;
  s->x3.s0 ^= s->x2.s0;
  s->x2.s0 = ~s->x2.s0;
  s->x1.s1 ^= s->x0.s1;
  s->x0.s1 ^= s->x4.s1;
  s->x3.s1 ^= s->x2.s1;
  s->x1.s2 ^= s->x0.s2;
  s->x0.s2 ^= s->x4.s2;
  s->x3.s2 ^= s->x2.s2;
  /* linear diffusion layer */
  s->x0.s2 ^= ROR64(s->x0.s2, 19) ^ ROR64(s->x0.s2, 28);
  s->x1.s2 ^= ROR64(s->x1.s2, 61) ^ ROR64(s->x1.s2, 39);
  s->x2.s2 ^= ROR64(s->x2.s2, 1) ^ ROR64(s->x2.s2, 6);
  s->x3.s2 ^= ROR64(s->x3.s2, 10) ^ ROR64(s->x3.s2, 17);
  s->x4.s2 ^= ROR64(s->x4.s2, 7) ^ ROR64(s->x4.s2, 41);
  s->x0.s1 ^= ROR64(s->x0.s1, 19) ^ ROR64(s->x0.s1, 28);
  s->x1.s1 ^= ROR64(s->x1.s1, 61) ^ ROR64(s->x1.s1, 39);
  s->x2.s1 ^= ROR64(s->x2.s1, 1) ^ ROR64(s->x2.s1, 6);
  s->x3.s1 ^= ROR64(s->x3.s1, 10) ^ ROR64(s->x3.s1, 17);
  s->x4.s1 ^= ROR64(s->x4.s1, 7) ^ ROR64(s->x4.s1, 41);
  s->x0.s0 ^= ROR64(s->x0.s0, 19) ^ ROR64(s->x0.s0, 28);
  s->x1.s0 ^= ROR64(s->x1.s0, 61) ^ ROR64(s->x1.s0, 39);
  s->x2.s0 ^= ROR64(s->x2.s0, 1) ^ ROR64(s->x2.s0, 6);
  s->x3.s0 ^= ROR64(s->x3.s0, 10) ^ ROR64(s->x3.s0, 17);
  s->x4.s0 ^= ROR64(s->x4.s0, 7) ^ ROR64(s->x4.s0, 41);
  printstate(" round output", s);
}

#endif /* ROUND_H_ */