Skip to main content

diskann_quantization/multi_vector/distance/
kernel.rs

1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT license.
3
4//! Object-safe kernel boundary trait plus BYOTE visitor trait.
5
6use crate::multi_vector::{MatRef, MaxSimError, Standard};
7
8/// Object-safe interface for computing per-query MaxSim scores.
9pub trait MaxSimKernel<T: Copy>: Send + Sync + std::fmt::Debug {
10    /// Number of query rows whose scores this kernel produces.
11    fn nrows(&self) -> usize;
12
13    /// Compute per-query MaxSim scores into `scores`. On zero docs, fills
14    /// every slot with `f32::MAX`.
15    ///
16    /// # Errors
17    ///
18    /// * [`MaxSimError::InvalidBufferLength`] if `scores.len() != self.nrows()`.
19    /// * [`MaxSimError::UnequalDim`] if `doc.vector_dim()` does not match the internal
20    ///   vector dim.
21    fn compute_max_sim(
22        &self,
23        doc: MatRef<'_, Standard<T>>,
24        scores: &mut [f32],
25    ) -> Result<(), MaxSimError>;
26}
27
28/// "Bring your own type erasure" visitor: the factory hands a concrete
29/// kernel to [`Erase::erase`], which decides how to package it (e.g. as
30/// `Box<dyn MaxSimKernel<T>>` via [`BoxErase`], a chamfer-only closure, a
31/// batched evaluator, …).
32pub trait Erase<T: Copy> {
33    type Output;
34    /// `K` is generic so the body sees its concrete type and the compiler
35    /// can inline it.
36    fn erase<K: MaxSimKernel<T> + 'static>(self, kernel: K) -> Self::Output;
37}
38
39/// Default boxing [`Erase`] impl.
40#[derive(Debug, Clone, Copy)]
41pub struct BoxErase;
42
43impl<T: Copy + 'static> Erase<T> for BoxErase {
44    type Output = Box<dyn MaxSimKernel<T>>;
45
46    fn erase<K: MaxSimKernel<T> + 'static>(self, kernel: K) -> Self::Output {
47        Box::new(kernel)
48    }
49}