Skip to main content

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 {
107                "features"
108            } else {
109                "feature"
110            }
111        }
112
113        match &self.0 {
114            FeaturesInner::Only(feature) => write!(f, "feature \"{}\"", feature),
115            FeaturesInner::Any(any) => {
116                if any.is_empty() {
117                    f.write_str("feature <missing>")
118                } else {
119                    write!(
120                        f,
121                        "{} {}",
122                        feature(any.len() != 1),
123                        fmt::Delimit::new(any.iter().map(fmt::Quote), ", ")
124                            .with_last(", or ")
125                            .with_pair(" or ")
126                    )
127                }
128            }
129            FeaturesInner::All(all) => {
130                if all.is_empty() {
131                    f.write_str("feature <missing>")
132                } else {
133                    write!(
134                        f,
135                        "{} {}",
136                        feature(all.len() != 1),
137                        fmt::Delimit::new(all.iter().map(fmt::Quote), ", ")
138                            .with_last(", and ")
139                            .with_pair(" and ")
140                    )
141                }
142            }
143            FeaturesInner::Custom { custom, plural } => {
144                write!(f, "{} {}", feature(*plural), custom)
145            }
146        }
147    }
148}
149
150///////////
151// Tests //
152///////////
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_rendering() {
160        let list = [
161            (Features::new("a"), "feature \"a\""),
162            (Features::from("b"), "feature \"b\""),
163            // Any
164            (Features::any([]), "feature <missing>"),
165            (Features::any(["a"]), "feature \"a\""),
166            (Features::any(["a", "b"]), "features \"a\" or \"b\""),
167            (
168                Features::any(["a", "b", "c"]),
169                "features \"a\", \"b\", or \"c\"",
170            ),
171            // All
172            (Features::all([]), "feature <missing>"),
173            (Features::all(["a"]), "feature \"a\""),
174            (Features::all(["a", "b"]), "features \"a\" and \"b\""),
175            (
176                Features::all(["a", "b", "c"]),
177                "features \"a\", \"b\", and \"c\"",
178            ),
179            // Custom
180            (
181                Features::custom("custom/stuff", false),
182                "feature custom/stuff",
183            ),
184            (
185                Features::custom("custom/stuff", true),
186                "features custom/stuff",
187            ),
188        ];
189
190        for (feature, expected) in list.iter() {
191            assert_eq!(feature.to_string(), *expected, "Failed for {:?}", feature);
192        }
193    }
194}