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

typedef const void *(*fn_t)(const void *, int, size_t);

#define ONES UINT64_C(0x0101010101010101)
#define HIGHS UINT64_C(0x8080808080808080)
#define HASZERO(x) (((x) - ONES) & ~(x) & HIGHS)

static inline uint64_t load64(const unsigned char *p) {
  uint64_t v;
  memcpy(&v, p, 8);
  return v;
}

__attribute__((noinline)) static const void *naive_memchr(const void *s, int c, size_t n) {
  const unsigned char *p = s;
  for (size_t i = 0; i < n; i++) if (p[i] == (unsigned char)c) return p + i;
  return NULL;
}

__attribute__((noinline)) static const void *swar_memchr(const void *s, int c, size_t n) {
  const unsigned char *p = s;
  uint64_t k = ONES * (unsigned char)c;
  while (n >= 8) {
    uint64_t m = HASZERO(load64(p) ^ k);
    if (m) return p + (__builtin_ctzll(m) >> 3);
    p += 8;
    n -= 8;
  }
  for (; n; p++, n--) if (*p == (unsigned char)c) return p;
  return NULL;
}

__attribute__((noinline)) static const void *swar2_memchr(const void *s, int c, size_t n) {
  const unsigned char *p = s;
  uint64_t k = ONES * (unsigned char)c;
  while (n >= 16) {
    uint64_t m0 = HASZERO(load64(p) ^ k);
    uint64_t m1 = HASZERO(load64(p + 8) ^ k);
    if (m0 | m1) {
      if (m0) return p + (__builtin_ctzll(m0) >> 3);
      return p + 8 + (__builtin_ctzll(m1) >> 3);
    }
    p += 16;
    n -= 16;
  }
  return swar_memchr(p, c, n);
}

__attribute__((noinline)) static const void *sse2_memchr(const void *s, int c, size_t n) {
  const unsigned char *p = s;
  __m128i needle = _mm_set1_epi8((char)c);
  while (n >= 16) {
    int mask = _mm_movemask_epi8(_mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)p), needle));
    if (mask) return p + __builtin_ctz(mask);
    p += 16;
    n -= 16;
  }
  for (; n; p++, n--) if (*p == (unsigned char)c) return p;
  return NULL;
}

__attribute__((noinline)) static const void *sse2x4_memchr(const void *s, int c, size_t n) {
  const unsigned char *p = s;
  __m128i needle = _mm_set1_epi8((char)c);
  while (n >= 64) {
    __m128i e0 = _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)p), needle);
    __m128i e1 = _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(p + 16)), needle);
    __m128i e2 = _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(p + 32)), needle);
    __m128i e3 = _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(p + 48)), needle);
    if (_mm_movemask_epi8(_mm_or_si128(_mm_or_si128(e0, e1), _mm_or_si128(e2, e3)))) {
      int m;
      if ((m = _mm_movemask_epi8(e0))) return p + __builtin_ctz(m);
      if ((m = _mm_movemask_epi8(e1))) return p + 16 + __builtin_ctz(m);
      if ((m = _mm_movemask_epi8(e2))) return p + 32 + __builtin_ctz(m);
      return p + 48 + __builtin_ctz(_mm_movemask_epi8(e3));
    }
    p += 64;
    n -= 64;
  }
  return sse2_memchr(p, c, n);
}

static double run(fn_t f, const unsigned char *buf, size_t n, long iters) {
  LARGE_INTEGER freq, t0, t1;
  size_t sum = 0;
  QueryPerformanceFrequency(&freq);
  QueryPerformanceCounter(&t0);
  for (long i = 0; i < iters; i++) {
    __asm__ volatile("" : : "r"(buf) : "memory");
    sum += (size_t)((const unsigned char *)f(buf, '\n', n) - buf);
  }
  QueryPerformanceCounter(&t1);
  if (sum != (size_t)iters * (n - 1)) printf("bad sum\n");
  return (double)(t1.QuadPart - t0.QuadPart) / freq.QuadPart;
}

int main(void) {
  const char *names[] = {"ucrtbase memchr", "naive", "swar", "swar x2", "sse2", "sse2 x4"};
  fn_t fns[6];
  fns[0] = (fn_t)GetProcAddress(LoadLibraryA("ucrtbase.dll"), "memchr");
  fns[1] = naive_memchr;
  fns[2] = swar_memchr;
  fns[3] = swar2_memchr;
  fns[4] = sse2_memchr;
  fns[5] = sse2x4_memchr;

  /* correctness against the naive loop */
  unsigned char rbuf[512];
  srand(1);
  for (int t = 0; t < 200000; t++) {
    size_t len = rand() % 300;
    size_t off = rand() % 64;
    for (size_t i = 0; i < len; i++) rbuf[off + i] = "ab\n\x80\xff"[rand() % 5];
    int c = (unsigned char)("\n" "\x80" "a"[rand() % 3]);
    const void *want = naive_memchr(rbuf + off, c, len);
    for (int k = 0; k < 6; k++) {
      if (fns[k](rbuf + off, c, len) != want) {
        printf("MISMATCH %s len=%zu off=%zu\n", names[k], len, off);
        return 1;
      }
    }
  }
  printf("correctness: ok (200000 random cases)\n");

  size_t sizes[] = {16, 64, 256, 1040, 65536, 1 << 20};
  size_t maxn = 1 << 20;
  unsigned char *buf = malloc(maxn);
  printf("%-16s", "size");
  for (int s = 0; s < 6; s++) printf("%14zu", sizes[s]);
  printf("\n");
  for (int k = 0; k < 6; k++) {
    printf("%-16s", names[k]);
    for (int s = 0; s < 6; s++) {
      size_t n = sizes[s];
      memset(buf, 'a', n - 1);
      buf[n - 1] = '\n';
      long iters = (long)((512ULL << 20) / n);
      double t = run(fns[k], buf, n, iters);
      printf("%9.1f GB/s", (double)n * iters / t / 1e9);
    }
    printf("\n");
  }

  /* edge cases: ns per call. pos < 0 means no match. */
  struct { const char *label; size_t n; long pos; } cases[] = {
    {"n=1, match", 1, 0},          {"n=7, match at end", 7, 6},
    {"n=15, match at end", 15, 14}, {"n=1040, match at 0", 1040, 0},
    {"n=16, no match", 16, -1},     {"n=64, no match", 64, -1},
    {"n=1040, no match", 1040, -1},
  };
  printf("\n%-22s", "ns per call");
  for (int k = 0; k < 6; k++) printf("%16s", names[k]);
  printf("\n");
  for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) {
    size_t n = cases[ci].n;
    memset(buf, 'a', n);
    if (cases[ci].pos >= 0) buf[cases[ci].pos] = '\n';
    const unsigned char *want = cases[ci].pos >= 0 ? buf + cases[ci].pos : NULL;
    long iters = 20000000;
    printf("%-22s", cases[ci].label);
    for (int k = 0; k < 6; k++) {
      LARGE_INTEGER freq, t0, t1;
      size_t bad = 0;
      QueryPerformanceFrequency(&freq);
      QueryPerformanceCounter(&t0);
      for (long i = 0; i < iters; i++) {
        __asm__ volatile("" : : "r"(buf) : "memory");
        bad += fns[k](buf, '\n', n) != want;
      }
      QueryPerformanceCounter(&t1);
      double ns = (double)(t1.QuadPart - t0.QuadPart) / freq.QuadPart / iters * 1e9;
      if (bad) printf("%16s", "WRONG");
      else printf("%16.2f", ns);
    }
    printf("\n");
  }
  return 0;
}