Skip to main content

diskann_tools/utils/
parameter_helper.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use num_cpus;
7
8pub fn get_num_threads(num_threads: Option<usize>) -> usize {
9    match num_threads {
10        Some(n) => n,
11        None => num_cpus::get(),
12    }
13}
14
15#[cfg(test)]
16mod tests {
17    use super::*;
18
19    #[test]
20    fn test_get_num_threads_with_some() {
21        assert_eq!(get_num_threads(Some(4)), 4);
22        assert_eq!(get_num_threads(Some(1)), 1);
23        assert_eq!(get_num_threads(Some(16)), 16);
24    }
25
26    #[test]
27    fn test_get_num_threads_with_none() {
28        let result = get_num_threads(None);
29        // Should return the number of CPUs, which is at least 1
30        assert!(result >= 1);
31        // Should match num_cpus::get()
32        assert_eq!(result, num_cpus::get());
33    }
34}