#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#include <time.h>
#include <sys/time.h>

static uint64_t
current_time(void)
{
        struct timeval tp;

        gettimeofday(&tp, NULL);
        return (uint64_t)tp.tv_sec * 1000000 + tp.tv_usec;
}

int
main(int argc, char *argv[])
{
        unsigned char init[256] __attribute__ ((aligned(16)));
        size_t i;
        uint64_t key, ref;
        uint64_t orig;

        if (argc < 2) {
                exit(EXIT_FAILURE);
        }
        ref = strtoull(argv[1], 0, 16);
        for (i = 0; i < 256; i ++) {
                init[i] = i;
        }
        orig = current_time();
        for (key = 0; key < ((uint64_t)1 << 40); key ++) {
                unsigned char buf[256] __attribute__ ((aligned(16)));
                unsigned c;
                unsigned k0, k1, k2, k3, k4;
                unsigned j;
                uint64_t out;

#define SWAP(a, b)   do { \
                unsigned char st = buf[(a)]; \
                buf[(a)] = buf[(b)]; \
                buf[(b)] = st; \
        } while (0)

                /*
                 * Key schedule.
                 */
                k0 = key & 0xFF;
                k1 = (key >> 8) & 0xFF;
                k2 = (key >> 16) & 0xFF;
                k3 = (key >> 24) & 0xFF;
                k4 = (key >> 32) & 0xFF;
                memcpy(buf, init, 256);
                c = 0;
                for (i = 0; i < 255; i += 5) {
                        c = (c + k0 + buf[i]) & 0xFF;
                        SWAP(i, c);
                        c = (c + k0 + buf[i + 1]) & 0xFF;
                        SWAP(i + 1, c);
                        c = (c + k0 + buf[i + 2]) & 0xFF;
                        SWAP(i + 2, c);
                        c = (c + k0 + buf[i + 3]) & 0xFF;
                        SWAP(i + 3, c);
                        c = (c + k0 + buf[i + 4]) & 0xFF;
                        SWAP(i + 4, c);
                }
                c = (c + k0 + buf[255]) & 0xFF;
                SWAP(255, c);

                /*
                 * Produce 64 bits of output.
                 */
                j = 0;
                out = 0;

#define STEP(k)   do { \
                unsigned b0, b1; \
                b0 = buf[k]; \
                j = (j + b0) & 0xFF; \
                b1 = buf[j]; \
                buf[j] = b0; \
                buf[k] = b1; \
                out = (out << 8) | buf[(b0 + b1) & 0xFF]; \
        } while (0)

                STEP(0);
                STEP(1);
                STEP(2);
                STEP(3);
                STEP(4);
                STEP(5);
                STEP(6);
                STEP(7);

                if (out == ref) {
                        printf("\nkey = %016llx\n", (unsigned long long)key);
                }
                if (((key + 1) & 0xFFFFF) == 0) {
                        uint64_t now;

                        now = current_time();
                        printf("\r%.2f keys/s    ",
                                (double)key * 1000000.0 / (now - orig));
                        fflush(stdout);
                }
        }
        return 0;
}