/* test_stub.c - Diffusion test stub - Don Yang (uguu.org)

   01/05/05
                                                                           */

#ifdef _WIN32
   #include<windows.h>
   #include<math.h>
   #include<stdio.h>
   #include<stdlib.h>
   #include<string.h>
   #include<time.h>

   typedef HANDLE Thread;
#else
   #include<math.h>
   #include<stdio.h>
   #include<stdlib.h>
   #include<string.h>
   #include<time.h>
   #include<unistd.h>
   #include<pthread.h>

   typedef pthread_t Thread;
#endif

#define BLOCK_COUNT     32
#define BLOCK_SIZE      (BLOCK_COUNT * 4)

#ifdef USE32
   #define COUNTER_FORMAT     "%u"
   typedef unsigned int Counter;
#else
   /* long long is not ISO C90 compatible, don't use "-ansi -pedantic" */
   #ifdef _WIN32
      #define COUNTER_FORMAT  "%I64u"
   #else
      #define COUNTER_FORMAT  "%llu"
   #endif
   typedef unsigned long long int Counter;
#endif

typedef struct
{
   Counter total, change;
   unsigned int freq[256], seed, *run;
   Thread handle;
} ThreadArg;

static void E(unsigned int *d)
{
   unsigned int i;
   for(i = 0; i < BLOCK_COUNT; i++)
      d[i] ^= ~0U;
}

static unsigned int BitCount(unsigned int x)
{
   x = (x & 0x55555555) + ((x & 0xaaaaaaaa) >> 1);
   x = (x & 0x33333333) + ((x & 0xcccccccc) >> 2);
   x = (x & 0x0f0f0f0f) + ((x & 0xf0f0f0f0) >> 4);
   x = (x & 0x00ff00ff) + ((x & 0xff00ff00) >> 8);
   return (x & 0x0000ffff) + ((x & 0xffff0000) >> 16);
}

static unsigned int CountDiff(unsigned int *a, unsigned int *b)
{
   unsigned int x, i;
   for(i = x = 0; i < BLOCK_COUNT; i++)
      x += BitCount(*(a++) ^ *(b++));
   return x;
}

static void InitRand(unsigned int *table, unsigned int seed)
{
   static const unsigned int seed0[55] =
   {
      509760043,  399328820, 99941072,   112282318, 611886020, 516451399,
      626288598,  337482183, 748548471,  808894867, 657927153, 386437385,
      42355480,   977713532, 311548488,  13857891,  307938721, 93724463,
      1041159001, 444711218, 1040610926, 233671814, 664494626, 1071756703,
      188709089,  420289414, 969883075,  513442196, 275039308, 918830973,
      598627151,  134083417, 823987070,  619204222, 81893604,  871834315,
      398384680,  475117924, 520153386,  324637501, 38588599,  435158812,
      168033706,  585877294, 328347186,  293179100, 671391820, 846150845,
      283985689,  502873302, 718642511,  938465128, 962756406, 107944131,
      192910970
   };
   unsigned int i;

   memcpy(table, seed0, 55 * sizeof(unsigned int));
   for(i = 0; i < 55; i++)
      table[i] ^= (seed = seed * 22695447 + 1);
}

#ifdef _WIN32
   static DWORD WINAPI ProcessThread(ThreadArg *a)
#else
   static unsigned int ProcessThread(ThreadArg *a)
#endif
{
   unsigned int rand_seed[55], rand_idx;
   unsigned int in[BLOCK_COUNT], out0[BLOCK_COUNT], out[BLOCK_COUNT], i, j, x;
   unsigned char *b;

   /* Initialize PRNG */
   InitRand(rand_seed, a->seed);
   rand_idx = a->seed % 55;
   while( *(a->run) != 0 )
   {
      /* Create block */
      for(i = 0; i < BLOCK_COUNT; i++)
      {
         rand_idx = (rand_idx + 1) % 55;
         in[i] = (rand_seed[rand_idx] += rand_seed[(rand_idx + 24) % 55]);
      }

      /* Encrypt reference block */
      memcpy(out0, in, BLOCK_SIZE);
      E(out0);

      /* Walk bits */
      for(i = 0; i < BLOCK_COUNT; i++)
      {
         for(j = 0; j < 32; j++)
         {
            in[i] ^= (1 << j);
            memcpy(out, in, BLOCK_SIZE);
            E(out);

            /* Count change in bits */
            a->total += 32 * BLOCK_COUNT;
            a->change += CountDiff(out0, out);

            /* Count frequency of output */
            b = (unsigned char*)out;
            for(x = 0; x < BLOCK_SIZE; x++)
               a->freq[*(b++)]++;

            in[i] ^= (1 << j);
         }
      }
   }

   return 0;
}

int main(int argc, char **argv)
{
   ThreadArg *a;
   unsigned int x, i, c, t, run;
   double e, p;

   if( argc < 3 )
      return printf("%s <thread count> <time>\n", *argv);
   if( (c = (unsigned int)atoi(argv[1])) < 1 )
      return printf("Invalid number of threads: %s\n", argv[1]);
   if( (t = (unsigned int)atoi(argv[2])) < 1 )
      return printf("Invalid time: %s\n", argv[2]);

   /* Create process threads */
   if( (a = (ThreadArg*)calloc((size_t)c, sizeof(ThreadArg))) == NULL )
      return puts("Out of memory");
   run = 1;
   srand((unsigned)time(NULL));
   for(i = 0; i < c; i++)
   {
      a[i].seed = (unsigned int)rand();
      a[i].run = &run;
      #ifdef _WIN32
         if( (a[i].handle = CreateThread(NULL, 0,
               (LPTHREAD_START_ROUTINE)ProcessThread,
               (void*)&a[i], 0, NULL)) == NULL )
            return printf("Error creating thread %d\n", i);
      #else
         if( pthread_create(&a[i].handle, NULL,
               (void*(*)(void*))ProcessThread, (void*)&a[i]) != 0 )
            return printf("Error creating thread %d\n", i);
      #endif
   }

   /* Run */
   #ifdef _WIN32
      Sleep(t * 1000);
      run = 0;
      for(i = 0; i < c; i++)
      {
         (void)WaitForSingleObject(a[i].handle, INFINITE);
         CloseHandle(a[i].handle);
      }
   #else
      sleep(t);
      run = 0;
      for(i = 0; i < c; i++)
         (void)pthread_join(a[i].handle, NULL);
   #endif

   /* Add statistics */
   for(i = 1; i < c; i++)
   {
      a->change += a[i].change;
      a->total += a[i].total;
      for(t = 0; t < 256; t++)
         a->freq[t] += a[i].freq[t];
   }
   if( a->total == 0 )
      return puts("Error?!");
   for(i = x = 0; i < 256; i++)
      x += a->freq[i];
   e = 0.0;
   for(i = 0; i < 256; i++)
   {
      if( a->freq[i] > 0 )
      {
         p = (double)(a->freq[i]) / (double)x;
         e -= p * log(p) / log(2.0);
      }
   }

   p = (double)(a->change) / (double)(a->total);
   printf(
      "Bits processed: "COUNTER_FORMAT" (%.3fMbit/s)\n"
      "Bits changed: "COUNTER_FORMAT" (%.3f%%)\n"
      "Change bias: %+f\n"
      "Entropy: %f\n",
      a->total, (double)(a->total) / (double)(1024 * 1024 * atoi(argv[2])),
      a->change, 100.0 * p, p - 0.5, e);

   free(a);
   return 0;
}