diskann_benchmark_runner/
features.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use crate::utils::fmt;
7
8/// A simple list of features for annotating gated benchmarks/inputs.
9///
10/// Features supports one of the following four patterns:
11///
12/// A singular feature:
13/// ```
14/// use diskann_benchmark_runner::Features;
15///
16/// let features = Features::new("a");
17/// assert_eq!(features.to_string(), "feature \"a\"");
18/// ```
19/// A conjunction:
20/// ```
21/// use diskann_benchmark_runner::Features;
22///
23/// let features = Features::all(["a", "b"]);
24/// assert_eq!(features.to_string(), "features \"a\" and \"b\"");
25/// ```
26/// A disjunction:
27/// ```
28/// use diskann_benchmark_runner::Features;
29///
30/// let features = Features::any(["a", "b"]);
31/// assert_eq!(features.to_string(), "features \"a\" or \"b\"");
32/// ```
33/// A custom layout:
34/// ```
35/// use diskann_benchmark_runner::Features;
36///
37/// let features = Features::custom("(\"a\" and \"b\") or \"c\"", true);
38/// assert_eq!(features.to_string(), "features (\"a\" and \"b\") or \"c\"");
39/// ```
40///
41/// See: [`crate::Registry::register_gated`] and [`crate::Registry::register_partially_gated`].
42#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
43pub struct Features(FeaturesInner);
44
45#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
46pub(crate) enum FeaturesInner {
47    Only(&'static str),
48    Any(Vec<&'static str>),
49    All(Vec<&'static str>),
50    Custom { custom: &'static str, plural: bool },
51}
52
53impl Features {
54    /// Construct a singular [`Features`].
55    pub fn new(feature: &'static str) -> Self {
56        Self(FeaturesInner::Only(feature))
57    }
58
59    /// Construct a [`Features`] as a disjunction of features.
60    pub fn any<I>(itr: I) -> Self
61    where
62        I: IntoIterator<Item = &'static str>,
63    {
64        Self(FeaturesInner::Any(itr.into_iter().collect()))
65    }
66
67    /// Construct a [`Features`] as a conjunction of features.
68    pub fn all<I>(itr: I) -> Self
69    where
70        I: IntoIterator<Item = &'static str>,
71    {
72        Self(FeaturesInner::All(itr.into_iter().collect()))
73    }
74
75    /// Construct a [`Features`] with a custom formatted string where `plural` can be used
76    /// to control whether the returned struct uses "feature" or "features" in its formatting.
77    pub fn custom(custom: &'static str, plural: bool) -> Self {
78        Self(FeaturesInner::Custom { custom, plural })
79    }
80
81    /// Return `true` if this feature requirement is satisfied by the set of `enabled` feature
82    /// names.
83    ///
84    /// [`FeaturesInner::Custom`] requirements are opaque to the registry and are therefore
85    /// never considered satisfied.
86    #[cfg(any(test, feature = "test-app"))]
87    pub(crate) fn satisfied_by(&self, enabled: &std::collections::HashSet<String>) -> bool {
88        match &self.0 {
89            FeaturesInner::Only(feature) => enabled.contains(*feature),
90            FeaturesInner::Any(any) => any.iter().any(|f| enabled.contains(*f)),
91            FeaturesInner::All(all) => all.iter().all(|f| enabled.contains(*f)),
92            FeaturesInner::Custom { .. } => false,
93        }
94    }
95}
96
97impl From<&'static str> for Features {
98    fn from(feature: &'static str) -> Features {
99        Features::new(feature)
100    }
101}
102
103impl std::fmt::Display for Features {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        fn feature(plural: bool) -> &'static str {
106            if plural { "features" } else { "feature" }
107        }
108
109        match &self.0 {
110            FeaturesInner::Only(feature) => write!(f, "feature \"{}\"", feature),
111            FeaturesInner::Any(any) => {
112                if any.is_empty() {
113                    f.write_str("feature <missing>")
114                } else {
115                    write!(
116                        f,
117                        "{} {}",
118                        feature(any.len() != 1),
119                        fmt::Delimit::new(any.iter().map(fmt::Quote), ", ")
120                            .with_last(", or ")
121                            .with_pair(" or ")
122                    )
123                }
124            }
125            FeaturesInner::All(all) => {
126                if all.is_empty() {
127                    f.write_str("feature <missing>")
128                } else {
129                    write!(
130                        f,
131                        "{} {}",
132                        feature(all.len() != 1),
133                        fmt::Delimit::new(all.iter().map(fmt::Quote), ", ")
134                            .with_last(", and ")
135                            .with_pair(" and ")
136                    )
137                }
138            }
139            FeaturesInner::Custom { custom, plural } => {
140                write!(f, "{} {}", feature(*plural), custom)
141            }
142        }
143    }
144}
145
146///////////
147// Tests //
148///////////
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn test_rendering() {
156        let list = [
157            (Features::new("a"), "feature \"a\""),
158            (Features::from("b"), "feature \"b\""),
159            // Any
160            (Features::any([]), "feature <missing>"),
161            (Features::any(["a"]), "feature \"a\""),
162            (Features::any(["a", "b"]), "features \"a\" or \"b\""),
163            (
164                Features::any(["a", "b", "c"]),
165                "features \"a\", \"b\", or \"c\"",
166            ),
167            // All
168            (Features::all([]), "feature <missing>"),
169            (Features::all(["a"]), "feature \"a\""),
170            (Features::all(["a", "b"]), "features \"a\" and \"b\""),
171            (
172                Features::all(["a", "b", "c"]),
173                "features \"a\", \"b\", and \"c\"",
174            ),
175            // Custom
176            (
177                Features::custom("custom/stuff", false),
178                "feature custom/stuff",
179            ),
180            (
181                Features::custom("custom/stuff", true),
182                "features custom/stuff",
183            ),
184        ];
185
186        for (feature, expected) in list.iter() {
187            assert_eq!(feature.to_string(), *expected, "Failed for {:?}", feature);
188        }
189    }
190}
diskann_benchmark_runner/features.rs

diskann_benchmark_runner/
features.rs
