diskann-benchmark-core 0.56.0

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use diskann::{
    graph::{SearchOutputBuffer, search_output_buffer::BufferState},
    neighbor::Neighbor,
};

/// A [`SearchOutputBuffer`] implementation that either references a slice in-place for
/// fixed sized outputs or references a growable vector.
///
/// We put this behind an enum to cut down on the amount of code that gets monomorphized.
///
/// If the underlying buffer is a slice, the number of written elements is tracked separately and
/// is available via [`Buffer::current_len`]. When the underlying buffer is a `Vec`, the length
/// of the vector is used.
#[derive(Debug)]
pub(crate) struct Buffer<'a, I>(Inner<'a, I>);

#[derive(Debug)]
enum Inner<'a, I> {
    Slice { slice: &'a mut [I], written: usize },
    Vec(&'a mut Vec<I>),
}

impl<'a, I> Buffer<'a, I> {
    /// Construct a buffer that writes into `slice`.
    pub(crate) fn slice(slice: &'a mut [I]) -> Self {
        Self(Inner::Slice { slice, written: 0 })
    }

    /// Construct a vector. This clear `v` when called.
    pub(crate) fn vector(v: &'a mut Vec<I>) -> Self {
        v.clear();
        Self(Inner::Vec(v))
    }

    /// Return the number of elements that have been written to the buffer.
    ///
    /// This is useful to track instances where searches return fewer than the
    /// requested number of ids.
    pub(crate) fn current_len(&self) -> usize {
        match &self.0 {
            Inner::Slice { written, .. } => *written,
            Inner::Vec(vec) => vec.len(),
        }
    }
}

impl<I, D> SearchOutputBuffer<I, D> for Buffer<'_, I> {
    fn size_hint(&self) -> Option<usize> {
        match &self.0 {
            Inner::Slice { slice, written } => Some(slice.len() - written),
            Inner::Vec(_) => None,
        }
    }

    fn current_len(&self) -> usize {
        <Buffer<I>>::current_len(self)
    }

    fn push(&mut self, neighbor: Neighbor<I, D>) -> BufferState {
        let (id, _) = neighbor.as_tuple();
        match &mut self.0 {
            Inner::Slice { slice, written } => match slice.get_mut(*written) {
                Some(slot) => {
                    *slot = id;
                    *written += 1;
                    if *written == slice.len() {
                        BufferState::Full
                    } else {
                        BufferState::Available
                    }
                }
                None => BufferState::Full,
            },
            Inner::Vec(vec) => {
                vec.push(id);
                BufferState::Available
            }
        }
    }

    fn extend<Itr>(&mut self, itr: Itr) -> usize
    where
        Itr: IntoIterator<Item = Neighbor<I, D>>,
    {
        match &mut self.0 {
            Inner::Slice { slice, written } => match slice.get_mut(*written..) {
                Some(left) => {
                    let count = std::iter::zip(left.iter_mut(), itr)
                        .map(|(dst, src)| {
                            *dst = src.as_tuple().0;
                        })
                        .count();
                    *written += count;
                    count
                }
                None => 0,
            },
            Inner::Vec(vec) => {
                let before = vec.len();
                vec.extend(itr.into_iter().map(|i| i.as_tuple().0));
                vec.len() - before
            }
        }
    }
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;

    fn size_hint<T>(buffer: &T) -> Option<usize>
    where
        T: SearchOutputBuffer<u32, f32>,
    {
        buffer.size_hint()
    }

    fn from_tuple<I, D>((id, distance): (I, D)) -> Neighbor<I, D> {
        Neighbor::new(id, distance)
    }

    #[test]
    fn test_slice_buffer_creation() {
        let mut data = [0u32; 5];
        let buffer = Buffer::slice(&mut data);

        assert_eq!(buffer.current_len(), 0);
        assert_eq!(size_hint(&buffer), Some(5));
    }

    #[test]
    fn test_vec_buffer_creation() {
        let mut vec = vec![1u32, 2, 3, 4];
        let buffer = Buffer::vector(&mut vec);

        assert_eq!(buffer.current_len(), 0);
        assert_eq!(size_hint(&buffer), None);
        assert!(vec.is_empty(), "vector should be cleared on construction");
    }

    #[test]
    fn test_slice_buffer_set_single_element() {
        let mut data = [0u32; 5];
        let mut buffer = Buffer::slice(&mut data);

        assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 1);
        assert_eq!(data, [42, 0, 0, 0, 0]);
    }

    #[test]
    fn test_slice_buffer_set_updates_written() {
        let mut data = [0u32; 5];
        let mut buffer = Buffer::slice(&mut data);

        assert_eq!(buffer.current_len(), 0);
        assert_eq!(size_hint(&buffer), Some(5));

        assert_eq!(buffer.push(Neighbor::new(100, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 1);
        assert_eq!(size_hint(&buffer), Some(4));

        assert_eq!(buffer.push(Neighbor::new(200, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 2);
        assert_eq!(size_hint(&buffer), Some(3));

        assert_eq!(buffer.push(Neighbor::new(300, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 3);
        assert_eq!(size_hint(&buffer), Some(2));

        assert_eq!(buffer.push(Neighbor::new(400, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 4);
        assert_eq!(size_hint(&buffer), Some(1));

        assert_eq!(buffer.push(Neighbor::new(500, 0.0)), BufferState::Full);
        assert_eq!(buffer.current_len(), 5);
        assert_eq!(size_hint(&buffer), Some(0));

        assert_eq!(buffer.push(Neighbor::new(600, 0.0)), BufferState::Full);
        assert_eq!(buffer.current_len(), 5);
        assert_eq!(size_hint(&buffer), Some(0));

        // Check that `data` was written properly.
        assert_eq!(data, [100, 200, 300, 400, 500]);
    }

    #[test]
    fn test_slice_buffer_empty() {
        let mut data = [0u32; 0];
        let mut buffer = Buffer::slice(&mut data);

        assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Full);
    }

    #[test]
    fn test_vec_buffer_push() {
        let mut vec = vec![0u32, 0, 0, 0, 0];
        let mut buffer = Buffer::vector(&mut vec);
        assert_eq!(
            size_hint(&buffer),
            None,
            "vector-type buffers have no upper bound"
        );

        assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 1);

        assert_eq!(buffer.push(Neighbor::new(50, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 2);

        assert_eq!(buffer.push(Neighbor::new(3, 0.0)), BufferState::Available);
        assert_eq!(buffer.current_len(), 3);

        assert_eq!(&vec, &[42, 50, 3]);
    }

    #[test]
    fn test_slice_buffer_set_from() {
        let mut data = [0u32; 5];
        let mut buffer = Buffer::slice(&mut data);

        let items = [(10, 1.0), (20, 2.0), (30, 3.0)].map(from_tuple);
        let count = buffer.extend(items);

        assert_eq!(count, 3);
        assert_eq!(buffer.current_len(), 3);
        assert_eq!(data, [10, 20, 30, 0, 0]);
    }

    #[test]
    fn test_slice_buffer_set_from_more_than_capacity() {
        let mut data = [0u32; 3];
        let mut buffer = Buffer::slice(&mut data);

        let items = [(10, 1.0), (20, 2.0), (30, 3.0), (40, 4.0), (50, 5.0)].map(from_tuple);
        let count = buffer.extend(items);

        // Only first 3 items should be written
        assert_eq!(count, 3);
        assert_eq!(buffer.current_len(), 3);
        assert_eq!(data, [10, 20, 30]);
    }

    #[test]
    fn test_vec_buffer_extend() {
        let mut vec = Vec::<u32>::new();
        let mut buffer = Buffer::vector(&mut vec);

        let items = [(100, 1.0), (200, 2.0)].map(from_tuple);
        let count = buffer.extend(items);

        assert_eq!(count, 2);
        assert_eq!(buffer.current_len(), 2);
        assert_eq!(&vec, &[100, 200]);
    }

    #[test]
    fn test_vec_buffer_extend_cascades() {
        let mut vec = vec![1u32, 2, 3, 4, 5];
        let mut buffer = Buffer::vector(&mut vec);

        let items = [(10, 1.0), (20, 2.0)].map(from_tuple);
        assert_eq!(buffer.extend(items), 2);

        let items = [(21, 1.0), (22, 2.0)].map(from_tuple);
        assert_eq!(buffer.extend(items), 2);

        assert_eq!(
            buffer.extend::<[Neighbor<u32, f32>; 0]>([]),
            0,
            "empty iterator should add nothing"
        );

        // Previous items should be cleared
        assert_eq!(&vec, &[10, 20, 21, 22]);
    }

    #[test]
    fn test_slice_push_and_extend_combinations() {
        // Push then extend
        let mut data = [0u32; 5];
        let mut buffer = Buffer::slice(&mut data);
        assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available);
        assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Available);
        assert_eq!(buffer.extend([(3, 0.0), (4, 0.0)].map(from_tuple)), 2);
        assert_eq!(data, [1, 2, 3, 4, 0]);

        // Extend then push to fill
        let mut data = [0u32; 5];
        let mut buffer = Buffer::slice(&mut data);
        buffer.extend([(10, 0.0), (20, 0.0)].map(from_tuple));
        assert_eq!(buffer.push(Neighbor::new(30, 0.0)), BufferState::Available);
        assert_eq!(buffer.push(Neighbor::new(40, 0.0)), BufferState::Available);
        assert_eq!(buffer.push(Neighbor::new(50, 0.0)), BufferState::Full);
        assert_eq!(data, [10, 20, 30, 40, 50]);

        // Interleaved operations
        let mut data = [0u32; 6];
        let mut buffer = Buffer::slice(&mut data);
        assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available);
        assert_eq!(buffer.extend([(2, 0.0), (3, 0.0)].map(from_tuple)), 2);
        assert_eq!(buffer.push(Neighbor::new(4, 0.0)), BufferState::Available);
        assert_eq!(
            buffer.extend([(5, 0.0), (6, 0.0), (7, 0.0)].map(from_tuple)),
            2
        );
        assert_eq!(data, [1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn test_slice_push_and_extend_at_capacity() {
        // Extend fills buffer, push returns Full
        let mut data = [0u32; 2];
        let mut buffer = Buffer::slice(&mut data);
        buffer.extend([(1, 0.0), (2, 0.0)].map(from_tuple));
        assert_eq!(buffer.push(Neighbor::new(99, 0.0)), BufferState::Full);
        assert_eq!(data, [1, 2]);

        // Push fills buffer, extend returns 0
        let mut data = [0u32; 2];
        let mut buffer = Buffer::slice(&mut data);
        assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available);
        assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Full);
        assert_eq!(buffer.extend([Neighbor::new(99, 0.0)]), 0);
        assert_eq!(data, [1, 2]);

        // Extend truncates when exceeding remaining capacity after push
        let mut data = [0u32; 4];
        let mut buffer = Buffer::slice(&mut data);
        assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available);
        assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Available);
        assert_eq!(
            buffer.extend([(3, 0.0), (4, 0.0), (5, 0.0), (6, 0.0)].map(from_tuple)),
            2
        );
        assert_eq!(data, [1, 2, 3, 4]);
    }

    #[test]
    fn test_vec_push_and_extend_combinations() {
        let mut vec = Vec::<u32>::new();
        let mut buffer = Buffer::vector(&mut vec);

        // Interleave push and extend - vec has no capacity limit
        assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available);
        assert_eq!(buffer.extend([(2, 0.0), (3, 0.0)].map(from_tuple)), 2);
        assert_eq!(buffer.push(Neighbor::new(4, 0.0)), BufferState::Available);
        assert_eq!(buffer.extend::<[Neighbor<u32, f32>; 0]>([]), 0); // empty extend
        assert_eq!(buffer.extend(vec![Neighbor::new(5, 0.0)]), 1);

        assert_eq!(&vec, &[1, 2, 3, 4, 5]);
    }
}