// Sorted-set intersection of a small array against a much larger array,
// using Knuth's optimal ("cmov-friendly") binary search inside fixed-size
// intervals of the large array.
//
// Design overview
// ---------------
// Both inputs are sorted ascending. We pick an interval size of roughly
// |large| / (B + 1) where B = |small| / ln(2). This splits `large` into
// B + 1 contiguous intervals. Each interval is delimited by a "tick"
// (its rightmost element); ticks let us decide which interval can possibly
// contain a given query from `small` with a single comparison per interval.
//
// For each query we (a) skip whole intervals by walking ticks until the
// query is bounded above by the interval's tick, then (b) run a branchless
// binary search inside that single interval. Because the linear walk is
// shared across consecutive queries (small is sorted, so successive queries
// only ever move forward), the amortized work per query is bounded.
//
// The width-wise variant (`LocalCmovSearchN<width>`) processes `width`
// consecutive queries in lock-step: their per-interval binary searches run
// inside a single loop over `n`, exposing instruction-level parallelism and
// keeping the inner loop free of unpredictable branches (the conditional
// updates lower to cmov on x86-64).
//
// Reference: Daniel Lemire, "Bounding the cost of the intersection between
// a small array and a large array," Daniel Lemire's blog, June 27, 2019.
// https://lemire.me/blog/2019/06/27/bounding-the-cost-of-the-intersection-between-a-small-array-and-a-large-array/
//
// Vaguely related: 
// Daniel Lemire, "Speeding up independent binary searches by interleaving them," in Daniel
// Lemire's blog, September 14, 2019, https://lemire.me/blog/2019/09/14/speeding-up-independent-binary-searches-by-interleaving-them/.
//
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <span>
#include <vector>

inline size_t simple_round(double x) {
  return static_cast<size_t>(x + 0.5);
}


  // Width-wise local cmov search: locates each of `width` sorted query values
  // inside [start, finish) of `haystack`, where the queries are known to be
  // partitioned across intervals of size `interval_size`, and writes any
  // matches to `out`.
  // Precondition: queries are sorted ascending and queries[width-1] is in
  // the final interval [finish - interval_size, finish).
  template <size_t width, class elem>
  elem * LocalCmovSearchN(std::span<elem> haystack, size_t start, size_t finish,
                          size_t interval_size,
                          std::span<const elem, width> queries, elem * out) {
    static_assert(width >= 1, "width must be at least 1");

    std::array<size_t, width> regionStart{};

    // Walk forward to find the interval for queries[0..width-2]. Since the
    // queries are sorted, each search resumes where the previous one left off.
    if constexpr (width >= 2) {
      size_t r = start;
      for (size_t i = 0; i + 1 < width; ++i) {
        while (queries[i] >= haystack[r + interval_size]) {
          r += interval_size;
        }
        regionStart[i] = r;
      }
    }
    // The largest query is guaranteed to fit in the final interval.
    regionStart[width - 1] = finish - interval_size;

    // Synchronized cmov-friendly binary search within each interval.
    size_t n = interval_size;
    while (n >= 2) {
      const size_t span = n / 2;
      for (size_t i = 0; i < width; ++i) {
        if (queries[i] >= haystack[regionStart[i] + span]) {
          regionStart[i] += span;
        }
      }
      n -= span;
    }

    for (size_t i = 0; i < width; ++i) {
      if (haystack[regionStart[i]] == queries[i]) {
        *out++ = queries[i];
      }
    }
    return out;
  }


template <class elem, size_t width>
size_t knuth_optimal_cmov(std::span<elem> small,  std::span<elem> large, elem* output) {
    static_assert(width >= 1, "width must be at least 1");
    const auto orig = output;
    if(small.empty() || large.empty()) {
      return 0; // These would mess up the math badly.
    }
    const auto largesize = large.size(); // >= 1
    // the number of subintervals should be B + 1, generated by B ticks.
    const double B = small.size() * 1.4426950408889634 - 1; // 1.44 is 1/ln(2)
    // B >= 0.444... so no negative numbers involved.
    // interval size must be at least one
    const size_t interval_size = std::max(size_t{1}, simple_round(largesize / ( B + 1 ))); // largesize >= 1
    size_t small_iter = 0;
    // We start with the first B intervals, always pointing at the end of the interval.
    size_t pointer = interval_size;
    elem max_val {0};
    if(pointer >= large.size())  {
       goto end;
    }
    max_val = large[pointer];

    for(; small_iter + width <= small.size() ; small_iter += width) {
       auto x = small[small_iter + width - 1];
       size_t pointer_start = pointer;
       while(x >= max_val) {
         pointer += interval_size;
         if(pointer >= largesize) {
           // reset
           pointer = pointer_start;
           max_val = large[pointer];
           goto nextpart;
         }
         max_val = large[pointer];
       }
       output = LocalCmovSearchN<width>(
                large, pointer_start - interval_size, pointer, interval_size,
                std::span<const elem, width>{small.data() + small_iter, width},
                output);
    }
    nextpart:
    for(; small_iter < small.size() ; small_iter++) {
       auto x = small[small_iter];
       while(x >= max_val) {
         pointer += interval_size;
         if(pointer >= largesize) goto end;
         max_val = large[pointer];
       }
       output = LocalCmovSearchN<1>(
                large, pointer - interval_size, pointer, interval_size,
                std::span<const elem, 1>{small.data() + small_iter, 1},
                output);
    }
    end:
    // we have one interval subinterval left, its size is (largesize - pointer + interval_size).
    // Strictly speaking, we don't need to do it separately.
    if (pointer - interval_size < largesize) {
      const size_t tail_size = largesize - (pointer - interval_size);
      while (small_iter < small.size()) {
        output = LocalCmovSearchN<1>(
                 large, pointer - interval_size, largesize, tail_size,
                 std::span<const elem, 1>{small.data() + small_iter, 1},
                 output);
        ++small_iter;
      }
    }
    return output - orig;
}


// Example usage: intersect a small sorted set against a much larger sorted
// set. Compile with a C++20 compiler, e.g.
//     c++ -std=c++20 -O3 sortedintersection.cpp -o sortedintersection
int main() {
  // The "small" array — sorted ascending.
  std::vector<uint32_t> small = {3, 8, 14, 22, 50, 91, 100, 137};

  // The "large" array — here the even numbers in [0, 200).
  std::vector<uint32_t> large;
  large.reserve(100);
  for (uint32_t v = 0; v < 200; v += 2) {
    large.push_back(v);
  }

  // Output buffer must be at least small.size().
  std::vector<uint32_t> output(small.size());

  // Process queries 4 at a time inside the wide search.
  constexpr size_t width = 4;
  const size_t count = knuth_optimal_cmov<uint32_t, width>(
      std::span<uint32_t>(small),
      std::span<uint32_t>(large),
      output.data());

  std::cout << "Intersection (" << count << " element"
            << (count == 1 ? "" : "s") << "):";
  for (size_t i = 0; i < count; ++i) {
    std::cout << ' ' << output[i];
  }
  std::cout << '\n';
  // Expected: 8 14 22 50 100 — the even values from `small` that appear in `large`.
  return 0;
}
