Initial commit

packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_handle.rs

@@ -0,0 +1,212 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use std::ffi::CString;
+use std::{io, ptr};
+
+use winapi::um::fileapi::OPEN_EXISTING;
+use winapi::um::winbase::{FILE_FLAG_NO_BUFFERING, FILE_FLAG_OVERLAPPED, FILE_FLAG_RANDOM_ACCESS};
+use winapi::um::winnt::{FILE_SHARE_DELETE, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE};
+
+use winapi::{
+    shared::minwindef::DWORD,
+    um::{
+        errhandlingapi::GetLastError,
+        fileapi::CreateFileA,
+        handleapi::{CloseHandle, INVALID_HANDLE_VALUE},
+        winnt::HANDLE,
+    },
+};
+
+pub const FILE_ATTRIBUTE_READONLY: DWORD = 0x00000001;
+
+/// `AccessMode` determines how a file can be accessed.
+/// These modes are used when creating or opening a file to decide what operations are allowed
+/// to be performed on the file.
+///
+/// # Variants
+///
+/// - `Read`: The file is opened in read-only mode.
+///
+/// - `Write`: The file is opened in write-only mode.
+///
+/// - `ReadWrite`: The file is opened for both reading and writing.
+pub enum AccessMode {
+    Read,
+    Write,
+    ReadWrite,
+}
+
+/// `ShareMode` determines how a file can be shared.
+///
+/// These modes are used when creating or opening a file to decide what operations other 
+/// opening instances of the file can perform on it.
+/// # Variants
+/// - `None`: Prevents other processes from opening a file if they request delete, 
+///   read, or write access.
+///
+/// - `Read`: Allows subsequent open operations on the same file to request read access.
+///
+/// - `Write`: Allows subsequent open operations on the same file file to request write access.
+///
+/// - `Delete`: Allows subsequent open operations on the same file file to request delete access.
+pub enum ShareMode {
+    None,
+    Read,
+    Write,
+    Delete,
+}
+
+/// # Windows File Handle Wrapper
+///
+/// Introduces a Rust-friendly wrapper around the native Windows `HANDLE` object, `FileHandle`.
+/// `FileHandle` provides safe creation and automatic cleanup of Windows file handles, leveraging Rust's ownership model.
+
+/// `FileHandle` struct that wraps a native Windows `HANDLE` object
+#[cfg(target_os = "windows")]
+pub struct FileHandle {
+    handle: HANDLE,
+}
+
+impl FileHandle {
+    /// Creates a new `FileHandle` by opening an existing file with the given access and shared mode.
+    ///
+    /// This function is marked unsafe because it creates a raw pointer to the filename and try to create
+    /// a Windows `HANDLE` object without checking if you have sufficient permissions.
+    ///
+    /// # Safety
+    ///
+    /// Ensure that the file specified by `file_name` is valid and the calling process has
+    /// sufficient permissions to perform the specified `access_mode` and `share_mode` operations.
+    ///
+    /// # Parameters
+    ///
+    /// - `file_name`: The name of the file.
+    /// - `access_mode`: The access mode to be used for the file.
+    /// - `share_mode`: The share mode to be used for the file
+    ///
+    /// # Errors
+    /// This function will return an error if the `file_name` is invalid or if the file cannot
+    /// be opened with the specified `access_mode` and `share_mode`.
+    pub unsafe fn new(
+        file_name: &str,
+        access_mode: AccessMode,
+        share_mode: ShareMode,
+    ) -> io::Result<Self> {
+        let file_name_c = CString::new(file_name).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("Invalid file name. {}", file_name),
+            )
+        })?;
+
+        let dw_desired_access = match access_mode {
+            AccessMode::Read => GENERIC_READ,
+            AccessMode::Write => GENERIC_WRITE,
+            AccessMode::ReadWrite => GENERIC_READ | GENERIC_WRITE,
+        };
+
+        let dw_share_mode = match share_mode {
+            ShareMode::None => 0,
+            ShareMode::Read => FILE_SHARE_READ,
+            ShareMode::Write => FILE_SHARE_WRITE,
+            ShareMode::Delete => FILE_SHARE_DELETE,
+        };
+
+        let dw_flags_and_attributes = FILE_ATTRIBUTE_READONLY
+            | FILE_FLAG_NO_BUFFERING
+            | FILE_FLAG_OVERLAPPED
+            | FILE_FLAG_RANDOM_ACCESS;
+
+        let handle = unsafe {
+            CreateFileA(
+                file_name_c.as_ptr(),
+                dw_desired_access,
+                dw_share_mode,
+                ptr::null_mut(),
+                OPEN_EXISTING,
+                dw_flags_and_attributes,
+                ptr::null_mut(),
+            )
+        };
+
+        if handle == INVALID_HANDLE_VALUE {
+            let error_code = unsafe { GetLastError() };
+            Err(io::Error::from_raw_os_error(error_code as i32))
+        } else {
+            Ok(Self { handle })
+        }
+    }
+
+    pub fn raw_handle(&self) -> HANDLE {
+        self.handle
+    }
+}
+
+impl Drop for FileHandle {
+    /// Automatically closes the `FileHandle` when it goes out of scope.
+    /// Any errors in closing the handle are logged, as `Drop` does not support returning `Result`.
+    fn drop(&mut self) {
+        let result = unsafe { CloseHandle(self.handle) };
+        if result == 0 {
+            let error_code = unsafe { GetLastError() };
+            let error = io::Error::from_raw_os_error(error_code as i32);
+
+            // Only log the error if dropping the handle fails, since Rust's Drop trait does not support returning Result types from the drop method,
+            // and panicking in the drop method is considered bad practice
+            log::warn!("Error when dropping IOCompletionPort: {:?}", error);
+        }
+    }
+}
+
+/// Returns a `FileHandle` with an `INVALID_HANDLE_VALUE`.
+impl Default for FileHandle {
+    fn default() -> Self {
+        Self {
+            handle: INVALID_HANDLE_VALUE,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs::File;
+    use std::path::Path;
+
+    #[test]
+    fn test_create_file() {
+        // Create a dummy file
+        let dummy_file_path = "dummy_file.txt";
+        {
+            let _file = File::create(dummy_file_path).expect("Failed to create dummy file.");
+        }
+
+        let path = Path::new(dummy_file_path);
+        {
+            let file_handle = unsafe {
+                FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read)
+            };
+
+            // Check that the file handle is valid
+            assert!(file_handle.is_ok());
+        }
+
+        // Try to delete the file. If the handle was correctly dropped, this should succeed.
+        match std::fs::remove_file(dummy_file_path) {
+            Ok(()) => (), // File was deleted successfully, which means the handle was closed.
+            Err(e) => panic!("Failed to delete file: {}", e), // Failed to delete the file, likely because the handle is still open.
+        }
+    }
+
+    #[test]
+    fn test_file_not_found() {
+        let path = Path::new("non_existent_file.txt");
+        let file_handle =
+            unsafe { FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read) };
+
+        // Check that opening a non-existent file returns an error
+        assert!(file_handle.is_err());
+    }
+}

packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/file_io.rs

@@ -0,0 +1,154 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+/// The module provides unsafe wrappers around two Windows API functions: `ReadFile` and `GetQueuedCompletionStatus`.
+///
+/// These wrappers aim to simplify and abstract the use of these functions, providing easier error handling and a safer interface.
+/// They return standard Rust `io::Result` types for convenience and consistency with the rest of the Rust standard library.
+use std::io;
+use std::ptr;
+
+use winapi::{
+    ctypes::c_void,
+    shared::{
+        basetsd::ULONG_PTR,
+        minwindef::{DWORD, FALSE},
+        winerror::{ERROR_IO_PENDING, WAIT_TIMEOUT},
+    },
+    um::{
+        errhandlingapi::GetLastError, fileapi::ReadFile, ioapiset::GetQueuedCompletionStatus,
+        minwinbase::OVERLAPPED,
+    },
+};
+
+use crate::FileHandle;
+use crate::IOCompletionPort;
+
+/// Asynchronously queue a read request from a file into a buffer slice.
+///
+/// Wraps the unsafe Windows API function `ReadFile`, making it safe to call only when the overlapped buffer
+/// remains valid and unchanged anywhere else during the entire async operation.
+///
+/// Returns a boolean indicating whether the read operation completed synchronously or is pending.
+///
+/// # Safety
+///
+/// This function is marked as `unsafe` because it uses raw pointers and requires the caller to ensure
+/// that the buffer slice and the overlapped buffer stay valid during the whole async operation.
+pub unsafe fn read_file_to_slice<T>(
+    file_handle: &FileHandle,
+    buffer_slice: &mut [T],
+    overlapped: *mut OVERLAPPED,
+    offset: u64,
+) -> io::Result<bool> {
+    let num_bytes = std::mem::size_of_val(buffer_slice);
+    unsafe {
+        ptr::write(overlapped, std::mem::zeroed());
+        (*overlapped).u.s_mut().Offset = offset as u32;
+        (*overlapped).u.s_mut().OffsetHigh = (offset >> 32) as u32;
+    }
+
+    let result = unsafe {
+        ReadFile(
+            file_handle.raw_handle(),
+            buffer_slice.as_mut_ptr() as *mut c_void,
+            num_bytes as DWORD,
+            ptr::null_mut(),
+            overlapped,
+        )
+    };
+
+    match result {
+        FALSE => {
+            let error = unsafe { GetLastError() };
+            if error != ERROR_IO_PENDING {
+                Err(io::Error::from_raw_os_error(error as i32))
+            } else {
+                Ok(false)
+            }
+        }
+        _ => Ok(true),
+    }
+}
+
+/// Retrieves the results of an asynchronous I/O operation on an I/O completion port.
+///
+/// Wraps the unsafe Windows API function `GetQueuedCompletionStatus`, making it safe to call only when the overlapped buffer
+/// remains valid and unchanged anywhere else during the entire async operation.
+///
+/// Returns a boolean indicating whether an I/O operation completed synchronously or is still pending.
+///
+/// # Safety
+///
+/// This function is marked as `unsafe` because it uses raw pointers and requires the caller to ensure
+/// that the overlapped buffer stays valid during the whole async operation.
+pub unsafe fn get_queued_completion_status(
+    completion_port: &IOCompletionPort,
+    lp_number_of_bytes: &mut DWORD,
+    lp_completion_key: &mut ULONG_PTR,
+    lp_overlapped: *mut *mut OVERLAPPED,
+    dw_milliseconds: DWORD,
+) -> io::Result<bool> {
+    let result = unsafe {
+        GetQueuedCompletionStatus(
+            completion_port.raw_handle(),
+            lp_number_of_bytes,
+            lp_completion_key,
+            lp_overlapped,
+            dw_milliseconds,
+        )
+    };
+
+    match result {
+        0 => {
+            let error = unsafe { GetLastError() };
+            if error == WAIT_TIMEOUT {
+                Ok(false)
+            } else {
+                Err(io::Error::from_raw_os_error(error as i32))
+            }
+        }
+        _ => Ok(true),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::file_handle::{AccessMode, ShareMode};
+
+    use super::*;
+    use std::fs::File;
+    use std::io::Write;
+    use std::path::Path;
+
+    #[test]
+    fn test_read_file_to_slice() {
+        // Create a temporary file and write some data into it
+        let path = Path::new("temp.txt");
+        {
+            let mut file = File::create(path).unwrap();
+            file.write_all(b"Hello, world!").unwrap();
+        }
+
+        let mut buffer: [u8; 512] = [0; 512];
+        let mut overlapped = unsafe { std::mem::zeroed::<OVERLAPPED>() };
+        {
+            let file_handle = unsafe {
+                FileHandle::new(path.to_str().unwrap(), AccessMode::Read, ShareMode::Read)
+            }
+            .unwrap();
+
+            // Call the function under test
+            let result =
+                unsafe { read_file_to_slice(&file_handle, &mut buffer, &mut overlapped, 0) };
+
+            assert!(result.is_ok());
+            let result_str = std::str::from_utf8(&buffer[.."Hello, world!".len()]).unwrap();
+            assert_eq!(result_str, "Hello, world!");
+        }
+
+        // Clean up
+        std::fs::remove_file("temp.txt").unwrap();
+    }
+}

packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/io_completion_port.rs

@@ -0,0 +1,142 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use std::io;
+
+use winapi::{
+    ctypes::c_void,
+    shared::{basetsd::ULONG_PTR, minwindef::DWORD},
+    um::{
+        errhandlingapi::GetLastError,
+        handleapi::{CloseHandle, INVALID_HANDLE_VALUE},
+        ioapiset::CreateIoCompletionPort,
+        winnt::HANDLE,
+    },
+};
+
+use crate::FileHandle;
+
+/// This module provides a safe and idiomatic Rust interface over the IOCompletionPort handle and associated Windows API functions.
+/// This struct represents an I/O completion port, which is an object used in asynchronous I/O operations on Windows.
+pub struct IOCompletionPort {
+    io_completion_port: HANDLE,
+}
+
+impl IOCompletionPort {
+    /// Create a new IOCompletionPort.
+    /// This function wraps the Windows CreateIoCompletionPort function, providing error handling and automatic resource management.
+    ///
+    /// # Arguments
+    ///
+    /// * `file_handle` - A reference to a FileHandle to associate with the IOCompletionPort.
+    /// * `existing_completion_port` - An optional reference to an existing IOCompletionPort. If provided, the new IOCompletionPort will be associated with it.
+    /// * `completion_key` - The completion key associated with the file handle.
+    /// * `number_of_concurrent_threads` - The maximum number of threads that the operating system can allow to concurrently process I/O completion packets for the I/O completion port.
+    ///
+    /// # Return
+    ///
+    /// Returns a Result with the new IOCompletionPort if successful, or an io::Error if the function fails.
+    pub fn new(
+        file_handle: &FileHandle,
+        existing_completion_port: Option<&IOCompletionPort>,
+        completion_key: ULONG_PTR,
+        number_of_concurrent_threads: DWORD,
+    ) -> io::Result<Self> {
+        let io_completion_port = unsafe {
+            CreateIoCompletionPort(
+                file_handle.raw_handle(),
+                existing_completion_port
+                    .map_or(std::ptr::null_mut::<c_void>(), |io_completion_port| {
+                        io_completion_port.raw_handle()
+                    }),
+                completion_key,
+                number_of_concurrent_threads,
+            )
+        };
+
+        if io_completion_port == INVALID_HANDLE_VALUE {
+            let error_code = unsafe { GetLastError() };
+            return Err(io::Error::from_raw_os_error(error_code as i32));
+        }
+
+        Ok(IOCompletionPort { io_completion_port })
+    }
+
+    pub fn raw_handle(&self) -> HANDLE {
+        self.io_completion_port
+    }
+}
+
+impl Drop for IOCompletionPort {
+    /// Drop method for IOCompletionPort.
+    /// This wraps the Windows CloseHandle function, providing automatic resource cleanup when the IOCompletionPort is dropped.
+    /// If an error occurs while dropping, it is logged and the drop continues. This is because panicking in Drop can cause unwinding issues.
+    fn drop(&mut self) {
+        let result = unsafe { CloseHandle(self.io_completion_port) };
+        if result == 0 {
+            let error_code = unsafe { GetLastError() };
+            let error = io::Error::from_raw_os_error(error_code as i32);
+
+            // Only log the error if dropping the handle fails, since Rust's Drop trait does not support returning Result types from the drop method,
+            // and panicking in the drop method is considered bad practice
+            log::warn!("Error when dropping IOCompletionPort: {:?}", error);
+        }
+    }
+}
+
+impl Default for IOCompletionPort {
+    /// Create a default IOCompletionPort, whose handle is set to INVALID_HANDLE_VALUE.
+    /// Returns a new IOCompletionPort with handle set to INVALID_HANDLE_VALUE.
+    fn default() -> Self {
+        Self {
+            io_completion_port: INVALID_HANDLE_VALUE,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::file_handle::{AccessMode, ShareMode};
+
+    #[test]
+    fn create_io_completion_port() {
+        let file_name = "../diskann/tests/data/delete_set_50pts.bin";
+        let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) }
+            .expect("Failed to create file handle.");
+
+        let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0);
+
+        assert!(
+            io_completion_port.is_ok(),
+            "Failed to create IOCompletionPort."
+        );
+    }
+
+    #[test]
+    fn drop_io_completion_port() {
+        let file_name = "../diskann/tests/data/delete_set_50pts.bin";
+        let file_handle = unsafe { FileHandle::new(file_name, AccessMode::Read, ShareMode::Read) }
+            .expect("Failed to create file handle.");
+
+        let io_completion_port = IOCompletionPort::new(&file_handle, None, 0, 0)
+            .expect("Failed to create IOCompletionPort.");
+
+        // After this line, io_completion_port goes out of scope and its Drop trait will be called.
+        let _ = io_completion_port;
+        // We have no easy way to test that the Drop trait works correctly, but if it doesn't,
+        // a resource leak or other problem may become apparent in later tests or in real use of the code.
+    }
+
+    #[test]
+    fn default_io_completion_port() {
+        let io_completion_port = IOCompletionPort::default();
+        assert_eq!(
+            io_completion_port.raw_handle(),
+            INVALID_HANDLE_VALUE,
+            "Default IOCompletionPort did not have INVALID_HANDLE_VALUE."
+        );
+    }
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/lib.rs

@@ -0,0 +1,20 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+#![cfg_attr(
+    not(test),
+    warn(clippy::panic, clippy::unwrap_used, clippy::expect_used)
+)]
+
+pub mod perf;
+pub use perf::{get_process_cycle_time, get_process_handle};
+
+pub mod file_io;
+pub use file_io::{get_queued_completion_status, read_file_to_slice};
+
+pub mod file_handle;
+pub use file_handle::FileHandle;
+
+pub mod io_completion_port;
+pub use io_completion_port::IOCompletionPort;

packages/leann-backend-diskann/third_party/DiskANN/rust/platform/src/perf.rs

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+#[cfg(target_os = "windows")]
+#[link(name = "kernel32")]
+extern "system" {
+    fn OpenProcess(dwDesiredAccess: u32, bInheritHandle: bool, dwProcessId: u32) -> usize;
+    fn QueryProcessCycleTime(hProcess: usize, lpCycleTime: *mut u64) -> bool;
+    fn GetCurrentProcessId() -> u32;
+}
+
+/// Get current process handle.
+pub fn get_process_handle() -> Option<usize> {
+    if cfg!(windows) {
+        const PROCESS_QUERY_INFORMATION: u32 = 0x0400;
+        const PROCESS_VM_READ: u32 = 0x0010;
+
+        unsafe {
+            let current_process_id = GetCurrentProcessId();
+            let handle = OpenProcess(
+                PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
+                false,
+                current_process_id,
+            );
+            if handle == 0 {
+                None
+            } else {
+                Some(handle)
+            }
+        }
+    } else {
+        None
+    }
+}
+
+pub fn get_process_cycle_time(process_handle: Option<usize>) -> Option<u64> {
+    let mut cycle_time: u64 = 0;
+    if cfg!(windows) {
+        if let Some(handle) = process_handle {
+            let result = unsafe { QueryProcessCycleTime(handle, &mut cycle_time as *mut u64) };
+            if result {
+                return Some(cycle_time);
+            }
+        }
+    }
+
+    None
+}
+