Repository: numpy/numpy
File: numpy/_core/src/npysort/binsearch.cpp (C++)

Code before:
         */
        if (cmp(last_key_val, key_val)) {
            max_idx = arr_len;
        }
        else {
            min_idx = 0;
            max_idx = (max_idx < arr_len) ? (max_idx + 1) : arr_len;
        }

        last_key_val = key_val;

        while (min_idx < max_idx) {
            const npy_intp mid_idx = min_idx + ((max_idx - min_idx) >> 1);
            const T mid_val = *(const T *)(arr + mid_idx * arr_str);
            if (cmp(mid_val, key_val)) {
                min_idx = mid_idx + 1;
            }
            else {
                max_idx = mid_idx;
            }
        }
        *(npy_intp *)ret = min_idx;
    }
}

template <class Tag, side_t side>
static int
argbinsearch(const char *arr, const char *key, const char *sort, char *ret,
             npy_intp arr_len, npy_intp key_len, npy_intp arr_str,
             npy_intp key_str, npy_intp sort_str, npy_intp ret_str,
             PyArrayObject *)
{
    using T = typename Tag::type;
    auto cmp = side_to_cmp<Tag, side>::value;
    npy_intp min_idx = 0;
    npy_intp max_idx = arr_len;
    T last_key_val;

    if (key_len == 0) {
        return 0;
    }
    last_key_val = *(const T *)key;

    for (; key_len > 0; key_len--, key += key_str, ret += ret_str) {
        const T key_val = *(const T *)key;
        /*
         * Updating only one of the indices based on the previous key
         * gives the search a big boost when keys are sorted, but slightly
         * slows down things for purely random ones.
         */
        if (cmp(last_key_val, key_val)) {

Review comment (suggestion):
Shouldn't this be `)` on the right? There also needs to be a distinction between range , and (base, length) pairs.

Code after:

    To batch multiple queries, we do one pivot pass for each different length
    for all keys, storing intermediate values of base_j for every key_j. To
    avoid consuming extra memory, we use the ret array to store intermediate
    values of each base until they become the final result in the last step.

    There are two benefits of this approach:

    1. Cache locallity of pivots. In early iterations each key is compared
    against the same set of pivots. For example, in the first iteration all
    keys are compared against the median. In the second iteration, all keys
    end up being compared against 1st and 3rd quartiles.

    2. Independent calculations for out-of-order execution. In the single-key
    version, step i+1 depends on computation of step i. Meaning that step i+1
    must wait for step i to complete before proceeding. When batching multiple
    keys, we compute each step for all keys before continuing on the next
    step. All the computations at a given step are independent across
    different keys. Meaning that the CPU can execute multiple keys
    out-of-order in parallel.

    Invariant:
    - cmp(arr[i], key_val) == true  for all i < base
    - cmp(arr[i], key_val) == false for all i >= base + length

    The insertion index candidates are in range [base, base+length] and
    on each iteration we shrink the range into either
        [base, ceil(length / 2)]
    or
        [base + floor(length / 2), ceil(length / 2)]

    Optimization: we unroll the first iteration for the following reasons:
        1. ret is not initialized with the bases, so we save |keys| writes
        by not having to intialize it with 0s.
        2. By assuming the initial base for every key is 0, we also save
        |keys| reads.
        3. In the first iteration, all elements are compared against the
        median. So we can store it in a variable and use it for all keys.

    This initial block replaces the initialization loop that is used for the
    arr_len==0 case. Note that when arr_len = 1, then half is 0 so the
    following block initializes the array as with 0s.
    */
    npy_intp interval_length = arr_len;
    npy_intp half = interval_length >> 1;
    interval_length -= half; // length -> ceil(length / 2)

    npy_intp base = 0;
    const T mid_val = *(const T *)(arr + (base + half) * arr_str);

    for (npy_intp i = 0; i < key_len; ++i) {