#ifdef __VSDSP__
#include <vstypes.h>
#else
#include <inttypes.h>
typedef uint32_t u_int32;
typedef uint16_t u_int16;
#define __mem_y
#define memcpyYX memcpy
#endif
#include <stdio.h>
#include <string.h>

void QuarterRoll(register u_int32 *a,
                 register u_int32 *b,
                 register u_int32 *c,
                 register u_int32 *d){
  *a += *b; *d ^= *a; *d = (*d << 16)|(*d >> (32 - 16));
  *c += *d; *b ^= *c; *b = (*b << 12)|(*b >> (32 - 12));
  *a += *b; *d ^= *a; *d = (*d <<  8)|(*d >> (32 -  8));
  *c += *d; *b ^= *c; *b = (*b <<  7)|(*b >> (32 -  7));
}

__mem_y u_int32 chachaState[32];
u_int32 chachaWork[32];

void Chacha20SetKey(register u_int32 *key,
            register u_int32 counter,
            register u_int32 *nonce ){
  static u_int32 keyConsts[] = {
    0x61707865, 0x3320646e, 0x79622d32, 0x6b206574
  };
  u_int32 *st = chachaState, *init = keyConsts;
  u_int16 i;
  for(i = 0; i < 4; i++){
    *st++ = *init++;
  }
  for(i = 0; i < 8; i++){
    *st++ = *key++;
  }
  *st++ = counter;
  for(i = 0; i < 3; i++){
    *st++ = *nonce++;
  }
}

void Chacha20Block(){
  u_int16 i;
  u_int32 *st = chachaState, *w = chachaWork;
  memcpyYX(chachaWork, chachaState, sizeof(chachaState));
  for(i = 0; i < 10; i++){
    QuarterRoll(chachaWork +  0, chachaWork +  4,
                chachaWork +  8, chachaWork + 12);
    QuarterRoll(chachaWork +  1, chachaWork +  5,
                chachaWork +  9, chachaWork + 13);
    QuarterRoll(chachaWork +  2, chachaWork +  6,
                chachaWork + 10, chachaWork + 14);
    QuarterRoll(chachaWork +  3, chachaWork +  7,
                chachaWork + 11, chachaWork + 15);
    QuarterRoll(chachaWork +  0, chachaWork +  5,
                chachaWork + 10, chachaWork + 15);
    QuarterRoll(chachaWork +  1, chachaWork +  6,
                chachaWork + 11, chachaWork + 12);
    QuarterRoll(chachaWork +  2, chachaWork +  7,
                chachaWork +  8, chachaWork + 13);
    QuarterRoll(chachaWork +  3, chachaWork +  4,
                chachaWork +  9, chachaWork + 14);
  }
  for(i = 0; i < 16; i++){
    /* Old version, writes over state. See test case and encryption
     * why state is wanted to preserved. */
    /* *st++ += *w++; */
    *w++ += *st++;
  }
}

void Chacha20Encrypt(char * msg, u_int16 len){
  u_int16 i, n;
  u_int32 *s;
  while(len >= 64){
    Chacha20Block();
    s = chachaWork;
    for(i = 0; i < 16; i++){
      *msg++ ^= (*s >> 0) & 0xff;
      *msg++ ^= (*s >> 8) & 0xff;
      *msg++ ^= (*s >> 16) & 0xff;
      *msg++ ^= (*s >> 24) & 0xff;
      s++;
    }
    /* Udpate counter */
    chachaState[12]++;
    len -= 64;
  }
  if(len){
    /* Last extra bytes */
    s = chachaWork;
    Chacha20Block();
    n = len;
    for(i = 0; i <= (len / 4); i++){
      if (n--){
	*msg++ ^= (*s >> 0) & 0xff;
	if (n--){
	  *msg++ ^= (*s >> 8) & 0xff;
	  if (n--){
	    *msg++ ^= (*s >> 16) & 0xff;
	    if (n--){
	      *msg++ ^= (*s >> 24) & 0xff;
	    }
	  }
	}
      }
      s++;
    }
  }
  return;
}
