vicinity 0.11.1

//! HNSW search algorithm with early termination optimizations.

use std::cell::RefCell;
use std::collections::{BinaryHeap, HashSet};

pub(crate) use crate::prefetch::prefetch_read_data;

#[cfg(feature = "benchmark")]
thread_local! {
    static SEARCH_COUNTERS: RefCell<HnswSearchCounters> = const {
        RefCell::new(HnswSearchCounters::new())
    };
}

/// Diagnostic counters for HNSW search benchmarks.
#[cfg(feature = "benchmark")]
#[derive(Clone, Copy, Debug, Default)]
pub struct HnswSearchCounters {
    /// Number of full vector distance evaluations.
    pub distance_evals: u64,
    /// Number of result-heap inserts while the heap is still filling.
    pub result_insertions: u64,
    /// Number of result-heap top replacements after the heap is full.
    pub result_replacements: u64,
    /// Number of distance-scored candidates rejected by the result heap.
    pub result_rejections: u64,
    /// Number of candidate-frontier heap pushes.
    pub candidate_pushes: u64,
    /// Number of candidate-frontier heap pops.
    pub candidate_pops: u64,
    /// Number of frontier-pruning retain calls.
    pub frontier_retain_calls: u64,
    /// Number of candidate entries removed by frontier pruning.
    pub frontier_pruned_candidates: u64,
    /// Maximum candidate-frontier heap length observed.
    pub max_frontier_len: usize,
}

#[cfg(feature = "benchmark")]
impl HnswSearchCounters {
    const fn new() -> Self {
        Self {
            distance_evals: 0,
            result_insertions: 0,
            result_replacements: 0,
            result_rejections: 0,
            candidate_pushes: 0,
            candidate_pops: 0,
            frontier_retain_calls: 0,
            frontier_pruned_candidates: 0,
            max_frontier_len: 0,
        }
    }
}

/// Reset the current thread's HNSW search benchmark counters.
#[cfg(feature = "benchmark")]
pub fn reset_search_counters() {
    SEARCH_COUNTERS.with(|cell| {
        *cell.borrow_mut() = HnswSearchCounters::new();
    });
}

/// Return and reset the current thread's HNSW search benchmark counters.
#[cfg(feature = "benchmark")]
pub fn take_search_counters() -> HnswSearchCounters {
    SEARCH_COUNTERS.with(|cell| {
        let mut counters = cell.borrow_mut();
        let out = *counters;
        *counters = HnswSearchCounters::new();
        out
    })
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_distance_evals(count: usize) {
    SEARCH_COUNTERS.with(|cell| {
        cell.borrow_mut().distance_evals += count as u64;
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_result_insertion() {
    SEARCH_COUNTERS.with(|cell| {
        cell.borrow_mut().result_insertions += 1;
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_result_replacement() {
    SEARCH_COUNTERS.with(|cell| {
        cell.borrow_mut().result_replacements += 1;
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_result_rejection() {
    SEARCH_COUNTERS.with(|cell| {
        cell.borrow_mut().result_rejections += 1;
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_candidate_push(candidates_len: usize) {
    SEARCH_COUNTERS.with(|cell| {
        let mut counters = cell.borrow_mut();
        counters.candidate_pushes += 1;
        counters.max_frontier_len = counters.max_frontier_len.max(candidates_len);
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_candidate_pop() {
    SEARCH_COUNTERS.with(|cell| {
        cell.borrow_mut().candidate_pops += 1;
    });
}

#[cfg(feature = "benchmark")]
#[inline]
fn record_frontier_retain(before: usize, after: usize) {
    SEARCH_COUNTERS.with(|cell| {
        let mut counters = cell.borrow_mut();
        counters.frontier_retain_calls += 1;
        counters.frontier_pruned_candidates += before.saturating_sub(after) as u64;
    });
}

// ─── Visited set ─────────────────────────────────────────────────────────────

/// Threshold below which we use a dense generation-counter array instead of HashSet.
/// 4M nodes = 4MB Vec<u8>, fits in L3 cache on most hardware (L3 is typically 8-32MB).
/// The dense array has O(1) clear (generation bump) vs O(capacity) for HashSet, and
/// avoids hashing overhead on every visited-node check.
const DENSE_VISITED_THRESHOLD: usize = 4_000_000;

/// Number of neighbor distances to compute before updating the HNSW heaps.
const DISTANCE_BATCH_SIZE: usize = 8;

/// Beam width where caching the current worst result starts paying for itself.
///
/// Criterion on the synthetic HNSW search bench showed the cached-worst loop
/// helps higher-ef search, but regresses very small beams where the result heap
/// changes often and `peek()` is cheap.
const CACHED_WORST_MIN_EF: usize = 64;

/// Beam width where periodically pruning stale frontier candidates may pay off.
const FRONTIER_PRUNE_MIN_EF: usize = 64;

/// Number of candidate pops between frontier-pruning attempts.
const FRONTIER_PRUNE_INTERVAL: usize = 64;

/// Fast visited-node tracker using the generation-counter pattern.
///
/// Dense variant: a `Vec<u8>` where `visited[id] == generation` means visited.
/// Incrementing `generation` logically clears the set in O(1). Only when the
/// u8 counter wraps (every 255 searches) does a memset occur.
///
/// Uses u8 (not u16) to minimize cache footprint: 1 byte/node = same as
/// the old Vec<bool>, but with O(1) clear instead of O(n).
///
/// Falls back to `HashSet<u32>` for large indexes where a full dense array
/// would waste memory.
pub(crate) enum VisitedSet {
    Dense { marks: Vec<u8>, generation: u8 },
    Sparse(HashSet<u32>),
}

impl VisitedSet {
    /// Create a visited set sized for `num_nodes` total nodes.
    fn new(num_nodes: usize, capacity_hint: usize) -> Self {
        if num_nodes <= DENSE_VISITED_THRESHOLD {
            Self::dense(num_nodes)
        } else {
            VisitedSet::Sparse(HashSet::with_capacity(capacity_hint))
        }
    }

    /// Create a dense generation-counter tracker for contiguous node IDs.
    pub(crate) fn dense(num_nodes: usize) -> Self {
        VisitedSet::Dense {
            marks: vec![0u8; num_nodes],
            generation: 1,
        }
    }

    /// Reset the visited set for a new search. O(1) amortized for the dense
    /// variant (increments generation; memsets only on u8 overflow).
    /// For the sparse variant, clears the HashSet.
    fn clear(&mut self) {
        match self {
            VisitedSet::Dense { marks, generation } => {
                if let Some(next) = generation.checked_add(1) {
                    *generation = next;
                } else {
                    // Overflow: reset all marks and restart at generation 1
                    marks.fill(0);
                    *generation = 1;
                }
            }
            VisitedSet::Sparse(s) => s.clear(),
        }
    }

    /// Check whether a node has been visited.
    #[cfg(test)]
    #[inline]
    fn contains(&self, id: u32) -> bool {
        match self {
            VisitedSet::Dense { marks, generation } => {
                let idx = id as usize;
                idx < marks.len() && marks[idx] == *generation
            }
            VisitedSet::Sparse(s) => s.contains(&id),
        }
    }

    /// Mark a node as visited. Returns `true` if the node was NOT previously visited.
    #[inline]
    pub(crate) fn insert(&mut self, id: u32) -> bool {
        match self {
            VisitedSet::Dense { marks, generation } => {
                let idx = id as usize;
                debug_assert!(
                    idx < marks.len(),
                    "VisitedSet::insert: id {} out of bounds (capacity {})",
                    id,
                    marks.len()
                );
                if idx < marks.len() {
                    if marks[idx] != *generation {
                        marks[idx] = *generation;
                        true
                    } else {
                        false
                    }
                } else {
                    // Out-of-bounds: treat as unvisited but don't track.
                    // This shouldn't happen if num_vectors is correct.
                    true
                }
            }
            VisitedSet::Sparse(s) => s.insert(id),
        }
    }

    /// Prepare for a new search with `num_nodes` total nodes. Reuses the
    /// existing allocation when possible, only reallocating if the index grew.
    fn prepare(&mut self, num_nodes: usize, capacity_hint: usize) {
        match self {
            VisitedSet::Dense { marks, .. } if num_nodes <= DENSE_VISITED_THRESHOLD => {
                if marks.len() < num_nodes {
                    // Index grew: resize and reset
                    marks.resize(num_nodes, 0);
                    // Force a full reset after resize since new slots are 0
                    // and we need generation != 0
                }
                self.clear();
            }
            VisitedSet::Sparse(s) if num_nodes > DENSE_VISITED_THRESHOLD => {
                s.clear();
            }
            _ => {
                // Variant mismatch (index crossed threshold): recreate
                *self = VisitedSet::new(num_nodes, capacity_hint);
            }
        }
    }
}

thread_local! {
    static THREAD_VISITED: RefCell<VisitedSet> = const { RefCell::new(
        VisitedSet::Dense { marks: Vec::new(), generation: 1 }
    ) };
    static THREAD_SEARCH_SCRATCH: RefCell<SearchScratch> = const { RefCell::new(SearchScratch::new()) };
}

struct SearchScratch {
    candidates: BinaryHeap<MinCandidate>,
    results: BinaryHeap<MaxResult>,
}

impl SearchScratch {
    const fn new() -> Self {
        Self {
            candidates: BinaryHeap::new(),
            results: BinaryHeap::new(),
        }
    }

    fn prepare(&mut self, ef: usize) {
        self.candidates.clear();
        self.results.clear();
        self.candidates.reserve(ef.saturating_mul(2));
        self.results.reserve(ef.saturating_add(1));
    }
}

/// Borrow the thread-local visited set, prepared for `num_nodes`.
/// The closure receives a mutable reference to the reused set.
pub(crate) fn with_visited_set<F, R>(num_nodes: usize, capacity_hint: usize, f: F) -> R
where
    F: FnOnce(&mut VisitedSet) -> R,
{
    THREAD_VISITED.with(|cell| {
        let mut visited = cell.borrow_mut();
        visited.prepare(num_nodes, capacity_hint);
        f(&mut visited)
    })
}

// ─── Candidate types ─────────────────────────────────────────────────────────

/// Candidate for min-heap (explore closest first).
#[derive(PartialEq)]
pub(crate) struct MinCandidate {
    pub(crate) id: u32,
    pub(crate) distance: f32,
}
impl Eq for MinCandidate {}
impl Ord for MinCandidate {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.distance.total_cmp(&self.distance)
    }
}
impl PartialOrd for MinCandidate {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

/// Result for max-heap (track worst result for pruning).
#[derive(PartialEq)]
pub(crate) struct MaxResult {
    pub(crate) id: u32,
    pub(crate) distance: f32,
}
impl Eq for MaxResult {}
impl Ord for MaxResult {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.distance.total_cmp(&other.distance)
    }
}
impl PartialOrd for MaxResult {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

// ─── Batched distance helpers ────────────────────────────────────────────────

#[inline]
fn insert_result_if_accepted(
    results: &mut BinaryHeap<MaxResult>,
    ef: usize,
    id: u32,
    distance: f32,
) -> bool {
    if results.len() < ef {
        results.push(MaxResult { id, distance });
        #[cfg(feature = "benchmark")]
        record_result_insertion();
        return true;
    }

    let Some(mut worst) = results.peek_mut() else {
        #[cfg(feature = "benchmark")]
        record_result_rejection();
        return false;
    };
    if distance >= worst.distance {
        #[cfg(feature = "benchmark")]
        record_result_rejection();
        return false;
    }

    *worst = MaxResult { id, distance };
    #[cfg(feature = "benchmark")]
    record_result_replacement();
    true
}

/// Process a batch of neighbor distance computations.
/// Computes distances for all batch entries, then inserts qualifying results
/// into the candidate and result heaps.
#[inline]
fn flush_batch(
    query: &[f32],
    batch_ids: &[u32; DISTANCE_BATCH_SIZE],
    count: usize,
    vectors: &[f32],
    dimension: usize,
    dist_fn: fn(&[f32], &[f32]) -> f32,
    candidates: &mut std::collections::BinaryHeap<MinCandidate>,
    results: &mut std::collections::BinaryHeap<MaxResult>,
    ef: usize,
) {
    // Compute all distances first (enables ILP -- CPU can overlap independent FP pipelines).
    let mut dists = [0.0f32; DISTANCE_BATCH_SIZE];
    for i in 0..count {
        let vec = get_vector(vectors, dimension, batch_ids[i] as usize);
        dists[i] = dist_fn(query, vec);
    }
    #[cfg(feature = "benchmark")]
    record_distance_evals(count);

    if ef < CACHED_WORST_MIN_EF {
        for i in 0..count {
            if insert_result_if_accepted(results, ef, batch_ids[i], dists[i]) {
                candidates.push(MinCandidate {
                    id: batch_ids[i],
                    distance: dists[i],
                });
                #[cfg(feature = "benchmark")]
                record_candidate_push(candidates.len());
            }
        }
        return;
    }

    // Process high-ef results with cached worst distance. Keep the current worst
    // distance locally so the inner loop does not repeatedly peek the result heap
    // when no update lands.
    let mut worst_dist = if results.len() < ef {
        f32::INFINITY
    } else {
        results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY)
    };
    for i in 0..count {
        if results.len() < ef || dists[i] < worst_dist {
            if insert_result_if_accepted(results, ef, batch_ids[i], dists[i]) {
                candidates.push(MinCandidate {
                    id: batch_ids[i],
                    distance: dists[i],
                });
                #[cfg(feature = "benchmark")]
                record_candidate_push(candidates.len());
            }
            worst_dist = if results.len() < ef {
                f32::INFINITY
            } else {
                results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY)
            };
        } else {
            #[cfg(feature = "benchmark")]
            record_result_rejection();
        }
    }
}

/// Same as [`flush_batch`] but for custom distance functions.
#[inline]
fn flush_batch_custom<F: Fn(&[f32], u32) -> f32>(
    query: &[f32],
    batch_ids: &[u32; DISTANCE_BATCH_SIZE],
    count: usize,
    dist_fn: &F,
    candidates: &mut std::collections::BinaryHeap<MinCandidate>,
    results: &mut std::collections::BinaryHeap<MaxResult>,
    ef: usize,
) {
    let mut dists = [0.0f32; DISTANCE_BATCH_SIZE];
    for i in 0..count {
        dists[i] = dist_fn(query, batch_ids[i]);
    }
    #[cfg(feature = "benchmark")]
    record_distance_evals(count);
    if ef < CACHED_WORST_MIN_EF {
        for i in 0..count {
            if insert_result_if_accepted(results, ef, batch_ids[i], dists[i]) {
                candidates.push(MinCandidate {
                    id: batch_ids[i],
                    distance: dists[i],
                });
                #[cfg(feature = "benchmark")]
                record_candidate_push(candidates.len());
            }
        }
        return;
    }

    let mut worst_dist = if results.len() < ef {
        f32::INFINITY
    } else {
        results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY)
    };
    for i in 0..count {
        if results.len() < ef || dists[i] < worst_dist {
            if insert_result_if_accepted(results, ef, batch_ids[i], dists[i]) {
                candidates.push(MinCandidate {
                    id: batch_ids[i],
                    distance: dists[i],
                });
                #[cfg(feature = "benchmark")]
                record_candidate_push(candidates.len());
            }
            worst_dist = if results.len() < ef {
                f32::INFINITY
            } else {
                results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY)
            };
        } else {
            #[cfg(feature = "benchmark")]
            record_result_rejection();
        }
    }
}

#[inline]
fn prune_unpromising_candidates(
    candidates: &mut BinaryHeap<MinCandidate>,
    results: &BinaryHeap<MaxResult>,
    ef: usize,
) {
    if ef < FRONTIER_PRUNE_MIN_EF || results.len() < ef || candidates.len() < ef {
        return;
    }
    let Some(worst) = results.peek() else {
        return;
    };
    let worst_dist = worst.distance;
    #[cfg(feature = "benchmark")]
    let before = candidates.len();
    candidates.retain(|candidate| candidate.distance <= worst_dist);
    #[cfg(feature = "benchmark")]
    record_frontier_retain(before, candidates.len());
}

// ─── Search functions ────────────────────────────────────────────────────────

/// Greedy search in a single layer using standard HNSW beam search.
///
/// Implements the correct HNSW search from Malkov & Yashunin (2016):
/// - Uses min-heap for candidates (explore closest first)
/// - Uses max-heap for results (track worst result for pruning)
/// - Continues until best unexplored candidate is worse than worst result
///
/// This is critical for achieving high recall (~98% on standard benchmarks).
#[cfg(feature = "hnsw")]
pub fn greedy_search_layer(
    query: &[f32],
    entry_point: u32,
    layer: &crate::hnsw::graph::Layer,
    vectors: &[f32],
    dimension: usize,
    ef: usize,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> Vec<(u32, f32)> {
    let num_vectors = vectors.len() / dimension;

    THREAD_SEARCH_SCRATCH.with(|scratch_cell| {
        let mut scratch = scratch_cell.borrow_mut();
        scratch.prepare(ef);
        let SearchScratch {
            candidates,
            results,
        } = &mut *scratch;

        with_visited_set(num_vectors, ef * 2, |visited| {
            // Start from entry point
            let entry_vector = get_vector(vectors, dimension, entry_point as usize);
            let entry_distance = dist_fn(query, entry_vector);
            #[cfg(feature = "benchmark")]
            record_distance_evals(1);
            candidates.push(MinCandidate {
                id: entry_point,
                distance: entry_distance,
            });
            #[cfg(feature = "benchmark")]
            record_candidate_push(candidates.len());
            results.push(MaxResult {
                id: entry_point,
                distance: entry_distance,
            });
            #[cfg(feature = "benchmark")]
            record_result_insertion();
            visited.insert(entry_point);

            // Standard HNSW beam search:
            // Continue while we have candidates that might improve results
            let should_prune_frontier = ef >= FRONTIER_PRUNE_MIN_EF;
            let mut pops_since_prune = 0usize;
            while let Some(candidate) = candidates.pop() {
                #[cfg(feature = "benchmark")]
                record_candidate_pop();

                // Stopping condition: if best candidate is worse than worst result
                // and we have enough results, we're done
                let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
                if candidate.distance > worst_dist && results.len() >= ef {
                    break;
                }
                if should_prune_frontier {
                    pops_since_prune += 1;
                }

                // Explore neighbors using batched distance computation.
                // Collect a small group of unvisited neighbors, prefetch their vectors,
                // compute distances together for better ILP and cache behavior.
                let neighbors = layer.get_neighbors(candidate.id);

                let mut batch_ids: [u32; DISTANCE_BATCH_SIZE] = [0; DISTANCE_BATCH_SIZE];
                let mut batch_count = 0usize;

                for &neighbor_id in neighbors.iter() {
                    if visited.insert(neighbor_id) {
                        batch_ids[batch_count] = neighbor_id;
                        batch_count += 1;

                        // Prefetch the vector for this neighbor (it will be read shortly).
                        if (neighbor_id as usize) < num_vectors {
                            let ptr = vectors
                                .as_ptr()
                                .wrapping_add(neighbor_id as usize * dimension);
                            prefetch_read_data(ptr);
                            if dimension > 16 {
                                prefetch_read_data(ptr.wrapping_add(16));
                            }
                        }

                        if batch_count == DISTANCE_BATCH_SIZE {
                            flush_batch(
                                query,
                                &batch_ids,
                                batch_count,
                                vectors,
                                dimension,
                                dist_fn,
                                candidates,
                                results,
                                ef,
                            );
                            batch_count = 0;
                        }
                    }
                }
                // Flush remaining (< 4).
                if batch_count > 0 {
                    flush_batch(
                        query,
                        &batch_ids,
                        batch_count,
                        vectors,
                        dimension,
                        dist_fn,
                        candidates,
                        results,
                        ef,
                    );
                }

                if should_prune_frontier && pops_since_prune >= FRONTIER_PRUNE_INTERVAL {
                    prune_unpromising_candidates(candidates, results, ef);
                    pops_since_prune = 0;
                }
            }

            let mut output: Vec<(u32, f32)> = results.drain().map(|r| (r.id, r.distance)).collect();
            candidates.clear();
            output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
            output
        })
    })
}

/// Greedy search seeded from multiple entry points.
///
/// Same beam search as [`greedy_search_layer`], but initialises the candidate
/// heap with all provided `entry_points`. Duplicate entry IDs are silently
/// deduplicated via the visited set. If `entry_points` is empty this returns
/// an empty result.
///
/// Used by `batch_search_mqo` to warm-start a query from the best result of a
/// nearby query already processed in the batch.
#[cfg(feature = "hnsw")]
pub fn greedy_search_layer_multi_entry(
    query: &[f32],
    entry_points: &[u32],
    layer: &crate::hnsw::graph::Layer,
    vectors: &[f32],
    dimension: usize,
    ef: usize,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> Vec<(u32, f32)> {
    if entry_points.is_empty() {
        return Vec::new();
    }

    let num_vectors = vectors.len() / dimension;

    with_visited_set(num_vectors, ef * 2, |visited| {
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::with_capacity(ef * 2);
        let mut results: BinaryHeap<MaxResult> = BinaryHeap::with_capacity(ef + 1);

        // Seed from every provided entry point.
        for &ep in entry_points {
            if visited.insert(ep) {
                let ep_vec = get_vector(vectors, dimension, ep as usize);
                let ep_dist = dist_fn(query, ep_vec);
                candidates.push(MinCandidate {
                    id: ep,
                    distance: ep_dist,
                });
                let _ = insert_result_if_accepted(&mut results, ef, ep, ep_dist);
            }
        }

        while let Some(candidate) = candidates.pop() {
            let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
            if candidate.distance > worst_dist && results.len() >= ef {
                break;
            }

            let neighbors = layer.get_neighbors(candidate.id);
            for (i, &neighbor_id) in neighbors.iter().enumerate() {
                if i + 1 < neighbors.len() {
                    let next_id = neighbors[i + 1] as usize;
                    if next_id < num_vectors {
                        let ptr = vectors.as_ptr().wrapping_add(next_id * dimension);
                        prefetch_read_data(ptr);
                        prefetch_read_data(ptr.wrapping_add(16));
                    }
                }
                if i + 4 < neighbors.len() {
                    let far_id = neighbors[i + 4] as usize;
                    if far_id < num_vectors {
                        prefetch_read_data(vectors.as_ptr().wrapping_add(far_id * dimension));
                    }
                }
                if visited.insert(neighbor_id) {
                    let neighbor_vector = get_vector(vectors, dimension, neighbor_id as usize);
                    let neighbor_distance = dist_fn(query, neighbor_vector);

                    if insert_result_if_accepted(&mut results, ef, neighbor_id, neighbor_distance) {
                        candidates.push(MinCandidate {
                            id: neighbor_id,
                            distance: neighbor_distance,
                        });
                    }
                }
            }
        }

        let mut output: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect();
        output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
        output
    })
}

/// Greedy search with a caller-provided distance function.
///
/// Unlike [`greedy_search_layer`] which computes `dist(query, stored_vector)`,
/// this variant takes a closure `dist_fn(query, internal_node_id) -> f32`.
/// The caller is responsible for resolving the node ID to whatever data they
/// need (box geometry, quantized codes, etc.).
///
/// This enables asymmetric distance computation: the graph is built with
/// center-to-center distance for navigability, but search uses any distance
/// function (box-to-point, quantized-to-float, etc.).
///
/// Prefetching is still performed on the stored vector array to keep the
/// center vectors warm for the common case where the custom distance function
/// reads them.
#[cfg(feature = "hnsw")]
pub fn greedy_search_layer_custom<F: Fn(&[f32], u32) -> f32>(
    query: &[f32],
    entry_point: u32,
    layer: &crate::hnsw::graph::Layer,
    vectors: &[f32],
    dimension: usize,
    ef: usize,
    dist_fn: &F,
) -> Vec<(u32, f32)> {
    let num_vectors = vectors.len() / dimension;

    with_visited_set(num_vectors, ef * 2, |visited| {
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::with_capacity(ef * 2);
        let mut results: BinaryHeap<MaxResult> = BinaryHeap::with_capacity(ef + 1);

        // Start from entry point
        let entry_distance = dist_fn(query, entry_point);
        candidates.push(MinCandidate {
            id: entry_point,
            distance: entry_distance,
        });
        results.push(MaxResult {
            id: entry_point,
            distance: entry_distance,
        });
        visited.insert(entry_point);

        while let Some(candidate) = candidates.pop() {
            let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
            if candidate.distance > worst_dist && results.len() >= ef {
                break;
            }

            let neighbors = layer.get_neighbors(candidate.id);
            let mut batch_ids: [u32; DISTANCE_BATCH_SIZE] = [0; DISTANCE_BATCH_SIZE];
            let mut batch_count = 0usize;

            for &neighbor_id in neighbors.iter() {
                if visited.insert(neighbor_id) {
                    batch_ids[batch_count] = neighbor_id;
                    batch_count += 1;

                    // Prefetch vector for this neighbor.
                    if (neighbor_id as usize) < num_vectors {
                        let ptr = vectors
                            .as_ptr()
                            .wrapping_add(neighbor_id as usize * dimension);
                        prefetch_read_data(ptr);
                        if dimension > 16 {
                            prefetch_read_data(ptr.wrapping_add(16));
                        }
                    }

                    if batch_count == DISTANCE_BATCH_SIZE {
                        flush_batch_custom(
                            query,
                            &batch_ids,
                            batch_count,
                            dist_fn,
                            &mut candidates,
                            &mut results,
                            ef,
                        );
                        batch_count = 0;
                    }
                }
            }
            if batch_count > 0 {
                flush_batch_custom(
                    query,
                    &batch_ids,
                    batch_count,
                    dist_fn,
                    &mut candidates,
                    &mut results,
                    ef,
                );
            }
        }

        let mut output: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect();
        output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
        output
    })
}

/// Greedy search where distance depends on both candidate (parent) and neighbor.
///
/// Used by vertex-relative quantization (SymphonyQGVR) where each edge (u -> v)
/// has a separate quantized code relative to u.
#[cfg(feature = "ivf_rabitq")]
///
/// The closure receives `(parent_id, neighbor_id, neighbor_slot)` where
/// `neighbor_slot` is the index within the parent's neighbor list (for looking
/// up per-edge codes). For the entry point, `parent_id == entry_point` (self-referential).
pub fn greedy_search_layer_edge_aware<F: Fn(u32, u32, usize) -> f32>(
    entry_point: u32,
    entry_dist: f32,
    layer: &crate::hnsw::graph::Layer,
    num_vectors: usize,
    ef: usize,
    dist_fn: &F,
) -> Vec<(u32, f32)> {
    with_visited_set(num_vectors, ef * 2, |visited| {
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::with_capacity(ef * 2);
        let mut results: BinaryHeap<MaxResult> = BinaryHeap::with_capacity(ef + 1);

        candidates.push(MinCandidate {
            id: entry_point,
            distance: entry_dist,
        });
        results.push(MaxResult {
            id: entry_point,
            distance: entry_dist,
        });
        visited.insert(entry_point);

        while let Some(candidate) = candidates.pop() {
            let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
            if candidate.distance > worst_dist && results.len() >= ef {
                break;
            }

            let neighbors = layer.get_neighbors(candidate.id);
            for (slot, &neighbor_id) in neighbors.iter().enumerate() {
                if visited.insert(neighbor_id) {
                    // The dist_fn closure accesses per-edge data; prefetching
                    // is the caller's responsibility via the closure itself.
                    let neighbor_distance = dist_fn(candidate.id, neighbor_id, slot);

                    if insert_result_if_accepted(&mut results, ef, neighbor_id, neighbor_distance) {
                        candidates.push(MinCandidate {
                            id: neighbor_id,
                            distance: neighbor_distance,
                        });
                    }
                }
            }
        }

        let mut output: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect();
        output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
        output
    })
}

/// Greedy search with adaptive early termination.
///
/// Same beam search as [`greedy_search_layer`], but uses an
/// `EarlyTerminationOracle` to stop once the distance
/// distribution suggests further exploration is unlikely to improve the top-k.
///
/// Returns `(results, num_evaluated)` so callers can inspect how many
/// distance computations were performed.
#[cfg(feature = "hnsw")]
pub fn greedy_search_layer_adaptive(
    query: &[f32],
    entry_point: u32,
    layer: &crate::hnsw::graph::Layer,
    vectors: &[f32],
    dimension: usize,
    ef: usize,
    k: usize,
    config: &crate::adaptive::AdaptiveConfig,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> (Vec<(u32, f32)>, usize) {
    use crate::adaptive::EarlyTerminationOracle;

    let num_vectors = vectors.len() / dimension;

    with_visited_set(num_vectors, ef * 2, |visited| {
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::with_capacity(ef * 2);
        let mut results: BinaryHeap<MaxResult> = BinaryHeap::with_capacity(ef + 1);
        let mut oracle = EarlyTerminationOracle::new(k, config.clone());

        // Start from entry point
        let entry_vector = get_vector(vectors, dimension, entry_point as usize);
        let entry_distance = dist_fn(query, entry_vector);
        oracle.observe(entry_distance);

        candidates.push(MinCandidate {
            id: entry_point,
            distance: entry_distance,
        });
        results.push(MaxResult {
            id: entry_point,
            distance: entry_distance,
        });
        visited.insert(entry_point);

        while let Some(candidate) = candidates.pop() {
            let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
            if candidate.distance > worst_dist && results.len() >= ef {
                break;
            }

            // Explore neighbors
            let neighbors = layer.get_neighbors(candidate.id);
            for (i, &neighbor_id) in neighbors.iter().enumerate() {
                // Prefetch upcoming neighbors' vectors to hide DRAM latency.
                // Two cache lines for vectors >64 bytes; 4-ahead for pipeline depth.
                if i + 1 < neighbors.len() {
                    let next_id = neighbors[i + 1] as usize;
                    if next_id < num_vectors {
                        let ptr = vectors.as_ptr().wrapping_add(next_id * dimension);
                        prefetch_read_data(ptr);
                        prefetch_read_data(ptr.wrapping_add(16));
                    }
                }
                if i + 4 < neighbors.len() {
                    let far_id = neighbors[i + 4] as usize;
                    if far_id < num_vectors {
                        prefetch_read_data(vectors.as_ptr().wrapping_add(far_id * dimension));
                    }
                }
                if visited.insert(neighbor_id) {
                    let neighbor_vector = get_vector(vectors, dimension, neighbor_id as usize);
                    let neighbor_distance = dist_fn(query, neighbor_vector);
                    oracle.observe(neighbor_distance);

                    if insert_result_if_accepted(&mut results, ef, neighbor_id, neighbor_distance) {
                        candidates.push(MinCandidate {
                            id: neighbor_id,
                            distance: neighbor_distance,
                        });
                    }
                }
            }

            // After exploring this candidate's neighbors, check early termination
            if oracle.should_terminate() && results.len() >= k {
                break;
            }
        }

        let num_evaluated = oracle.num_evaluated();
        let mut output: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect();
        output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
        (output, num_evaluated)
    })
}

/// Get vector from SoA storage.
#[inline]
fn get_vector(vectors: &[f32], dimension: usize, idx: usize) -> &[f32] {
    let start = idx * dimension;
    let end = start + dimension;
    &vectors[start..end]
}

/// Greedy search with PRT (Probabilistic Routing Test) pre-filtering.
///
/// Before computing the full O(d) distance to a candidate, checks the
/// O(k)-dimensional projected distance. Candidates whose projected distance
/// exceeds the current worst result (scaled by the TFB ratio) are skipped
/// entirely, saving the full distance computation.
///
/// Returns `(results, full_distance_count)`.
#[cfg(feature = "hnsw")]
pub fn greedy_search_layer_prt(
    query: &[f32],
    entry_point: u32,
    layer: &crate::hnsw::graph::Layer,
    vectors: &[f32],
    dimension: usize,
    ef: usize,
    dist_fn: fn(&[f32], &[f32]) -> f32,
    prt: &crate::prt::ProbabilisticRoutingTest,
    query_proj: &[f32],
    tfb: &mut crate::prt::TestFeedbackBuffer,
) -> (Vec<(u32, f32)>, usize) {
    let num_vectors = vectors.len() / dimension;
    let mut full_dist_count: usize = 0;

    let results = with_visited_set(num_vectors, ef * 2, |visited| {
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::with_capacity(ef * 2);
        let mut results: BinaryHeap<MaxResult> = BinaryHeap::with_capacity(ef + 1);

        let entry_vector = get_vector(vectors, dimension, entry_point as usize);
        let entry_distance = dist_fn(query, entry_vector);
        full_dist_count += 1;

        candidates.push(MinCandidate {
            id: entry_point,
            distance: entry_distance,
        });
        results.push(MaxResult {
            id: entry_point,
            distance: entry_distance,
        });
        visited.insert(entry_point);

        while let Some(candidate) = candidates.pop() {
            let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
            if candidate.distance > worst_dist && results.len() >= ef {
                break;
            }

            let neighbors = layer.get_neighbors(candidate.id);
            for (i, &neighbor_id) in neighbors.iter().enumerate() {
                // Prefetch.
                if i + 1 < neighbors.len() {
                    let next_id = neighbors[i + 1] as usize;
                    if next_id < num_vectors {
                        let ptr = vectors.as_ptr().wrapping_add(next_id * dimension);
                        prefetch_read_data(ptr);
                    }
                }
                if !visited.insert(neighbor_id) {
                    continue;
                }

                // PRT pre-filter: skip candidates whose projected distance
                // exceeds the worst result distance (scaled by TFB ratio).
                let worst_dist = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY);
                if results.len() >= ef
                    && !prt.should_compute_full_distance(query_proj, neighbor_id, worst_dist, tfb)
                {
                    continue;
                }

                // Full distance computation.
                let neighbor_vector = get_vector(vectors, dimension, neighbor_id as usize);
                let neighbor_distance = dist_fn(query, neighbor_vector);
                full_dist_count += 1;

                if insert_result_if_accepted(&mut results, ef, neighbor_id, neighbor_distance) {
                    tfb.record_true_positive();
                    candidates.push(MinCandidate {
                        id: neighbor_id,
                        distance: neighbor_distance,
                    });
                } else {
                    tfb.record_false_positive();
                }
            }
        }

        let mut output: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect();
        output.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
        output
    });

    (results, full_dist_count)
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn test_candidate_ordering() {
        let mut heap = BinaryHeap::new();
        heap.push(MinCandidate {
            id: 0,
            distance: 0.5,
        });
        heap.push(MinCandidate {
            id: 1,
            distance: 0.1,
        });
        heap.push(MinCandidate {
            id: 2,
            distance: 0.3,
        });

        // Should pop in order: 0.1, 0.3, 0.5 (min-heap)
        assert_eq!(heap.pop().unwrap().distance, 0.1);
        assert_eq!(heap.pop().unwrap().distance, 0.3);
        assert_eq!(heap.pop().unwrap().distance, 0.5);
    }

    #[test]
    fn test_visited_set_dense() {
        let mut v = VisitedSet::new(100, 10);
        assert!(!v.contains(5));
        assert!(v.insert(5));
        assert!(v.contains(5));
        assert!(!v.insert(5)); // already visited
    }

    #[test]
    fn test_visited_set_dense_clear() {
        let mut v = VisitedSet::new(100, 10);
        assert!(v.insert(5));
        assert!(v.contains(5));
        v.clear();
        // After clear, previously visited nodes are no longer marked
        assert!(!v.contains(5));
        assert!(v.insert(5));
    }

    #[test]
    fn test_visited_set_dense_generation_overflow() {
        // Create a dense set and force generation to u8::MAX
        let mut v = VisitedSet::new(100, 10);
        if let VisitedSet::Dense {
            ref mut generation, ..
        } = v
        {
            *generation = u8::MAX;
        }
        assert!(v.insert(5));
        assert!(v.contains(5));
        // This clear triggers the overflow path (memset)
        v.clear();
        assert!(!v.contains(5));
        assert!(v.insert(5));
        assert!(v.contains(5));
    }

    #[test]
    fn test_visited_set_sparse() {
        // Force sparse by using a large num_nodes
        let mut v = VisitedSet::new(DENSE_VISITED_THRESHOLD + 1, 10);
        assert!(!v.contains(42));
        assert!(v.insert(42));
        assert!(v.contains(42));
        assert!(!v.insert(42));
    }

    #[test]
    fn test_visited_set_sparse_clear() {
        let mut v = VisitedSet::new(DENSE_VISITED_THRESHOLD + 1, 10);
        assert!(v.insert(42));
        v.clear();
        assert!(!v.contains(42));
        assert!(v.insert(42));
    }
}