#include "params.h"
#include "poly.h"
#include "polyvec.h"
#include "packing.h"

/*************************************************
* Name:        pack_pk
* 
* Description: Bit-pack public key pk = (rho, t1).
*
* Arguments:   - unsigned char pk[]: output byte array
*              - const unsigned char rho[]: byte array containing rho
*              - const polyveck *t1: pointer to vector t1
**************************************************/
void pack_pk(unsigned char pk[PK_SIZE_PACKED],
             const unsigned char rho[SEEDBYTES],
             const polyveck *t1)
{
  unsigned int i;

  for(i = 0; i < SEEDBYTES; ++i)
    pk[i] = rho[i];
  pk += SEEDBYTES;

  for(i = 0; i < K; ++i)
    polyt1_pack(pk + i*POLT1_SIZE_PACKED, t1->vec+i);
}

/*************************************************
* Name:        unpack_pk
* 
* Description: Unpack public key pk = (rho, t1).
*
* Arguments:   - const unsigned char rho[]: output byte array for rho
*              - const polyveck *t1: pointer to output vector t1
*              - unsigned char pk[]: byte array containing bit-packed pk
**************************************************/
void unpack_pk(unsigned char rho[SEEDBYTES],
               polyveck *t1,
               const unsigned char pk[PK_SIZE_PACKED])
{
  unsigned int i;

  for(i = 0; i < SEEDBYTES; ++i)
    rho[i] = pk[i];
  pk += SEEDBYTES;

  for(i = 0; i < K; ++i)
    polyt1_unpack(t1->vec+i, pk + i*POLT1_SIZE_PACKED);
}

/*************************************************
* Name:        pack_sk
* 
* Description: Bit-pack secret key sk = (rho, key, tr, s1, s2, t0).
*
* Arguments:   - unsigned char sk[]: output byte array
*              - const unsigned char rho[]: byte array containing rho
*              - const unsigned char key[]: byte array containing key
*              - const unsigned char tr[]: byte array containing tr
*              - const polyvecl *s1: pointer to vector s1
*              - const polyveck *s2: pointer to vector s2
*              - const polyveck *t0: pointer to vector t0
**************************************************/
void pack_sk(unsigned char sk[SK_SIZE_PACKED],
             const unsigned char rho[SEEDBYTES],
             const unsigned char key[SEEDBYTES],
             const unsigned char tr[CRHBYTES],
             const polyvecl *s1,
             const polyveck *s2,
             const polyveck *t0)
{
  unsigned int i;

  for(i = 0; i < SEEDBYTES; ++i)
    sk[i] = rho[i];
  sk += SEEDBYTES;

  for(i = 0; i < SEEDBYTES; ++i)
    sk[i] = key[i];
  sk += SEEDBYTES;

  for(i = 0; i < CRHBYTES; ++i)
    sk[i] = tr[i];
  sk += CRHBYTES;

  for(i = 0; i < L; ++i)
    polyeta_pack(sk + i*POLETA_SIZE_PACKED, s1->vec+i);
  sk += L*POLETA_SIZE_PACKED;

  for(i = 0; i < K; ++i)
    polyeta_pack(sk + i*POLETA_SIZE_PACKED, s2->vec+i);
  sk += K*POLETA_SIZE_PACKED;

  for(i = 0; i < K; ++i)
    polyt0_pack(sk + i*POLT0_SIZE_PACKED, t0->vec+i);
}

/*************************************************
* Name:        unpack_sk
* 
* Description: Unpack secret key sk = (rho, key, tr, s1, s2, t0).
*
* Arguments:   - const unsigned char rho[]: output byte array for rho
*              - const unsigned char key[]: output byte array for key
*              - const unsigned char tr[]: output byte array for tr
*              - const polyvecl *s1: pointer to output vector s1
*              - const polyveck *s2: pointer to output vector s2
*              - const polyveck *r0: pointer to output vector t0
*              - unsigned char sk[]: byte array containing bit-packed sk
**************************************************/
void unpack_sk(unsigned char rho[SEEDBYTES],
               unsigned char key[SEEDBYTES],
               unsigned char tr[CRHBYTES],
               polyvecl *s1,
               polyveck *s2,
               polyveck *t0,
               const unsigned char sk[SK_SIZE_PACKED])
{
  unsigned int i;

  for(i = 0; i < SEEDBYTES; ++i)
    rho[i] = sk[i];
  sk += SEEDBYTES;

  for(i = 0; i < SEEDBYTES; ++i)
    key[i] = sk[i];
  sk += SEEDBYTES;

  for(i = 0; i < CRHBYTES; ++i)
    tr[i] = sk[i];
  sk += CRHBYTES;

  for(i=0; i < L; ++i)
    polyeta_unpack(s1->vec+i, sk + i*POLETA_SIZE_PACKED);
  sk += L*POLETA_SIZE_PACKED;
  
  for(i=0; i < K; ++i)
    polyeta_unpack(s2->vec+i, sk + i*POLETA_SIZE_PACKED);
  sk += K*POLETA_SIZE_PACKED;

  for(i=0; i < K; ++i)
    polyt0_unpack(t0->vec+i, sk + i*POLT0_SIZE_PACKED);
}

/*************************************************
* Name:        pack_sig
* 
* Description: Bit-pack signature sig = (z, h, c).
*
* Arguments:   - unsigned char sig[]: output byte array
*              - const polyvecl *z: pointer to vector z
*              - const polyveck *h: pointer to hint vector h
*              - const poly *c: pointer to challenge polynomial
**************************************************/
void pack_sig(unsigned char sig[SIG_SIZE_PACKED],
              const polyvecl *z,
              const polyveck *h,
              const poly *c)
{
  unsigned int i, j, k;
  uint64_t signs, mask;

  for(i = 0; i < L; ++i)
    polyz_pack(sig + i*POLZ_SIZE_PACKED, z->vec+i);
  sig += L*POLZ_SIZE_PACKED;

  /* Encode h */
  k = 0;
  for(i = 0; i < K; ++i) {
    for(j = 0; j < N; ++j)
      if(h->vec[i].coeffs[j] == 1)
        sig[k++] = j;

    sig[OMEGA + i] = k;
  }
  while(k < OMEGA) sig[k++] = 0;
  sig += OMEGA + K;
  
  /* Encode c */
  signs = 0;
  mask = 1;
  for(i = 0; i < N/8; ++i) {
    sig[i] = 0;
    for(j = 0; j < 8; ++j) {
      if(c->coeffs[8*i+j] != 0) {
        sig[i] |= (1 << j);
        if(c->coeffs[8*i+j] == (Q - 1)) signs |= mask;
        mask <<= 1;
      }
    }
  }
  sig += N/8;
  for(i = 0; i < 8; ++i)
    sig[i] = signs >> 8*i;
}

/*************************************************
* Name:        unpack_sig
* 
* Description: Unpack signature sig = (z, h, c).
*
* Arguments:   - const polyvecl *z: pointer to output vector z
*              - const polyveck *h: pointer to output hint vector h
*              - const poly *c: pointer to output challenge polynomial
*              - unsigned char sig[]: byte array containing bit-packed signature
*
* Returns 1 in case of malformed signature; otherwise 0
**************************************************/
int unpack_sig(polyvecl *z,
               polyveck *h,
               poly *c,
               const unsigned char sig[SIG_SIZE_PACKED])
{
  unsigned int i, j, k;
  uint64_t signs, mask;

  for(i = 0; i < L; ++i)
    polyz_unpack(z->vec+i, sig + i*POLZ_SIZE_PACKED);
  sig += L*POLZ_SIZE_PACKED;

  /* Enforce standard representatives for strong unforgeability */
  for(i = 0; i < L; ++i)
    for(j = 0; j < N; ++j)
      if(z->vec[i].coeffs[j] >= Q)
        return 1;

  /* Decode h */
  k = 0;
  for(i = 0; i < K; ++i) {
    for(j = 0; j < N; ++j)
      h->vec[i].coeffs[j] = 0;

    if(sig[OMEGA + i] < k || sig[OMEGA + i] > OMEGA)
      return 1;

    for(j = k; j < sig[OMEGA + i]; ++j) {
      /* Coefficients are ordered for strong unforgeability */
      if(j > k && sig[j] < sig[j-1]) return 1;
      h->vec[i].coeffs[sig[j]] = 1;
    }

    k = sig[OMEGA + i];
  }

  /* Extra indices are zero for strong unforgeability */
  for(j = k; j < OMEGA; ++j)
    if(sig[j])
      return 1;

  sig += OMEGA + K;

  /* Decode c */
  for(i = 0; i < N; ++i)
    c->coeffs[i] = 0;

  signs = 0;
  for(i = 0; i < 8; ++i)
    signs |= (uint64_t)sig[N/8+i] << 8*i;

  /* Extra sign bits are zero for strong unforgeability */
  if(signs >> 60)
    return 1;

  mask = 1;
  for(i = 0; i < N/8; ++i) {
    for(j = 0; j < 8; ++j) {
      if((sig[i] >> j) & 0x01) {
        c->coeffs[8*i+j] = (signs & mask) ? Q - 1 : 1;
        mask <<= 1;
      }
    }
  }

  return 0;
}
