#ifndef WORD_H_
#define WORD_H_

#include <stdint.h>

#include "config.h"
#include "random.h"

typedef struct {
  uint32_t e;
  uint32_t o;
} share_t;

typedef struct {
  share_t s0;
  share_t s1;
  share_t s2;
} word_t;

__forceinline word_t WORD_T(uint64_t x) {
  word_t w;
  w.s0.e = (uint32_t)x;
  w.s0.o = x >> 32;
  w.s1.e = 0;
  w.s1.o = 0;
  w.s2.e = 0;
  w.s2.o = 0;
  return w;
}

__forceinline uint64_t UINT64_T(word_t w) {
  return (uint64_t)w.s0.o << 32 | w.s0.e;
}

share_t TOBI32(share_t in);

share_t FROMBI32(share_t in);

__forceinline word_t TOSHARES(uint64_t in) {
  uint32_t r0 = rand32();
  uint32_t r1 = rand32();
  uint32_t r2 = rand32();
  uint32_t r3 = rand32();
  word_t w;
  w.s0.e = (uint32_t)in ^ r0 ^ r2;
  w.s0.o = (in >> 32) ^ r1 ^ r3;
  w.s1.e = r0;
  w.s1.o = r1;
  w.s2.e = r2;
  w.s2.o = r3;
  return w;
}

__forceinline uint64_t FROMSHARES(word_t in) {
  return (uint64_t)(in.s0.o ^ in.s1.o ^ in.s2.o) << 32 |
         (in.s0.e ^ in.s1.e ^ in.s2.e);
}

__forceinline word_t U64TOWORD(uint64_t x) {
#if ASCON_MASK_LOADS
  word_t w = TOSHARES(x);
  w.s0 = TOBI32(w.s0);
  w.s1 = TOBI32(w.s1);
  w.s2 = TOBI32(w.s2);
  return w;
#else
  word_t w = WORD_T(x);
  w.s0 = TOBI32(w.s0);
  return w;
#endif
}

__forceinline uint64_t WORDTOU64(word_t w) {
  w.s0 = FROMBI32(w.s0);
  w.s1 = FROMBI32(w.s1);
  w.s2 = FROMBI32(w.s2);
  return (uint64_t)FROMSHARES(w);
}

#define XOR(a, b)        \
  do {                   \
    word_t tb = b;       \
    (a).s0.e ^= tb.s0.e; \
    (a).s0.o ^= tb.s0.o; \
    (a).s1.e ^= tb.s1.e; \
    (a).s1.o ^= tb.s1.o; \
    (a).s2.e ^= tb.s2.e; \
    (a).s2.o ^= tb.s2.o; \
  } while (0)

#define AND(a, b)                                                        \
  do {                                                                   \
    word_t ta = a;                                                       \
    word_t tb = b;                                                       \
    (a).s0.e =                                                           \
        (ta.s0.e & tb.s0.e) ^ (ta.s0.e & tb.s1.e) ^ (ta.s0.e & tb.s2.e); \
    (a).s0.o =                                                           \
        (ta.s0.o & tb.s0.o) ^ (ta.s0.o & tb.s1.o) ^ (ta.s0.o & tb.s2.o); \
    (a).s1.e =                                                           \
        (ta.s1.e & tb.s0.e) ^ (ta.s1.e & tb.s1.e) ^ (ta.s1.e & tb.s2.e); \
    (a).s1.o =                                                           \
        (ta.s1.o & tb.s0.o) ^ (ta.s1.o & tb.s1.o) ^ (ta.s1.o & tb.s2.o); \
    (a).s2.e =                                                           \
        (ta.s2.e & tb.s0.e) ^ (ta.s2.e & tb.s1.e) ^ (ta.s2.e & tb.s2.e); \
    (a).s2.o =                                                           \
        (ta.s2.o & tb.s0.o) ^ (ta.s2.o & tb.s1.o) ^ (ta.s2.o & tb.s2.o); \
  } while (0)

__forceinline uint32_t ROR32(uint32_t x, int n) {
  return x >> n | x << (32 - n);
}

__forceinline word_t KEYROT(word_t lo2hi, word_t hi2lo) {
  word_t r;
  r.s0.e = lo2hi.s0.e << 16 | hi2lo.s0.e >> 16;
  r.s0.o = lo2hi.s0.o << 16 | hi2lo.s0.o >> 16;
  r.s1.e = lo2hi.s1.e << 16 | hi2lo.s1.e >> 16;
  r.s1.o = lo2hi.s1.o << 16 | hi2lo.s1.o >> 16;
  r.s2.e = lo2hi.s2.e << 16 | hi2lo.s2.e >> 16;
  r.s2.o = lo2hi.s2.o << 16 | hi2lo.s2.o >> 16;
  return r;
}

__forceinline int NOTZERO(word_t a, word_t b) {
  int result = 0;
  for (int i = 0; i < 8; ++i)
    result |= ((uint8_t*)&(a.s0))[i] ^ ((uint8_t*)&(a.s1))[i] ^
              ((uint8_t*)&(a.s2))[i];
  for (int i = 0; i < 8; ++i)
    result |= ((uint8_t*)&(b.s0))[i] ^ ((uint8_t*)&(b.s1))[i] ^
              ((uint8_t*)&(b.s2))[i];
  return result;
}

/* set padding byte in 64-bit Ascon word */
__forceinline word_t PAD(int i) {
  return WORD_T((uint64_t)(0x08 << (28 - 4 * i)) << 32);
}

/* byte mask for 64-bit Ascon word (1 <= n <= 8) */
__forceinline word_t XMASK(int n) {
  uint32_t mask = 0x0fffffff >> (n * 4 - 4);
  return WORD_T((uint64_t)mask << 32 | mask);
}

#endif /* WORD_H_ */