// gcc -O2  -o haystack haystack.c  -march=native
#include <stdint.h>
#include <stdio.h>
#include <x86intrin.h>
#include <stddef.h>
#include <assert.h>
#include <stdbool.h>

#define RDTSC_START(cycles)                                             \
    do {                                                                \
        uint32_t cyc_high, cyc_low;                                     \
        __asm volatile("cpuid\n"                                        \
                       "rdtsc\n"                                        \
                       "mov %%edx, %0\n"                                \
                       "mov %%eax, %1" :                                \
                       "=r" (cyc_high),                                 \
                       "=r"(cyc_low) :                                  \
                       : /* no read only */                             \
                       "%rax", "%rbx", "%rcx", "%rdx" /* clobbers */    \
                       );                                               \
        (cycles) = ((uint64_t)cyc_high << 32) | cyc_low;                \
    } while (0)

#define RDTSC_STOP(cycles)                                              \
    do {                                                                \
        uint32_t cyc_high, cyc_low;                                     \
        __asm volatile("rdtscp\n"                                       \
                       "mov %%edx, %0\n"                                \
                       "mov %%eax, %1\n"                                \
                       "cpuid" :                                        \
                       "=r"(cyc_high),                                  \
                       "=r"(cyc_low) :                                  \
                       /* no read only registers */ :                   \
                       "%rax", "%rbx", "%rcx", "%rdx" /* clobbers */    \
                       );                                               \
        (cycles) = ((uint64_t)cyc_high << 32) | cyc_low;                \
    } while (0)

static __attribute__ ((noinline))
uint64_t rdtsc_overhead_func(uint64_t dummy) {
    return dummy;
}

uint64_t global_rdtsc_overhead = (uint64_t) UINT64_MAX;

#define RDTSC_SET_OVERHEAD(test, repeat)			      \
  do {								      \
    uint64_t cycles_start, cycles_final, cycles_diff;		      \
    uint64_t min_diff = UINT64_MAX;				      \
    for (int i = 0; i < repeat; i++) {			      \
      __asm volatile("" ::: /* pretend to clobber */ "memory");	      \
      RDTSC_START(cycles_start);				      \
      test;							      \
      RDTSC_STOP(cycles_final);                                       \
      cycles_diff = (cycles_final - cycles_start);		      \
      if (cycles_diff < min_diff) min_diff = cycles_diff;	      \
    }								      \
    global_rdtsc_overhead = min_diff;				      \
    printf("rdtsc_overhead set to %d\n", (int)global_rdtsc_overhead);     \
  } while (0)							      \


#define BEST_TIME(test, expected, pre, repeat, size, verbose)                         \
        do {                                                              \
            if (global_rdtsc_overhead == UINT64_MAX) {                    \
               RDTSC_SET_OVERHEAD(rdtsc_overhead_func(1), repeat);        \
            }                                                             \
            if(verbose) printf("%-60s\t: ", #test);                                        \
            fflush(NULL);                                                 \
            uint64_t cycles_start, cycles_final, cycles_diff;             \
            uint64_t min_diff = (uint64_t)-1;                             \
            uint64_t sum_diff = 0;                                        \
            for (int i = 0; i < repeat; i++) {                            \
                pre;                                                      \
                __asm volatile("" ::: /* pretend to clobber */ "memory"); \
                RDTSC_START(cycles_start);                                \
                if(test != expected) {printf("not expected (%d , %d )",(int)test,(int)expected);break;}                     \
                RDTSC_STOP(cycles_final);                                \
                cycles_diff = (cycles_final - cycles_start - global_rdtsc_overhead);           \
                if (cycles_diff < min_diff) min_diff = cycles_diff;       \
                sum_diff += cycles_diff;                                  \
            }                                                             \
            uint64_t S = size;                                            \
            float cycle_per_op = (min_diff) / (double)S;                  \
            float avg_cycle_per_op = (sum_diff) / ((double)S * repeat);   \
            if(verbose) printf(" %.3f cycles per operation (best) ", cycle_per_op);   \
            if(verbose) printf("\t%.3f cycles per operation (avg) ", avg_cycle_per_op);   \
            if(verbose) printf("\n");                                                 \
            if(!verbose) printf(" %.3f ",cycle_per_op);                   \
            fflush(NULL);                                                 \
 } while (0)


#define _mm512_set1_epu8(c) _mm512_set1_epi32((uint32_t)(c) * 0x01010101u)

__mmask16 zero_byte_mask(const __m512i v) {

    const __m512i v01  = _mm512_set1_epi32(0x01010101u);
    const __m512i v80  = _mm512_set1_epi32(0x80808080u);

    const __m512i v1   = _mm512_sub_epi32(v, v01);
    // tmp1 = (v - 0x01010101) & ~v & 0x80808080
    const __m512i tmp1 = _mm512_ternarylogic_epi32(v1, v, v80, 0x20);

    return _mm512_test_epi32_mask(tmp1, tmp1);
}


size_t avx512f_strstr_v2_anysize(const char* string, size_t n, const char* needle, size_t k) {

    assert(n > 0);
    assert(k > 0);

    const __m512i first = _mm512_set1_epu8(needle[0]);
    const __m512i last  = _mm512_set1_epu8(needle[k - 1]);

    char* haystack = (char*)(string);
    char* end      = haystack + n;
    size_t counter = 0;
    for (/**/; haystack < end; haystack += 64) {

        const __m512i block_first = _mm512_loadu_si512(haystack + 0);
        const __m512i block_last  = _mm512_loadu_si512(haystack + k - 1);

#if 0
        const __m512i first_zeros = _mm512_xor_si512(block_first, first);
        const __m512i last_zeros  = _mm512_xor_si512(block_last, last);
        const __m512i zeros       = _mm512_or_si512(first_zeros, last_zeros);
#else
        const __m512i first_zeros = _mm512_xor_si512(block_first, first);
        /*
            first_zeros | block_last | last |  first_zeros | (block_last ^ last)
            ------------+------------+------+------------------------------------
                 0      |      0     |   0  |      0
                 0      |      0     |   1  |      1
                 0      |      1     |   0  |      1
                 0      |      1     |   1  |      0
                 1      |      0     |   0  |      1
                 1      |      0     |   1  |      1
                 1      |      1     |   0  |      1
                 1      |      1     |   1  |      1
        */
        const __m512i zeros       = _mm512_ternarylogic_epi32(first_zeros, block_last, last, 0xf6);
#endif

        uint32_t mask = zero_byte_mask(zeros);
        if(mask) counter ++;
   }

    return counter;
}


void demo() {
  const size_t N = 1000 * 64 ;
  printf("Prints the number of cycles per byte.\n");
  char * blah = malloc(N + 1);
  for(size_t i = 0; i < N; i++) blah[i] = (char) i;
  blah[N] = '\0';
  char * needle = "abcd";
  size_t expected =  avx512f_strstr_v2_anysize(blah, N, needle,4);
  const int repeat = 1;
  BEST_TIME(strstr (blah,needle), NULL, , repeat, N, true);
  BEST_TIME(avx512f_strstr_v2_anysize(blah, N, needle,4), expected, , repeat, N, true);
  free(blah);
}

int main() {
  demo();
}

