diff --git a/Cargo.lock b/Cargo.lock index 76aeda3862..41b7858263 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -518,7 +518,10 @@ version = "0.60.0" dependencies = [ "anyhow", "clap", + "diskann-macros", + "diskann-utils", "half", + "hashbrown 0.16.1", "indicatif", "serde", "serde_json", @@ -685,6 +688,17 @@ dependencies = [ "thiserror 2.0.21", ] +[[package]] +name = "diskann-macros" +version = "0.60.0" +dependencies = [ + "diskann-benchmark-runner", + "proc-macro2", + "quote", + "syn 2.0.119", + "trybuild", +] + [[package]] name = "diskann-providers" version = "0.60.0" diff --git a/Cargo.toml b/Cargo.toml index 69ba9ff3ff..f28a55cd2c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ members = [ "diskann-record", "diskann-tools", "diskann-bftree", + "diskann-macros", ] default-members = [ @@ -66,6 +67,7 @@ diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0 diskann-disk = { path = "diskann-disk", version = "0.60.0" } diskann-label-filter = { path = "diskann-label-filter", version = "0.60.0" } # Infra +diskann-macros = { path = "diskann-macros", version = "0.60.0" } diskann-benchmark-runner = { path = "diskann-benchmark-runner", version = "0.60.0" } diskann-benchmark-core = { path = "diskann-benchmark-core", version = "0.60.0" } diskann-tools = { path = "diskann-tools", version = "0.60.0" } diff --git a/diskann-benchmark-runner/Cargo.toml b/diskann-benchmark-runner/Cargo.toml index 041f16df33..43baa244ad 100644 --- a/diskann-benchmark-runner/Cargo.toml +++ b/diskann-benchmark-runner/Cargo.toml @@ -10,7 +10,9 @@ edition = "2024" [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } +diskann-macros = { workspace = true } half = { workspace = true } +hashbrown = { workspace = true } indicatif = "0.18.3" serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } @@ -25,6 +27,7 @@ expect_used = "warn" panic = "warn" [dev-dependencies] +diskann-utils = { workspace = true, default-features = false, features = ["testing"] } tempfile.workspace = true [features] diff --git a/diskann-benchmark-runner/dev/main.rs b/diskann-benchmark-runner/dev/main.rs index 788e9016b3..3e9db47e3f 100644 --- a/diskann-benchmark-runner/dev/main.rs +++ b/diskann-benchmark-runner/dev/main.rs @@ -25,7 +25,13 @@ fn main() -> anyhow::Result<()> { #[derive(Debug, clap::Parser)] struct Cli { /// Emulate enabling various features for feature gated functionality. - #[arg(long, value_delimiter = ',', num_args = 0..)] + /// + /// Currently available features: + /// * `gated-feature-0` + /// * `gated-feature-1` + /// * `gated-feature-2` + /// * `gated-feature-3` + #[arg(long, value_delimiter = ',', num_args = 1)] features: Vec, /// The actual application. diff --git a/diskann-benchmark-runner/src/app.rs b/diskann-benchmark-runner/src/app.rs index 442732068b..c8178d5b22 100644 --- a/diskann-benchmark-runner/src/app.rs +++ b/diskann-benchmark-runner/src/app.rs @@ -121,6 +121,11 @@ pub enum Commands { }, #[command(subcommand)] Check(Check), + /// Provide information about all registered types. + TypeInfo { + /// Provide information for the given type. + describe: Option, + }, } /// Subcommands for regression check operations. @@ -207,12 +212,28 @@ impl App { if let Some(describe) = describe { if let Some(input) = registry.input(describe) { let repr = jobs::Unprocessed::format_input(input)?; + + // Render JSON. writeln!( output, - "The example JSON representation for \"{}\" is:", + "The example JSON representation for \"{}\" is:\n", describe )?; writeln!(output, "{}", serde_json::to_string_pretty(&repr)?)?; + + // Render Type Info. + match input.raw_reflection() { + Some(reflection) => { + writeln!(output, "\nType Information:\n\n{}", reflection.render())?; + + writeln!( + output, + "More type information available using `type-info`" + )?; + } + None => writeln!(output, "\n\nNo Type Information Available")?, + } + return Ok(()); } else { writeln!(output, "No input found for \"{}\"", describe)?; @@ -386,6 +407,11 @@ impl App { } // Extensions Commands::Check(check) => return self.check(check, registry, output), + + // Types + Commands::TypeInfo { describe } => { + self.type_info(describe.as_deref(), registry, output)? + } }; Ok(()) } @@ -482,6 +508,29 @@ impl App { } } } + + fn type_info( + &self, + describe: Option<&str>, + registry: ®istry::Registry, + mut output: &mut dyn Output, + ) -> anyhow::Result<()> { + match describe { + Some(type_name) => match registry.type_info(type_name) { + Some(reflection) => writeln!(output, "{}", reflection.render())?, + None => anyhow::bail!("No type information for \"{}\"", type_name), + }, + None => { + let mut all_types: Vec<_> = registry.type_names().collect(); + all_types.sort_unstable(); + writeln!(output, "All registered types:")?; + for type_name in all_types { + writeln!(output, " {}", type_name)?; + } + } + } + Ok(()) + } } /////////// @@ -541,8 +590,6 @@ mod tests { use crate::{registry, test::TestConfig, ux}; - const ENV: &str = "DISKANN_TEST"; - // Expected I/O files. const STDIN: &str = "stdin.txt"; const STDOUT: &str = "stdout.txt"; @@ -560,42 +607,6 @@ mod tests { const ALL_GENERATED_OUTPUTS: [&str; 2] = [OUTPUT_FILE, CHECK_OUTPUT_FILE]; - // Read the entire contents of a file to a string. - fn read_to_string>(path: P, ctx: &str) -> String { - match std::fs::read_to_string(path.as_ref()) { - Ok(s) => ux::normalize(s), - Err(err) => panic!( - "failed to read {} {:?} with error: {}", - ctx, - path.as_ref(), - err - ), - } - } - - // Check if `DISKANN_TEST=overwrite` is configured. Return `true` if so - otherwise - // return `false`. - // - // If `DISKANN_TEST` is set but its value is not `overwrite` - panic. - fn overwrite() -> bool { - match std::env::var(ENV) { - Ok(v) => { - if v == "overwrite" { - true - } else { - panic!( - "Unknown value for {}: \"{}\". Expected \"overwrite\"", - ENV, v - ); - } - } - Err(std::env::VarError::NotPresent) => false, - Err(std::env::VarError::NotUnicode(_)) => { - panic!("Value for {} is not unicode", ENV); - } - } - } - // Test Runner struct Test { dir: PathBuf, @@ -606,7 +617,7 @@ mod tests { fn new(dir: &Path) -> Self { Self { dir: dir.into(), - overwrite: overwrite(), + overwrite: ux::test::overwrite(), } } @@ -614,7 +625,7 @@ mod tests { let path = self.dir.join(STDIN); // Read the standard input file to a string. - let stdin = read_to_string(&path, "standard input"); + let stdin = ux::test::read_to_string(&path, "standard input"); let output: Vec = stdin .lines() @@ -723,7 +734,7 @@ mod tests { if self.overwrite { std::fs::write(output, stdout).unwrap(); } else { - let expected = read_to_string(&output, "expected standard output"); + let expected = ux::test::read_to_string(&output, "expected standard output"); if stdout != expected { panic!("Got:\n--\n{}\n--\nExpected:\n--\n{}\n--", stdout, expected); } @@ -765,9 +776,9 @@ mod tests { } else { match (was_generated, is_expected) { (true, true) => { - let output_contents = read_to_string(generated_path, "generated"); + let output_contents = ux::test::read_to_string(generated_path, "generated"); - let expected_contents = read_to_string(expected_path, "expected"); + let expected_contents = ux::test::read_to_string(expected_path, "expected"); if output_contents != expected_contents { panic!( @@ -777,7 +788,7 @@ mod tests { } } (true, false) => { - let output_contents = read_to_string(generated_path, "generated"); + let output_contents = ux::test::read_to_string(generated_path, "generated"); panic!( "{} was generated when none was expected. Contents:\n\n{}", diff --git a/diskann-benchmark-runner/src/files.rs b/diskann-benchmark-runner/src/files.rs index 02b0c97979..0206dcd31d 100644 --- a/diskann-benchmark-runner/src/files.rs +++ b/diskann-benchmark-runner/src/files.rs @@ -7,7 +7,7 @@ use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; -use super::Checker; +use crate::{Checker, Reflect}; /// A file that is used as an input to for a benchmark. /// @@ -27,6 +27,16 @@ pub struct InputFile { path: PathBuf, } +impl Reflect for InputFile { + fn ty() -> crate::reflect::Type { + ::ty() + } + + fn format_type_name(f: &mut dyn std::fmt::Write) -> std::fmt::Result { + f.write_str("benchmark::InputFile") + } +} + impl InputFile { /// Create a new input file from the path-like `path`. pub fn new

(path: P) -> Self diff --git a/diskann-benchmark-runner/src/input.rs b/diskann-benchmark-runner/src/input.rs index 20ba42dd0a..8b4495a7d7 100644 --- a/diskann-benchmark-runner/src/input.rs +++ b/diskann-benchmark-runner/src/input.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use crate::Checker; +use crate::{Checker, Reflect}; /// Inputs to [`Benchmarks`](crate::Benchmark). /// @@ -16,7 +16,7 @@ pub trait Input: Sized + std::fmt::Debug + 'static { /// [`Deserialize`](serde::Deserialize) implementation. /// /// Final object validation is performed via [`from_raw`](Self::from_raw). - type Raw: serde::de::DeserializeOwned + serde::Serialize; + type Raw: serde::de::DeserializeOwned + serde::Serialize + Reflect; /// Return the discriminant associated with this type. /// @@ -39,7 +39,7 @@ pub trait Input: Sized + std::fmt::Debug + 'static { } pub(crate) mod internal { - use crate::{Checker, Features, internal::visibility::Visibility}; + use crate::{Checker, Features, Reflection, internal::visibility::Visibility}; /// Runtime representation of a deserialized [`Input`]. #[derive(Debug)] @@ -145,6 +145,10 @@ pub(crate) mod internal { self.0.visibility() } + pub(crate) fn raw_reflection(&self) -> Option { + self.0.raw_reflection() + } + pub(crate) fn as_any(&self) -> &dyn std::any::Any { self.0.as_any() } @@ -221,6 +225,7 @@ pub(crate) mod internal { fn visibility(&self) -> Visibility<'_>; // reflection + fn raw_reflection(&self) -> Option; fn as_any(&self) -> &dyn std::any::Any; fn type_name(&self) -> &'static str; } @@ -261,6 +266,9 @@ pub(crate) mod internal { fn visibility(&self) -> Visibility<'_> { Visibility::Available } + fn raw_reflection(&self) -> Option { + Some(Reflection::new::()) + } fn as_any(&self) -> &dyn std::any::Any { self } @@ -305,6 +313,9 @@ pub(crate) mod internal { features: &self.features, } } + fn raw_reflection(&self) -> Option { + None + } fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/diskann-benchmark-runner/src/internal/mod.rs b/diskann-benchmark-runner/src/internal/mod.rs index 21e53228fa..698bce6288 100644 --- a/diskann-benchmark-runner/src/internal/mod.rs +++ b/diskann-benchmark-runner/src/internal/mod.rs @@ -6,6 +6,7 @@ use anyhow::Context; pub(crate) mod regression; +pub(crate) mod typemap; pub(crate) mod visibility; /// Attempt to load and deserialize from a JSON file on disk identified with `path`. diff --git a/diskann-benchmark-runner/src/internal/typemap.rs b/diskann-benchmark-runner/src/internal/typemap.rs new file mode 100644 index 0000000000..d3d9b62156 --- /dev/null +++ b/diskann-benchmark-runner/src/internal/typemap.rs @@ -0,0 +1,407 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + any::TypeId, + collections::{HashMap, hash_map::Entry}, +}; + +use hashbrown::{HashSet, hash_set}; + +use crate::Reflection; + +/// Collection of [`Reflection`] types. +#[derive(Debug)] +pub(crate) struct TypeMap { + /// Registered types for documentation. + /// + /// The keys are generated by [`Reflection::type_name`]. + name_map: HashMap, + + /// The type IDs present in `name_map`. + type_ids: HashSet, +} + +impl TypeMap { + /// Create a new empty type-map. + pub(crate) fn new() -> Self { + Self { + name_map: HashMap::new(), + type_ids: HashSet::new(), + } + } + + /// Return the [`Reflection`] type information for the provided type-name. + pub(crate) fn type_info(&self, type_name: &str) -> Option { + self.name_map.get(type_name).copied() + } + + /// Return an iterator over all the registered type-names. + pub(crate) fn type_names(&self) -> impl ExactSizeIterator { + self.name_map.keys().map(|v| &**v) + } + + pub(crate) fn register( + &mut self, + reflections: &[Reflection], + ) -> Result, Conflict> { + let mut t = Transaction::new(self); + + reflections.iter().try_for_each(|r| { + r.visit_with(|r: Reflection| -> Result { + let type_id = r.type_id(); + + // If the type doesn't exist - try to add it via type-name. + // Only if all that succeeds to we commit the type-id. + if let hash_set::Entry::Vacant(vacant_type_id) = t.map.type_ids.entry(type_id) { + let type_name = r.type_name().to_string(); + match t.map.name_map.entry(type_name) { + // Everything checks out, we're good to go. + Entry::Vacant(vacant_type_name) => { + vacant_type_name.insert(r); + vacant_type_id.insert(); + t.added.push(type_id); + } + Entry::Occupied(occupied_type_name) => { + // This branch is only reachable if `type_id` did not exist + // in `type_ids`, which means it never could have entered into + // `name_map` anyways. + debug_assert_ne!(occupied_type_name.get().type_id(), type_id); + + return Err(Conflict { + type_name: occupied_type_name.key().clone(), + }); + } + } + + // Continue exploring. + Ok(true) + } else { + // Already seen this node. No need to recurse. + Ok(false) + } + }) + })?; + + Ok(t) + } +} + +#[derive(Debug)] +pub(crate) struct Conflict { + type_name: String, +} + +impl std::fmt::Display for Conflict { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "A different type with the type name \"{}\" was already registered", + self.type_name + ) + } +} + +impl std::error::Error for Conflict {} + +#[derive(Debug)] +#[must_use = "transactions must be explicitly completed - otherwise they automatically abort"] +pub(crate) struct Transaction<'a> { + map: &'a mut TypeMap, + added: Vec, +} + +impl<'a> Transaction<'a> { + fn new(map: &'a mut TypeMap) -> Self { + Self { + map, + added: Vec::new(), + } + } + + fn roll_back(&mut self) { + // Nothing to roll back. + if self.added.is_empty() { + return; + } + + // We first roll-back the type IDs that were added, which is relatively easy. + // Then, we cannot rely on `Reflect::type_name` returning the same name every time. + // So instead, we have to loop over all entries in `name_map` and inspect the + // `TypeId`s directly (which we control) to determing removal. + // + // Fortunately, `extract_if` does exactly what we want. + let mut added_type_ids = HashSet::new(); + for type_id in std::mem::take(&mut self.added) { + self.map.type_ids.remove(&type_id); + added_type_ids.insert(type_id); + } + + self.map + .name_map + .retain(|_, v| !added_type_ids.contains(&v.type_id())); + } + + pub(crate) fn complete(mut self) { + self.added.clear(); + } +} + +impl Drop for Transaction<'_> { + fn drop(&mut self) { + self.roll_back() + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use crate::Reflect; + + // The types `A`, `B`, `C`, etc. are the valid types for testing registration. + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct A { + a: usize, + b: Option, + } + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct B(A, f32); + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct C { + a: f64, + b: isize, + } + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + enum D { + V0(A), + V1 { first: B, second: B }, + V2, + V3(Vec>), + } + + // The type tree `Bad0`, `Bad1`, etc. are the types we use to test the erroring path. + // + // The idea is that somewhere at the bottom of this path is a type with a type-name + // that conflicts with one of the "good" types above. + #[derive(Reflect)] + #[reflect(type_name = "usize")] + struct Bad0; + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct Bad1 { + a: usize, + b: D, + c: i32, + d: Option, + } + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct Bad2(u32, Vec); + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + enum Bad3 { + Unit, + S(String), + B { field: Bad2 }, + } + + fn register(map: &mut TypeMap) -> Result, Conflict> + where + T: Reflect, + { + map.register(&[Reflection::new::()]) + } + + fn type_name_for() -> String + where + T: Reflect, + { + Reflection::new::().type_name().to_string() + } + + fn assert_is_present(map: &TypeMap) + where + T: Reflect, + { + let type_name = type_name_for::(); + let r = map.type_info(&type_name).unwrap(); + + assert_eq!(r.type_id(), TypeId::of::()); + assert_eq!(r.type_name().to_string(), type_name); + } + + fn assert_is_absent(map: &TypeMap) + where + T: Reflect, + { + let type_name = type_name_for::(); + if let Some(r) = map.type_info(&type_name) { + assert_ne!( + r.type_id(), + Reflection::new::().type_id(), + "Did not expect \"{}\" to be registered", + std::any::type_name::(), + ); + } + } + + // We expect the following types to be present. + // + // * usize + // * String + // * Option + // * f32 + // * f64 + // * isize + // * A + // * B + // * C + // * Option + // * Vec> + // * D + // + // The following types should *not* be present. + // + // * i32 + // * Bad0 + // * Option + // * Bad1 + // * u32 + // * Vec + // * Bad2 + // * Bad3 + fn check_types_are_present(r: &TypeMap) { + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::>(r); + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::(r); + assert_is_present::>(r); + assert_is_present::>>(r); + assert_is_present::(r); + + assert_is_absent::(r); + assert_is_absent::(r); + assert_is_absent::>(r); + assert_is_absent::(r); + assert_is_absent::(r); + assert_is_absent::>(r); + assert_is_absent::(r); + assert_is_absent::(r); + } + + // Register the entire type-tree at once. + #[test] + fn test_registration_all() { + let mut map = TypeMap::new(); + register::(&mut map).unwrap().complete(); + check_types_are_present(&map); + + // Registering again should be fine. + register::(&mut map).unwrap().complete(); + check_types_are_present(&map); + } + + // Multiple total types can be registered in one call. + #[test] + fn test_register_multiple() { + let mut map = TypeMap::new(); + map.register(&[Reflection::new::(), Reflection::new::()]) + .unwrap() + .complete(); + + assert_is_present::(&map); + assert_is_present::(&map); + } + + // Roll-back undoes multiple types. + #[test] + fn test_multiple_roll_back() { + let mut map = TypeMap::new(); + + // Drop aborts the transaction + let _ = map + .register(&[Reflection::new::(), Reflection::new::()]) + .unwrap(); + + assert!(map.name_map.is_empty()); + assert!(map.type_ids.is_empty()); + } + + // Register portions of the type-tree. + #[test] + fn test_incremental_registration() { + let mut map = TypeMap::new(); + register::(&mut map).unwrap().complete(); + register::(&mut map).unwrap().complete(); + register::(&mut map).unwrap().complete(); + register::(&mut map).unwrap().complete(); + + check_types_are_present(&map); + + register::(&mut map).unwrap().complete(); + check_types_are_present(&map); + } + + // Immediate errors do not leave invalid state in the map. + #[test] + fn test_errors_are_recoverable() { + let mut map = TypeMap::new(); + + // Registering the "bad" type should eventually hit a naming conflict. + // Everything should be rolled back and adding the "good" type should still work. + let err = register::(&mut map).unwrap_err(); + let msg = err.to_string(); + assert_eq!( + msg, + "A different type with the type name \"usize\" was already registered" + ); + + assert!( + map.name_map.is_empty(), + "invalid registration should not commit registered items" + ); + + assert!( + map.type_ids.is_empty(), + "invalid registration should not commit registered items" + ); + + let err = map + .register(&[Reflection::new::(), Reflection::new::()]) + .unwrap_err(); + assert_eq!(err.to_string(), msg); + assert!(map.name_map.is_empty()); + assert!(map.type_ids.is_empty()); + + // After this, we should succeed in adding more types correctly. + register::(&mut map).unwrap().complete(); + check_types_are_present(&map); + + // Again, bad registration should not corrupt the map. + register::(&mut map).unwrap_err(); + check_types_are_present(&map); + } +} diff --git a/diskann-benchmark-runner/src/lib.rs b/diskann-benchmark-runner/src/lib.rs index 403f15fc4e..c6eb596aa8 100644 --- a/diskann-benchmark-runner/src/lib.rs +++ b/diskann-benchmark-runner/src/lib.rs @@ -5,6 +5,9 @@ //! A moderately functional utility for making simple benchmarking CLI applications. +#[doc(hidden)] +extern crate self as diskann_benchmark_runner; + pub mod benchmark; mod checker; mod features; @@ -16,6 +19,7 @@ pub mod app; pub mod files; pub mod input; pub mod output; +pub mod reflect; pub mod registry; pub mod utils; @@ -25,6 +29,7 @@ pub use checker::Checker; pub use features::Features; pub use input::Input; pub use output::Output; +pub use reflect::{Reflect, Reflection}; pub use registry::{Registry, RegistryError}; pub use result::Checkpoint; diff --git a/diskann-benchmark-runner/src/reflect/mod.rs b/diskann-benchmark-runner/src/reflect/mod.rs new file mode 100644 index 0000000000..3ca62fd6fa --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/mod.rs @@ -0,0 +1,416 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Run time inspection of types. + +use std::{ + any::TypeId, + fmt::{self, Write}, +}; + +pub use diskann_macros::Reflect; + +mod render; +pub mod tree; +pub use tree::Type; + +#[cfg(test)] +mod test; + +/// Provide run-time information about compile-time types, including documentation. +/// +/// This trait is derivable and supports a subset of `serde` attributes. +/// +/// ``` +/// use diskann_benchmark_runner::{Reflect, Reflection}; +/// +/// /// An example struct. +/// #[derive(Reflect)] +/// struct Foo { +/// /// An awesome field. +/// #[serde(rename = "bar")] +/// foo: usize, +/// baz: usize, +/// } +/// +/// // Get information about `Foo`. +/// let ty = Foo::ty(); +/// +/// // The documentation of `Foo` will automatically be extracted from the docstrings. +/// assert_eq!(ty.doc().unwrap(), "An example struct."); +/// +/// // The type name can be extracted as well. +/// let mut name = String::new(); +/// Foo::format_type_name(&mut name); +/// assert_eq!(name, "Foo"); +/// +/// // To make typenames easier to extract, a `Reflection` can be use. +/// let reflection = Reflection::new::(); +/// assert_eq!(reflection.type_name().to_string(), "Foo"); +/// ``` +/// +/// # Attributes +/// +/// For derivation, a combination of `serde` attributes and custom `reflect` attributes +/// are supported. +/// +/// ## Serde Attributes +/// +/// Like the [serde crate](https://serde.rs/attributes.html), attributes are categorized by +/// container, variant, or field. +/// +/// ### Container Attributes +/// +/// * `#[serde(rename_all = "...")]`: Rename all fields (if a struct) or variants (if enum) +/// arrocding to the given case. +/// +/// Possible values are "lowercase", "snake_case", and "kebab-case". +/// +/// * `#[serde(tag = "type")]`: Used for internally tagged enums. +/// +/// * `#[serde(tag = "t", content = "c")]`: Used for adjacently tagged enums. +/// +/// ### Variant Attributes +/// +/// * `#[serde(rename = "name")]`: Describe with the given name instead of its Rust name. +/// +/// * `#[serde(rename_all = "...)]`: Rename all fiels of this struct variant with the given +/// case convention. +/// +/// Possible values are "lowercase", "snake_case", and "kebab-case". +/// +/// ### Field Attributes +/// +/// * `#[serde(rename = "name")]`: Describe with the given name instead of its Rust name. +/// +/// ## Reflect Attributes +/// +/// ### Container Attributes +/// +/// * `#[reflect(type_name = "name")]`: Use the given name to describe a struct instead of +/// its Rust name. This is used to avoid naming conflicts within a [`Registry`], which +/// enforces that type names are unique. +/// +/// Because of this, this attribute cannot be used on generic structs. +/// +/// This is mutually exclusive with the `prefix` attribute. +/// +/// ``` +/// use diskann_benchmark_runner::{Reflect, Reflection}; +/// +/// #[derive(Reflect)] +/// #[reflect(type_name = "Bar")] +/// struct Foo; +/// +/// assert_eq!(Reflection::new::().type_name().to_string(), "Bar"); +/// ``` +/// +/// * `#[reflect(prefix = "...")]`: Prefix the Rust name with the provided prefix. Like the +/// `type_name` attribute, this can be used to create name spaces to help generate unique +/// type names. +/// +/// ``` +/// use diskann_benchmark_runner::{Reflect, Reflection}; +/// +/// #[derive(Reflect)] +/// #[reflect(prefix = "mod::")] +/// struct Foo; +/// +/// assert_eq!(Reflection::new::().type_name().to_string(), "mod::Foo"); +/// ``` +pub trait Reflect: 'static { + /// Return the [`Type`] containing the information about `self`. + fn ty() -> Type; + + /// Write the type-name for `self` into the buffer. + fn format_type_name(f: &mut dyn Write) -> fmt::Result; +} + +/// A [`Reflect`]ed type. +#[derive(Clone, Copy)] +pub struct Reflection { + reflection: &'static internal::VTable, +} + +impl Reflection { + /// Construct a new [`Reflection`] for `T`. + pub const fn new() -> Self + where + T: Reflect, + { + Self { + reflection: internal::VTable::new::(), + } + } + + /// Return the [`Type`] for the type being reflected. + pub fn ty(&self) -> Type { + (self.reflection.ty)() + } + + /// Return a [`std::fmt::Display`] compatible struct for rendering the name of the type + /// being reflected. + pub fn type_name(&self) -> TypeName { + TypeName(*self) + } + + /// Return the [`TypeId`] of the + pub fn type_id(&self) -> TypeId { + (self.reflection.type_id)() + } + + /// Return a [`std::fmt::Display`] comaptible struct for rendering the reflected type. + pub(crate) fn render(&self) -> Render { + Render(*self) + } + + /// Visit all types reachable from the reflected type. + /// + /// This will traverse through all structs, enum variants, container types etc. + /// + /// The closure `f` can be used to direct the exploration by returning the following values: + /// + /// * `Ok(true)`: Continue exploring through the argument [`Reflection`]. + /// * `Ok(false)`: Do not continue exploring through the argument [`Reflection`]. + /// * `Err(E)`: Immediately stop exploring and return the error `E`. + pub(crate) fn visit_with(&self, f: F) -> Result<(), E> + where + F: FnMut(Reflection) -> Result, + { + visit_with(*self, f) + } +} + +impl fmt::Debug for Reflection { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Reflection") + .field("type_name", &self.type_name()) + .finish_non_exhaustive() + } +} + +/// A [`std::fmt::Display`] compatible type for [`Reflection`]. +/// +/// See: [`Reflection::type_name`]. +pub struct TypeName(Reflection); + +impl TypeName { + fn format_type_name(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + (self.0.reflection.format_type_name)(f) + } +} + +impl std::fmt::Debug for TypeName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.format_type_name(f) + } +} + +impl std::fmt::Display for TypeName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.format_type_name(f) + } +} + +/// A [`std::fmt::Display`] compatible type for rendering a [`Reflection`]. +/// +/// See: [`Reflection::render`]. +pub(crate) struct Render(Reflection); + +impl std::fmt::Display for Render { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut r = render::Renderer::new(f, 2); + r.render_subject(self.0) + } +} + +//------------// +// Algorithms // +//------------// + +fn visit_with(mut reflection: Reflection, mut f: F) -> Result<(), E> +where + F: FnMut(Reflection) -> Result, +{ + use tree::Fields; + + let mut stack = Vec::new(); + loop { + // visit: expand this node if the closure returns `true`. + if f(reflection)? { + let mut push = |r: Reflection| stack.push(r); + let mut push_fields = |fields: &Fields| match fields { + Fields::Named(named) => named.iter().for_each(|field| push(field.field())), + Fields::Unnamed(unnamed) => unnamed.iter().for_each(|field| push(field.field())), + Fields::NewType(newtype) => push(newtype.field()), + Fields::Unit => {} + }; + + // explore + match reflection.ty() { + // Nothing to do for primitives as there is no other object that can be reached. + Type::Primitive(_) => {} + Type::Aggregate(aggregate) => push_fields(aggregate.fields()), + Type::Enum(enum_) => enum_ + .variants() + .iter() + .for_each(|variant| push_fields(variant.fields())), + Type::Sequence(seq) => push(seq.element()), + Type::Optional(opt) => push(opt.value()), + } + } + + // loop + if let Some(r) = stack.pop() { + reflection = r; + } else { + break Ok(()); + } + } +} + +/////////////// +// Bootstrap // +/////////////// + +impl Reflect for std::marker::PhantomData +where + T: Reflect, +{ + fn ty() -> Type { + Type::primitive(tree::PrimitiveKind::Null, None) + } + + fn format_type_name(f: &mut dyn Write) -> fmt::Result { + write!(f, "PhantomData<{}>", Reflection::new::().type_name()) + } +} + +macro_rules! primitive { + ($T:ty, $kind:ident, $doc:literal, $type_name:literal) => { + impl Reflect for $T { + fn ty() -> Type { + Type::primitive(tree::PrimitiveKind::$kind, Some($doc.into())) + } + + fn format_type_name(f: &mut dyn Write) -> fmt::Result { + f.write_str($type_name) + } + } + }; +} + +primitive!((), Null, "empty", "()"); +primitive!( + usize, + Number, + "A system dependent unsigned integer", + "usize" +); +primitive!(isize, Number, "A system dependent signed integer", "isize"); + +primitive!(u8, Number, "An 8-bit unsigned integer", "u8"); +primitive!(u16, Number, "A 16-bit unsigned integer", "u16"); +primitive!(u32, Number, "A 32-bit unsigned integer", "u32"); +primitive!(u64, Number, "A 64-bit unsigned integer", "u64"); + +primitive!(i8, Number, "An 8-bit signed integer", "i8"); +primitive!(i16, Number, "A 16-bit signed integer", "i16"); +primitive!(i32, Number, "A 32-bit signed integer", "i32"); +primitive!(i64, Number, "A 64-bit signed integer", "i64"); + +primitive!(f32, Number, "An 32-bit floating-point number", "f32"); +primitive!(f64, Number, "An 64-bit floating-point number", "f64"); + +primitive!( + std::num::NonZeroU32, + Number, + "A system dependent, 32-bit unsigned integer", + "NonZero" +); +primitive!( + std::num::NonZeroUsize, + Number, + "A system dependent, non-zero, unsigned integer", + "NonZero" +); + +primitive!(bool, Boolean, "A value of \"true\" or \"false\"", "bool"); +primitive!(String, String, "A string", "string"); +primitive!(std::path::PathBuf, String, "A file path", "PathBuf"); + +impl Reflect for Option +where + T: Reflect, +{ + fn ty() -> Type { + Type::optional::(Some("An optional type".into())) + } + + fn format_type_name(f: &mut dyn Write) -> fmt::Result { + f.write_str("Option<")?; + T::format_type_name(f)?; + f.write_str(">") + } +} + +impl Reflect for Vec +where + T: Reflect, +{ + fn ty() -> Type { + Type::sequence::(Some("An ordered collection of elements".into())) + } + + fn format_type_name(f: &mut dyn Write) -> fmt::Result { + write!(f, "Vec<{}>", Reflection::new::().type_name()) + } +} + +////////////// +// Internal // +////////////// + +pub mod internal { + pub(super) struct VTable { + pub(super) ty: fn() -> super::Type, + pub(super) format_type_name: fn(&mut dyn std::fmt::Write) -> std::fmt::Result, + pub(super) type_id: fn() -> std::any::TypeId, + } + + impl VTable { + pub(super) const fn new() -> &'static Self + where + T: super::Reflect, + { + &Self { + ty: ty::, + format_type_name: format_type_name::, + type_id: type_id::, + } + } + } + + fn ty() -> super::Type + where + T: super::Reflect, + { + ::ty() + } + + fn format_type_name(f: &mut dyn std::fmt::Write) -> std::fmt::Result + where + T: super::Reflect, + { + ::format_type_name(f) + } + + fn type_id() -> std::any::TypeId + where + T: super::Reflect, + { + std::any::TypeId::of::() + } +} diff --git a/diskann-benchmark-runner/src/reflect/render.rs b/diskann-benchmark-runner/src/reflect/render.rs new file mode 100644 index 0000000000..91d611be01 --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/render.rs @@ -0,0 +1,658 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::fmt::{self, Write}; + +use crate::utils::fmt::Quote; + +use super::{ + Reflection, + tree::{ + Aggregate, Enum, EnumRepr, Fields, NamedField, Optional, Sequence, Type, UnnamedField, + Variant, + }, +}; + +const INDENT: usize = 2; + +#[derive(Debug)] +struct Tagged { + ty: Type, + reflection: Reflection, +} + +impl Tagged { + fn new(reflection: Reflection) -> Self { + Self { + ty: reflection.ty(), + reflection, + } + } + + fn ty(&self) -> &Type { + &self.ty + } + + fn reflection(&self) -> Reflection { + self.reflection + } +} + +pub(super) struct Renderer<'a> { + output: &'a mut dyn Write, + indent: usize, + depth: usize, + max_depth: usize, +} + +impl<'a> Renderer<'a> { + pub(super) fn new(output: &'a mut dyn Write, max_depth: usize) -> Self { + Self { + output, + indent: 0, + depth: 0, + max_depth, + } + } + + fn at_bottom(&self) -> bool { + self.depth >= self.max_depth + } + + fn line(&mut self, display: D) -> fmt::Result + where + D: std::fmt::Display, + { + let indent = INDENT * self.indent; + writeln!(self.output, "{: >indent$}{}", "", display) + } + + fn blank(&mut self) -> fmt::Result { + self.output.write_char('\n') + } + + fn maybe_indent(&mut self, indent: bool, f: F) -> Result + where + F: FnOnce(&mut Self) -> Result, + { + if indent { + self.indent += 1; + } + let result = f(self); + if indent { + self.indent -= 1; + } + result + } + + fn indent(&mut self, f: F) -> Result + where + F: FnOnce(&mut Self) -> Result, + { + self.maybe_indent(true, f) + } + + fn next_with( + &mut self, + pre: impl FnOnce(&mut Self) -> fmt::Result, + body: impl FnOnce(&mut Self) -> fmt::Result, + post: impl FnOnce(&mut Self) -> fmt::Result, + ) -> fmt::Result { + if self.at_bottom() { + Ok(()) + } else { + pre(self)?; + self.depth += 1; + let result = self.indent(body); + self.depth -= 1; + post(self)?; + result + } + } + + fn next(&mut self, f: F) -> fmt::Result + where + F: FnOnce(&mut Self) -> fmt::Result, + { + self.next_with(|_| Ok(()), f, |_| Ok(())) + } + + fn render_doc(&mut self, s: Option<&str>) -> Result { + if let Some(s) = s { + let mut rendered = false; + for ln in s.lines() { + if ln.is_empty() { + self.blank()?; + } else { + self.line(ln)?; + } + + rendered = true; + } + Ok(rendered) + } else { + Ok(false) + } + } + + /// Return `true` if a nested type will be rendered. + fn will_render(&self, ty: &Type) -> bool { + !self.at_bottom() && ty.has_body() + } + + //-------// + // Types // + //-------// + + pub(super) fn render_subject(&mut self, reflection: Reflection) -> fmt::Result { + let tagged = Tagged::new(reflection); + if let Some(kind) = tagged.ty().json_kind() { + self.line(format_args!("{} (json {})", reflection.type_name(), kind))?; + } else { + self.line(reflection.type_name())?; + } + + self.indent(|r| { + let wrote_doc = r.render_doc(tagged.ty().doc())?; + + if wrote_doc && r.will_render(tagged.ty()) { + r.blank()?; + } + + r.render_body(&tagged) + }) + } + + fn render_body(&mut self, tagged: &Tagged) -> fmt::Result { + match tagged.ty() { + Type::Primitive(_) => Ok(()), + Type::Aggregate(aggregate) => self.render_aggregate(aggregate), + Type::Enum(enum_) => self.render_enum(enum_), + Type::Sequence(sequence) => self.render_sequence(sequence), + Type::Optional(opt) => self.render_optional(opt), + } + } + + fn render_aggregate(&mut self, aggregate: &Aggregate) -> fmt::Result { + self.render_fields(aggregate.fields()) + } + + fn render_enum(&mut self, enum_: &Enum) -> fmt::Result { + match enum_.repr() { + EnumRepr::External => self.line("Representation: externally tagged")?, + EnumRepr::Internal { tag } => { + self.line(format_args!("Discriminant field: {}", Quote(tag)))? + } + EnumRepr::Adjacent { tag, content } => { + self.line(format_args!("Discriminant field: {}", Quote(tag)))?; + self.line(format_args!("Content field: {}", Quote(content)))?; + } + } + + self.blank()?; + self.line("Options:")?; + self.indent(|r| { + let mut previous: Option<&Variant> = None; + for variant in enum_.variants().iter() { + // Decide whether or not to put a space before this variant. + // Spaces and be skipped if the previous one was a `Unit` with no docs. + if let Some(previous) = previous { + let skip = previous.fields().is_unit() && previous.doc().is_none(); + if !skip { + r.blank()?; + } + } + + r.render_variant(variant)?; + previous = Some(variant) + } + + Ok(()) + }) + } + + fn render_sequence(&mut self, sequence: &Sequence) -> fmt::Result { + self.next_with( + |r| r.line(format_args!("Elements: {}", sequence.element().type_name())), + |r| r.render_body(&Tagged::new(sequence.element())), + |_| Ok(()), + ) + } + + fn render_optional(&mut self, op: &Optional) -> fmt::Result { + self.line("May be `null`.")?; + self.next(|r| r.render_body(&Tagged::new(op.value()))) + } + + //--------// + // Fields // + //--------// + + fn render_fields(&mut self, fields: &Fields) -> fmt::Result { + match fields { + Fields::Named(named) => { + let mut first = true; + for field in named.iter() { + if !first { + self.blank()?; + } + self.render_named_field(field)?; + first = false; + } + } + Fields::Unnamed(unnamed) => { + for (i, field) in unnamed.iter().enumerate() { + if i != 0 { + self.blank()?; + } + self.render_unnamed_field(Some(i), field)?; + } + } + Fields::NewType(newtype) => self.render_unnamed_field(None, newtype)?, + Fields::Unit => {} + } + + Ok(()) + } + + fn render_named_field(&mut self, field: &NamedField) -> fmt::Result { + let tagged = Tagged::new(field.field()); + let will_render_body = self.will_render(tagged.ty()); + + self.line(format_args!( + "{}: {}", + Quote(field.name()), + tagged.reflection().type_name() + ))?; + + let rendered_doc = self.indent(|r| r.render_doc(field.doc()))?; + if rendered_doc && will_render_body { + self.blank()?; + } + + if will_render_body { + self.next(|r| r.render_body(&tagged))?; + } + Ok(()) + } + + fn render_unnamed_field(&mut self, index: Option, field: &UnnamedField) -> fmt::Result { + let tagged = Tagged::new(field.field()); + let will_render_body = self.will_render(tagged.ty()); + + match index { + Some(index) => self.line(format_args!( + "{}: {}", + index, + tagged.reflection().type_name() + ))?, + None => self.line(tagged.reflection().type_name())?, + } + + let rendered_doc = self.indent(|r| r.render_doc(field.doc()))?; + if rendered_doc && will_render_body { + self.blank()?; + } + + if will_render_body { + self.next(|r| r.render_body(&tagged))?; + } + Ok(()) + } + + //---------// + // Variant // + //---------// + + fn render_variant(&mut self, variant: &Variant) -> fmt::Result { + self.line(Quote(variant.name()))?; + + self.indent(|r| { + let rendered_doc = r.render_doc(variant.doc())?; + + if rendered_doc && variant.fields().has_body() { + r.blank()?; + } + + r.render_fields(variant.fields()) + }) + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use std::{ + fs::File, + io::{BufRead, BufReader, Write}, + path::{Path, PathBuf}, + }; + + use serde::{Deserialize, Serialize}; + + use crate::{Reflect, ux}; + + // For these tests, we use a variation of baseline tests where all the expected results + // are put into a single file, mainly to keep from generating a bunch of files for the + // relatively small tests. + + fn baseline_path() -> PathBuf { + format!( + "{}/tests/rendered_reflections.txt", + env!("CARGO_MANIFEST_DIR") + ) + .into() + } + + fn overwrite_hint() -> &'static str { + "Baselines can be regenerated by running tests with `DISKANN_TEST=overwrite`" + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + struct Content { + name: String, + description: String, + } + + const CASE_SEPARATOR: &str = "========"; + const OUTPUT_SEPARATOR: &str = "--------"; + + #[derive(Debug)] + struct Case { + content: Content, + r: Reflection, + } + + impl Case { + fn new(name: &str, description: &str) -> Self + where + T: Reflect, + { + Self { + content: Content { + name: name.into(), + description: description.into(), + }, + r: Reflection::new::(), + } + } + } + + #[derive(Debug)] + struct Rendered { + content: Content, + rendered: String, + } + + #[derive(Default, Debug, PartialEq)] + enum State { + #[default] + Content, + Rendered, + Done, + } + + impl State { + fn set_to_rendered(&mut self) { + assert_eq!(*self, Self::Content); + *self = Self::Rendered; + } + + fn set_to_done(&mut self) { + assert_eq!(*self, Self::Rendered); + *self = Self::Done; + } + + fn assert_is_done(self) { + assert_eq!(self, Self::Done); + } + } + + impl Rendered { + fn parse_if_present(lines: &mut std::io::Lines) -> Option + where + B: std::io::BufRead, + { + assert_eq!(lines.next()?.unwrap(), CASE_SEPARATOR); + + let mut content = String::new(); + let mut rendered = String::new(); + + let mut state = State::default(); + + for ln in lines { + let ln: &str = &ln.unwrap(); + match ln { + CASE_SEPARATOR => { + state.set_to_done(); + break; + } + OUTPUT_SEPARATOR => { + state.set_to_rendered(); + } + ln => match state { + State::Content => content.push_str(ln), + State::Rendered => { + rendered.push('\n'); + rendered.push_str(ln); + } + State::Done => panic!("invalid state"), + }, + } + } + + state.assert_is_done(); + + Some(Self { + content: serde_json::from_str(&content).unwrap(), + rendered: rendered.trim().to_string(), + }) + } + } + + fn parse_baselines(path: &Path) -> Vec { + let file = match File::open(path) { + Ok(file) => file, + Err(err) => panic!( + "Opening path \"{}\" failed with {}. {}", + path.display(), + err, + overwrite_hint() + ), + }; + let mut lines = BufReader::new(file).lines(); + let mut baselines = Vec::new(); + while let Some(baseline) = Rendered::parse_if_present(&mut lines) { + baselines.push(baseline); + } + + baselines + } + + fn write_baselines(rendered: &[Rendered], path: &Path) { + let mut io = File::create(path).unwrap(); + + for r in rendered.iter() { + let Rendered { content, rendered } = r; + writeln!(io, "{}", CASE_SEPARATOR).unwrap(); + writeln!(io, "{}", serde_json::to_string_pretty(content).unwrap()).unwrap(); + writeln!(io, "{}", OUTPUT_SEPARATOR).unwrap(); + writeln!(io, "{}", ux::normalize(rendered.clone())).unwrap(); + writeln!(io, "{}", CASE_SEPARATOR).unwrap(); + } + } + + fn run_tests_inner(cases: &[Case], path: &Path, overwrite: bool) { + // Generate the current set of baseline. + let current: Vec<_> = cases + .iter() + .map(|case| { + let rendered = ux::normalize(case.r.render().to_string()); + Rendered { + content: case.content.clone(), + rendered, + } + }) + .collect(); + + if overwrite { + write_baselines(¤t, path); + } else { + let expected = parse_baselines(path); + assert_eq!( + current.len(), + expected.len(), + "Number of baseline cases differs. {}", + overwrite_hint(), + ); + + for (current, expected) in std::iter::zip(current.iter(), expected.iter()) { + assert_eq!( + current.content, + expected.content, + "Baseline headers differ. {}", + overwrite_hint(), + ); + + if current.rendered != expected.rendered { + panic!( + "Difference for case name {}\n\nEXPECTED\n\n{}\n\nGOT\n\n{}\n\n{}", + expected.content.name, + expected.rendered, + current.rendered, + overwrite_hint() + ); + } + } + } + } + + /// Select the data type please. + #[derive(Reflect)] + #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "render::")] + #[expect(unused, reason = "testing")] + enum SimpleDataType { + Float32, + Float16, + } + + /// Select the data type please. + #[derive(Reflect)] + #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "render::")] + #[expect(unused, reason = "testing")] + enum AnnotatedDataType { + /// Use high-precision. + Float32, + /// Use lower precision. + Float16, + Int8, + } + + #[derive(Reflect)] + #[serde(rename_all = "snake_case")] + #[reflect(prefix = "render::")] + #[expect(unused, reason = "testing")] + enum Source { + /// Build from scratch. + Build { + /// The type of the input data. + data_type: SimpleDataType, + + /// Input data in the `.bin` binary format. + file: String, + + /// Output file. + /// + /// If provided, saved data will go here. + output: Option, + }, + /// Run from a previously generated output. + FromPrevious { + data_type: AnnotatedDataType, + /// The previously generated output. + file: String, + }, + } + + /// A new-type wrapper around `Source`. + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct SourceWrapper(Source); + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct Id(usize); + + /// A top level config. + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct Config { + source: SourceWrapper, + /// This does one thing. + param1: String, + /// This does another. + param2: Vec, + } + + /// A tuple like type! + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + struct Tuples( + /// Some docs, + usize, + String, + ); + + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + #[serde(tag = "internally-tagged")] + enum InternallyTagged { + /// Config level docs. + Config( + /// Field docs. + Config, + ), + Other { + foo: usize, + /// Hello world + bar: String, + }, + Unit, + } + + /// An adjacently tagged enum. + #[derive(Reflect)] + #[expect(unused, reason = "testing")] + #[serde(tag = "tag", content = "content")] + enum AdjacentlyTagged { + Config(Config), + Other { foo: usize, bar: String }, + Unit, + Tuple(usize, usize), + } + + #[test] + fn run_tests() { + let cases = [ + Case::new::("usize", "a simple test"), + Case::new::("source", "a configurable source enum"), + Case::new::("source wrapper", "render a newtype wrapper"), + Case::new::("id", "render a primitive newtype wrapper"), + Case::new::("config", "a sample struct config."), + Case::new::("tuples", "a sample tuple-like struct."), + Case::new::("internally-tagged", "an internally tagged enum"), + Case::new::("adjacenctly-tagged", "an adjacently tagged enum"), + ]; + + run_tests_inner(&cases, &baseline_path(), ux::test::overwrite()); + } +} diff --git a/diskann-benchmark-runner/src/reflect/test/mod.rs b/diskann-benchmark-runner/src/reflect/test/mod.rs new file mode 100644 index 0000000000..359f96c9af --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/test/mod.rs @@ -0,0 +1,10 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +// Directed macro unit test. +mod unit; + +// Serde compatibility. +mod serde; diff --git a/diskann-benchmark-runner/src/reflect/test/serde.rs b/diskann-benchmark-runner/src/reflect/test/serde.rs new file mode 100644 index 0000000000..19646e67b2 --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/test/serde.rs @@ -0,0 +1,803 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`Reflect`] macro serde-compatibility tests. +//! +//! Serde does allow some differences between serialization and deserialization. We don't +//! care to support such fanciness in our benchmark inputs and outputs, especially since we +//! need example JSONs to be round-trippable back to their serialized representations. +//! +//! As such, the main entry points ([`check_struct`] and [`check_enums`]) just take [`Serialize`] +//! bounds. We do not check round-trippability here. That's the user's problem. + +use std::{assert_matches, borrow::Cow}; + +use hashbrown::HashSet; +use serde::Serialize; +use serde_json::Value; + +use crate::reflect::{Reflect, Reflection, Type, tree}; + +/// Check that the serialized representation of `s` in JSON matches the [`Reflection`] +/// generated for this type. +fn check_struct(s: T) +where + T: Serialize + Reflect, +{ + let r = Reflection::new::(); + let val = serde_json::to_value(s).unwrap(); + check_reflection( + r, + &val, + Context::new(&val, &val, format_args!("struct: {}", r.type_name())), + ); +} + +/// Check that the serialized representations of all the variants of the enum `T` match the +/// [`Reflection`] for `T`. +/// +/// Note that `examples` is expected to contain examples of all variants in declaration order +/// and this function will panic if this is not the case. +fn check_enums(examples: &[T], ctx: std::fmt::Arguments<'_>) +where + T: Serialize + Reflect, +{ + let r = Reflection::new::(); + let serialized: Vec<_> = examples + .iter() + .map(|e| serde_json::to_value(e).unwrap()) + .collect(); + + let Type::Enum(e) = r.ty() else { + panic!("expected an enum for type {}", r.type_name()); + }; + + check_enum_variants(&e, &serialized, ctx); +} + +//////////////////// +// Implementation // +//////////////////// + +/// A context for displaying where we are in the type tree. +/// +/// Contains the top level JSON we're working, the current JSON, and a stack of +/// [`std::fmt::Arguments`] that describe the sequence of operations that led us into a mess. +/// +/// Use the [`context`] macro for creating nested contexts. +#[derive(Debug, Clone, Copy)] +struct Context<'a> { + top: &'a Value, + current: &'a Value, + stack: std::fmt::Arguments<'a>, +} + +impl<'a> Context<'a> { + fn new(top: &'a Value, current: &'a Value, stack: std::fmt::Arguments<'a>) -> Self { + Self { + top, + current, + stack, + } + } +} + +impl std::fmt::Display for Context<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Full JSON:\n\n{}\n\nCurrent JSON:\n\n{}\n\nContext: {}", + serde_json::to_string_pretty(self.top).unwrap(), + serde_json::to_string_pretty(self.current).unwrap(), + self.stack + ) + } +} + +macro_rules! context { + ($current:ident, $next:ident, $fmt:expr) => { + Context { + top: $current.top, + current: $next, + stack: format_args!( + "{}\n- {}", + $current.stack, + format_args!($fmt), + ), + } + }; + ($current:ident, $next:ident, $fmt:expr, $($args:tt)*) => { + Context { + top: $current.top, + current: $next, + stack: format_args!( + "{}\n- {}", + $current.stack, + format_args!($fmt, $($args)*), + ), + } + }; +} + +fn check_reflection(r: Reflection, s: &Value, ctx: Context<'_>) { + check_type( + &r.ty(), + s, + context!(ctx, s, "type_name = {}", r.type_name()), + ); +} + +fn check_type(ty: &Type, s: &Value, ctx: Context<'_>) { + match ty { + Type::Primitive(p) => check_primitive(p, s, context!(ctx, s, "primitive")), + Type::Aggregate(a) => check_fields(a.fields(), s, context!(ctx, s, "aggregate")), + Type::Enum(e) => check_enum(e, s, context!(ctx, s, "enum")), + Type::Sequence(seq) => check_sequence(seq, s, context!(ctx, s, "sequence")), + Type::Optional(opt) => check_optional(opt, s, context!(ctx, s, "optional")), + } +} + +fn check_primitive(p: &tree::Primitive, s: &Value, ctx: Context<'_>) { + // Verify that `s` has the correct `PrimitiveKind`. + match p.kind() { + tree::PrimitiveKind::Null => { + if !s.is_null() { + panic!("expected `Null`\n\n{}", ctx); + } + } + tree::PrimitiveKind::Boolean => { + if !s.is_boolean() { + panic!("expected `Bool`\n\n{}", ctx); + } + } + tree::PrimitiveKind::Number => { + if !s.is_number() { + panic!("expected `Number`\n\n{}", ctx); + } + } + tree::PrimitiveKind::String => { + if !s.is_string() { + panic!("expected `String`\n\n{}", ctx); + } + } + } +} + +fn check_fields(fields: &tree::Fields, s: &Value, ctx: Context<'_>) { + match fields { + tree::Fields::Named(named_fields) => { + // All names must be unique. + assert_all_unique(named_fields.iter().map(|f| f.name()), ctx); + let map = value_as_map(s, ctx); + assert_eq!( + map.len(), + named_fields.len(), + "mismatch in number of fields\n\n{}", + ctx, + ); + + for f in named_fields { + let val = match map.get(f.name()) { + Some(val) => val, + None => panic!("Could not match field name \"{}\"\n\n{}", f.name(), ctx), + }; + + check_reflection( + f.field(), + val, + context!(ctx, val, "named field \"{}\"", f.name()), + ); + } + } + tree::Fields::Unnamed(unnamed_fields) => { + let array = value_as_array(s, ctx); + assert_eq!( + array.len(), + unnamed_fields.len(), + "mismatch in number of fields\n\n{}", + ctx + ); + + for (i, (f, val)) in std::iter::zip(unnamed_fields.iter(), array.iter()).enumerate() { + check_reflection(f.field(), val, context!(ctx, val, "unnamed field {}", i)); + } + } + tree::Fields::NewType(newtype) => { + check_reflection(newtype.field(), s, context!(ctx, s, "newtype")) + } + tree::Fields::Unit => { + assert_matches!(s, Value::Null, "expected a unit\n{}", ctx,); + } + } +} + +fn check_enum_variants(e: &tree::Enum, s: &[Value], ctx: std::fmt::Arguments<'_>) { + let num_variants = e.variants().len(); + assert_eq!(num_variants, s.len(), "expected one example per variant"); + + for (i, (variant, example)) in std::iter::zip(e.variants().iter(), s.iter()).enumerate() { + let ctx = Context::new(example, example, ctx); + + // Verify that the extracted tag matches. + let (tag, _content) = extract_tag_and_content(e.repr(), example, ctx); + assert_eq!(variant.name(), tag, "{}", ctx); + + // Use `check_enum`. + // + // This is a little wasteful because we've already extracted the tag and the enum, + // but this is test-only code so we can afford it. + check_enum( + e, + example, + context!( + ctx, + example, + "variant = {} ({} of {})", + variant.name(), + i + 1, + num_variants + ), + ); + } +} + +/// To summarize. For a struct like +/// ``` +/// enum Foo { +/// Baz { +/// a: usize, +/// }, +/// } +/// ``` +/// +/// External looks like +/// ```json +/// { +/// "baz": { +/// "a": 10 +/// } +/// } +/// ``` +/// Internal (say "tag = "mytag") looks like +/// ```json +/// { +/// "mytag": "baz", +/// "a": 10, +/// } +/// ``` +/// Adjacency (say "tag = "mytag", "content" = "mycontent") looks like +/// ```json +/// { +/// "mytag": "baz", +/// "mycontent": { +/// "a": 10 +/// } +/// } +/// ``` +fn extract_tag_and_content<'a>( + repr: &tree::EnumRepr, + s: &'a Value, + ctx: Context<'_>, +) -> (&'a str, Option>) { + match repr { + // When using external tagging, the associated value can look like one of two things: + // + // 1. A raw string. This is only valid if the associated payload is a unit type. + // 2. A map consisting of a single key-value pair. The `key` is the enum tag. + tree::EnumRepr::External => match s { + Value::String(s) => (s, None), + Value::Object(m) => { + assert_eq!( + m.len(), + 1, + "externally tagged enums should only have a single key-value pair\n\n{}", + ctx + ); + + let kv = m.iter().next().unwrap(); + (kv.0, Some(Cow::Borrowed(kv.1))) + } + _ => panic!("invalid representation\n\n{}", ctx), + }, + + // For internally tagged enums - we remove the tag after retrieval. + // + // If the remaining dictionary is empty, then we change it to `None`. + // This is technically a little ambiguous between unit variants and empty struct + // variants (e.g. `Enum::Unit` and `Enum::Empty {}`. + // + // However, we teach the latter case to expect empty results with internal tagging + // for purposes of the check. + tree::EnumRepr::Internal { tag } => { + let map = value_as_map(s, ctx); + + let t = match map.get(*tag) { + Some(value) => value_as_str(value, context!(ctx, value, "tag extraction")), + None => panic!("Could not find tag \"{}\"\n\n{}", tag, ctx), + }; + + // Delete the tag field to reuse the rest of the checking infrastructure. + let mut map = map.clone(); + map.remove(*tag); + (t, Some(Cow::Owned(Value::Object(map)))) + } + + // For adjacent tagging, the "content" field is omitted when the corresponding + // variant is a unit variant. + tree::EnumRepr::Adjacent { tag, content } => { + let outer = value_as_map(s, ctx); + assert!(outer.len() == 1 || outer.len() == 2); + + let t = match outer.get(*tag) { + Some(value) => value_as_str(value, context!(ctx, value, "tag extraction")), + None => panic!("expected tag field \"{}\"", tag), + }; + + if outer.len() == 1 { + (t, None) + } else { + let c = match outer.get(*content) { + Some(c) => c, + None => panic!("expected content field \"{}\"", content), + }; + + (t, Some(Cow::Borrowed(c))) + } + } + } +} + +fn check_enum(e: &tree::Enum, s: &Value, ctx: Context<'_>) { + let (tag, content) = extract_tag_and_content(e.repr(), s, ctx); + assert_all_unique(e.variants().iter().map(|f| f.name()), ctx); + + let variant = match e.variants().iter().find(|v| v.name() == tag) { + Some(variant) => variant, + None => panic!("could not find variant \"{}\"\n\n{}", tag, ctx), + }; + + let next: &Value = match (e.repr(), &content) { + (tree::EnumRepr::External, None) => { + assert!( + variant.fields().is_unit(), + "content may only be excluded for unit variants\n\n{}", + ctx + ); + return; + } + (tree::EnumRepr::External, Some(c)) => c, + (tree::EnumRepr::Internal { .. }, None) => unreachable!("internal always returns content"), + (tree::EnumRepr::Internal { .. }, Some(c)) => { + if variant.fields().is_unit() { + let map = value_as_map(c, ctx); + assert!(map.is_empty(), "unit enums should have no remaining values"); + return; + } else { + c + } + } + (tree::EnumRepr::Adjacent { .. }, None) => { + assert!( + variant.fields().is_unit(), + "content may only be excluded for unit variants\n\n{}", + ctx + ); + return; + } + (tree::EnumRepr::Adjacent { .. }, Some(c)) => c, + }; + + check_fields( + variant.fields(), + next, + context!(ctx, next, "variant \"{}\"", variant.name()), + ) +} + +//----------// +// sequence // +//----------// + +fn check_sequence(seq: &tree::Sequence, s: &Value, ctx: Context<'_>) { + let a = value_as_array(s, ctx); + for (i, v) in a.iter().enumerate() { + check_reflection( + seq.element(), + v, + context!(ctx, v, "element {} of {}", i + 1, a.len()), + ); + } +} + +//----------// +// optional // +//----------// + +fn check_optional(opt: &tree::Optional, s: &Value, ctx: Context<'_>) { + if !s.is_null() { + check_reflection(opt.value(), s, context!(ctx, s, "present optional")) + } +} + +//---------// +// Helpers // +//---------// + +fn assert_all_unique(itr: I, ctx: Context<'_>) +where + I: IntoIterator, + I::Item: std::hash::Hash + Eq + std::fmt::Debug + Clone, +{ + let mut seen = HashSet::new(); + for i in itr { + if !seen.insert(i.clone()) { + panic!("item {:?} seen multiple times\n\n{}", i, ctx); + } + } +} + +fn value_as_str<'a>(v: &'a Value, ctx: Context<'_>) -> &'a str { + if let Value::String(s) = v { + s + } else { + panic!("expected value to be a string\n\n{}", ctx); + } +} + +fn value_as_map<'a>(v: &'a Value, ctx: Context<'_>) -> &'a serde_json::value::Map { + if let Value::Object(m) = v { + m + } else { + panic!("expected value to be an object\n\n{}", ctx); + } +} + +fn value_as_array<'a>(v: &'a Value, ctx: Context<'_>) -> &'a [Value] { + if let Value::Array(a) = v { + a + } else { + panic!("expected value to be an array\n\n{}", ctx); + } +} + +/////////// +// Tests // +/////////// + +#[test] +fn test_primitives() { + check_struct(()); + + check_struct(false); + + check_struct(0u8); + check_struct(0u16); + check_struct(0u32); + check_struct(0u64); + check_struct(0usize); + + check_struct(0i8); + check_struct(0i16); + check_struct(0i32); + check_struct(0i64); + check_struct(0isize); + + check_struct(0f32); + check_struct(0f64); + + check_struct(std::marker::PhantomData::); + + check_struct(String::from("hello")); +} + +#[test] +fn simple_aggregate() { + #[derive(Serialize, Reflect, Clone)] + struct Simple { + a: usize, + b: usize, + r#type: (), + } + + let s = Simple { + a: 10, + b: 20, + r#type: (), + }; + check_struct(s.clone()); + + #[derive(Serialize, Reflect)] + #[serde(rename_all = "kebab-case")] + struct Nested { + some_long_field: usize, + #[serde(rename = "s")] + manual_rename: Simple, + } + + check_struct(Nested { + some_long_field: 4, + manual_rename: s, + }); +} + +#[test] +fn simple_tuple() { + // A simple newtype tuple. + #[derive(Serialize, Reflect, Clone)] + struct NewType(String); + + let newtype = NewType(String::from("hello")); + check_struct(newtype.clone()); + + // A tuple struct of length 2 + #[derive(Serialize, Reflect, Clone)] + struct Tuple2(usize, NewType); + + let tuple2 = Tuple2(100, newtype.clone()); + check_struct(tuple2.clone()); + + // A tuple struct of length 3 with a generic + #[derive(Serialize, Reflect)] + struct Tuple3(usize, NewType, T); + + let tuple3 = Tuple3(4, NewType(String::from("foo")), tuple2.clone()); + check_struct(tuple3); + + // A tuple struct of length 3 with a generic and zero-sized type. + #[derive(Serialize, Reflect)] + struct Unit; + + #[derive(Serialize, Reflect)] + struct Tuple4(usize, Unit, NewType, T); + + let tuple3 = Tuple4(4, Unit, newtype, tuple2); + check_struct(tuple3); +} + +#[test] +fn nested_newtype() { + #[derive(Serialize, Reflect)] + struct NewType0(usize); + + #[derive(Serialize, Reflect)] + struct NewType1(NewType0); + + #[derive(Serialize, Reflect)] + struct NewType2(NewType1); + + check_struct(NewType2(NewType1(NewType0(10)))); +} + +#[test] +fn enum_externally_tagged() { + #[derive(Serialize, Reflect)] + struct Aggregate { + foo: usize, + bar: usize, + } + + #[derive(Serialize, Reflect)] + #[serde(rename_all = "kebab-case")] + enum Enum { + Unit, + NewType(T), + Tuple0(), + Tuple2(usize, usize), + Tuple3(usize, usize, Aggregate), + #[serde(rename = "longer-unit-2")] + Unit2 {}, + + #[serde(rename_all = "kebab-case")] + Struct1 { + hello_world: usize, + }, + + #[serde(rename_all = "kebab-case")] + Struct2 { + hello_world: usize, + #[serde(rename = "baz")] + foo: usize, + }, + } + + check_enums( + &[ + Enum::Unit, + Enum::NewType(0usize), + Enum::Tuple0(), + Enum::Tuple2(1, 2), + Enum::Tuple3(1, 2, Aggregate { foo: 10, bar: 20 }), + Enum::Unit2 {}, + Enum::Struct1 { hello_world: 3 }, + Enum::Struct2 { + hello_world: 4, + foo: 5, + }, + ], + format_args!("externally tagged enums"), + ); +} + +#[test] +fn acronym_heavy_variant_name() { + #[derive(Serialize, Reflect)] + #[serde(rename_all = "snake_case")] + enum Enum { + XMLHttpRequest, + } + + check_enums( + &[Enum::XMLHttpRequest], + format_args!("acronym-heavy variant names"), + ); +} + +#[test] +fn enum_internally_tagged() { + #[derive(Serialize, Reflect)] + struct NewTypePayload { + value: usize, + } + + #[derive(Serialize, Reflect)] + struct NewTypeEmpty {} + + #[derive(Serialize, Reflect)] + #[serde(rename_all = "kebab-case")] + #[serde(tag = "fizzle")] + enum Enum { + Unit, + EmptyStruct {}, + NewType(NewTypePayload), + NewTypeEmpty(NewTypeEmpty), + Struct1 { a: usize }, + Struct2 { a: usize, b: usize }, + Struct3 { a: usize, b: usize, c: usize }, + } + + check_enums( + &[ + Enum::Unit, + Enum::EmptyStruct {}, + Enum::NewType(NewTypePayload { value: 10 }), + Enum::NewTypeEmpty(NewTypeEmpty {}), + Enum::Struct1 { a: 0 }, + Enum::Struct2 { a: 0, b: 1 }, + Enum::Struct3 { a: 0, b: 1, c: 3 }, + ], + format_args!("internally tagged enums"), + ); +} + +#[test] +fn enum_adjacently_tagged() { + #[derive(Serialize, Reflect)] + struct Nested { + a: usize, + b: usize, + } + + #[derive(Serialize, Reflect)] + #[serde(rename_all = "kebab-case")] + #[serde(tag = "fizzle", content = "fuzzle")] + enum Enum { + Unit, + r#RawUnit, + Empty {}, + NewType(()), + NewType2(usize), + NewType3(Nested), + Tuple0(), + Tuple1(usize), + Tuple2(usize, ()), + Struct1 { a: usize }, + Struct2 { a: usize, b: usize }, + Struct3 { a: usize, b: usize, c: usize }, + } + + check_enums( + &[ + Enum::Unit, + Enum::r#RawUnit, + Enum::Empty {}, + Enum::NewType(()), + Enum::NewType2(10), + Enum::NewType3(Nested { a: 0, b: 1 }), + Enum::Tuple0(), + Enum::Tuple1(0), + Enum::Tuple2(0, ()), + Enum::Struct1 { a: 0 }, + Enum::Struct2 { a: 0, b: 1 }, + Enum::Struct3 { a: 0, b: 1, c: 3 }, + ], + format_args!("adjacently tagged enums"), + ); +} + +#[test] +fn nested_enum() { + #[derive(Serialize, Reflect)] + enum Nested { + Unit, + NewType(usize), + Struct { value: usize }, + } + + #[derive(Serialize, Reflect)] + struct Outer { + nested: Nested, + } + + check_struct(Outer { + nested: Nested::Unit, + }); + check_struct(Outer { + nested: Nested::NewType(10), + }); + check_struct(Outer { + nested: Nested::Struct { value: 20 }, + }); +} + +#[test] +fn sequence_of_newtypes() { + #[derive(Serialize, Reflect)] + struct NewType(usize); + + #[derive(Serialize, Reflect)] + struct Outer { + values: Vec, + } + + check_struct(Outer { + values: vec![NewType(10), NewType(20)], + }); +} + +#[test] +fn test_optionals() { + #[derive(Serialize, Reflect)] + #[serde(tag = "tag", content = "content", rename_all = "snake_case")] + enum CasesAdjacent { + NewType(Option), + Struct { val: Option }, + } + + #[derive(Serialize, Reflect)] + struct NewType(Option); + + #[derive(Serialize, Reflect)] + struct More { + a: usize, + b: usize, + } + + #[derive(Serialize, Reflect)] + struct Struct { + e0: CasesAdjacent, + e1: CasesAdjacent, + more: Option, + opt: Option, + newtype: NewType, + } + + check_struct(Struct { + e0: CasesAdjacent::NewType(None), + e1: CasesAdjacent::Struct { val: None }, + more: None, + opt: None, + newtype: NewType(None), + }); + + let some = Some(1); + + check_struct(Struct { + e0: CasesAdjacent::NewType(some), + e1: CasesAdjacent::Struct { val: some }, + more: Some(More { a: 0, b: 1 }), + opt: some, + newtype: NewType(some), + }); +} diff --git a/diskann-benchmark-runner/src/reflect/test/unit.rs b/diskann-benchmark-runner/src/reflect/test/unit.rs new file mode 100644 index 0000000000..6dd149082d --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/test/unit.rs @@ -0,0 +1,395 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`Reflect`] macro unit tests. + +use std::assert_matches; + +use crate::reflect::{Reflect, Reflection, tree::Fields}; + +#[test] +fn test_unit() { + /// A unit struct. + #[derive(Reflect)] + struct Unit; + + let r = Reflection::new::(); + let ty = r.ty(); + + assert_eq!(r.ty().doc().unwrap(), "A unit struct."); + assert_eq!(r.type_name().to_string(), "Unit"); + assert!(!ty.has_body()); + + // Unit structs should be unit aggregates. + let agg = ty.as_aggregate().unwrap(); + assert!(agg.fields().is_unit()); +} + +#[test] +fn test_unit_rename() { + /// A unit struct. + #[derive(Reflect)] + #[reflect(type_name = "something-else")] + struct Unit; + + let r = Reflection::new::(); + let ty = r.ty(); + + assert_eq!(r.ty().doc().unwrap(), "A unit struct."); + assert_eq!(r.type_name().to_string(), "something-else"); + assert!(!ty.has_body()); + + // Unit structs should be unit aggregates. + let agg = ty.as_aggregate().unwrap(); + assert!(agg.fields().is_unit()); +} + +#[test] +fn test_unit_prefix() { + /// A unit struct. + #[derive(Reflect)] + #[reflect(prefix = "module::")] + struct Unit; + + let r = Reflection::new::(); + let ty = r.ty(); + + assert_eq!(r.ty().doc().unwrap(), "A unit struct."); + assert_eq!(r.type_name().to_string(), "module::Unit"); + assert!(!ty.has_body()); +} + +#[test] +fn test_unit_const_generic() { + /// A unit struct. + #[derive(Reflect)] + struct Unit; + + let r = Reflection::new::>(); + let ty = r.ty(); + + assert_eq!(ty.doc().unwrap(), "A unit struct."); + assert_eq!(r.type_name().to_string(), "Unit<10>"); + assert!(!ty.has_body()); + + // Unit structs should be unit aggregates. + let agg = ty.as_aggregate().unwrap(); + assert!(agg.fields().is_unit()); +} + +#[test] +fn test_unit_const_generic_prefix() { + /// A unit struct. + #[derive(Reflect)] + #[reflect(prefix = "module::")] + struct Unit; + + let r = Reflection::new::>(); + let ty = r.ty(); + + assert_eq!(ty.doc().unwrap(), "A unit struct."); + assert_eq!(r.type_name().to_string(), "module::Unit<10>"); + assert!(!ty.has_body()); +} + +#[test] +fn test_unit_const_generic_2() { + #[derive(Reflect)] + struct Unit; + + let r = Reflection::new::>(); + let ty = r.ty(); + + assert!(ty.doc().is_none()); + assert_eq!(r.type_name().to_string(), "Unit<10, 20>"); + assert!(!ty.has_body()); +} + +#[test] +fn test_empty_tuple_like() { + /// An empty tuple-like struct. + #[derive(Reflect)] + struct Empty(); + + let r = Reflection::new::(); + let ty = r.ty(); + + assert_eq!(ty.doc().unwrap(), "An empty tuple-like struct."); + assert_eq!(r.type_name().to_string(), "Empty"); + assert!(!ty.has_body()); + + // Unit tuple structs should be aggregates with zero unnamed fields. + let agg = ty.as_aggregate().unwrap(); + let fields = agg.fields().as_unnamed().unwrap(); + assert!(fields.is_empty()); +} + +#[test] +fn test_empty_struct_like() { + /// An empty struct. + #[derive(Reflect)] + struct Empty {} + + let r = Reflection::new::(); + let ty = r.ty(); + + assert_eq!(ty.doc().unwrap(), "An empty struct."); + assert_eq!(r.type_name().to_string(), "Empty"); + assert!(!ty.has_body()); + + // Unit structs should be aggregates with zero unnamed fields. + let agg = ty.as_aggregate().unwrap(); + let fields = agg.fields().as_named().unwrap(); + assert!(fields.is_empty()); +} + +#[test] +fn test_struct() { + /// A struct with two fields. + #[expect(unused)] + #[derive(Reflect)] + struct Woo { + /// Foo + foo: usize, + /// Bar + bar: usize, + } + + let r = Reflection::new::(); + let ty = r.ty(); + assert_eq!(ty.doc().unwrap(), "A struct with two fields."); + assert_eq!(r.type_name().to_string(), "Woo"); + assert!(ty.has_body()); + + let f = ty.as_aggregate().unwrap().fields().as_named().unwrap(); + assert_eq!(f.len(), 2); + + assert_eq!(f[0].name(), "foo"); + assert_eq!(f[0].doc().unwrap(), "Foo"); + assert_eq!(f[0].field().type_name().to_string(), "usize"); + + assert_eq!(f[1].name(), "bar"); + assert_eq!(f[1].doc().unwrap(), "Bar"); + assert_eq!(f[1].field().type_name().to_string(), "usize"); +} + +#[test] +fn test_struct_rename() { + /// A struct with three fields. + #[expect(unused)] + #[derive(Reflect)] + #[serde(rename_all = "kebab-case")] + struct Woo { + /// Foo + foo_bar: usize, + /// Bar + baz: usize, + + #[serde(rename = "oops")] + biz: usize, + } + + let r = Reflection::new::(); + let ty = r.ty(); + assert_eq!(ty.doc().unwrap(), "A struct with three fields."); + assert_eq!(r.type_name().to_string(), "Woo"); + assert!(ty.has_body()); + + let f = ty.as_aggregate().unwrap().fields().as_named().unwrap(); + assert_eq!(f.len(), 3); + + assert_eq!( + f[0].name(), + "foo-bar", + "fields should be renamed by `rename_all`" + ); + assert_eq!(f[0].doc().unwrap(), "Foo"); + assert_eq!(f[0].field().type_name().to_string(), "usize"); + + assert_eq!(f[1].name(), "baz"); + assert_eq!(f[1].doc().unwrap(), "Bar"); + assert_eq!(f[1].field().type_name().to_string(), "usize"); + + assert_eq!( + f[2].name(), + "oops", + "explicit rename should take precedence" + ); + assert!(f[2].doc().is_none()); + assert_eq!(f[2].field().type_name().to_string(), "usize"); +} + +#[test] +fn test_tuple1() { + /// A tuple with one field. This should be a "newtype". + #[expect(unused)] + #[derive(Reflect)] + struct Tuple1( + /// Field 0. + usize, + ); + + let r = Reflection::new::(); + let ty = r.ty(); + assert_eq!( + ty.doc().unwrap(), + "A tuple with one field. This should be a \"newtype\"." + ); + assert_eq!(r.type_name().to_string(), "Tuple1"); + + assert!(ty.has_body()); + + let f = ty.as_aggregate().unwrap().fields().as_newtype().unwrap(); + assert_eq!(f.doc().unwrap(), "Field 0."); + assert_eq!(f.field().type_name().to_string(), "usize"); +} + +#[test] +fn test_tuple2() { + #[expect(unused)] + #[derive(Reflect)] + struct Tuple2( + T, + /// Field 1. + Vec, + ); + + let r = Reflection::new::>(); + let ty = r.ty(); + assert!(ty.doc().is_none()); + assert_eq!(r.type_name().to_string(), "Tuple2"); + assert!(ty.has_body()); + + let f = ty.as_aggregate().unwrap().fields().as_unnamed().unwrap(); + assert_eq!(f.len(), 2); + + assert!(f[0].doc().is_none()); + assert_eq!(f[0].field().type_name().to_string(), "usize"); + + assert_eq!(f[1].doc().unwrap(), "Field 1."); + assert_eq!(f[1].field().type_name().to_string(), "Vec"); +} + +#[test] +fn test_empty_enum() { + /// An empty enum. + #[derive(Reflect)] + enum Empty {} + + let r = Reflection::new::(); + let ty = r.ty(); + assert_eq!(ty.doc().unwrap(), "An empty enum."); + assert_eq!(r.type_name().to_string(), "Empty"); + assert!(!ty.has_body()); +} + +#[test] +fn test_enum_with_generics() { + /// An enum with generics. + #[expect(unused)] + #[derive(Reflect)] + enum Either { + A(Vec), + /// It's a bee! + B( + /// Buzz buzz + B, + ), + } + + let r = Reflection::new::>(); + let ty = r.ty(); + assert_eq!(ty.doc().unwrap(), "An enum with generics."); + assert_eq!(r.type_name().to_string(), "Either"); + assert!(ty.has_body()); + + let variants = ty.as_enum().unwrap().variants(); + assert_eq!(variants.len(), 2); + + // Variant 0 + assert!(variants[0].doc().is_none()); + assert_eq!(variants[0].name(), "A"); + assert!(variants[0].fields().has_body()); + + let f = variants[0].fields().as_newtype().unwrap(); + assert!(f.doc().is_none()); + assert_eq!(f.field().type_name().to_string(), "Vec"); + + // Variant 1 + assert_eq!(variants[1].doc().unwrap(), "It's a bee!"); + assert_eq!(variants[1].name(), "B"); + assert!(variants[1].fields().has_body()); + + let f = variants[1].fields().as_newtype().unwrap(); + assert_eq!(f.doc().unwrap(), "Buzz buzz"); + assert_eq!(f.field().type_name().to_string(), "string"); +} + +#[test] +fn test_enum_variants() { + /// All the enums. + #[expect(unused)] + #[derive(Reflect)] + #[serde(tag = "tag", content = "content")] + enum All { + /// A unit variant. + Unit, + /// A tuple-like variant. + Tuple( + /// Field 0. + usize, + String, + ), + /// Struct-like. + Struct { + foo: usize, + /// All the bars! + bar: u32, + }, + } + + let r = Reflection::new::(); + let ty = r.ty(); + assert_eq!(ty.doc().unwrap(), "All the enums."); + assert_eq!(r.type_name().to_string(), "All"); + assert!(ty.has_body()); + + let variants = ty.as_enum().unwrap().variants(); + assert_eq!(variants.len(), 3); + + // Variant 0 + assert_eq!(variants[0].doc().unwrap(), "A unit variant."); + assert_eq!(variants[0].name(), "Unit"); + assert!(!variants[0].fields().has_body()); + assert_matches!(variants[0].fields(), Fields::Unit); + + // Variant 1 + assert_eq!(variants[1].doc().unwrap(), "A tuple-like variant."); + assert_eq!(variants[1].name(), "Tuple"); + assert!(variants[1].fields().has_body()); + + let f = variants[1].fields().as_unnamed().unwrap(); + assert_eq!(f.len(), 2); + + assert_eq!(f[0].doc().unwrap(), "Field 0."); + assert_eq!(f[0].field().type_name().to_string(), "usize"); + + assert!(f[1].doc().is_none()); + assert_eq!(f[1].field().type_name().to_string(), "string"); + + // Variant 2 + assert_eq!(variants[2].name(), "Struct"); + assert_eq!(variants[2].doc().unwrap(), "Struct-like."); + let f = variants[2].fields().as_named().unwrap(); + assert_eq!(f.len(), 2); + + assert_eq!(f[0].name(), "foo"); + assert!(f[0].doc().is_none()); + assert_eq!(f[0].field().type_name().to_string(), "usize"); + + assert_eq!(f[1].name(), "bar"); + assert_eq!(f[1].doc().unwrap(), "All the bars!"); + assert_eq!(f[1].field().type_name().to_string(), "u32"); +} diff --git a/diskann-benchmark-runner/src/reflect/tree.rs b/diskann-benchmark-runner/src/reflect/tree.rs new file mode 100644 index 0000000000..1543bee896 --- /dev/null +++ b/diskann-benchmark-runner/src/reflect/tree.rs @@ -0,0 +1,500 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Model of the Rust type system. +//! +//! This is largely based on the structure of types in [`syn`](https://docs.rs/syn/latest/syn/) +//! with some extra entries (e.g. [`Type::Optional`]) for a closer representation with how +//! types are serialized by [`serde`]. + +use super::{Reflect, Reflection}; + +pub type Doc = std::borrow::Cow<'static, str>; + +/// Classification of types. +#[derive(Debug)] +pub enum Type { + Primitive(Primitive), + Aggregate(Aggregate), + Enum(Enum), + Sequence(Sequence), + Optional(Optional), +} + +impl Type { + /// Construc a new [`Type::Primitive`]. + pub fn primitive(kind: PrimitiveKind, doc: Option) -> Self { + Self::Primitive(Primitive::new(kind, doc)) + } + + /// Construc a new [`Type::Aggregate`]. + pub fn aggregate(fields: Fields, doc: Option) -> Self { + Self::Aggregate(Aggregate::new(fields, doc)) + } + + /// Construc a new [`Type::Enum`]. + pub fn enum_( + repr: EnumRepr, + variants: impl IntoIterator, + doc: Option, + ) -> Self { + Self::Enum(Enum::new(repr, variants, doc)) + } + + /// Construc a new [`Type::Sequence`]. + pub fn sequence(doc: Option) -> Self + where + T: Reflect, + { + Self::Sequence(Sequence::new::(doc)) + } + + /// Keep the constructor private since we don't want users constructing the very special + /// `Optional` type for their own types. + pub(crate) fn optional(doc: Option) -> Self + where + T: Reflect, + { + Self::Optional(Optional::new::(doc)) + } + + /// Return the struct level documentation if available. + pub fn doc(&self) -> Option<&str> { + match self { + Self::Primitive(p) => p.doc(), + Self::Aggregate(a) => a.doc(), + Self::Enum(e) => e.doc(), + Self::Sequence(s) => s.doc(), + Self::Optional(o) => o.doc(), + } + } + + /// Return the primitive JSON kind (if there is one). + pub(super) fn json_kind(&self) -> Option<&str> { + match self { + Self::Primitive(p) => Some(p.kind().json_kind()), + Self::Aggregate(_) | Self::Enum(_) | Self::Sequence(_) | Self::Optional(_) => None, + } + } + + /// Return `true` if there is field level information of some kind to render. + pub(super) fn has_body(&self) -> bool { + match self { + Self::Primitive(_) => false, + Self::Aggregate(a) => a.has_body(), + Self::Enum(e) => e.has_body(), + Self::Sequence(_) => true, + Self::Optional(_) => true, + } + } + + #[cfg(test)] + pub(super) fn as_aggregate(&self) -> Option<&Aggregate> { + if let Self::Aggregate(aggregate) = self { + Some(aggregate) + } else { + None + } + } + + #[cfg(test)] + pub(super) fn as_enum(&self) -> Option<&Enum> { + if let Self::Enum(enum_) = self { + Some(enum_) + } else { + None + } + } +} + +//-----------// +// Primitive // +//-----------// + +/// The native JSON representation for a primitive. +#[derive(Debug, Clone, Copy)] +pub enum PrimitiveKind { + Null, + Boolean, + Number, + String, +} + +impl PrimitiveKind { + pub(crate) fn json_kind(&self) -> &'static str { + match self { + Self::Null => "null", + Self::Boolean => "bool", + Self::Number => "number", + Self::String => "string", + } + } +} + +/// A primitive type that maps closely to a native JSON type. +#[derive(Debug)] +pub struct Primitive { + kind: PrimitiveKind, + doc: Option, +} + +impl Primitive { + pub(crate) fn new(kind: PrimitiveKind, doc: Option) -> Self { + Self { kind, doc } + } + + pub(super) fn kind(&self) -> PrimitiveKind { + self.kind + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} + +//-----------// +// Aggregate // +//-----------// + +/// A representation of aggregates like normal structs, unit structs, and tuple-like structs. +#[derive(Debug)] +pub struct Aggregate { + fields: Fields, + doc: Option, +} + +impl Aggregate { + pub(crate) fn new(fields: Fields, doc: Option) -> Self { + Self { fields, doc } + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } + + pub(super) fn fields(&self) -> &Fields { + &self.fields + } + + pub(super) fn has_body(&self) -> bool { + self.fields.has_body() + } +} + +/// Represent the fields of a struct. +#[derive(Debug)] +pub enum Fields { + /// Standard Rust structs. + Named(Vec), + + /// Tuple-like structs. + Unnamed(Vec), + + /// Tuple-like structs with a single named field. + /// + /// These are treated specially by `serde` and thus get their own variant. + NewType(UnnamedField), + + /// Unit structs. + Unit, +} + +impl Fields { + pub(super) fn has_body(&self) -> bool { + match self { + Self::Named(fields) => !fields.is_empty(), + Self::Unnamed(fields) => !fields.is_empty(), + Self::NewType(_) => true, + Self::Unit => false, + } + } + + /// Construct [`Fields::Named`] from the iterator. + pub fn named(itr: impl IntoIterator) -> Self { + Self::Named(itr.into_iter().collect()) + } + + /// Construct [`Fields::Unamed`] from the iterator. + pub fn unnamed(itr: impl IntoIterator) -> Self { + Self::Unnamed(itr.into_iter().collect()) + } + + /// Construct [`Fields::NewType`] from the iterator. + pub fn newtype(field: UnnamedField) -> Self { + Self::NewType(field) + } + + #[cfg(test)] + pub(super) fn as_named(&self) -> Option<&[NamedField]> { + if let Self::Named(fields) = self { + Some(fields) + } else { + None + } + } + + #[cfg(test)] + pub(super) fn as_unnamed(&self) -> Option<&[UnnamedField]> { + if let Self::Unnamed(fields) = self { + Some(fields) + } else { + None + } + } + + #[cfg(test)] + pub(super) fn as_newtype(&self) -> Option<&UnnamedField> { + if let Self::NewType(field) = self { + Some(field) + } else { + None + } + } + + /// Return `true` if `self` is [`Self::Unit`]. + pub(super) fn is_unit(&self) -> bool { + matches!(self, Self::Unit) + } +} + +/// A struct field with a name. +#[derive(Debug)] +pub struct NamedField { + name: &'static str, + field: Reflection, + doc: Option, +} + +impl NamedField { + /// Construct a new [`NamedField`] for `T`. + pub fn new(name: &'static str, doc: Option) -> Self + where + T: Reflect, + { + Self { + name, + field: Reflection::new::(), + doc, + } + } + + pub(super) fn name(&self) -> &str { + self.name + } + + pub(super) fn field(&self) -> Reflection { + self.field + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} + +/// An unnamed field. +#[derive(Debug)] +pub struct UnnamedField { + field: Reflection, + doc: Option, +} + +impl UnnamedField { + /// Construct a new [`UnnamedField`] for `T`. + pub fn new(doc: Option) -> Self + where + T: Reflect, + { + Self { + field: Reflection::new::(), + doc, + } + } + + pub(super) fn field(&self) -> Reflection { + self.field + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} + +//------// +// Enum // +//------// + +/// A representation for enums. +#[derive(Debug)] +pub struct Enum { + repr: EnumRepr, + variants: Vec, + doc: Option, +} + +impl Enum { + /// Construct a new [`Enum`]. + pub fn new( + repr: EnumRepr, + variants: impl IntoIterator, + doc: Option, + ) -> Self { + Self { + repr, + variants: variants.into_iter().collect(), + doc, + } + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } + + pub(super) fn repr(&self) -> &EnumRepr { + &self.repr + } + + pub(super) fn variants(&self) -> &[Variant] { + &self.variants + } + + pub(super) fn has_body(&self) -> bool { + !self.variants.is_empty() + } +} + +/// Describe how an enum is being represented by `serde`. +#[derive(Debug)] +#[non_exhaustive] +pub enum EnumRepr { + /// Enums are tagged as the key in a collection. + External, + + /// Enums are tagged by an internal field. + /// + /// This contains the name of the field that is used as the tag. + /// + /// ```json + /// { + /// "tag": "some-enum-tag", + /// "value": 10, + /// "members": [ + /// 1, + /// "world" + /// ] + /// } + /// ``` + Internal { tag: &'static str }, + + /// Enums have a separate tag and content payload. + /// + /// ```json + /// { + /// "tag": "some-enum-tag", + /// "content": { + /// "value": 10, + /// "members": [ + /// 1, + /// "world" + /// ] + /// } + /// } + /// ``` + Adjacent { + tag: &'static str, + content: &'static str, + }, +} + +/// A variant of an [`Enum`]. +#[derive(Debug)] +pub struct Variant { + name: &'static str, + fields: Fields, + doc: Option, +} + +impl Variant { + /// Construct a new [`Variant`]. + pub fn new(name: &'static str, fields: Fields, doc: Option) -> Self { + Self { name, fields, doc } + } + + pub(super) fn name(&self) -> &'static str { + self.name + } + + pub(super) fn fields(&self) -> &Fields { + &self.fields + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} + +//----------// +// Sequence // +//----------// + +/// A homogeneous sequence of values. +#[derive(Debug)] +pub struct Sequence { + element: Reflection, + doc: Option, +} + +impl Sequence { + /// Create a new [`Sequence`] containing `T`. + pub fn new(doc: Option) -> Self + where + T: Reflect, + { + Self { + element: Reflection::new::(), + doc, + } + } + + pub(super) fn element(&self) -> Reflection { + self.element + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} + +//----------// +// Optional // +//----------// + +/// An [`Option`]. +#[derive(Debug)] +pub struct Optional { + element: Reflection, + doc: Option, +} + +impl Optional { + /// Keep the constructor private since we don't want users constructing the very special + /// `Optional` type for their own types. + pub(super) fn new(doc: Option) -> Self + where + T: Reflect, + { + Self { + element: Reflection::new::(), + doc, + } + } + + pub(super) fn value(&self) -> Reflection { + self.element + } + + pub(super) fn doc(&self) -> Option<&str> { + self.doc.as_deref() + } +} diff --git a/diskann-benchmark-runner/src/registry.rs b/diskann-benchmark-runner/src/registry.rs index 66a9f9468f..3b2c893d05 100644 --- a/diskann-benchmark-runner/src/registry.rs +++ b/diskann-benchmark-runner/src/registry.rs @@ -11,16 +11,21 @@ use std::{ use thiserror::Error; use crate::{ - Checkpoint, Features, Input, Output, + Checkpoint, Features, Input, Output, Reflection, benchmark::{self, Benchmark, MatchContext, Regression, Score, internal::AnnotatedMatch}, input, - internal::visibility::Visibility, + internal::{typemap::TypeMap, visibility::Visibility}, }; /// A collection of registered inputs and benchmarks. pub struct Registry { // Inputs keyed by their tag type. inputs: HashMap<&'static str, input::internal::Input>, + + /// Collection of input types. + types: TypeMap, + + /// The registered benchmarks. benchmarks: Vec, } @@ -29,10 +34,15 @@ impl Registry { pub fn new() -> Self { Self { inputs: HashMap::new(), + types: TypeMap::new(), benchmarks: Vec::new(), } } + //-------// + // Input // + //-------// + /// Return the input with the registered `tag` if present. Otherwise, return `None`. /// /// Inputs are automatically registered as a side-effect of: @@ -47,6 +57,20 @@ impl Registry { self.inputs.values() } + //-----------// + // Type Info // + //-----------// + + /// Return the [`Reflection`] type information for the provided type-name. + pub fn type_info(&self, type_name: &str) -> Option { + self.types.type_info(type_name) + } + + /// Return an iterator over all the registered type-names. + pub fn type_names(&self) -> impl ExactSizeIterator { + self.types.type_names() + } + //--------------// // Registration // //--------------// @@ -64,7 +88,10 @@ impl Registry { where T: Benchmark, { - self.register_input(input::internal::Input::new::())?; + self.register_input( + input::internal::Input::new::(), + &[Reflection::new::<::Raw>()], + )?; self.benchmarks.push(RegisteredBenchmark { name: name.into(), @@ -94,7 +121,7 @@ impl Registry { features: Features, description: impl Into, ) -> Result<(), RegistryError> { - self.register_input(input::internal::Input::gated(tag, features.clone()))?; + self.register_input(input::internal::Input::gated(tag, features.clone()), &[])?; self.benchmarks.push(RegisteredBenchmark { name: name.into(), @@ -125,7 +152,10 @@ impl Registry { where I: Input, { - self.register_input(input::internal::Input::new::())?; + self.register_input( + input::internal::Input::new::(), + &[Reflection::new::()], + )?; self.benchmarks.push(RegisteredBenchmark { name: name.into(), @@ -220,12 +250,31 @@ impl Registry { find_best_match(job, self.benchmarks()) } - fn register_input(&mut self, input: input::internal::Input) -> Result<(), RegistryError> { + fn register_input( + &mut self, + input: input::internal::Input, + types: &[Reflection], + ) -> Result<(), RegistryError> { + // Register all types. Only complete if type input registration succeeds. + // + // This allows multiple rounds of `Reflections` to be added even if the input type + // is the same (which can occur when collecting type information for regression + // inputs). + // + // Defer evaluating the `Result` until we check the type. This keeps incorrect + // inputs the first cause of error instead of a more obscure type conflict. + let result = self.types.register(types); + let tag = input.tag(); - match self.inputs.entry(tag) { + let input_result = match self.inputs.entry(tag) { Entry::Vacant(v) => { + // Make sure type insertion doesn't conflict. + match result { + Ok(transaction) => transaction.complete(), + Err(err) => return Err(RegistryError::type_name_conflict(err)), + }; v.insert(input); - Ok(()) + return Ok(()); } Entry::Occupied(o) => { use Visibility::{Available, Gated}; @@ -236,36 +285,45 @@ impl Registry { if o.as_any().type_id() == input.as_any().type_id() { Ok(()) } else { - Err(RegistryError { + Err(RegistryError::input_conflict( tag, - existing: Kind::Available(o.type_name()), - new: Kind::Available(input.type_name()), - }) + Kind::Available(o.type_name()), + Kind::Available(input.type_name()), + )) } } - (Available, Gated { features }) => Err(RegistryError { + (Available, Gated { features }) => Err(RegistryError::input_conflict( tag, - existing: Kind::Available(o.type_name()), - new: Kind::Gated(features.to_string()), - }), - (Gated { features }, Available) => Err(RegistryError { + Kind::Available(o.type_name()), + Kind::Gated(features.to_string()), + )), + (Gated { features }, Available) => Err(RegistryError::input_conflict( tag, - existing: Kind::Gated(features.to_string()), - new: Kind::Available(input.type_name()), - }), + Kind::Gated(features.to_string()), + Kind::Available(input.type_name()), + )), (Gated { features: existing }, Gated { features: new }) => { if existing == new { Ok(()) } else { - Err(RegistryError { + Err(RegistryError::input_conflict( tag, - existing: Kind::Gated(existing.to_string()), - new: Kind::Gated(new.to_string()), - }) + Kind::Gated(existing.to_string()), + Kind::Gated(new.to_string()), + )) } } } } + }; + + input_result?; + match result { + Ok(transaction) => { + transaction.complete(); + Ok(()) + } + Err(err) => Err(RegistryError::type_name_conflict(err)), } } @@ -289,12 +347,19 @@ impl Registry { where T: Regression, { - self.register_input(input::internal::Input::new::())?; + self.register_input( + input::internal::Input::new::(), + &[ + Reflection::new::<::Raw>(), + Reflection::new::<::Raw>(), + ], + )?; let registered = benchmark::internal::Wrapper::::new( benchmark, benchmark::internal::WithRegression, ); + self.benchmarks.push(RegisteredBenchmark { name: name.into(), benchmark: Box::new(registered), @@ -454,16 +519,40 @@ impl std::fmt::Display for Display<'_> { /// Error for [`Registry::register`] or [`Registry::register_regression`]. #[derive(Debug, Error)] -#[error( - "A different input with tag \"{}\" was already registered. Existing {}. New {}", - self.tag, - self.existing, - self.new, -)] +#[error(transparent)] pub struct RegistryError { - tag: &'static str, - existing: Kind, - new: Kind, + inner: RegistryErrorInner, +} + +impl RegistryError { + fn input_conflict(tag: &'static str, existing: Kind, new: Kind) -> Self { + Self { + inner: RegistryErrorInner::InputConflict { tag, existing, new }, + } + } + + fn type_name_conflict(conflict: crate::internal::typemap::Conflict) -> Self { + Self { + inner: RegistryErrorInner::TypeNameConflict(conflict), + } + } +} + +#[derive(Debug, Error)] +enum RegistryErrorInner { + #[error( + "A different input with tag \"{}\" was already registered. Existing {}. New {}", + tag, + existing, + new + )] + InputConflict { + tag: &'static str, + existing: Kind, + new: Kind, + }, + #[error(transparent)] + TypeNameConflict(crate::internal::typemap::Conflict), } #[derive(Debug)] @@ -584,7 +673,10 @@ pub(crate) struct RegisteredTolerance<'a> { mod tests { use super::*; - use crate::{Checker, input}; + use diskann_utils::assert_contains; + use serde::{Deserialize, Serialize}; + + use crate::{Checker, Reflect, input}; macro_rules! input { ($T:ident, $tag:literal) => { @@ -618,10 +710,10 @@ mod tests { fn test_tag_conflicts() { let mut registry = Registry::new(); registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap(); registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap(); let mut tags: Vec<_> = registry.inputs().map(|i| i.tag()).collect(); @@ -645,7 +737,7 @@ mod tests { } let err = registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap_err(); let msg = err.to_string(); assert!( @@ -655,18 +747,68 @@ mod tests { ); } + #[test] + fn test_duplicate_input_registers_additional_types() { + #[derive(Reflect)] + struct First; + + #[derive(Reflect)] + struct Second; + + #[derive(Reflect)] + struct RolledBack; + + let mut registry = Registry::new(); + registry + .register_input( + input::internal::Input::new::(), + &[Reflection::new::()], + ) + .unwrap(); + registry + .register_input( + input::internal::Input::new::(), + &[Reflection::new::()], + ) + .unwrap(); + + assert!( + registry + .type_info(&Reflection::new::().type_name().to_string()) + .is_some() + ); + assert!( + registry + .type_info(&Reflection::new::().type_name().to_string()) + .is_some() + ); + + registry + .register_input( + input::internal::Input::new::(), + &[Reflection::new::()], + ) + .unwrap_err(); + + assert!( + registry + .type_info(&Reflection::new::().type_name().to_string()) + .is_none() + ); + } + #[test] fn test_gated_to_real_conflict() { let mut registry = Registry::new(); registry - .register_input(input::internal::Input::gated( - "type-a", - Features::new("feature"), - )) + .register_input( + input::internal::Input::gated("type-a", Features::new("feature")), + &[], + ) .unwrap(); registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap(); let mut tags: Vec<_> = registry.inputs().map(|i| i.tag()).collect(); @@ -679,7 +821,7 @@ mod tests { } let err = registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap_err(); let msg = err.to_string(); assert!( @@ -694,17 +836,17 @@ mod tests { let mut registry = Registry::new(); registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap(); registry - .register_input(input::internal::Input::new::()) + .register_input(input::internal::Input::new::(), &[]) .unwrap(); let err = registry - .register_input(input::internal::Input::gated( - "type-a", - Features::new("feature"), - )) + .register_input( + input::internal::Input::gated("type-a", Features::new("feature")), + &[], + ) .unwrap_err(); let msg = err.to_string(); @@ -720,26 +862,26 @@ mod tests { let mut registry = Registry::new(); registry - .register_input(input::internal::Input::gated( - "type-a", - Features::new("feature"), - )) + .register_input( + input::internal::Input::gated("type-a", Features::new("feature")), + &[], + ) .unwrap(); // If we register with the same feature set, there is no error. registry - .register_input(input::internal::Input::gated( - "type-a", - Features::new("feature"), - )) + .register_input( + input::internal::Input::gated("type-a", Features::new("feature")), + &[], + ) .unwrap(); // If we register with a different feature-set, that is an error. let err = registry - .register_input(input::internal::Input::gated( - "type-a", - Features::any(["feature", "another"]), - )) + .register_input( + input::internal::Input::gated("type-a", Features::any(["feature", "another"])), + &[], + ) .unwrap_err(); let msg = err.to_string(); @@ -749,4 +891,100 @@ mod tests { msg ); } + + // Test that type registration happens before input registration, and that if type + // registration fails, the input is not registered. + #[test] + fn test_input_registration_aborts_correctly() { + #[derive(Serialize, Deserialize, Reflect, Debug)] + struct GoodInput { + a: usize, + b: isize, + } + + impl Input for GoodInput { + type Raw = Self; + fn tag() -> &'static str { + "good-input" + } + fn from_raw(_raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + unimplemented!("this struct is for test only"); + } + fn serialize(&self) -> anyhow::Result { + unimplemented!("this struct is for test only"); + } + fn example() -> Self::Raw { + unimplemented!("this struct is for test only"); + } + } + + #[derive(Serialize, Deserialize, Reflect, Debug)] + #[reflect(type_name = "isize")] + struct Boom; + + #[derive(Serialize, Deserialize, Reflect, Debug)] + struct BadInput { + /// This type should not be registered on failure. + a: f32, + b: Boom, + /// Put another `f32` on the other side of `Boom` so no matter the expansion + /// order, a `f32` is registered before `Boom`. + c: f32, + } + + impl Input for BadInput { + type Raw = Self; + fn tag() -> &'static str { + "bad-input" + } + fn from_raw(_raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + unimplemented!("this struct is for test only"); + } + fn serialize(&self) -> anyhow::Result { + unimplemented!("this struct is for test only"); + } + fn example() -> Self::Raw { + unimplemented!("this struct is for test only"); + } + } + + let mut registry = Registry::new(); + registry + .register_input( + input::internal::Input::new::(), + &[Reflection::new::()], + ) + .unwrap(); + + assert!(registry.input("good-input").is_some()); + assert!(registry.input("bad-input").is_none()); + + let mut type_names: Vec<_> = registry.type_names().collect(); + type_names.sort(); + assert_eq!(type_names, &["GoodInput", "isize", "usize"]); + + // This should hit a type-conflict. + // + // Also include an extra `f32` so we can test that all the reflections are checked + // before finalization. + let err = registry + .register_input( + input::internal::Input::new::(), + &[Reflection::new::(), Reflection::new::()], + ) + .unwrap_err(); + let msg = err.to_string(); + assert_contains!(msg, "isize"); + + assert!(registry.input("good-input").is_some()); + assert!( + registry.input("bad-input").is_none(), + "bad input should not be registered on failure" + ); + + // Registered type names + let mut type_names: Vec<_> = registry.type_names().collect(); + type_names.sort(); + assert_eq!(type_names, &["GoodInput", "isize", "usize"]); + } } diff --git a/diskann-benchmark-runner/src/test/dim.rs b/diskann-benchmark-runner/src/test/dim.rs index 329cb64a34..f7787b77a8 100644 --- a/diskann-benchmark-runner/src/test/dim.rs +++ b/diskann-benchmark-runner/src/test/dim.rs @@ -8,7 +8,7 @@ use std::io::Write; use serde::{Deserialize, Serialize}; use crate::{ - Benchmark, Checker, Checkpoint, Input, Output, + Benchmark, Checker, Checkpoint, Input, Output, Reflect, benchmark::{MatchContext, PassFail, Regression, Score}, }; @@ -16,7 +16,8 @@ use crate::{ // Input // /////////// -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::")] pub(super) struct DimInput { dim: Option, } @@ -55,7 +56,8 @@ impl Input for DimInput { // Tolerance // /////////////// -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::dim::")] pub(super) struct Tolerance { succeed: bool, error_in_check: bool, diff --git a/diskann-benchmark-runner/src/test/gated.rs b/diskann-benchmark-runner/src/test/gated.rs index 033fab55db..064c051566 100644 --- a/diskann-benchmark-runner/src/test/gated.rs +++ b/diskann-benchmark-runner/src/test/gated.rs @@ -16,7 +16,8 @@ use std::io::Write; use serde::{Deserialize, Serialize}; use crate::{ - Benchmark, Checker, Checkpoint, Input, Output, benchmark::MatchContext, benchmark::Score, + Benchmark, Checker, Checkpoint, Input, Output, Reflect, benchmark::MatchContext, + benchmark::Score, }; use super::{dim::DimInput, typed::TypeInput}; @@ -85,7 +86,8 @@ impl Benchmark for AnotherGatedBench { // Partially Gated with Input Always Registered // ////////////////////////////////////////////////// -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::")] pub(super) struct SampleInput { value: String, } @@ -145,7 +147,8 @@ impl Benchmark for GatedWithIndependentInput { // The input backing the fully-gated benchmark. This is only compiled and registered when the // controlling feature set is enabled, standing in for an input whose validation would otherwise // pull in a heavy optional dependency. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::")] pub(super) struct PhantomInput { value: usize, } diff --git a/diskann-benchmark-runner/src/test/typed.rs b/diskann-benchmark-runner/src/test/typed.rs index 3ab0e7a891..e48960cc69 100644 --- a/diskann-benchmark-runner/src/test/typed.rs +++ b/diskann-benchmark-runner/src/test/typed.rs @@ -8,7 +8,7 @@ use std::io::Write; use serde::{Deserialize, Serialize}; use crate::{ - Benchmark, Checker, Checkpoint, Input, Output, + Benchmark, Checker, Checkpoint, Input, Output, Reflect, benchmark::{MatchContext, PassFail, Regression, Score}, utils::datatype::{AsDataType, DataType}, }; @@ -24,11 +24,13 @@ pub(crate) struct TypeInput { error_when_checked: bool, } -#[derive(Serialize, Deserialize)] +/// This is a test input for testing corner cases in the benchmark runner. +#[derive(Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::")] pub(crate) struct TypeInputRaw { data_type: DataType, dim: usize, - // Should we return an error when deserializing? + /// Should we return an error when deserializing? error_when_checked: bool, } @@ -78,9 +80,10 @@ impl Input for TypeInput { // Tolerance // /////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "benchmark::test::typed::")] pub(super) struct Tolerance { - // Should we return an error when `from_raw` is called? + /// Should we return an error when `from_raw` is called? pub(super) error_when_checked: bool, } diff --git a/diskann-benchmark-runner/src/utils/datatype.rs b/diskann-benchmark-runner/src/utils/datatype.rs index 470251af01..bce0d70288 100644 --- a/diskann-benchmark-runner/src/utils/datatype.rs +++ b/diskann-benchmark-runner/src/utils/datatype.rs @@ -6,11 +6,14 @@ use half::f16; use serde::{Deserialize, Serialize}; +use crate::Reflect; + /// An enum representation for common DiskANN data types. /// /// See also: [`AsDataType`]. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "lowercase")] +#[reflect(prefix = "benchmark::")] pub enum DataType { Float64, Float32, diff --git a/diskann-benchmark-runner/src/utils/mod.rs b/diskann-benchmark-runner/src/utils/mod.rs index 9f48f9413e..24caccbd92 100644 --- a/diskann-benchmark-runner/src/utils/mod.rs +++ b/diskann-benchmark-runner/src/utils/mod.rs @@ -8,5 +8,7 @@ pub mod fmt; pub mod microseconds; pub mod num; pub mod percentiles; +mod required; pub use microseconds::MicroSeconds; +pub use required::RequiredOption; diff --git a/diskann-benchmark-runner/src/utils/num.rs b/diskann-benchmark-runner/src/utils/num.rs index 049c2800f9..a2766a60c2 100644 --- a/diskann-benchmark-runner/src/utils/num.rs +++ b/diskann-benchmark-runner/src/utils/num.rs @@ -8,6 +8,8 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; use thiserror::Error; +use crate::Reflect; + /// Compute the relative change from `before` to `after`. /// /// This helper is intentionally opinionated for benchmark-style metrics: @@ -54,7 +56,8 @@ pub enum RelativeChangeError { } /// A finite floating-point value that is greater than or equal to zero. -#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)] +#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Reflect)] +#[reflect(prefix = "benchmark::")] pub struct NonNegativeFinite(f64); impl NonNegativeFinite { diff --git a/diskann-benchmark-runner/src/utils/required.rs b/diskann-benchmark-runner/src/utils/required.rs new file mode 100644 index 0000000000..39198cf4d9 --- /dev/null +++ b/diskann-benchmark-runner/src/utils/required.rs @@ -0,0 +1,159 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use serde::{Deserialize, Serialize}; + +use crate::{Reflect, reflect::tree}; + +/// Like `Option`, but requires the containing field to be present in the input JSON. +/// +/// To represent [`None`], the field must be explicitly set to `null`. +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +pub struct RequiredOption(Option); + +impl RequiredOption { + /// Wrap an optional value. + pub fn new(opt: Option) -> Self { + Self(opt) + } + + /// Wrap a present value. + pub fn some(value: T) -> Self { + Self::new(Some(value)) + } + + /// Construct an absent value. + pub fn none() -> Self { + Self::new(None) + } + + /// See: [`Option::unwrap_or`]. + pub fn unwrap_or(self, default: T) -> T { + self.0.unwrap_or(default) + } + + /// See: [`Option::get_or_insert`]. + pub fn get_or_insert(&mut self, v: T) -> &mut T { + self.0.get_or_insert(v) + } + + /// Return the wrapped optional value. + pub fn into_inner(self) -> Option { + self.0 + } + + /// See: [`Option::as_ref`]. + pub fn as_ref(&self) -> Option<&T> { + self.0.as_ref() + } + + /// See: [`Option::as_mut`]. + pub fn as_mut(&mut self) -> Option<&mut T> { + self.0.as_mut() + } + + /// See: [`Option::as_deref`]. + pub fn as_deref(&self) -> Option<&T::Target> + where + T: std::ops::Deref, + { + self.0.as_deref() + } + + /// See: [`Option::as_deref_mut`]. + pub fn as_deref_mut(&mut self) -> Option<&mut T::Target> + where + T: std::ops::DerefMut, + { + self.0.as_deref_mut() + } +} + +impl Reflect for RequiredOption +where + T: Reflect, +{ + fn ty() -> tree::Type { + let doc = "A required optional type.\n\n\ + Unlike `Option`, where an omitted field implies `None`, the field must be \ + present and explicitly set to `null` to represent `None`."; + + // We're special - we get to instantiate the `Optional` type. + tree::Type::optional::(Some(doc.into())) + } + + fn format_type_name(f: &mut dyn std::fmt::Write) -> std::fmt::Result { + f.write_str("RequiredOption<")?; + T::format_type_name(f)?; + f.write_str(">") + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[derive(Debug, PartialEq, Serialize, Deserialize)] + struct Input { + value: RequiredOption, + } + + #[test] + fn rejects_missing_field() { + let error = serde_json::from_value::(json!({})).unwrap_err(); + + assert!(error.to_string().contains("missing field `value`")); + } + + #[test] + fn accepts_null_field() { + let input: Input = serde_json::from_value(json!({ "value": null })).unwrap(); + + assert_eq!(input.value, RequiredOption::none()); + } + + #[test] + fn accepts_present_value() { + let input: Input = serde_json::from_value(json!({ "value": 42 })).unwrap(); + + assert_eq!(input.value, RequiredOption::some(42)); + } + + #[test] + fn serializes_transparently() { + let none = Input { + value: RequiredOption::none(), + }; + let some = Input { + value: RequiredOption::some(42), + }; + + assert_eq!( + serde_json::to_value(none).unwrap(), + json!({ "value": null }) + ); + assert_eq!(serde_json::to_value(some).unwrap(), json!({ "value": 42 })); + } + + #[test] + fn methods_smoke_test() { + let value = RequiredOption::new(Some(1)); + assert_eq!(value.into_inner(), Some(1)); + assert_eq!(RequiredOption::some(2).unwrap_or(3), 2); + + let mut value = RequiredOption::none(); + assert_eq!(value.as_ref(), None); + assert_eq!(value.get_or_insert(String::from("hello")), "hello"); + assert_eq!(value.as_deref(), Some("hello")); + value.as_mut().unwrap().push('!'); + value.as_deref_mut().unwrap().make_ascii_uppercase(); + assert_eq!(value.into_inner(), Some(String::from("HELLO!"))); + } +} diff --git a/diskann-benchmark-runner/src/ux.rs b/diskann-benchmark-runner/src/ux.rs index 9f862d5619..7f761f5578 100644 --- a/diskann-benchmark-runner/src/ux.rs +++ b/diskann-benchmark-runner/src/ux.rs @@ -90,3 +90,52 @@ pub fn strip_backtrace(s: String) -> String { lines.join("\n") } + +//--------------// +// Crate Shared // +//--------------// + +#[cfg(test)] +pub(crate) mod test { + use super::*; + + use std::path::Path; + + // Read the entire contents of a file to a string. + pub(crate) fn read_to_string>(path: P, ctx: &str) -> String { + match std::fs::read_to_string(path.as_ref()) { + Ok(s) => normalize(s), + Err(err) => panic!( + "failed to read {} {:?} with error: {}", + ctx, + path.as_ref(), + err + ), + } + } + + const ENV: &str = "DISKANN_TEST"; + + // Check if `DISKANN_TEST=overwrite` is configured. Return `true` if so - otherwise + // return `false`. + // + // If `DISKANN_TEST` is set but its value is not `overwrite` - panic. + pub(crate) fn overwrite() -> bool { + match std::env::var(ENV) { + Ok(v) => { + if v == "overwrite" { + true + } else { + panic!( + "Unknown value for {}: \"{}\". Expected \"overwrite\"", + ENV, v + ); + } + } + Err(std::env::VarError::NotPresent) => false, + Err(std::env::VarError::NotUnicode(_)) => { + panic!("Value for {} is not unicode", ENV); + } + } + } +} diff --git a/diskann-benchmark-runner/tests/benchmark/test-2/stdout.txt b/diskann-benchmark-runner/tests/benchmark/test-2/stdout.txt index 0c2c55cb3e..0440dc40af 100644 --- a/diskann-benchmark-runner/tests/benchmark/test-2/stdout.txt +++ b/diskann-benchmark-runner/tests/benchmark/test-2/stdout.txt @@ -1,7 +1,16 @@ The example JSON representation for "test-input-dim" is: + { "type": "test-input-dim", "content": { "dim": 128 } -} \ No newline at end of file +} + +Type Information: + +benchmark::test::DimInput + "dim": Option + May be `null`. + +More type information available using `type-info` \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/benchmark/test-3/stdout.txt b/diskann-benchmark-runner/tests/benchmark/test-3/stdout.txt index 06d12d7086..dcf65baba1 100644 --- a/diskann-benchmark-runner/tests/benchmark/test-3/stdout.txt +++ b/diskann-benchmark-runner/tests/benchmark/test-3/stdout.txt @@ -1,4 +1,5 @@ The example JSON representation for "test-input-types" is: + { "type": "test-input-types", "content": { @@ -6,4 +7,33 @@ The example JSON representation for "test-input-types" is: "dim": 128, "error_when_checked": false } -} \ No newline at end of file +} + +Type Information: + +benchmark::test::TypeInputRaw + This is a test input for testing corner cases in the benchmark runner. + + "data_type": benchmark::DataType + Representation: externally tagged + + Options: + "float64" + "float32" + "float16" + "uint8" + "uint16" + "uint32" + "uint64" + "int8" + "int16" + "int32" + "int64" + "bool" + + "dim": usize + + "error_when_checked": bool + Should we return an error when deserializing? + +More type information available using `type-info` \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdin.txt b/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdin.txt new file mode 100644 index 0000000000..8e5adac4a3 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdin.txt @@ -0,0 +1 @@ +type-info benchmark::test::TypeInputRaw diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdout.txt b/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdout.txt new file mode 100644 index 0000000000..34b5423302 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-describe/stdout.txt @@ -0,0 +1,24 @@ +benchmark::test::TypeInputRaw + This is a test input for testing corner cases in the benchmark runner. + + "data_type": benchmark::DataType + Representation: externally tagged + + Options: + "float64" + "float32" + "float16" + "uint8" + "uint16" + "uint32" + "uint64" + "int8" + "int16" + "int32" + "int64" + "bool" + + "dim": usize + + "error_when_checked": bool + Should we return an error when deserializing? \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/features.txt b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/features.txt new file mode 100644 index 0000000000..da13a750a7 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/features.txt @@ -0,0 +1,4 @@ +gated-feature-0 +gated-feature-1 +gated-feature-2 +gated-feature-3 diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdin.txt b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdin.txt new file mode 100644 index 0000000000..3ae7b5f6ff --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdin.txt @@ -0,0 +1 @@ +type-info diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdout.txt b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdout.txt new file mode 100644 index 0000000000..8908687320 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-list-all-features/stdout.txt @@ -0,0 +1,12 @@ +All registered types: + Option + benchmark::DataType + benchmark::test::DimInput + benchmark::test::PhantomInput + benchmark::test::SampleInput + benchmark::test::TypeInputRaw + benchmark::test::dim::Tolerance + benchmark::test::typed::Tolerance + bool + string + usize \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-list/stdin.txt b/diskann-benchmark-runner/tests/benchmark/type-info-list/stdin.txt new file mode 100644 index 0000000000..3ae7b5f6ff --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-list/stdin.txt @@ -0,0 +1 @@ +type-info diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-list/stdout.txt b/diskann-benchmark-runner/tests/benchmark/type-info-list/stdout.txt new file mode 100644 index 0000000000..d405fb9879 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-list/stdout.txt @@ -0,0 +1,11 @@ +All registered types: + Option + benchmark::DataType + benchmark::test::DimInput + benchmark::test::SampleInput + benchmark::test::TypeInputRaw + benchmark::test::dim::Tolerance + benchmark::test::typed::Tolerance + bool + string + usize \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdin.txt b/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdin.txt new file mode 100644 index 0000000000..a17e298506 --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdin.txt @@ -0,0 +1 @@ +type-info benchmark::test::Missing diff --git a/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdout.txt b/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdout.txt new file mode 100644 index 0000000000..3ae52a773c --- /dev/null +++ b/diskann-benchmark-runner/tests/benchmark/type-info-missing/stdout.txt @@ -0,0 +1 @@ +No type information for "benchmark::test::Missing" \ No newline at end of file diff --git a/diskann-benchmark-runner/tests/rendered_reflections.txt b/diskann-benchmark-runner/tests/rendered_reflections.txt new file mode 100644 index 0000000000..f64dbde9eb --- /dev/null +++ b/diskann-benchmark-runner/tests/rendered_reflections.txt @@ -0,0 +1,259 @@ +======== +{ + "name": "usize", + "description": "a simple test" +} +-------- +usize (json number) + A system dependent unsigned integer +======== +======== +{ + "name": "source", + "description": "a configurable source enum" +} +-------- +render::Source + Representation: externally tagged + + Options: + "build" + Build from scratch. + + "data_type": render::SimpleDataType + The type of the input data. + + Representation: externally tagged + + Options: + "float32" + "float16" + + "file": string + Input data in the `.bin` binary format. + + "output": Option + Output file. + + If provided, saved data will go here. + + May be `null`. + + "from_previous" + Run from a previously generated output. + + "data_type": render::AnnotatedDataType + Representation: externally tagged + + Options: + "float32" + Use high-precision. + + "float16" + Use lower precision. + + "int8" + + "file": string + The previously generated output. +======== +======== +{ + "name": "source wrapper", + "description": "render a newtype wrapper" +} +-------- +SourceWrapper + A new-type wrapper around `Source`. + + render::Source + Representation: externally tagged + + Options: + "build" + Build from scratch. + + "data_type": render::SimpleDataType + The type of the input data. + + Representation: externally tagged + + Options: + "float32" + "float16" + + "file": string + Input data in the `.bin` binary format. + + "output": Option + Output file. + + If provided, saved data will go here. + + May be `null`. + + "from_previous" + Run from a previously generated output. + + "data_type": render::AnnotatedDataType + Representation: externally tagged + + Options: + "float32" + Use high-precision. + + "float16" + Use lower precision. + + "int8" + + "file": string + The previously generated output. +======== +======== +{ + "name": "id", + "description": "render a primitive newtype wrapper" +} +-------- +Id + usize +======== +======== +{ + "name": "config", + "description": "a sample struct config." +} +-------- +Config + A top level config. + + "source": SourceWrapper + render::Source + Representation: externally tagged + + Options: + "build" + Build from scratch. + + "data_type": render::SimpleDataType + The type of the input data. + + "file": string + Input data in the `.bin` binary format. + + "output": Option + Output file. + + If provided, saved data will go here. + + "from_previous" + Run from a previously generated output. + + "data_type": render::AnnotatedDataType + + "file": string + The previously generated output. + + "param1": string + This does one thing. + + "param2": Vec + This does another. + + Elements: render::AnnotatedDataType + Representation: externally tagged + + Options: + "float32" + Use high-precision. + + "float16" + Use lower precision. + + "int8" +======== +======== +{ + "name": "tuples", + "description": "a sample tuple-like struct." +} +-------- +Tuples + A tuple like type! + + 0: usize + Some docs, + + 1: string +======== +======== +{ + "name": "internally-tagged", + "description": "an internally tagged enum" +} +-------- +InternallyTagged + Discriminant field: "internally-tagged" + + Options: + "Config" + Config level docs. + + Config + Field docs. + + "source": SourceWrapper + render::Source + + "param1": string + This does one thing. + + "param2": Vec + This does another. + + + "Other" + "foo": usize + + "bar": string + Hello world + + "Unit" +======== +======== +{ + "name": "adjacenctly-tagged", + "description": "an adjacently tagged enum" +} +-------- +AdjacentlyTagged + An adjacently tagged enum. + + Discriminant field: "tag" + Content field: "content" + + Options: + "Config" + Config + "source": SourceWrapper + render::Source + + "param1": string + This does one thing. + + "param2": Vec + This does another. + + + "Other" + "foo": usize + + "bar": string + + "Unit" + "Tuple" + 0: usize + + 1: usize +======== diff --git a/diskann-benchmark-simd/src/lib.rs b/diskann-benchmark-simd/src/lib.rs index 353c8504c5..29057addeb 100644 --- a/diskann-benchmark-simd/src/lib.rs +++ b/diskann-benchmark-simd/src/lib.rs @@ -26,7 +26,7 @@ use diskann_benchmark_runner::{ num::{relative_change, NonNegativeFinite}, percentiles, MicroSeconds, }, - Benchmark, Checker, Input, Registry, + Benchmark, Checker, Input, Reflect, Registry, }; //////////////// @@ -55,8 +55,9 @@ impl std::ops::Deref for DisplayWrapper<'_, T> { // Inputs // //////////// -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "simd::")] pub enum SimilarityMeasure { SquaredL2, InnerProduct, @@ -74,17 +75,41 @@ impl std::fmt::Display for SimilarityMeasure { } } -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +/// DiskANN provides specialization for different micro-architectures, allowing binaries +/// compiled for an older machine to still execute accelerated kernels when the runtime +/// CPU allows. +/// +/// This enum selects the target micro-architecture's implementation to benchmark. +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] +#[reflect(prefix = "simd::")] enum Arch { + /// Target AVX-512 with additional VNNI and population-count instructions. + /// + /// This is generally IceLake server or better. + /// + /// Only usable when compiling for x86-64. #[serde(rename = "x86-64-v4")] #[expect(non_camel_case_types)] X86_64_V4, + + /// Target AVX2. + /// + /// Only usable when compiling for x86-64. #[serde(rename = "x86-64-v3")] #[expect(non_camel_case_types)] X86_64_V3, + /// Target the Aarch64 Neon instruction set. + /// + /// Only usable when compiling for aarch64. Neon, + + /// Use the `diskann-vector` scalar fallbacks. + /// + /// These use a mixture of actual scalar code and auto-vectorization friendly loops. Scalar, + + /// Reference operations using purely scalar loops. Reference, } @@ -101,20 +126,39 @@ impl std::fmt::Display for Arch { } } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +/// Parameters controlling the benchmark for an individual kernel. +/// +/// Kernels work by using a fixed `query` vector and a variable of `data` vectors. +/// Internal timers measure how long it takes to compute all distances from `query` to each +/// data vector. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "simd::")] struct Run { + /// The kernel type to test. distance: SimilarityMeasure, + /// The number of dimensions in each vector. dim: NonZeroUsize, + /// The number of data points to loop over. num_points: NonZeroUsize, + /// The number of loops to run between timing samples. loops_per_measurement: NonZeroUsize, + /// The total number of measurements to take. num_measurements: NonZeroUsize, } -#[derive(Debug, Serialize, Deserialize)] +/// A full-precision SIMD Accelerated kernel. +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "simd::")] pub struct SimdOp { + /// The data type for the left-hand side. + /// + /// This is typically the query in KNN scenarios. query_type: DataType, + /// The data type for the right-hand side. data_type: DataType, + /// The micro-architecture specialization. arch: Arch, + /// Kernel configurations. runs: Vec, } @@ -205,7 +249,8 @@ impl Input for SimdOp { /// /// Each field specifies the maximum allowed relative increase in the corresponding metric. /// For example, a value of `0.10` means a 10% increase is tolerated. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "simd::")] struct SimdTolerance { min_time_regression: NonNegativeFinite, } diff --git a/diskann-benchmark/example/graph-index-bftree-stream.json b/diskann-benchmark/example/graph-index-bftree-stream.json index 2fbfd2e165..9d275f936e 100644 --- a/diskann-benchmark/example/graph-index-bftree-stream.json +++ b/diskann-benchmark/example/graph-index-bftree-stream.json @@ -47,7 +47,8 @@ "l_value": 40 } }, - "ip_delete_num_to_replace": 3 + "ip_delete_num_to_replace": 3, + "consolidate_threshold": 0.2 }, "vector_store_config": { "cb_size_byte": 67108864, diff --git a/diskann-benchmark/src/disk_index/benchmarks.rs b/diskann-benchmark/src/disk_index/benchmarks.rs index 1acc56abc9..7f56e0c778 100644 --- a/diskann-benchmark/src/disk_index/benchmarks.rs +++ b/diskann-benchmark/src/disk_index/benchmarks.rs @@ -15,7 +15,7 @@ use diskann_benchmark_runner::{ fmt::Table, num::{relative_change, NonNegativeFinite}, }, - Benchmark, Checker, Checkpoint, Input, Registry, + Benchmark, Checker, Checkpoint, Input, Reflect, Registry, }; use diskann_providers::storage::FileStorageProvider; use half::f16; @@ -130,7 +130,8 @@ pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> /// or decrease (for "higher is better" metrics) before a regression is flagged. /// /// For example, `recall_regression: 0.01` means recall must not drop by more than 1%. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "disk::")] pub(super) struct DiskIndexTolerance { /// Max allowed relative increase in build time (e.g., 0.10 = 10%). build_time_regression: NonNegativeFinite, diff --git a/diskann-benchmark/src/index/benchmarks.rs b/diskann-benchmark/src/index/benchmarks.rs index 3aad4d8275..4db862daf8 100644 --- a/diskann-benchmark/src/index/benchmarks.rs +++ b/diskann-benchmark/src/index/benchmarks.rs @@ -226,7 +226,7 @@ where build::set_start_points( index.provider(), data.as_view(), - *build.start_point_strategy(), + build.start_point_strategy(), )?; Ok(index) }, @@ -304,14 +304,7 @@ where ) -> anyhow::Result>> { writeln!(output, "{}", input)?; - let groundtruth_directory = input - .runbook_params - .resolved_gt_directory - .as_ref() - .ok_or_else(|| { - anyhow::anyhow!("Ground truth directory path was not resolved during validation") - })?; - + let groundtruth_directory = &input.runbook_params.gt_directory; let mut runbook = bigann::RunBook::load( &input.runbook_params.runbook_path, &input.runbook_params.dataset_name, @@ -843,9 +836,10 @@ where { let topk = input.search_phase.as_topk()?; - let consolidate_threshold: f32 = input + let consolidate_threshold: f32 = *input .runbook_params .consolidate_threshold + .as_ref() .ok_or_else(|| anyhow::anyhow!("consolidate_threshold is required for inmem streaming"))?; let data = datafiles::load_dataset::(datafiles::BinFile(input.build.data()))?; @@ -865,7 +859,7 @@ where build::set_start_points( index.provider(), data.as_view(), - *input.build.start_point_strategy(), + input.build.start_point_strategy(), )?; let num_threads_and_tasks = NonZeroUsize::new(input.build.num_threads()).unwrap(); diff --git a/diskann-benchmark/src/index/bftree/streaming_utils.rs b/diskann-benchmark/src/index/bftree/streaming_utils.rs index d505eeae76..20ba612e5c 100644 --- a/diskann-benchmark/src/index/bftree/streaming_utils.rs +++ b/diskann-benchmark/src/index/bftree/streaming_utils.rs @@ -28,13 +28,7 @@ where T: 'static, F: FnOnce(usize) -> anyhow::Result>>, { - let groundtruth_directory = runbook_params - .resolved_gt_directory - .as_ref() - .ok_or_else(|| { - anyhow::anyhow!("Ground truth directory path was not resolved during validation") - })?; - + let groundtruth_directory = &runbook_params.gt_directory; let mut runbook = bigann::RunBook::load( &runbook_params.runbook_path, &runbook_params.dataset_name, diff --git a/diskann-benchmark/src/index/inmem/product.rs b/diskann-benchmark/src/index/inmem/product.rs index 5956d8ebe3..305df3e496 100644 --- a/diskann-benchmark/src/index/inmem/product.rs +++ b/diskann-benchmark/src/index/inmem/product.rs @@ -202,7 +202,7 @@ mod imp { build::set_start_points( index.provider(), data_view, - *build.start_point_strategy(), + build.start_point_strategy(), )?; Ok(index) }; diff --git a/diskann-benchmark/src/index/inmem/scalar.rs b/diskann-benchmark/src/index/inmem/scalar.rs index da0c8a3c17..2a20268c67 100644 --- a/diskann-benchmark/src/index/inmem/scalar.rs +++ b/diskann-benchmark/src/index/inmem/scalar.rs @@ -258,7 +258,7 @@ mod imp { inmem::WithBits::<$N>::new(quantizer), common::NoDeletes, )?; - build::set_start_points(index.provider(), data_view, *build.start_point_strategy())?; + build::set_start_points(index.provider(), data_view, build.start_point_strategy())?; Ok(index) }; diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index 0d574b16d8..63259a8df0 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -26,6 +26,7 @@ use diskann_benchmark_runner::{ utils::{ datatype::{AsDataType, DataType}, fmt::{Delimit, KeyValue, Quote}, + RequiredOption, }, Benchmark, Checker, Checkpoint, Input, Registry, }; @@ -69,14 +70,18 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> mod dto { use super::*; - #[derive(Debug, Serialize, Deserialize)] + use diskann_benchmark_runner::Reflect; + + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct KnnSweep { pub(super) search_n: usize, pub(super) search_l: Vec, pub(super) recall_k: usize, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct KnnSearch { pub(super) queries: InputFile, pub(super) groundtruth: InputFile, @@ -85,14 +90,16 @@ mod dto { pub(super) runs: Vec, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct Data { pub(super) data_type: DataType, pub(super) data: InputFile, pub(super) distance: SimilarityMeasure, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct BuildParams { pub(super) pruned_degree: usize, pub(super) max_degree: usize, @@ -105,13 +112,13 @@ mod dto { // Quantization // //--------------// - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] pub(super) enum Quantization { Spherical(Spherical), } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] pub(super) enum SphericalBits { One, @@ -119,7 +126,7 @@ mod dto { Four, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] pub(super) struct Spherical { pub(super) bits: SphericalBits, } @@ -128,7 +135,8 @@ mod dto { // Streaming // //-----------// - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct StreamingKnnSearch { pub(super) queries: InputFile, pub(super) reps: NonZeroUsize, @@ -136,7 +144,8 @@ mod dto { pub(super) runs: Vec, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct RunBook { pub(super) path: InputFile, pub(super) dataset: String, @@ -149,15 +158,17 @@ mod dto { // Top Level Inputs // //------------------// - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct StaticBuild { pub(super) data: Data, pub(super) build: BuildParams, pub(super) search: KnnSearch, - pub(super) quantization: Option, + pub(super) quantization: RequiredOption, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "inmem::")] pub(super) struct BigANNStreaming { pub(super) data: Data, pub(super) build: BuildParams, @@ -507,7 +518,7 @@ impl StaticBuild { let data = Data::from_raw(data, checker.as_deref_mut())?; let build = BuildParams::from_raw(build, data.distance)?; let search = KnnSearch::from_raw(search, checker)?; - let quantization = Quantization::from_raw(quantization); + let quantization = Quantization::from_raw(quantization.into_inner()); Ok(Self { data, @@ -600,7 +611,7 @@ impl Input for StaticBuild { recall_k: 10, }], }, - quantization: None, + quantization: RequiredOption::none(), } } } diff --git a/diskann-benchmark/src/inputs/bftree.rs b/diskann-benchmark/src/inputs/bftree.rs index 32cc7adb16..b214a6b68d 100644 --- a/diskann-benchmark/src/inputs/bftree.rs +++ b/diskann-benchmark/src/inputs/bftree.rs @@ -10,7 +10,10 @@ use crate::inputs::{ write_field, Example, PRINT_WIDTH, }; use diskann::graph::config; -use diskann_benchmark_runner::{utils::datatype::DataType, Checker}; +use diskann_benchmark_runner::{ + utils::{datatype::DataType, RequiredOption}, + Checker, Reflect, +}; use diskann_bftree::BfTreeProviderParameters; use serde::{Deserialize, Serialize}; @@ -20,7 +23,8 @@ use serde::{Deserialize, Serialize}; /// /// Required fields control memory sizing before data spills to disk. /// Optional fields tune internal behavior and default to bf_tree's defaults. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "bftree::")] pub(crate) struct BfTreeStoreConfig { /// Size of the circular buffer (in-memory write cache) in bytes. pub(crate) cb_size_byte: usize, @@ -29,32 +33,25 @@ pub(crate) struct BfTreeStoreConfig { pub(crate) leaf_page_size: usize, /// Maximum record size that can be stored in the circular buffer. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) cb_max_record_size: Option, + pub(crate) cb_max_record_size: RequiredOption, /// Minimum record size for the circular buffer. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) cb_min_record_size: Option, + pub(crate) cb_min_record_size: RequiredOption, /// Probability (0-100) of promoting a read record to the front of the buffer. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) read_promotion_rate: Option, + pub(crate) read_promotion_rate: RequiredOption, /// Probability (0-100) of promoting a scanned record to the front of the buffer. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) scan_promotion_rate: Option, + pub(crate) scan_promotion_rate: RequiredOption, /// Ratio of buffer used before copy-on-access kicks in. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) cb_copy_on_access_ratio: Option, + pub(crate) cb_copy_on_access_ratio: RequiredOption, /// Whether to cache full pages on read. - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) read_record_cache: Option, + pub(crate) read_record_cache: RequiredOption, /// If true, only use the in-memory circular buffer (no disk pages). - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) cache_only: Option, + pub(crate) cache_only: RequiredOption, /// Whether to enable CPR snapshot support for this store. pub(crate) use_snapshot: bool, @@ -86,26 +83,26 @@ impl BfTreeStoreConfig { let mut c = bf_tree::Config::default(); c.cb_size_byte(self.cb_size_byte); c.leaf_page_size(self.leaf_page_size); - if let Some(v) = self.cb_max_record_size { - c.cb_max_record_size(v); + if let Some(v) = self.cb_max_record_size.as_ref() { + c.cb_max_record_size(*v); } - if let Some(v) = self.cb_min_record_size { - c.cb_min_record_size(v); + if let Some(v) = self.cb_min_record_size.as_ref() { + c.cb_min_record_size(*v); } - if let Some(v) = self.read_promotion_rate { - c.read_promotion_rate(v); + if let Some(v) = self.read_promotion_rate.as_ref() { + c.read_promotion_rate(*v); } - if let Some(v) = self.scan_promotion_rate { - c.scan_promotion_rate(v); + if let Some(v) = self.scan_promotion_rate.as_ref() { + c.scan_promotion_rate(*v); } - if let Some(v) = self.cb_copy_on_access_ratio { - c.cb_copy_on_access_ratio(v); + if let Some(v) = self.cb_copy_on_access_ratio.as_ref() { + c.cb_copy_on_access_ratio(*v); } - if let Some(v) = self.read_record_cache { - c.read_record_cache(v); + if let Some(v) = self.read_record_cache.as_ref() { + c.read_record_cache(*v); } - if let Some(v) = self.cache_only { - c.cache_only(v); + if let Some(v) = self.cache_only.as_ref() { + c.cache_only(*v); } c.use_snapshot(self.use_snapshot); c @@ -139,13 +136,13 @@ impl Default for BfTreeStoreConfig { Self { cb_size_byte: 32 * 1024 * 1024, // 32MB leaf_page_size: 4096, - cb_max_record_size: None, - cb_min_record_size: None, - read_promotion_rate: None, - scan_promotion_rate: None, - cb_copy_on_access_ratio: None, - read_record_cache: None, - cache_only: None, + cb_max_record_size: RequiredOption::none(), + cb_min_record_size: RequiredOption::none(), + read_promotion_rate: RequiredOption::none(), + scan_promotion_rate: RequiredOption::none(), + cb_copy_on_access_ratio: RequiredOption::none(), + read_record_cache: RequiredOption::none(), + cache_only: RequiredOption::none(), use_snapshot: false, } } @@ -161,25 +158,25 @@ impl std::fmt::Display for BfTreeStoreConfig { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write_field!(f, "cb_size_byte", self.cb_size_byte)?; write_field!(f, "leaf_page_size", self.leaf_page_size)?; - if let Some(v) = self.cb_max_record_size { + if let Some(v) = self.cb_max_record_size.as_ref() { write_field!(f, "cb_max_record_size", v)?; } - if let Some(v) = self.cb_min_record_size { + if let Some(v) = self.cb_min_record_size.as_ref() { write_field!(f, "cb_min_record_size", v)?; } - if let Some(v) = self.read_promotion_rate { + if let Some(v) = self.read_promotion_rate.as_ref() { write_field!(f, "read_promotion_rate", v)?; } - if let Some(v) = self.scan_promotion_rate { + if let Some(v) = self.scan_promotion_rate.as_ref() { write_field!(f, "scan_promotion_rate", v)?; } - if let Some(v) = self.cb_copy_on_access_ratio { + if let Some(v) = self.cb_copy_on_access_ratio.as_ref() { write_field!(f, "cb_copy_on_access_ratio", v)?; } - if let Some(v) = self.read_record_cache { + if let Some(v) = self.read_record_cache.as_ref() { write_field!(f, "read_record_cache", v)?; } - if let Some(v) = self.cache_only { + if let Some(v) = self.cache_only.as_ref() { write_field!(f, "cache_only", v)?; } Ok(()) @@ -190,7 +187,7 @@ impl std::fmt::Display for BfTreeStoreConfig { /// /// Returns the agreed-upon value. If no configs are present or none set /// `use_snapshot`, returns `false`. If configs disagree, returns an error. -fn reconcile_use_snapshot(configs: &[(&str, &Option)]) -> anyhow::Result { +fn reconcile_use_snapshot(configs: &[(&str, Option<&BfTreeStoreConfig>)]) -> anyhow::Result { let mut resolved: Option = None; for (name, config) in configs { @@ -218,9 +215,9 @@ fn bftree_parameters_from( build: &IndexBuild, num_points: usize, dim: usize, - vector_store_config: &Option, - neighbor_store_config: &Option, - quant_store_config: &Option, + vector_store_config: Option<&BfTreeStoreConfig>, + neighbor_store_config: Option<&BfTreeStoreConfig>, + quant_store_config: Option<&BfTreeStoreConfig>, ) -> anyhow::Result { let use_snapshot = reconcile_use_snapshot(&[ ("vector_store_config", vector_store_config), @@ -236,14 +233,17 @@ fn bftree_parameters_from( dim, metric: build.distance().into(), vector_provider_config: vector_store_config - .clone() + .cloned() .unwrap_or_default() .into_config(), neighbor_list_provider_config: neighbor_store_config - .clone() + .cloned() + .unwrap_or_default() + .into_config(), + quant_vector_provider_config: quant_store_config + .cloned() .unwrap_or_default() .into_config(), - quant_vector_provider_config: quant_store_config.clone().unwrap_or_default().into_config(), graph_params: None, use_snapshot, }) @@ -251,14 +251,13 @@ fn bftree_parameters_from( as_input!(BfTreeFullPrecisionBuild); -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "bftree::")] pub(crate) struct BfTreeFullPrecisionBuild { build: IndexBuild, search_phase: SearchPhase, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, + vector_store_config: RequiredOption, + neighbor_store_config: RequiredOption, } impl BfTreeFullPrecisionBuild { @@ -291,19 +290,19 @@ impl BfTreeFullPrecisionBuild { &self.build, num_points, dim, - &self.vector_store_config, - &self.neighbor_store_config, - &None, + self.vector_store_config.as_ref(), + self.neighbor_store_config.as_ref(), + None, ) } pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.build.validate(checker)?; self.search_phase.validate(checker)?; - if let Some(cfg) = &mut self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } @@ -318,8 +317,8 @@ impl Example for BfTreeFullPrecisionBuild { Self { build, search_phase: SearchPhase::Topk(TopkSearchPhase::example()), - vector_store_config: None, - neighbor_store_config: None, + vector_store_config: RequiredOption::none(), + neighbor_store_config: RequiredOption::none(), } } } @@ -335,11 +334,11 @@ impl std::fmt::Display for BfTreeFullPrecisionBuild { writeln!(f)?; self.build.summarize_fields(f)?; - if let Some(ref cfg) = self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_ref() { writeln!(f, "\n Vector Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_ref() { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } @@ -350,15 +349,14 @@ impl std::fmt::Display for BfTreeFullPrecisionBuild { as_input!(BfTreeDynamicRun); -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "bftree::")] pub(crate) struct BfTreeDynamicRun { build: IndexBuild, search_phase: SearchPhase, runbook_params: DynamicRunbookParams, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, + vector_store_config: RequiredOption, + neighbor_store_config: RequiredOption, } impl BfTreeDynamicRun { @@ -395,20 +393,20 @@ impl BfTreeDynamicRun { &self.build, num_points, dim, - &self.vector_store_config, - &self.neighbor_store_config, - &None, + self.vector_store_config.as_ref(), + self.neighbor_store_config.as_ref(), + None, ) } pub(crate) fn validate(&mut self, checker: &mut Checker) -> Result<(), anyhow::Error> { self.build.validate(checker)?; self.search_phase.validate(checker)?; self.runbook_params.validate(checker)?; - if let Some(cfg) = &mut self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } @@ -424,8 +422,8 @@ impl Example for BfTreeDynamicRun { build, search_phase: SearchPhase::Topk(TopkSearchPhase::example()), runbook_params: DynamicRunbookParams::example_immediate(), - vector_store_config: None, - neighbor_store_config: None, + vector_store_config: RequiredOption::none(), + neighbor_store_config: RequiredOption::none(), } } } @@ -441,11 +439,11 @@ impl std::fmt::Display for BfTreeDynamicRun { writeln!(f)?; self.build.summarize_fields(f)?; - if let Some(ref cfg) = self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_ref() { writeln!(f, "\n Vector Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_ref() { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } @@ -458,7 +456,8 @@ impl std::fmt::Display for BfTreeDynamicRun { as_input!(BfTreeSphericalBuild); -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "bftree::")] pub(crate) struct BfTreeSphericalBuild { build: IndexBuild, search_phase: SearchPhase, @@ -466,12 +465,9 @@ pub(crate) struct BfTreeSphericalBuild { transform_kind: exhaustive::TransformKind, num_bits: NonZeroUsize, pre_scale: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - quant_store_config: Option, + vector_store_config: RequiredOption, + neighbor_store_config: RequiredOption, + quant_store_config: RequiredOption, } impl BfTreeSphericalBuild { @@ -492,9 +488,9 @@ impl BfTreeSphericalBuild { &self.build, num_points, dim, - &self.vector_store_config, - &self.neighbor_store_config, - &self.quant_store_config, + self.vector_store_config.as_ref(), + self.neighbor_store_config.as_ref(), + self.quant_store_config.as_ref(), ) } @@ -529,15 +525,15 @@ impl BfTreeSphericalBuild { pub(crate) fn validate(&mut self, checker: &mut Checker) -> anyhow::Result<()> { self.build.validate(checker)?; self.search_phase.validate(checker)?; - if let Some(cfg) = &mut self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.quant_store_config { + if let Some(cfg) = self.quant_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } @@ -560,9 +556,9 @@ impl Example for BfTreeSphericalBuild { transform_kind: exhaustive::TransformKind::Null, num_bits: NonZeroUsize::new(1).unwrap(), pre_scale: None, - vector_store_config: None, - neighbor_store_config: None, - quant_store_config: None, + vector_store_config: RequiredOption::none(), + neighbor_store_config: RequiredOption::none(), + quant_store_config: RequiredOption::none(), } } } @@ -581,15 +577,15 @@ impl std::fmt::Display for BfTreeSphericalBuild { writeln!(f)?; self.build.summarize_fields(f)?; - if let Some(ref cfg) = self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_ref() { writeln!(f, "\n Vector Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_ref() { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.quant_store_config { + if let Some(cfg) = self.quant_store_config.as_ref() { writeln!(f, "\n Quant Store:")?; write!(f, "{}", cfg)?; } @@ -602,7 +598,8 @@ impl std::fmt::Display for BfTreeSphericalBuild { as_input!(BfTreeSphericalDynamicRun); -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "bftree::")] pub(crate) struct BfTreeSphericalDynamicRun { build: IndexBuild, search_phase: SearchPhase, @@ -611,12 +608,9 @@ pub(crate) struct BfTreeSphericalDynamicRun { transform_kind: exhaustive::TransformKind, num_bits: NonZeroUsize, pre_scale: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - vector_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - neighbor_store_config: Option, - #[serde(deserialize_with = "Deserialize::deserialize")] - quant_store_config: Option, + vector_store_config: RequiredOption, + neighbor_store_config: RequiredOption, + quant_store_config: RequiredOption, } impl BfTreeSphericalDynamicRun { @@ -669,9 +663,9 @@ impl BfTreeSphericalDynamicRun { &self.build, num_points, dim, - &self.vector_store_config, - &self.neighbor_store_config, - &self.quant_store_config, + self.vector_store_config.as_ref(), + self.neighbor_store_config.as_ref(), + self.quant_store_config.as_ref(), ) } @@ -679,15 +673,15 @@ impl BfTreeSphericalDynamicRun { self.build.validate(checker)?; self.search_phase.validate(checker)?; self.runbook_params.validate(checker)?; - if let Some(cfg) = &mut self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } - if let Some(cfg) = &mut self.quant_store_config { + if let Some(cfg) = self.quant_store_config.as_mut() { cfg.fill_defaults(); cfg.validate()?; } @@ -711,9 +705,9 @@ impl Example for BfTreeSphericalDynamicRun { transform_kind: exhaustive::TransformKind::Null, num_bits: NonZeroUsize::new(1).unwrap(), pre_scale: None, - vector_store_config: None, - neighbor_store_config: None, - quant_store_config: None, + vector_store_config: RequiredOption::none(), + neighbor_store_config: RequiredOption::none(), + quant_store_config: RequiredOption::none(), } } } @@ -732,15 +726,15 @@ impl std::fmt::Display for BfTreeSphericalDynamicRun { writeln!(f)?; self.build.summarize_fields(f)?; - if let Some(ref cfg) = self.vector_store_config { + if let Some(cfg) = self.vector_store_config.as_ref() { writeln!(f, "\n Vector Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.neighbor_store_config { + if let Some(cfg) = self.neighbor_store_config.as_ref() { writeln!(f, "\n Neighbor Store:")?; write!(f, "{}", cfg)?; } - if let Some(ref cfg) = self.quant_store_config { + if let Some(cfg) = self.quant_store_config.as_ref() { writeln!(f, "\n Quant Store:")?; write!(f, "{}", cfg)?; } diff --git a/diskann-benchmark/src/inputs/disk.rs b/diskann-benchmark/src/inputs/disk.rs index 032665cd7a..7fe5797f23 100644 --- a/diskann-benchmark/src/inputs/disk.rs +++ b/diskann-benchmark/src/inputs/disk.rs @@ -11,7 +11,9 @@ use std::collections::HashSet; #[cfg(feature = "disk-index")] use diskann::graph; -use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker}; +use diskann_benchmark_runner::{ + files::InputFile, reflect, utils::datatype::DataType, Checker, Reflect, +}; #[cfg(feature = "disk-index")] use diskann_disk::search::search_mode::SearchMode; #[cfg(feature = "disk-index")] @@ -191,6 +193,19 @@ impl DiskIndexOperation { } } +impl Reflect for DiskIndexOperation { + fn ty() -> reflect::tree::Type { + reflect::tree::Type::aggregate( + reflect::tree::Fields::Unit, + Some("The disk index is currently incompatible with type information, so none is available.".into()) + ) + } + + fn format_type_name(f: &mut dyn std::fmt::Write) -> std::fmt::Result { + write!(f, "disk::DiskIndexOperation") + } +} + impl DiskIndexLoad { pub(crate) fn validate(&mut self, _checker: &mut Checker) -> anyhow::Result<()> { let files = [ diff --git a/diskann-benchmark/src/inputs/exhaustive.rs b/diskann-benchmark/src/inputs/exhaustive.rs index 5331e71d68..5849dc1ed5 100644 --- a/diskann-benchmark/src/inputs/exhaustive.rs +++ b/diskann-benchmark/src/inputs/exhaustive.rs @@ -6,7 +6,7 @@ use std::num::NonZeroUsize; use anyhow::{anyhow, Context}; -use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker}; +use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker, Reflect}; use serde::{Deserialize, Serialize}; use crate::{ @@ -33,7 +33,8 @@ as_input!(MinMax); // Search // //////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "exhaustive::")] pub(crate) struct SearchValues { pub(crate) recall_k: Vec, pub(crate) recall_n: Vec, @@ -85,7 +86,8 @@ impl SearchValues { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "exhaustive::")] pub(crate) struct SearchPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile, @@ -124,8 +126,9 @@ impl Example for SearchPhase { //////////////////////////////// // Transforms related methods // /////////////////////////////// -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "transform::")] pub(crate) enum TargetDim { Same, Natural, @@ -152,8 +155,9 @@ impl From for diskann_quantization::algorithms::transforms::TargetDim } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "transform::")] pub(crate) enum TransformKind { PaddingHadamard(TargetDim), RandomRotation(TargetDim), @@ -200,8 +204,9 @@ impl From<&TransformKind> for diskann_quantization::algorithms::transforms::Tran // Product Quantization Methods // ////////////////////////////////// -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] +#[reflect(prefix = "product::")] pub(crate) enum PQTableStyle { FixedChunk, Padded, @@ -218,7 +223,7 @@ impl PQTableStyle { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct Product { pub(crate) data: InputFile, pub(crate) data_type: DataType, @@ -293,8 +298,9 @@ impl std::fmt::Display for Product { // Spherical-quantization-based methods // ////////////////////////////////////////// -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "spherical::")] pub(crate) enum SphericalQuery { SameAsData, FourBitTransposed, @@ -354,8 +360,9 @@ pub(super) fn check_compatibility(num_bits: usize, query: SphericalQuery) -> any } } -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "spherical::")] pub(crate) enum PreScale { None, Some(f32), @@ -399,7 +406,8 @@ impl PreScale { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "exhaustive::spherical::")] pub(crate) struct Spherical { pub(crate) data: InputFile, pub(crate) data_type: DataType, @@ -482,8 +490,9 @@ impl std::fmt::Display for Spherical { // MinMax-quantization-based methods // /////////////////////////////////////// -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "minmax::")] pub(crate) enum MinMaxQuery { SameAsData, FullPrecision, @@ -501,7 +510,8 @@ impl std::fmt::Display for MinMaxQuery { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "minmax::exhaustive::")] pub(crate) struct MinMax { pub(crate) data: InputFile, pub(crate) data_type: DataType, diff --git a/diskann-benchmark/src/inputs/filters.rs b/diskann-benchmark/src/inputs/filters.rs index 942c6da12b..40b9685139 100644 --- a/diskann-benchmark/src/inputs/filters.rs +++ b/diskann-benchmark/src/inputs/filters.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_benchmark_runner::{files::InputFile, Checker}; +use diskann_benchmark_runner::{files::InputFile, Checker, Reflect}; use serde::{Deserialize, Serialize}; use crate::inputs::{as_input, Example}; @@ -18,7 +18,7 @@ as_input!(MetadataIndexBuild); // Metadata-only Index Build // /////////////////////////////// -#[derive(Default, Debug, Serialize, Deserialize, Clone, Copy)] +#[derive(Default, Debug, Serialize, Deserialize, Reflect, Clone, Copy)] pub(crate) enum InvertedIndexKind { #[serde(rename = "bftree")] #[default] @@ -32,22 +32,20 @@ impl std::fmt::Display for InvertedIndexKind { } } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct FilterParams { pub(crate) query_predicates: InputFile, pub(crate) data_labels: InputFile, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct MetadataIndexBuild { /// Filter parameters describing predicate and label file locations. The /// actual label file used to build the inverted index is taken from /// `filter_params.data_labels`. pub(crate) filter_params: FilterParams, - /// Which inverted-index implementation to use when building/evaluating - /// bitmap filters. If omitted in input files, defaults to `fast`. - #[serde(default)] + /// Which inverted-index implementation to use when building/evaluating bitmap filters. pub(crate) inverted_index_type: InvertedIndexKind, } diff --git a/diskann-benchmark/src/inputs/flat.rs b/diskann-benchmark/src/inputs/flat.rs index d25e77cf76..b998aceae8 100644 --- a/diskann-benchmark/src/inputs/flat.rs +++ b/diskann-benchmark/src/inputs/flat.rs @@ -6,7 +6,7 @@ use std::num::NonZeroUsize; use anyhow::Context; -use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker}; +use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker, Reflect}; use serde::{Deserialize, Serialize}; use crate::{ @@ -25,7 +25,7 @@ as_input!(FlatSearch); /////////// /// Input specification for a flat-index (brute-force kNN) benchmark. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct FlatSearch { /// Path to the dataset vectors (`.bin` format). pub(crate) data: InputFile, @@ -91,7 +91,7 @@ impl Example for FlatSearch { /////////////////// /// Parameters controlling the search phase of a flat benchmark. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct SearchPhase { /// Path to the query vectors (`.bin` format). pub(crate) queries: InputFile, diff --git a/diskann-benchmark/src/inputs/graph_index.rs b/diskann-benchmark/src/inputs/graph_index.rs index c7ed6c1410..e548f63e40 100644 --- a/diskann-benchmark/src/inputs/graph_index.rs +++ b/diskann-benchmark/src/inputs/graph_index.rs @@ -3,15 +3,20 @@ * Licensed under the MIT license. */ -use std::num::{NonZero, NonZeroU32, NonZeroUsize}; +use std::{ + num::{NonZero, NonZeroU32, NonZeroUsize}, + path::PathBuf, +}; use anyhow::{anyhow, Context}; use diskann::{ - graph::{self, config, search::Range, RangeSearchError, StartPointStrategy}, + graph::{self, config, search::Range, RangeSearchError}, utils::IntoUsize, }; use diskann_benchmark_core::streaming::executors::bigann; -use diskann_benchmark_runner::{files::InputFile, utils::datatype::DataType, Checker}; +use diskann_benchmark_runner::{ + files::InputFile, utils::datatype::DataType, utils::RequiredOption, Checker, Reflect, +}; use diskann_providers::{ model::{ configuration::IndexConfiguration, @@ -41,7 +46,8 @@ as_input!(DynamicIndexRun); // Search // //////////// -#[derive(Debug, Serialize, Deserialize, Clone)] +#[derive(Debug, Serialize, Deserialize, Reflect, Clone)] +#[reflect(prefix = "graph::")] pub(crate) struct GraphSearch { pub(crate) search_n: usize, pub(crate) search_l: Vec, @@ -65,7 +71,8 @@ impl GraphSearch { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct GraphRangeSearch { pub(crate) initial_search_l: Vec, pub(crate) radius: f32, @@ -102,7 +109,8 @@ impl GraphRangeSearch { } } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct TopkSearchPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile, @@ -156,7 +164,8 @@ impl Example for TopkSearchPhase { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct RangeSearchPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile, @@ -166,7 +175,8 @@ pub(crate) struct RangeSearchPhase { pub(crate) runs: Vec, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct FilteredRangeSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, @@ -206,7 +216,8 @@ impl RangeSearchPhase { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct BetaSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, @@ -242,7 +253,8 @@ impl BetaSearchPhase { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct MultihopFilterSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, @@ -269,7 +281,8 @@ impl MultihopFilterSearchPhase { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct AdaptiveL { pub(crate) sample_count: NonZeroUsize, pub(crate) scale_factor: f64, @@ -283,7 +296,8 @@ impl AdaptiveL { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct InlineFilterSearchPhase { pub(crate) queries: InputFile, pub(crate) query_predicates: InputFile, @@ -292,8 +306,7 @@ pub(crate) struct InlineFilterSearchPhase { pub(crate) data_labels: InputFile, pub(crate) num_threads: Vec, pub(crate) runs: Vec, - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(crate) adaptive_l: Option, + pub(crate) adaptive_l: RequiredOption, } impl InlineFilterSearchPhase { @@ -306,7 +319,7 @@ impl InlineFilterSearchPhase { run.validate(checker) .with_context(|| format!("search run {}", i))?; } - if let Some(ref adaptive_l) = self.adaptive_l { + if let Some(adaptive_l) = self.adaptive_l.as_ref() { adaptive_l.validate(checker)?; } @@ -314,7 +327,7 @@ impl InlineFilterSearchPhase { } pub(crate) fn adaptive_l(&self) -> Result, anyhow::Error> { - if let Some(ref adaptive_l) = self.adaptive_l { + if let Some(adaptive_l) = self.adaptive_l.as_ref() { let adaptive_l = graph::search::AdaptiveL::new( adaptive_l.sample_count.into(), adaptive_l.scale_factor, @@ -327,8 +340,9 @@ impl InlineFilterSearchPhase { } /// A one-to-one correspondence with [`diskann::graph::config::IntraBatchCandidates`]. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] +#[reflect(prefix = "graph::")] pub(crate) enum IntraBatchCandidates { /// No intra-batch candidates will be considered. None, @@ -359,7 +373,8 @@ impl From for config::IntraBatchCandidates { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct MultiInsert { pub(crate) batch_size: NonZeroUsize, pub(crate) batch_parallelism: NonZeroUsize, @@ -379,7 +394,8 @@ impl Example for MultiInsert { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct TopkDeterminantDiversityPhase { pub(crate) queries: InputFile, pub(crate) groundtruth: InputFile, @@ -426,8 +442,9 @@ impl Example for TopkDeterminantDiversityPhase { } } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Debug, Deserialize, Serialize, Reflect)] #[serde(tag = "search-type", rename_all = "kebab-case")] +#[reflect(prefix = "graph::")] pub(crate) enum SearchPhase { Topk(TopkSearchPhase), Range(RangeSearchPhase), @@ -596,7 +613,8 @@ impl std::fmt::Display for SearchPhaseKind { // Build - Full Precision // //////////////////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct IndexLoad { pub(crate) data_type: DataType, pub(crate) distance: SimilarityMeasure, @@ -673,17 +691,18 @@ impl std::fmt::Display for IndexLoad { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct InsertRetry { num_insert_attempts: NonZeroU32, retry_threshold: f32, saturate_inserts: bool, } -#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq)] -#[serde(remote = "StartPointStrategy")] +#[derive(Debug, Serialize, Deserialize, Reflect, Clone, Copy, PartialEq)] #[serde(rename_all = "snake_case")] -pub enum StartPointStrategyRef { +#[reflect(prefix = "graph::")] +pub enum StartPointStrategy { /// Randomly select vector(s) with given norm as starting points with seed provided. /// Requires the norm (f32), number of samples (usize), and random seed (u64) to be provided. RandomVectors { @@ -707,7 +726,32 @@ pub enum StartPointStrategyRef { FirstVector, } -#[derive(Debug, Serialize, Deserialize)] +impl StartPointStrategy { + fn as_diskann(&self) -> graph::StartPointStrategy { + match *self { + Self::RandomVectors { + norm, + nsamples, + seed, + } => graph::StartPointStrategy::RandomVectors { + norm, + nsamples, + seed, + }, + Self::RandomSamples { nsamples, seed } => { + graph::StartPointStrategy::RandomSamples { nsamples, seed } + } + Self::Medoid => graph::StartPointStrategy::Medoid, + Self::LatinHyperCube { nsamples, seed } => { + graph::StartPointStrategy::LatinHyperCube { nsamples, seed } + } + Self::FirstVector => graph::StartPointStrategy::FirstVector, + } + } +} + +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct IndexBuild { data_type: DataType, data: InputFile, @@ -715,7 +759,6 @@ pub(crate) struct IndexBuild { max_degree: usize, l_build: usize, insert_retry: Option, - #[serde(with = "StartPointStrategyRef")] start_point_strategy: StartPointStrategy, alpha: f32, backedge_ratio: f32, @@ -789,7 +832,7 @@ impl IndexBuild { ) -> DefaultProviderParameters { DefaultProviderParameters { max_points: num_points, - frozen_points: NonZero::new(self.start_point_strategy.count()).unwrap(), + frozen_points: NonZero::new(self.start_point_strategy().count()).unwrap(), metric: self.distance.into(), dim, max_degree: self.exact_max_degree() as u32, @@ -804,7 +847,7 @@ impl IndexBuild { write_field!(f, "max degree", self.max_degree)?; write_field!(f, "L-build", self.l_build)?; write_field!(f, "alpha", self.alpha)?; - write_field!(f, "start point strategy", self.start_point_strategy)?; + write_field!(f, "start point strategy", self.start_point_strategy())?; write_field!(f, "backedge ratio", self.backedge_ratio)?; match &self.multi_insert { None => write_field!(f, "Using Multi Insert", "NO")?, @@ -814,7 +857,7 @@ impl IndexBuild { write_field!(f, "Intra Batch Candidates", mi.intra_batch_candidates)?; } } - write_field!(f, "start_point_strategy", self.start_point_strategy)?; + write_field!(f, "start_point_strategy", self.start_point_strategy())?; write_field!(f, "build threads", self.num_threads)?; match &self.save_path { None => write_field!(f, "Save Path", "None")?, @@ -864,8 +907,8 @@ impl IndexBuild { &self.data } - pub(crate) fn start_point_strategy(&self) -> &StartPointStrategy { - &self.start_point_strategy + pub(crate) fn start_point_strategy(&self) -> graph::StartPointStrategy { + self.start_point_strategy.as_diskann() } pub(crate) fn multi_insert(&self) -> Option<&MultiInsert> { @@ -906,8 +949,9 @@ impl std::fmt::Display for IndexBuild { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(tag = "index-source")] // Use tagged enums for JSON +#[reflect(prefix = "graph::")] pub enum IndexSource { Load(IndexLoad), Build(IndexBuild), @@ -936,7 +980,8 @@ impl IndexSource { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct IndexOperation { pub(crate) source: IndexSource, // either load or build pub(crate) search_phase: SearchPhase, @@ -978,7 +1023,8 @@ impl std::fmt::Display for IndexOperation { // Graph Index Build PQ // ////////////////////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct IndexPQOperation { pub(crate) index_operation: IndexOperation, // either load or build pub(crate) num_pq_chunks: usize, @@ -1068,7 +1114,8 @@ impl std::fmt::Display for IndexPQOperation { // Graph Index Build SQ // ////////////////////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct IndexSQOperation { pub(crate) index_operation: IndexOperation, pub(crate) num_bits: usize, @@ -1158,7 +1205,8 @@ impl std::fmt::Display for IndexSQOperation { // Graph Index Build Spherical // ///////////////////////////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct SphericalQuantBuild { pub(crate) build: IndexBuild, // spherical does not support saving and loading pub(crate) search_phase: SearchPhase, @@ -1276,8 +1324,9 @@ impl std::fmt::Display for SphericalQuantBuild { // Dynamic Runbook Params // //////////////////////////// -#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize)] +#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize, Reflect)] #[serde(tag = "method", content = "params")] +#[reflect(prefix = "graph::")] pub enum InplaceDeleteMethod { #[serde(rename = "visited_and_top_k")] VisitedAndTopK { k_value: usize, l_value: usize }, @@ -1300,19 +1349,17 @@ impl From for graph::InplaceDeleteMethod { } /// Runbook loading and phase type definitions are in utils.datafiles -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct DynamicRunbookParams { pub(crate) runbook_path: InputFile, pub(crate) dataset_name: String, - pub(crate) gt_directory: String, + pub(crate) gt_directory: PathBuf, pub(crate) ip_delete_method: InplaceDeleteMethod, pub(crate) ip_delete_num_to_replace: usize, /// Threshold for deferred consolidation. Required for soft-delete providers (inmem). /// Hard-delete providers (bf-tree) ignore this field. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub(crate) consolidate_threshold: Option, - #[serde(skip)] - pub(crate) resolved_gt_directory: Option, + pub(crate) consolidate_threshold: RequiredOption, } // Validates: @@ -1324,7 +1371,7 @@ impl DynamicRunbookParams { self.runbook_path.resolve(checker)?; // Validate consolidate_threshold if provided - if let Some(threshold) = self.consolidate_threshold { + if let Some(&threshold) = self.consolidate_threshold.as_ref() { if threshold <= 0.0 { return Err(anyhow::anyhow!( "consolidate_threshold must be greater than 0, but got {}", @@ -1334,23 +1381,23 @@ impl DynamicRunbookParams { } let final_gt_directory = checker.find_input_dir(self.gt_directory.as_ref())?; - self.resolved_gt_directory = Some(final_gt_directory.clone()); // Run pre-flight checks for the runbook. let _runbook = bigann::RunBook::load( &self.runbook_path, &self.dataset_name, - &mut bigann::ScanDirectory::new(final_gt_directory)?, + &mut bigann::ScanDirectory::new(&final_gt_directory)?, ) .with_context(|| { format!( "Failed to validate runbook '{}' with dataset '{}' and gt_directory '{}'", self.runbook_path.display(), self.dataset_name, - self.gt_directory + self.gt_directory.display(), ) })?; + self.gt_directory = final_gt_directory; Ok(()) } } @@ -1366,8 +1413,7 @@ impl Example for DynamicRunbookParams { l_value: 64, }, ip_delete_num_to_replace: 3, - consolidate_threshold: Some(0.2), - resolved_gt_directory: None, + consolidate_threshold: RequiredOption::some(0.2), } } } @@ -1377,7 +1423,7 @@ impl DynamicRunbookParams { #[cfg(feature = "bftree")] pub(crate) fn example_immediate() -> Self { Self { - consolidate_threshold: None, + consolidate_threshold: RequiredOption::none(), ..Self::example() } } @@ -1388,13 +1434,7 @@ impl std::fmt::Display for DynamicRunbookParams { writeln!(f, "Dynamic Runbook Parameters")?; write_field!(f, "Runbook Path", self.runbook_path.display())?; write_field!(f, "Dataset Name", self.dataset_name)?; - - // Show resolved path if available, otherwise show original - let gt_dir_display = match &self.resolved_gt_directory { - Some(resolved) => resolved.display().to_string(), - None => self.gt_directory.clone(), - }; - write_field!(f, "Ground Truth Directory", gt_dir_display)?; + write_field!(f, "Ground Truth Directory", self.gt_directory.display())?; match self.ip_delete_method { InplaceDeleteMethod::VisitedAndTopK { k_value, l_value } => { @@ -1410,7 +1450,7 @@ impl std::fmt::Display for DynamicRunbookParams { } } write_field!(f, "IP Delete Num to Replace", self.ip_delete_num_to_replace)?; - if let Some(threshold) = self.consolidate_threshold { + if let Some(threshold) = self.consolidate_threshold.as_ref() { write_field!(f, "Consolidate Threshold", threshold)?; } @@ -1422,7 +1462,8 @@ impl std::fmt::Display for DynamicRunbookParams { // Graph Index Dynamic // /////////////////////////// -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "graph::")] pub(crate) struct DynamicIndexRun { pub(crate) build: IndexBuild, pub(crate) search_phase: SearchPhase, diff --git a/diskann-benchmark/src/inputs/multi_vector.rs b/diskann-benchmark/src/inputs/multi_vector.rs index 6f1ec9dddd..2231c5c752 100644 --- a/diskann-benchmark/src/inputs/multi_vector.rs +++ b/diskann-benchmark/src/inputs/multi_vector.rs @@ -5,7 +5,7 @@ use std::num::NonZeroUsize; -use diskann_benchmark_runner::{utils::datatype::DataType, Checker, Input}; +use diskann_benchmark_runner::{utils::datatype::DataType, Checker, Input, Reflect}; use diskann_quantization::multi_vector::MaxSimIsa; use serde::{Deserialize, Serialize}; @@ -16,7 +16,7 @@ use serde::{Deserialize, Serialize}; /// JSON-facing shadow of [`MaxSimIsa`]. The library's enum is deliberately /// serde-free; this owns the kebab-case JSON shape and converts via `From`. /// Stays variant-for-variant in sync with `MaxSimIsa` manually. -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] #[non_exhaustive] pub(crate) enum BenchIsa { @@ -60,7 +60,7 @@ impl From for MaxSimIsa { } /// One benchmark configuration: a single shape measurement. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Reflect)] pub(crate) struct Run { pub(crate) num_query_vectors: NonZeroUsize, pub(crate) num_doc_vectors: NonZeroUsize, @@ -74,7 +74,7 @@ pub(crate) struct Run { /////////////////////// /// A complete multi-vector benchmark job. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Reflect)] pub(crate) struct MultiVectorOp { pub(crate) element_type: DataType, pub(crate) isa: BenchIsa, diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index d25a56fc92..9644043da7 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -49,6 +49,8 @@ impl Cli { fn run(&self, output: &mut dyn runner::Output) -> anyhow::Result<()> { self.check_target(output)?; + let now = std::time::Instant::now(); + // Collect benchmarks. let mut registry = runner::Registry::new(); exhaustive::register_benchmarks(&mut registry)?; @@ -58,6 +60,8 @@ impl Cli { filters::register_benchmarks(&mut registry)?; multi_vector::register_benchmarks(&mut registry)?; + println!("registration took {}us", now.elapsed().as_micros()); + self.app.run(®istry, output) } diff --git a/diskann-benchmark/src/multi_vector/driver.rs b/diskann-benchmark/src/multi_vector/driver.rs index 5cec88bc02..44011504f9 100644 --- a/diskann-benchmark/src/multi_vector/driver.rs +++ b/diskann-benchmark/src/multi_vector/driver.rs @@ -12,7 +12,7 @@ use diskann_benchmark_runner::{ num::{relative_change, NonNegativeFinite}, percentiles, MicroSeconds, }, - Checker, Input, + Checker, Input, Reflect, }; use diskann_quantization::multi_vector::{ Mat, MatRef, MaxSimElement, MaxSimKernel, Overflow, Standard, @@ -32,7 +32,8 @@ use crate::utils::DisplayWrapper; ////////////////////// /// Tolerance thresholds for multi-vector benchmark regression detection. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "multi_vector::")] pub(super) struct MultiVectorTolerance { pub(super) min_time_regression: NonNegativeFinite, } diff --git a/diskann-benchmark/src/utils/mod.rs b/diskann-benchmark/src/utils/mod.rs index cd8e6510cf..71d6ec671d 100644 --- a/diskann-benchmark/src/utils/mod.rs +++ b/diskann-benchmark/src/utils/mod.rs @@ -6,6 +6,7 @@ use diskann_benchmark_runner::{ benchmark::Score, utils::datatype::{AsDataType, DataType}, + Reflect, }; use serde::{Deserialize, Serialize}; @@ -26,8 +27,9 @@ where } } -#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "snake_case")] +#[reflect(prefix = "benchmark::")] pub(crate) enum SimilarityMeasure { SquaredL2, InnerProduct, diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index 0e5d1e95a8..3a9d89272f 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -8,10 +8,10 @@ use std::{io::Write, sync::Arc}; use anyhow::Context; use diskann::graph::{DiskANNIndex, search::Knn}; use diskann_benchmark_runner::{ - Checker, Checkpoint, Output, Registry, RegistryError, + Checker, Checkpoint, Output, Reflect, Registry, RegistryError, benchmark::{MatchContext, PassFail, Regression, Score}, files::InputFile, - utils::fmt::Indent, + utils::{RequiredOption, fmt::Indent}, }; use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; @@ -43,8 +43,9 @@ mod dto { use serde::{Deserialize, Serialize}; - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "graph::")] pub(super) enum SerdeMetric { L2, InnerProduct, @@ -73,8 +74,9 @@ mod dto { } } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "graph::")] pub(super) enum Preprocess { Halve, Floor, @@ -98,7 +100,8 @@ mod dto { } } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "graph::")] pub(super) struct Data { pub(super) data: InputFile, pub(super) queries: InputFile, @@ -115,24 +118,27 @@ mod dto { pub(super) mod spherical { use super::*; - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "graph::spherical::")] pub(in crate::index::runner) enum Bits { One, Two, Four, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "graph::spherical::")] pub(in crate::index::runner) enum Rerank { None, F16, } } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] + #[reflect(prefix = "graph::")] pub(super) enum Representation { FullPrecision { data_type: DataType, @@ -143,7 +149,8 @@ mod dto { }, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "graph::")] pub(super) struct Build { pub(super) pruned_degree: usize, pub(super) max_degree: usize, @@ -151,20 +158,22 @@ mod dto { pub(super) alpha: f32, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "graph::")] pub(super) struct KnnSearch { pub(super) knn: usize, pub(super) search_l: usize, - #[serde(deserialize_with = "Deserialize::deserialize")] - pub(super) beam_width: Option, + pub(super) beam_width: RequiredOption, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "graph::")] pub(super) struct Search { pub(super) knn: Vec, } - #[derive(Debug, Serialize, Deserialize)] + #[derive(Debug, Serialize, Deserialize, Reflect)] + #[reflect(prefix = "graph::")] pub(super) struct Test { pub(super) data: Data, pub(super) representation: Representation, @@ -400,7 +409,10 @@ struct Search { impl Search { fn from_raw(raw: dto::Search) -> anyhow::Result { fn make_knn(raw: &dto::KnnSearch) -> anyhow::Result<(usize, Knn)> { - Ok((raw.knn, Knn::new(raw.search_l, raw.beam_width)?)) + Ok(( + raw.knn, + Knn::new(raw.search_l, raw.beam_width.as_ref().copied())?, + )) } Ok(Self { @@ -417,7 +429,7 @@ impl Search { dto::KnnSearch { knn: *k, search_l: knn.l_value().get(), - beam_width: Some(knn.beam_width().get()), + beam_width: RequiredOption::some(knn.beam_width().get()), } } @@ -635,17 +647,17 @@ impl diskann_benchmark_runner::Input for Test { dto::KnnSearch { knn: 10, search_l: 50, - beam_width: None, + beam_width: RequiredOption::none(), }, dto::KnnSearch { knn: 10, search_l: 50, - beam_width: Some(3), + beam_width: RequiredOption::some(3), }, dto::KnnSearch { knn: 20, search_l: 100, - beam_width: Some(3), + beam_width: RequiredOption::some(3), }, ], }, diff --git a/diskann-inmem/integration/store/checked.rs b/diskann-inmem/integration/store/checked.rs index 4bc5f8a982..63ddd1b4c5 100644 --- a/diskann-inmem/integration/store/checked.rs +++ b/diskann-inmem/integration/store/checked.rs @@ -14,7 +14,8 @@ pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::Registry } /// Configuration for a [`Stress`] run. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, dbr::Reflect)] +#[reflect(prefix = "store::checked::")] struct Input { /// Shared stress test setup. setup: super::Setup, diff --git a/diskann-inmem/integration/store/intrusive.rs b/diskann-inmem/integration/store/intrusive.rs index 9e5882ed1c..bcd4158b22 100644 --- a/diskann-inmem/integration/store/intrusive.rs +++ b/diskann-inmem/integration/store/intrusive.rs @@ -14,7 +14,8 @@ pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::Registry } /// Configuration for a [`Stress`] run. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, dbr::Reflect)] +#[reflect(prefix = "store::intrusive::")] struct Input { /// Shared stress test setup. setup: super::Setup, diff --git a/diskann-inmem/integration/store/mod.rs b/diskann-inmem/integration/store/mod.rs index 893133b63a..3120a0f5e0 100644 --- a/diskann-inmem/integration/store/mod.rs +++ b/diskann-inmem/integration/store/mod.rs @@ -27,7 +27,7 @@ use std::{ time::{Duration, Instant}, }; -use diskann_benchmark_runner::{Registry, RegistryError, utils::fmt::KeyValue}; +use diskann_benchmark_runner::{Reflect, Registry, RegistryError, utils::fmt::KeyValue}; use rand::{Rng, SeedableRng, distr::Uniform, rngs::StdRng}; use serde::{Deserialize, Serialize}; @@ -58,7 +58,8 @@ pub(super) fn register(registry: &mut Registry) -> Result<(), RegistryError> { // Input // /////////// -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Reflect)] +#[reflect(prefix = "store::")] pub struct Setup { /// Number of reader threads. Must be below `epoch_guard_slots`. readers: usize, diff --git a/diskann-inmem/integration/store/simple.rs b/diskann-inmem/integration/store/simple.rs index 11095a67db..aa82486dde 100644 --- a/diskann-inmem/integration/store/simple.rs +++ b/diskann-inmem/integration/store/simple.rs @@ -14,7 +14,8 @@ pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::Registry } /// Configuration for a [`Stress`] run. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, dbr::Reflect)] +#[reflect(prefix = "store::simple::")] struct Input { /// Shared stress test setup. setup: super::Setup, diff --git a/diskann-inmem/integration/support/datatype.rs b/diskann-inmem/integration/support/datatype.rs index b420453317..ad6b8ba703 100644 --- a/diskann-inmem/integration/support/datatype.rs +++ b/diskann-inmem/integration/support/datatype.rs @@ -3,6 +3,7 @@ * Licensed under the MIT license. */ +use diskann_benchmark_runner::Reflect; use diskann_utils::{ sampling::medoid::ComputeMedoid, views::rowmajor::{self, Matrix, MatrixMut}, @@ -16,7 +17,7 @@ use thiserror::Error; // DataType // ////////////// -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)] #[serde(rename_all = "kebab-case")] pub(crate) enum DataType { F32, diff --git a/diskann-inmem/integration/support/tolerance.rs b/diskann-inmem/integration/support/tolerance.rs index 7f8efa8709..2f6555325d 100644 --- a/diskann-inmem/integration/support/tolerance.rs +++ b/diskann-inmem/integration/support/tolerance.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_benchmark_runner::{Checker, Input}; +use diskann_benchmark_runner::{Checker, Input, Reflect}; use serde::{Deserialize, Serialize}; /// A tolerance [`Input`] for [`diskann_benchmark_runner::benchmark::Regression`]s that /// do not need any external tolerances. -#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Reflect)] pub(crate) struct Empty; impl Input for Empty { diff --git a/diskann-macros/Cargo.toml b/diskann-macros/Cargo.toml new file mode 100644 index 0000000000..e85dcef53b --- /dev/null +++ b/diskann-macros/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "diskann-macros" +version.workspace = true +description.workspace = true +authors.workspace = true +license.workspace = true +edition = "2024" + +[lib] +proc-macro = true + +[dependencies] +syn = { version = "2", features = ["full"] } +quote = "1" +proc-macro2 = "1" + +[dev-dependencies] +diskann-benchmark-runner = { workspace = true } +trybuild = "1.0.120" + +[lints] +workspace = true diff --git a/diskann-macros/src/lib.rs b/diskann-macros/src/lib.rs new file mode 100644 index 0000000000..ac159d8651 --- /dev/null +++ b/diskann-macros/src/lib.rs @@ -0,0 +1,44 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use syn::{DeriveInput, parse_macro_input}; + +mod reflect; + +/// Derive macro for the `Reflect` trait. +/// +/// Supports named structs, tuple structs, and enums. Doc comments on the item +/// and its fields/variants are captured as reflection metadata. +/// +/// # Example +/// +/// ```ignore +/// use diskann_benchmark_runner::Reflect; +/// +/// /// A test aggregate. +/// #[derive(Reflect)] +/// struct MyInput { +/// /// The number of threads. +/// threads: usize, +/// } +/// ``` +/// +/// # Serde Compatibility +/// +/// This is meant to be used in conjunction with `serde` for documenting benchmark inputs +/// and as such, it respects several common `serde` attributes: +/// +/// * rename_all = "lowercase" | "snake_case" | "kebab-case" (on structs, enums, and variants) +/// * rename = "..." (on fields) +/// * tag = "..." (on enums) +/// * tag = "...", content = "..." (on enums) +#[proc_macro_derive(Reflect, attributes(reflect, serde))] +pub fn derive_reflect(input: proc_macro::TokenStream) -> proc_macro::TokenStream { + let input = parse_macro_input!(input as DeriveInput); + + reflect::expand(&input) + .unwrap_or_else(syn::Error::into_compile_error) + .into() +} diff --git a/diskann-macros/src/reflect/attributes.rs b/diskann-macros/src/reflect/attributes.rs new file mode 100644 index 0000000000..9b5d0c41f9 --- /dev/null +++ b/diskann-macros/src/reflect/attributes.rs @@ -0,0 +1,654 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! These are modeled after the attributes documented in . + +use proc_macro2::{Span, TokenStream}; +use quote::quote; + +#[must_use] +fn is_serde_attr(attr: &syn::Attribute) -> bool { + attr.path().is_ident("serde") +} + +#[must_use] +fn is_reflect_attr(attr: &syn::Attribute) -> bool { + attr.path().is_ident("reflect") +} + +#[must_use] +fn is_doc_attr(attr: &syn::Attribute) -> bool { + attr.path().is_ident("doc") +} + +fn identity(x: T) -> T { + x +} + +fn set_unique(opt: &mut Option, value: syn::LitStr, attr: &str) -> syn::Result<()> { + if opt.is_some() { + Err(syn::Error::new_spanned( + value, + format!("serde attribute `{}` found multiple times", attr), + )) + } else { + *opt = Some(value); + Ok(()) + } +} + +fn reject_reflect_attributes(attrs: &[syn::Attribute], error: &str) -> syn::Result<()> { + if let Some(attr) = attrs.iter().find(|a| is_reflect_attr(a)) { + return Err(syn::Error::new_spanned(attr, error)); + } + Ok(()) +} + +#[derive(Default)] +pub(crate) struct Doc { + doc: Option, +} + +impl Doc { + /// Parse and concatenate with newlines all doc attributes. + pub(crate) fn append(&mut self, meta: &syn::Meta) { + if let syn::Meta::NameValue(name) = meta + && let syn::Expr::Lit(literal) = &name.value + && let syn::Lit::Str(s) = &literal.lit + { + let value = s.value(); + + // Doc attributes are automatically prefixed with a single space. + // + // This removes that. + let processed = value.strip_prefix(" ").unwrap_or(&value); + match self.doc.as_mut() { + None => self.doc = Some(processed.into()), + Some(doc) => { + doc.push('\n'); + doc.push_str(processed); + } + } + } + } + + /// Emite the parsed docs as a borrowed `Option>`. + pub(crate) fn emit(&self) -> TokenStream { + match self.doc.as_deref() { + None => quote!(::std::option::Option::None), + Some(doc) => quote! { + ::std::option::Option::Some(::std::borrow::Cow::Borrowed(#doc)) + }, + } + } +} + +/// Attributes applicable to struct definitions. +pub(crate) struct Struct { + /// A univeral rename rule for all fields. + /// + /// [`Field`] specific renames take precedence. + pub(crate) rename_all: RenameAll, + + /// The documentation for this struct. + pub(crate) doc: Doc, +} + +/// Attributes applicable to enum definitions. +pub(crate) struct Enum { + /// A universal rename rule for all variants. + /// + /// [`Variant`] specific renames take precedence. + pub(crate) rename_all: RenameAll, + + /// Enum's representation. + pub(crate) enum_repr: EnumRepr, + + /// The documentation for this struct. + pub(crate) doc: Doc, +} + +/// Attributes on the top-level [`syn::DeriveInput`]. +/// +/// Uses should go through [`Container::into_enum`] or [`Container::try_as_struct`] to ensure +/// the attributes are appropriate for the actual type. +pub(crate) struct Container { + rename_all: RenameAll, + enum_repr: EnumRepr, + type_name: TypeName, + doc: Doc, +} + +impl Container { + /// Parse container level attributes. Accepts: + /// + /// * `doc` + /// * `serde(rename_all = "")` + /// * `serde(tag = "")` + /// * `serde(content = "")` + /// * `reflect(prefix = "")` + /// * `reflect(type_name = "")` + /// + /// `reflect(prefix)` and `reflect(type_name)` are mutually exclusive. + /// + /// All other serde and reflect attributes are rejected. + pub(crate) fn parse(attrs: &[syn::Attribute]) -> syn::Result { + let mut rename_all = RenameAll::default(); + let mut tag = Option::None; + let mut content = Option::None; + let mut type_name = TypeName::None; + let mut doc = Doc::default(); + + // Parse serde attributes. + for attr in attrs.iter().filter(|a| is_serde_attr(a)) { + attr.parse_nested_meta(|meta| { + // serde(rename_all = "...") + if meta.path.is_ident("rename_all") { + let value: syn::LitStr = meta.value()?.parse()?; + rename_all.parse_once(&value)?; + return Ok(()); + } + + // serde(tag = "...") + if meta.path.is_ident("tag") { + let value: syn::LitStr = meta.value()?.parse()?; + set_unique(&mut tag, value, "tag")?; + return Ok(()); + } + + // serde(content = "...") + if meta.path.is_ident("content") { + let value: syn::LitStr = meta.value()?.parse()?; + set_unique(&mut content, value, "content")?; + return Ok(()); + } + + Err(meta.error("unsupported Serde attribute for Reflect")) + })?; + } + + // Parse reflect attributes + for attr in attrs.iter().filter(|a| is_reflect_attr(a)) { + attr.parse_nested_meta(|meta| { + // reflect(prefix = "...") + if meta.path.is_ident("prefix") { + let value: syn::LitStr = meta.value()?.parse()?; + type_name.set_unique(TypeNameKind::Prefix, value)?; + return Ok(()); + } + + // reflect(type_name = "...") + if meta.path.is_ident("type_name") { + let value: syn::LitStr = meta.value()?.parse()?; + type_name.set_unique(TypeNameKind::Rename, value)?; + return Ok(()); + } + + Err(meta.error("unsupported attribute for Reflect")) + })?; + } + + // Parse docs + for attr in attrs.iter().filter(|a| is_doc_attr(a)) { + doc.append(&attr.meta); + } + + Ok(Self { + rename_all, + enum_repr: EnumRepr::from_parsed(tag, content)?, + type_name, + doc, + }) + } + + /// Verify the parsed attributes are compatible with an aggregate definition. + pub(crate) fn try_as_struct(self) -> syn::Result { + self.enum_repr.assert_struct_compatible()?; + Ok(Struct { + rename_all: self.rename_all, + doc: self.doc, + }) + } + + /// Verify the parsed attributes are compatible with an enum definition. + pub(crate) fn into_enum(self) -> Enum { + Enum { + rename_all: self.rename_all, + enum_repr: self.enum_repr, + doc: self.doc, + } + } + + /// Extract the type-name attributes. + pub(crate) fn type_name(&self) -> TypeName { + self.type_name.clone() + } +} + +/// The `serde` representation for an enum. +pub(crate) enum EnumRepr { + External, + Internal { + tag: syn::LitStr, + }, + Adjacent { + tag: syn::LitStr, + content: syn::LitStr, + }, +} + +impl EnumRepr { + /// Verify that the parsed `tag` and `content` fields are coherent. + /// + /// This checks that `content` cannot be applied without a `tag`. + pub(crate) fn from_parsed( + tag: Option, + content: Option, + ) -> syn::Result { + match (tag, content) { + (None, None) => Ok(EnumRepr::External), + (Some(tag), None) => Ok(EnumRepr::Internal { tag }), + (Some(tag), Some(content)) => { + if tag.value() == content.value() { + Err(syn::Error::new_spanned( + &content, + "adjacently tagged `content` conflicts with `tag`", + )) + } else { + Ok(EnumRepr::Adjacent { tag, content }) + } + } + (None, Some(content)) => Err(syn::Error::new_spanned( + content, + "serde attribute `content` provided without a `tag`", + )), + } + } + + /// Verify that no `enum` tag attributes are present. + /// + /// These do not apply to structs, so we give a compile error with a diagnostic if they + /// are observed. + pub(crate) fn assert_struct_compatible(&self) -> syn::Result<()> { + match self { + Self::External => Ok(()), + Self::Internal { tag } => Err(syn::Error::new_spanned( + tag, + "serde attribute `tag` provided on a non-enum", + )), + Self::Adjacent { tag, .. } => Err(syn::Error::new_spanned( + tag, + "serde attributes `tag` and `content` provided on a non-enum", + )), + } + } +} + +/// Rename rule for fields and variants. +#[derive(Debug, Default, Clone, Copy)] +pub(crate) struct RenameAll { + inner: RenameRule, + span: Option, +} + +impl RenameAll { + /// Attempt to parse `s`, returning an error if `self` is already parsed. + fn parse_once(&mut self, s: &syn::LitStr) -> syn::Result<()> { + if self.inner != RenameRule::None { + Err(syn::Error::new_spanned( + s, + "serde attribute `rename_all` found multiple times", + )) + } else { + let value = s.value(); + match RenameRule::parse(&value) { + Some(inner) => { + self.inner = inner; + self.span = Some(s.span()); + Ok(()) + } + None => Err(syn::Error::new_spanned( + s, + format!( + "unsupported serde `rename_all` rule \"{}\" - expected one of {}", + value, + RenameRule::supported() + ), + )), + } + } + } + + /// Return the [`Span`] for the rename all rule - returning `None` if no rename-all + /// rule is in effect. + pub(crate) fn span_if_present(&self) -> Option { + self.span + } + + /// Apply the rename rule to `variant`. + pub(crate) fn apply_to_variant(&self, variant: syn::LitStr) -> syn::LitStr { + if self.inner == RenameRule::None { + variant + } else { + syn::LitStr::new( + &self.inner.apply_to_variant_str(&variant.value()), + variant.span(), + ) + } + } + + /// Apply the rename rule to `field`. + pub(crate) fn apply_to_field(&self, field: syn::LitStr) -> syn::LitStr { + if self.inner == RenameRule::None { + field + } else { + syn::LitStr::new(&self.inner.apply_to_field_str(&field.value()), field.span()) + } + } +} + +/// Supported subset of `serde(rename_all = "...")` +#[derive(Default, Debug, Clone, Copy, PartialEq)] +pub(crate) enum RenameRule { + #[default] + None, + Lower, + Snake, + Kebab, +} + +impl RenameRule { + fn supported() -> &'static str { + "\"lowercase\", \"snake_case\", or \"kebab-case\"" + } + + fn parse(s: &str) -> Option { + match s { + "lowercase" => Some(Self::Lower), + "snake_case" => Some(Self::Snake), + "kebab-case" => Some(Self::Kebab), + _ => None, + } + } + + /// These methods are taken from the `serde_derive` internals as they need to match. + /// + /// See: + fn apply_to_variant_str(&self, variant: &str) -> String { + match self { + Self::None => variant.to_owned(), + Self::Lower => variant.to_ascii_lowercase(), + Self::Snake => { + let mut snake = String::new(); + for (i, ch) in variant.char_indices() { + if i > 0 && ch.is_uppercase() { + snake.push('_'); + } + snake.push(ch.to_ascii_lowercase()); + } + snake + } + Self::Kebab => (Self::Snake) + .apply_to_variant_str(variant) + .replace('_', "-"), + } + } + + /// These methods are taken from the `serde_derive` internals as they need to match. + /// + /// Since Rust field are generally in lower snake case, there's less work to do. + /// + /// See: + fn apply_to_field_str(&self, field: &str) -> String { + match self { + Self::None | Self::Lower | Self::Snake => field.to_owned(), + Self::Kebab => field.replace('_', "-"), + } + } +} + +/// A one-type variant or field renamer. +#[derive(Default)] +pub(crate) struct RenameOnce { + rename: Option, +} + +impl RenameOnce { + /// Return the source span for the rename attribute, or `None` if no attribute was present. + pub(crate) fn span(&self) -> Option { + self.rename.as_ref().map(|lit| lit.span()) + } + + /// Apply the configured rename to `variant`. If no rename is configured, apply `or_else`. + pub(crate) fn apply_to_variant(&self, variant: syn::LitStr, or_else: RenameAll) -> syn::LitStr { + self.rename + .clone() + .map_or_else(|| or_else.apply_to_variant(variant), identity) + } + + /// Apply the configured rename to `field`. If no rename is configured, apply `or_else`. + pub(crate) fn apply_to_field(&self, field: syn::LitStr, or_else: RenameAll) -> syn::LitStr { + self.rename + .clone() + .map_or_else(|| or_else.apply_to_field(field), identity) + } +} + +/// Strategy for generating type-names. +#[derive(Default, Clone)] +pub(crate) enum TypeName { + #[default] + None, + Prefix(syn::LitStr), + Rename(syn::LitStr), +} + +enum TypeNameKind { + Prefix, + Rename, +} + +impl TypeName { + fn set_unique(&mut self, kind: TypeNameKind, value: syn::LitStr) -> syn::Result<()> { + let error = match (&*self, &kind) { + (Self::Prefix(_), TypeNameKind::Prefix) => { + Some("reflect attribute `prefix` found multiple times") + } + (Self::Rename(_), TypeNameKind::Rename) => { + Some("reflect attribute `type_name` found multiple times") + } + (Self::Prefix(_), TypeNameKind::Rename) | (Self::Rename(_), TypeNameKind::Prefix) => { + Some("reflect attributes `prefix` and `type_name` are mutually exclusive") + } + (Self::None, _) => None, + }; + + if let Some(error) = error { + return Err(syn::Error::new_spanned(value, error)); + } + + match kind { + TypeNameKind::Prefix => *self = Self::Prefix(value), + TypeNameKind::Rename => *self = Self::Rename(value), + } + Ok(()) + } +} + +//---------// +// Variant // +//---------// + +/// Variant level attributes. +#[derive(Default)] +pub(crate) struct Variant { + pub(crate) rename_variant: RenameOnce, + pub(crate) rename_variant_fields: RenameAll, + pub(crate) doc: Doc, +} + +impl Variant { + /// Parse variant level attributes. Accepts: + /// + /// * `doc` + /// * `serde(rename_all)` + /// * `serde(rename)` + /// + /// All other `serde` and `reflect` attributes are rejected. + pub(crate) fn parse(attrs: &[syn::Attribute]) -> syn::Result { + let mut me = Self::default(); + + for attr in attrs.iter().filter(|a| is_serde_attr(a)) { + attr.parse_nested_meta(|meta| { + // serde(rename_all = "...") + if meta.path.is_ident("rename_all") { + let value: syn::LitStr = meta.value()?.parse()?; + me.rename_variant_fields.parse_once(&value)?; + return Ok(()); + } + + // serde(rename = "...") + if meta.path.is_ident("rename") { + let value: syn::LitStr = meta.value()?.parse()?; + set_unique(&mut me.rename_variant.rename, value, "rename")?; + return Ok(()); + } + + Err(meta.error("unsupported Serde attribute for Reflect")) + })?; + } + + // Reject any `reflect` attributes. + reject_reflect_attributes(attrs, "Reflect does not support variant-level attributes")?; + + // Parse docs + for attr in attrs.iter().filter(|a| is_doc_attr(a)) { + me.doc.append(&attr.meta); + } + + Ok(me) + } +} + +//-------// +// Field // +//-------// + +/// Field level attributes. +#[derive(Default)] +pub(crate) struct Field { + pub(crate) rename_field: RenameOnce, + pub(crate) doc: Doc, +} + +impl Field { + /// Parse field level attributes. Accepts: + /// + /// * `doc` + /// * `serde(rename)` + /// + /// All other serde and reflect attributes are rejected. + pub(crate) fn parse(attrs: &[syn::Attribute]) -> syn::Result { + let mut me = Self::default(); + + for attr in attrs.iter().filter(|a| is_serde_attr(a)) { + attr.parse_nested_meta(|meta| { + // serde(rename = "...") + if meta.path.is_ident("rename") { + let value: syn::LitStr = meta.value()?.parse()?; + set_unique(&mut me.rename_field.rename, value, "rename")?; + return Ok(()); + } + + Err(meta.error("unsupported Serde attribute for Reflect")) + })?; + } + + // Reject any `reflect` attributes. + reject_reflect_attributes(attrs, "Reflect does not support field-level attributes")?; + + // Parse docs + for attr in attrs.iter().filter(|a| is_doc_attr(a)) { + me.doc.append(&attr.meta); + } + + Ok(me) + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_rename_rule_parse() { + assert!(RenameRule::parse("none").is_none()); + assert_eq!(RenameRule::parse("lowercase").unwrap(), RenameRule::Lower); + assert_eq!(RenameRule::parse("snake_case").unwrap(), RenameRule::Snake); + assert_eq!(RenameRule::parse("kebab-case").unwrap(), RenameRule::Kebab); + + assert!(RenameRule::parse("foo").is_none()); + assert!(RenameRule::parse("bar").is_none()); + } + + #[test] + fn test_apply_to_variant() { + assert_eq!( + RenameRule::None.apply_to_variant_str("MiXeDUpper_Case"), + "MiXeDUpper_Case" + ); + + assert_eq!( + RenameRule::Lower.apply_to_variant_str("MiXeDUpper_Case"), + "mixedupper_case" + ); + assert_eq!( + RenameRule::Lower.apply_to_variant_str("all_lower"), + "all_lower" + ); + + assert_eq!( + RenameRule::Snake.apply_to_variant_str("MixedUpperCase"), + "mixed_upper_case" + ); + assert_eq!( + RenameRule::Snake.apply_to_variant_str("X86_64_V4"), + "x86_64__v4" + ); + + assert_eq!( + RenameRule::Kebab.apply_to_variant_str("MixedUpperCase"), + "mixed-upper-case" + ); + assert_eq!( + RenameRule::Kebab.apply_to_variant_str("X86_64_V4"), + "x86-64--v4" + ); + } + + #[test] + fn test_apply_to_field() { + assert_eq!( + RenameRule::None.apply_to_field_str("a_standard_field"), + "a_standard_field" + ); + assert_eq!( + RenameRule::Lower.apply_to_field_str("a_standard_field"), + "a_standard_field" + ); + assert_eq!( + RenameRule::Snake.apply_to_field_str("a_standard_field"), + "a_standard_field" + ); + assert_eq!( + RenameRule::Kebab.apply_to_field_str("a_standard_field"), + "a-standard-field" + ); + } +} diff --git a/diskann-macros/src/reflect/mod.rs b/diskann-macros/src/reflect/mod.rs new file mode 100644 index 0000000000..d9435f59ef --- /dev/null +++ b/diskann-macros/src/reflect/mod.rs @@ -0,0 +1,18 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use proc_macro2::TokenStream; +use syn::DeriveInput; + +mod attributes; +mod repr; + +/// The implementation of `diskann_benchmark_runner::Reflect`. +pub(crate) fn expand(input: &DeriveInput) -> syn::Result { + let input = repr::Input::parse(input)?; + let crate_name = syn::parse_quote!(::diskann_benchmark_runner::reflect); + let ts = input.emit(&crate_name); + Ok(ts) +} diff --git a/diskann-macros/src/reflect/repr.rs b/diskann-macros/src/reflect/repr.rs new file mode 100644 index 0000000000..16ed16a0a2 --- /dev/null +++ b/diskann-macros/src/reflect/repr.rs @@ -0,0 +1,754 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! This macro is made annoying complicated by the need to be compatible with `serde` and +//! more-or-less consistent with its corner cases. Particularly tricky or non-obvious (to +//! the author at time of writing) aspects of macro writing are recorded here for posterity. +//! +//! ## Generic Bounds +//! +//! If we have a struct like +//! ``` +//! struct Container(T); +//! +//! struct Foo<'a, T, U> { +//! baz: Container, +//! bar: String, +//! buz: U, +//! } +//! ``` +//! And we want to derive `Reflect`, we do so by recursing through all fields. For fields +//! like `buz`, this will fail unless we add a bound `U: Reflect` to `Foo`'s implementation +//! of `Reflect`. In addition, we don't know that types like `Container` implement +//! `Reflect`. A simple bound like `T: Reflect` doesn't work because `Container`'s +//! implementation may not follow from `T`. So we need to add the following +//! ```text +//! impl Reflect for Foo<'_, T, U> +//! where +//! T: Reflect, +//! U: Reflect, +//! Containter: Reflect, +//! String: Reflect, +//! ``` +//! We can do this by walking the generics for `Foo` as well as all field types and adding +//! bounds for all discovered items. Note that this can result in duplicates, but that's +//! fine. +//! +//! The struct [`Generics`] is responsible for this aggregation and for creating the expanded +//! where-clause. +//! +//! ## Renaming +//! +//! The attributes `rename_all` and `rename` are supported. `rename_all` can go on: +//! +//! * Structs: All fields are renamed according to the rule. Tuple-like and unit structs +//! do not support this attribute and should return an error. +//! +//! [`Fields::parse`] enforces this. +//! +//! * Enums: All variants are renamed according to the rule. +//! +//! * Variants: Struct-like will have their fields renamed. Tuple-like and unit structs +//! do not support this attribute and should return an error. +//! +//! Attribute `rename` can go on struct fields, variants, and fields of struct-like variants. +//! When present, it supersedes an enclosing `rename_all`. +//! +//! ## New Types +//! +//! Serde treats tuple-like structs with a single member as special new-types. These are not +//! modeled as single-element containers and instead modeled as transparent types. +//! +//! We don't have to do anything special during parsing to handle this. However, during code +//! emission, [`Fields::Unnamed`] needs to emit a `NewType` when there is just a single field. +//! +//! ## Enum Woes +//! +//! Serde supports four different enum tagging mechanism. We support the following three: +//! +//! ### Externally tagged +//! +//! Externally tagged enums are modeled as +//! ```json +//! "discriminant": { +//! "key0": 0, +//! "key1": 1, +//! } +//! ``` +//! where "discriminant" is the name for the active enum variant. This is relatively easy +//! to support and doesn't require additional legalization checks. +//! +//! External tagging is the default for enum when neither `serde(tag = "")` nor +//! `serde(content = "") attributes are provided. +//! +//! ### Adjacently Tagged +//! +//! Adjacently tagged enums look like +//! ```json +//! { +//! "tag": "discriminant", +//! "content": { +//! "key0": 0, +//! "key1": 1, +//! } +//! } +//! ``` +//! +//! This gets selected when both the "tag" and "content" serde attributes are provided. The +//! values for these attributes are used as the actual strings in the "tag" and "content" +//! positions in the example JSON. +//! +//! Like externally tagged enums, there are no additional legalization steps. +//! +//! ### Internally Tagged +//! +//! These are evil. These look like +//! ```text +//! { +//! "tag": "discriminant", <--- Embedded with the content +//! "key0": 0, +//! "key1": 1, +//! } +//! ``` +//! This gets selected when just the `serde(tag = "")` attribute is applied to an enum, with +//! the actual "tag" string in the sample replaced with the value of said attribute. +//! +//! There are several additional checks that internally tagged enums need beyond the other two: +//! +//! * Internal Tagging can only be used if *all* variants are either new-types, struct-like, +//! or units. Tuple-like variants don't have any named fields, so a tag doesn't work. +//! +//! Serde supports new-types but throws a runtime error if the new-type is not struct-like. +//! We don't throw this run time error - we just assume it's going to work. +//! +//! * We need to check that the "tag" string does not conflict with any field of the +//! struct-like variant. This needs to happen *after* all renames have been applied. +//! +//! ### Untagged +//! +//! We don't support untagged enums. +//! +//! ## Name Uniqueness +//! +//! We also need to ensure that after renaming, all final field names are unique and all +//! final variant names are unique. +//! +//! ## Type Names +//! +//! Type names are what we use to provide unique names for reflected types. Two attributes +//! are supported: +//! +//! * `reflect(prefix = "")`: Apply the prefix to the type-name. +//! * `reflect(type_name = "")`: Completely rewrite the type name. +//! +//! Since type registries assume that type names are unique, we don't allow `type_name` to +//! be applied to types with generics. Otherwise, each instantiation of the generic would +//! have the same name, limiting usefulness. +//! + +use std::collections::HashSet; + +use proc_macro2::{Span, TokenStream}; +use quote::{quote, quote_spanned}; +use syn::{parse_quote, spanned::Spanned}; + +use super::attributes; + +pub(crate) struct Input<'a> { + type_name: &'a syn::Ident, + generics: Generics<'a>, + format_type_name: TypeName<'a>, + container: Container<'a>, +} + +impl<'a> Input<'a> { + pub(crate) fn parse(input: &'a syn::DeriveInput) -> syn::Result { + let attrs = attributes::Container::parse(&input.attrs)?; + + let type_name = &input.ident; + let format_type_name = TypeName::parse(type_name, &input.generics, attrs.type_name())?; + + // Parse the type tree for the container. + let container = Container::parse(&input.data, attrs, input.span())?; + let generics = Generics::parse(&input.generics, &container); + + Ok(Self { + type_name, + generics, + format_type_name, + container, + }) + } + + pub(crate) fn emit(&self, path: &syn::Path) -> TokenStream { + let type_name = self.type_name; + let generics = self.generics.emit(path); + let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); + let format_type_name = self.format_type_name.emit(path); + let container = self.container.emit(path); + + quote! { + impl #impl_generics #path::Reflect for #type_name #ty_generics #where_clause { + fn ty() -> #path::Type { + #container + } + + fn format_type_name(f: &mut dyn ::std::fmt::Write) -> ::std::fmt::Result { + #format_type_name + } + } + } + } +} + +/// Emit a `where` clause predicate for all generics and field types. +struct Generics<'a> { + generics: &'a syn::Generics, + /// The type parameters of the original object. + params: Vec<&'a syn::Ident>, + /// All field types. + types: Vec<&'a syn::Type>, +} + +impl<'a> Generics<'a> { + fn parse(generics: &'a syn::Generics, container: &Container<'a>) -> Self { + let params = generics.type_params().map(|p| &p.ident).collect(); + let mut types = Vec::new(); + container.for_each_type(|ty| types.push(ty)); + Self { + generics, + params, + types, + } + } + + fn emit(&self, path: &syn::Path) -> syn::Generics { + let mut generics = self.generics.clone(); + + let predicates = &mut generics.make_where_clause().predicates; + for ident in self.params.iter() { + predicates.push(parse_quote!(#ident: #path::Reflect)); + } + + for ty in self.types.iter() { + predicates.push(parse_quote!(#ty: #path::Reflect)); + } + + generics + } +} + +enum Container<'a> { + Struct(Struct<'a>), + Enum(Enum<'a>), +} + +impl<'a> Container<'a> { + fn parse(data: &'a syn::Data, attrs: attributes::Container, span: Span) -> syn::Result { + match data { + syn::Data::Struct(s) => Ok(Self::Struct(Struct::parse(s, attrs.try_as_struct()?)?)), + syn::Data::Enum(e) => Ok(Self::Enum(Enum::parse(e, attrs.into_enum())?)), + syn::Data::Union(_) => Err(syn::Error::new( + span, + "Reflect cannot be derived for unions", + )), + } + } + + pub(crate) fn for_each_type(&self, f: F) + where + F: FnMut(&'a syn::Type), + { + match self { + Self::Struct(s) => s.for_each_type(f), + Self::Enum(e) => e.for_each_type(f), + } + } + + fn emit(&self, path: &syn::Path) -> TokenStream { + match self { + Self::Struct(s) => s.emit(path), + Self::Enum(e) => e.emit(path), + } + } +} + +enum TypeName<'a> { + Override(syn::LitStr), + Generate { + prefix: Option, + raw_name: &'a syn::Ident, + generics: Vec>, + }, +} + +impl<'a> TypeName<'a> { + fn parse( + raw_name: &'a syn::Ident, + generics: &'a syn::Generics, + type_name: attributes::TypeName, + ) -> syn::Result { + let generics: Vec<_> = generics + .params + .iter() + .filter_map(|p| match p { + syn::GenericParam::Type(p) => Some(Generic::Type(&p.ident)), + syn::GenericParam::Const(p) => Some(Generic::Const(&p.ident)), + syn::GenericParam::Lifetime(_) => None, + }) + .collect(); + + let prefix = match type_name { + attributes::TypeName::Rename(rename) => { + // Reject `reflect(type_name = "...")` on types with generic parameters. + if !generics.is_empty() { + return Err(syn::Error::new_spanned( + rename, + "The `type_name` attribute cannot be applied to types with generics", + )); + } else { + return Ok(Self::Override(rename)); + }; + } + attributes::TypeName::Prefix(prefix) => Some(prefix), + attributes::TypeName::None => None, + }; + + Ok(Self::Generate { + prefix, + raw_name, + generics, + }) + } + + /// When generics are absent, do the easy thing and apply the prefix with the raw type + /// name from the [`syn::DeriveInput`]. + /// + /// When generics are present (for example) + /// ``` + /// struct Foo { + /// // ... + /// } + /// ``` + /// use the type-names for each generic to get the final type. The generated code roughly + /// ```text + /// f.write_str("Foo")?; + /// f.write_str("<")?; + /// ::type_name(f)?; // The type name for the generic + /// f.write_str(", ")?; + /// ::type_name(f)?; // The type name for the generic + /// f.write_str(", ")?; + /// f.write!("{}", N)?; + /// f.write_str(">")?; + /// ``` + fn emit(&self, path: &syn::Path) -> TokenStream { + match self { + Self::Override(type_name) => quote!(f.write_str(#type_name)), + Self::Generate { + prefix, + raw_name, + generics, + } => { + let raw_name = raw_name.to_string(); + let prefix = prefix.as_ref().map(|p| quote!(f.write_str(#p)?;)); + if generics.is_empty() { + quote! { + #prefix + f.write_str(#raw_name)?; + Ok(()) + } + } else { + let generics = generics.iter().enumerate().map(|(i, g)| { + let separator = if i == 0 { + None + } else { + Some(quote!(f.write_str(", ")?;)) + }; + let ts = g.emit(path); + quote! { + #separator + #ts + } + }); + + quote! { + #prefix + f.write_str(#raw_name)?; + f.write_str("<")?; + #(#generics)* + f.write_str(">")?; + Ok(()) + } + } + } + } + } +} + +enum Generic<'a> { + Type(&'a syn::Ident), + Const(&'a syn::Ident), +} + +impl Generic<'_> { + fn emit(&self, path: &syn::Path) -> TokenStream { + match self { + Self::Type(ident) => quote! { + ::std::write!(f, "{}", #path::Reflection::new::<#ident>().type_name())?; + }, + Self::Const(ident) => quote!(::std::write!(f, "{}", #ident)?;), + } + } +} + +//////////// +// Struct // +//////////// + +pub(crate) struct Struct<'a> { + doc: attributes::Doc, + fields: Fields<'a>, +} + +impl<'a> Struct<'a> { + pub(crate) fn parse(s: &'a syn::DataStruct, attrs: attributes::Struct) -> syn::Result { + let attributes::Struct { rename_all, doc } = attrs; + let fields = Fields::parse(&s.fields, rename_all)?; + Ok(Self { doc, fields }) + } + + pub(crate) fn for_each_type(&self, f: F) + where + F: FnMut(&'a syn::Type), + { + self.fields.for_each_type(f) + } + + pub(crate) fn emit(&self, path: &syn::Path) -> TokenStream { + let doc = self.doc.emit(); + let fields = self.fields.emit(path); + quote! { + #path::Type::aggregate( + #fields, + #doc + ) + } + } +} + +////////// +// Enum // +////////// + +pub(crate) struct Enum<'a> { + doc: attributes::Doc, + enum_repr: attributes::EnumRepr, + variants: Vec>, +} + +impl<'a> Enum<'a> { + pub(crate) fn parse(enum_: &'a syn::DataEnum, attr: attributes::Enum) -> syn::Result { + let attributes::Enum { + rename_all, + enum_repr, + doc, + } = attr; + + // Parse all variants. + let variants = enum_ + .variants + .iter() + .map(|v| Variant::parse(v, rename_all)) + .collect::>>()?; + + // Ensure all variant names are unique. + let mut seen_variants = HashSet::new(); + for v in variants.iter() { + let variant_name = v.name.value(); + if !seen_variants.insert(variant_name.clone()) { + return Err(syn::Error::new_spanned( + &v.name, + format!("Variant name \"{}\" seen more than once", variant_name), + )); + } + } + + // Ensure that: + // + // 1. No field name conflicts with the tag. + // 2. Non-newtype tuple variants are rejected. + if let attributes::EnumRepr::Internal { tag } = &enum_repr { + let tag = tag.value(); + for v in variants.iter() { + match &v.fields { + Fields::Named(named) => { + for field in named.iter() { + if tag == field.name.value() { + return Err(syn::Error::new_spanned( + &field.name, + format!( + "field \"{}\" conflicts with internal discriminant tag", + tag + ), + )); + } + } + } + Fields::Unnamed(unnamed) => { + if unnamed.len() != 1 { + return Err(syn::Error::new_spanned( + &v.name, + "non-newtype tuple type variants are not allowed with internal tagging", + )); + } + } + Fields::Unit => {} + } + } + } + + Ok(Self { + doc, + enum_repr, + variants, + }) + } + + pub(crate) fn for_each_type(&self, mut f: F) + where + F: FnMut(&'a syn::Type), + { + self.variants.iter().for_each(|v| v.for_each_type(&mut f)) + } + + pub(crate) fn emit(&self, path: &syn::Path) -> TokenStream { + let enum_repr = match &self.enum_repr { + attributes::EnumRepr::External => quote!(#path::tree::EnumRepr::External), + attributes::EnumRepr::Internal { tag } => { + quote!(#path::tree::EnumRepr::Internal { tag: #tag }) + } + attributes::EnumRepr::Adjacent { tag, content } => { + quote!(#path::tree::EnumRepr::Adjacent { tag: #tag, content: #content }) + } + }; + + let variants = self.variants.iter().map(|v| v.emit(path)); + let doc = self.doc.emit(); + quote! { + #path::Type::enum_( + #enum_repr, + [#(#variants),*], + #doc + ) + } + } +} + +//////////////////// +// Implementation // +//////////////////// + +pub(crate) struct UnnamedField<'a> { + doc: attributes::Doc, + ty: &'a syn::Type, +} + +impl<'a> UnnamedField<'a> { + fn parse(field: &'a syn::Field) -> syn::Result { + assert!( + field.ident.is_none(), + "expected `syn` to not attach names to `FieldsNamed` entries" + ); + + let attributes::Field { rename_field, doc } = attributes::Field::parse(&field.attrs)?; + + // Rejact renames on unnamed fields. + if let Some(span) = rename_field.span() { + return Err(syn::Error::new( + span, + "rename attributes are not supported on unnamed fields", + )); + } + + Ok(Self { doc, ty: &field.ty }) + } + + fn emit(&self, path: &syn::Path) -> TokenStream { + let ty = self.ty; + let doc = self.doc.emit(); + quote_spanned! { ty.span()=> #path::tree::UnnamedField::new::<#ty>(#doc) } + } +} + +pub(crate) struct NamedField<'a> { + doc: attributes::Doc, + name: syn::LitStr, + ty: &'a syn::Type, +} + +impl<'a> NamedField<'a> { + fn parse(field: &'a syn::Field, rename_all: attributes::RenameAll) -> syn::Result { + let ident = field + .ident + .as_ref() + .expect("named fields should have identifiers"); + + let name = syn::LitStr::new(strip_raw_prefix(&ident.to_string()), ident.span()); + let attributes::Field { rename_field, doc } = attributes::Field::parse(&field.attrs)?; + let name = rename_field.apply_to_field(name, rename_all); + + Ok(Self { + doc, + name, + ty: &field.ty, + }) + } + + fn emit(&self, path: &syn::Path) -> TokenStream { + let ty = self.ty; + let doc = self.doc.emit(); + let name = &self.name; + quote_spanned! { + ty.span()=> #path::tree::NamedField::new::<#ty>(#name, #doc) + } + } +} + +pub(crate) enum Fields<'a> { + Named(Vec>), + Unnamed(Vec>), + Unit, +} + +impl<'a> Fields<'a> { + fn parse(fields: &'a syn::Fields, rename_all: attributes::RenameAll) -> syn::Result { + let me = match fields { + syn::Fields::Named(fields) => { + let named = fields + .named + .iter() + .map(|f| NamedField::parse(f, rename_all)) + .collect::>>()?; + + let mut seen = HashSet::new(); + for f in named.iter() { + let field_name = f.name.value(); + if !seen.insert(field_name.clone()) { + return Err(syn::Error::new_spanned( + &f.name, + format!("field name \"{}\" seen more than once", field_name), + )); + } + } + + Self::Named(named) + } + syn::Fields::Unnamed(fields) => { + if let Some(span) = rename_all.span_if_present() { + return Err(syn::Error::new( + span, + "`rename_all` cannot be applied to tuple structs", + )); + } + + Self::Unnamed( + fields + .unnamed + .iter() + .map(UnnamedField::parse) + .collect::>()?, + ) + } + syn::Fields::Unit => { + if let Some(span) = rename_all.span_if_present() { + return Err(syn::Error::new( + span, + "`rename_all` cannot be applied to unit structs", + )); + } + Self::Unit + } + }; + Ok(me) + } + + fn for_each_type(&self, mut f: F) + where + F: FnMut(&'a syn::Type), + { + match self { + Self::Named(named) => named.iter().for_each(|field| f(field.ty)), + Self::Unnamed(unnamed) => unnamed.iter().for_each(|field| f(field.ty)), + Self::Unit => {} + } + } + + fn emit(&self, path: &syn::Path) -> TokenStream { + match self { + Self::Named(named) => { + let ts = named.iter().map(|f| f.emit(path)); + quote!(#path::tree::Fields::Named(vec![#(#ts),*])) + } + Self::Unnamed(unnamed) => { + if unnamed.len() == 1 { + let ts = unnamed[0].emit(path); + quote!(#path::tree::Fields::NewType(#ts)) + } else { + let ts = unnamed.iter().map(|f| f.emit(path)); + quote!(#path::tree::Fields::Unnamed(vec![#(#ts),*])) + } + } + Self::Unit => quote!(#path::tree::Fields::Unit), + } + } +} + +struct Variant<'a> { + doc: attributes::Doc, + name: syn::LitStr, + fields: Fields<'a>, +} + +impl<'a> Variant<'a> { + fn parse(variant: &'a syn::Variant, rename_all: attributes::RenameAll) -> syn::Result { + let attributes::Variant { + rename_variant, + rename_variant_fields, + doc, + } = attributes::Variant::parse(&variant.attrs)?; + + let fields = Fields::parse(&variant.fields, rename_variant_fields)?; + + let name = syn::LitStr::new( + strip_raw_prefix(&variant.ident.to_string()), + variant.ident.span(), + ); + let name = rename_variant.apply_to_variant(name, rename_all); + + Ok(Self { doc, name, fields }) + } + + fn for_each_type(&self, f: F) + where + F: FnMut(&'a syn::Type), + { + self.fields.for_each_type(f) + } + + fn emit(&self, path: &syn::Path) -> TokenStream { + let name = &self.name; + let fields = self.fields.emit(path); + let doc = self.doc.emit(); + quote!(#path::tree::Variant::new(#name, #fields, #doc)) + } +} + +//-----// +// raw // +//-----// + +fn strip_raw_prefix(s: &str) -> &str { + s.strip_prefix("r#").unwrap_or(s) +} diff --git a/diskann-macros/tests/compile.rs b/diskann-macros/tests/compile.rs new file mode 100644 index 0000000000..f79da56766 --- /dev/null +++ b/diskann-macros/tests/compile.rs @@ -0,0 +1,11 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +#[test] +fn compile_tests() { + let t = trybuild::TestCases::new(); + t.pass("tests/ui/pass/*.rs"); + t.compile_fail("tests/ui/fail/*.rs"); +} diff --git a/diskann-macros/tests/ui/fail/name_collisions.rs b/diskann-macros/tests/ui/fail/name_collisions.rs new file mode 100644 index 0000000000..efe7018b5a --- /dev/null +++ b/diskann-macros/tests/ui/fail/name_collisions.rs @@ -0,0 +1,50 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +#[serde(rename_all = "lowercase")] +enum DuplicateVariant { + Foo, + #[serde(rename = "foo")] + Other, +} + +#[derive(Reflect)] +struct DuplicateStructField { + value: usize, + #[serde(rename = "value")] + other: usize, +} + +#[derive(Reflect)] +#[serde(rename_all = "kebab-case")] +struct DuplicateStructFieldAfterRenameAll { + #[serde(rename = "foo-bar")] + other: usize, + foo_bar: usize, +} + +#[derive(Reflect)] +enum DuplicateVariantField { + Value { + value: usize, + #[serde(rename = "value")] + other: usize, + }, +} + +#[derive(Reflect)] +enum DuplicateVariantFieldAfterRenameAll { + #[serde(rename_all = "kebab-case")] + Value { + #[serde(rename = "foo-bar")] + other: usize, + foo_bar: usize, + }, +} + +fn main() {} diff --git a/diskann-macros/tests/ui/fail/name_collisions.stderr b/diskann-macros/tests/ui/fail/name_collisions.stderr new file mode 100644 index 0000000000..4258b9dc6d --- /dev/null +++ b/diskann-macros/tests/ui/fail/name_collisions.stderr @@ -0,0 +1,29 @@ +error: Variant name "foo" seen more than once + --> tests/ui/fail/name_collisions.rs:12:22 + | +12 | #[serde(rename = "foo")] + | ^^^^^ + +error: field name "value" seen more than once + --> tests/ui/fail/name_collisions.rs:19:22 + | +19 | #[serde(rename = "value")] + | ^^^^^^^ + +error: field name "foo-bar" seen more than once + --> tests/ui/fail/name_collisions.rs:28:5 + | +28 | foo_bar: usize, + | ^^^^^^^ + +error: field name "value" seen more than once + --> tests/ui/fail/name_collisions.rs:35:26 + | +35 | #[serde(rename = "value")] + | ^^^^^^^ + +error: field name "foo-bar" seen more than once + --> tests/ui/fail/name_collisions.rs:46:9 + | +46 | foo_bar: usize, + | ^^^^^^^ diff --git a/diskann-macros/tests/ui/fail/reflect_attributes.rs b/diskann-macros/tests/ui/fail/reflect_attributes.rs new file mode 100644 index 0000000000..c7487e7e08 --- /dev/null +++ b/diskann-macros/tests/ui/fail/reflect_attributes.rs @@ -0,0 +1,54 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +#[reflect(prefix = "one::", prefix = "two::")] +struct DuplicatePrefix; + +#[derive(Reflect)] +#[reflect(type_name = "One", type_name = "Two")] +struct DuplicateTypeName; + +#[derive(Reflect)] +#[reflect(prefix = "benchmark::", type_name = "Renamed")] +struct PrefixAndTypeName; + +#[derive(Reflect)] +#[reflect(type_name = "Generic")] +struct GenericTypeName { + value: T, +} + +#[derive(Reflect)] +#[reflect(rename = "Unsupported")] +struct UnsupportedContainerReflectAttribute; + +#[derive(Reflect)] +struct UnsupportedFieldReflectAttribute { + #[reflect(rename = "value")] + value: usize, +} + +#[derive(Reflect)] +struct UnsupportedFieldReflectAttributeEmpty { + #[reflect] + value: usize, +} + +#[derive(Reflect)] +enum UnsupportedVariantReflectAttribute { + #[reflect(rename = "value")] + Value, +} + +#[derive(Reflect)] +enum UnsupportedVariantReflectAttributeEmpty { + #[reflect] + Value, +} + +fn main() {} diff --git a/diskann-macros/tests/ui/fail/reflect_attributes.stderr b/diskann-macros/tests/ui/fail/reflect_attributes.stderr new file mode 100644 index 0000000000..ad16b9926c --- /dev/null +++ b/diskann-macros/tests/ui/fail/reflect_attributes.stderr @@ -0,0 +1,53 @@ +error: reflect attribute `prefix` found multiple times + --> tests/ui/fail/reflect_attributes.rs:9:38 + | +9 | #[reflect(prefix = "one::", prefix = "two::")] + | ^^^^^^^ + +error: reflect attribute `type_name` found multiple times + --> tests/ui/fail/reflect_attributes.rs:13:42 + | +13 | #[reflect(type_name = "One", type_name = "Two")] + | ^^^^^ + +error: reflect attributes `prefix` and `type_name` are mutually exclusive + --> tests/ui/fail/reflect_attributes.rs:17:47 + | +17 | #[reflect(prefix = "benchmark::", type_name = "Renamed")] + | ^^^^^^^^^ + +error: The `type_name` attribute cannot be applied to types with generics + --> tests/ui/fail/reflect_attributes.rs:21:23 + | +21 | #[reflect(type_name = "Generic")] + | ^^^^^^^^^ + +error: unsupported attribute for Reflect + --> tests/ui/fail/reflect_attributes.rs:27:11 + | +27 | #[reflect(rename = "Unsupported")] + | ^^^^^^ + +error: Reflect does not support field-level attributes + --> tests/ui/fail/reflect_attributes.rs:32:5 + | +32 | #[reflect(rename = "value")] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: Reflect does not support field-level attributes + --> tests/ui/fail/reflect_attributes.rs:38:5 + | +38 | #[reflect] + | ^^^^^^^^^^ + +error: Reflect does not support variant-level attributes + --> tests/ui/fail/reflect_attributes.rs:44:5 + | +44 | #[reflect(rename = "value")] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +error: Reflect does not support variant-level attributes + --> tests/ui/fail/reflect_attributes.rs:50:5 + | +50 | #[reflect] + | ^^^^^^^^^^ diff --git a/diskann-macros/tests/ui/fail/serde_attributes.rs b/diskann-macros/tests/ui/fail/serde_attributes.rs new file mode 100644 index 0000000000..e3728f3fe4 --- /dev/null +++ b/diskann-macros/tests/ui/fail/serde_attributes.rs @@ -0,0 +1,81 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +#[serde(rename_all = "snake_case", rename_all = "kebab-case")] +struct DuplicateRenameAll { + field_name: usize, +} + +#[derive(Reflect)] +struct DuplicateFieldRename { + #[serde(rename = "first", rename = "second")] + value: usize, +} + +#[derive(Reflect)] +enum DuplicateVariantRename { + #[serde(rename = "first", rename = "second")] + Unit, +} + +#[derive(Reflect)] +#[serde(rename_all = "camelCase")] +struct UnsupportedRenameAll { + field_name: usize, +} + +#[derive(Reflect)] +#[serde(default)] +struct UnsupportedContainerAttribute { + value: usize, +} + +#[derive(Reflect)] +struct UnsupportedFieldAttribute { + #[serde(default)] + value: usize, +} + +#[derive(Reflect)] +enum UnsupportedVariantAttribute { + #[serde(skip)] + Unit, +} + +#[derive(Reflect)] +struct RenameUnnamedField(#[serde(rename = "value")] usize); + +#[derive(Reflect)] +#[serde(rename_all = "snake_case")] +struct RenameAllTupleStruct(usize); + +#[derive(Reflect)] +#[serde(rename_all = "snake_case")] +struct RenameAllUnitStruct; + +#[derive(Reflect)] +enum DuplicateVariantRenameAll { + #[serde(rename_all = "snake_case", rename_all = "snake_case")] + Value { + foo: usize, + } +} + +#[derive(Reflect)] +enum RenameAllTupleVariant { + #[serde(rename_all = "snake_case")] + Value(usize), +} + +#[derive(Reflect)] +enum RenameAllUnitVariant { + #[serde(rename_all = "snake_case")] + Value, +} + +fn main() {} diff --git a/diskann-macros/tests/ui/fail/serde_attributes.stderr b/diskann-macros/tests/ui/fail/serde_attributes.stderr new file mode 100644 index 0000000000..b82392d40c --- /dev/null +++ b/diskann-macros/tests/ui/fail/serde_attributes.stderr @@ -0,0 +1,77 @@ +error: serde attribute `rename_all` found multiple times + --> tests/ui/fail/serde_attributes.rs:9:49 + | +9 | #[serde(rename_all = "snake_case", rename_all = "kebab-case")] + | ^^^^^^^^^^^^ + +error: serde attribute `rename` found multiple times + --> tests/ui/fail/serde_attributes.rs:16:40 + | +16 | #[serde(rename = "first", rename = "second")] + | ^^^^^^^^ + +error: serde attribute `rename` found multiple times + --> tests/ui/fail/serde_attributes.rs:22:40 + | +22 | #[serde(rename = "first", rename = "second")] + | ^^^^^^^^ + +error: unsupported serde `rename_all` rule "camelCase" - expected one of "lowercase", "snake_case", or "kebab-case" + --> tests/ui/fail/serde_attributes.rs:27:22 + | +27 | #[serde(rename_all = "camelCase")] + | ^^^^^^^^^^^ + +error: unsupported Serde attribute for Reflect + --> tests/ui/fail/serde_attributes.rs:33:9 + | +33 | #[serde(default)] + | ^^^^^^^ + +error: unsupported Serde attribute for Reflect + --> tests/ui/fail/serde_attributes.rs:40:13 + | +40 | #[serde(default)] + | ^^^^^^^ + +error: unsupported Serde attribute for Reflect + --> tests/ui/fail/serde_attributes.rs:46:13 + | +46 | #[serde(skip)] + | ^^^^ + +error: rename attributes are not supported on unnamed fields + --> tests/ui/fail/serde_attributes.rs:51:44 + | +51 | struct RenameUnnamedField(#[serde(rename = "value")] usize); + | ^^^^^^^ + +error: `rename_all` cannot be applied to tuple structs + --> tests/ui/fail/serde_attributes.rs:54:22 + | +54 | #[serde(rename_all = "snake_case")] + | ^^^^^^^^^^^^ + +error: `rename_all` cannot be applied to unit structs + --> tests/ui/fail/serde_attributes.rs:58:22 + | +58 | #[serde(rename_all = "snake_case")] + | ^^^^^^^^^^^^ + +error: serde attribute `rename_all` found multiple times + --> tests/ui/fail/serde_attributes.rs:63:53 + | +63 | #[serde(rename_all = "snake_case", rename_all = "snake_case")] + | ^^^^^^^^^^^^ + +error: `rename_all` cannot be applied to tuple structs + --> tests/ui/fail/serde_attributes.rs:71:26 + | +71 | #[serde(rename_all = "snake_case")] + | ^^^^^^^^^^^^ + +error: `rename_all` cannot be applied to unit structs + --> tests/ui/fail/serde_attributes.rs:77:26 + | +77 | #[serde(rename_all = "snake_case")] + | ^^^^^^^^^^^^ diff --git a/diskann-macros/tests/ui/fail/tagging.rs b/diskann-macros/tests/ui/fail/tagging.rs new file mode 100644 index 0000000000..4845206d1b --- /dev/null +++ b/diskann-macros/tests/ui/fail/tagging.rs @@ -0,0 +1,78 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +#[serde(content = "value")] +enum ContentWithoutTag { + Unit, +} + +#[derive(Reflect)] +#[serde(tag = "kind", content = "first", content = "second")] +enum DuplicateContent { + Unit, +} + +#[derive(Reflect)] +#[serde(tag = "first", tag = "second")] +enum DuplicateTag { + Unit, +} + +#[derive(Reflect)] +#[serde(tag = "same", content = "same")] +enum SameTagAndContent { + Unit, +} + +#[derive(Reflect)] +#[serde(tag = "kind")] +struct TagOnStruct { + value: usize, +} + +#[derive(Reflect)] +#[serde(tag = "kind", content = "value")] +struct TagAndContentOnStruct { + value: usize, +} + +#[derive(Reflect)] +#[serde(tag = "tag")] +enum InternallyTaggedEmptyTuple { + Tuple(), +} + +#[derive(Reflect)] +#[serde(tag = "tag")] +enum InternallyTaggedTuple { + Tuple(usize, usize), +} + +#[derive(Reflect)] +#[serde(tag = "kind")] +enum DirectTagConflict { + Value { kind: usize }, +} + +#[derive(Reflect)] +#[serde(tag = "kind")] +enum RenamedTagConflict { + Value { + #[serde(rename = "kind")] + value: usize, + }, +} + +#[derive(Reflect)] +#[serde(tag = "foo-bar")] +enum RenameAllTagConflict { + #[serde(rename_all = "kebab-case")] + Value { foo_bar: usize }, +} + +fn main() {} diff --git a/diskann-macros/tests/ui/fail/tagging.stderr b/diskann-macros/tests/ui/fail/tagging.stderr new file mode 100644 index 0000000000..f7392f043c --- /dev/null +++ b/diskann-macros/tests/ui/fail/tagging.stderr @@ -0,0 +1,65 @@ +error: serde attribute `content` provided without a `tag` + --> tests/ui/fail/tagging.rs:9:19 + | +9 | #[serde(content = "value")] + | ^^^^^^^ + +error: serde attribute `content` found multiple times + --> tests/ui/fail/tagging.rs:15:52 + | +15 | #[serde(tag = "kind", content = "first", content = "second")] + | ^^^^^^^^ + +error: serde attribute `tag` found multiple times + --> tests/ui/fail/tagging.rs:21:30 + | +21 | #[serde(tag = "first", tag = "second")] + | ^^^^^^^^ + +error: adjacently tagged `content` conflicts with `tag` + --> tests/ui/fail/tagging.rs:27:33 + | +27 | #[serde(tag = "same", content = "same")] + | ^^^^^^ + +error: serde attribute `tag` provided on a non-enum + --> tests/ui/fail/tagging.rs:33:15 + | +33 | #[serde(tag = "kind")] + | ^^^^^^ + +error: serde attributes `tag` and `content` provided on a non-enum + --> tests/ui/fail/tagging.rs:39:15 + | +39 | #[serde(tag = "kind", content = "value")] + | ^^^^^^ + +error: non-newtype tuple type variants are not allowed with internal tagging + --> tests/ui/fail/tagging.rs:47:5 + | +47 | Tuple(), + | ^^^^^ + +error: non-newtype tuple type variants are not allowed with internal tagging + --> tests/ui/fail/tagging.rs:53:5 + | +53 | Tuple(usize, usize), + | ^^^^^ + +error: field "kind" conflicts with internal discriminant tag + --> tests/ui/fail/tagging.rs:59:13 + | +59 | Value { kind: usize }, + | ^^^^ + +error: field "kind" conflicts with internal discriminant tag + --> tests/ui/fail/tagging.rs:66:26 + | +66 | #[serde(rename = "kind")] + | ^^^^^^ + +error: field "foo-bar" conflicts with internal discriminant tag + --> tests/ui/fail/tagging.rs:75:13 + | +75 | Value { foo_bar: usize }, + | ^^^^^^^ diff --git a/diskann-macros/tests/ui/fail/unsupported_union.rs b/diskann-macros/tests/ui/fail/unsupported_union.rs new file mode 100644 index 0000000000..3b0e43c96e --- /dev/null +++ b/diskann-macros/tests/ui/fail/unsupported_union.rs @@ -0,0 +1,13 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +union UnsupportedUnion { + value: usize, +} + +fn main() {} diff --git a/diskann-macros/tests/ui/fail/unsupported_union.stderr b/diskann-macros/tests/ui/fail/unsupported_union.stderr new file mode 100644 index 0000000000..e0d1b0acec --- /dev/null +++ b/diskann-macros/tests/ui/fail/unsupported_union.stderr @@ -0,0 +1,5 @@ +error: Reflect cannot be derived for unions + --> tests/ui/fail/unsupported_union.rs:9:1 + | +9 | union UnsupportedUnion { + | ^^^^^ diff --git a/diskann-macros/tests/ui/pass/supported.rs b/diskann-macros/tests/ui/pass/supported.rs new file mode 100644 index 0000000000..b668392561 --- /dev/null +++ b/diskann-macros/tests/ui/pass/supported.rs @@ -0,0 +1,31 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_runner::Reflect; + +#[derive(Reflect)] +#[reflect(prefix = "benchmark::")] +#[serde(rename_all = "kebab-case", tag = "kind", content = "value")] +enum Generic { + Unit, + #[serde(rename = "tuple")] + Tuple(T, Vec), + #[serde(rename_all = "snake_case")] + Struct { + #[serde(rename = "renamed")] + field_name: T, + }, +} + +#[derive(Reflect)] +#[reflect(type_name = "CompleteOverride")] +struct Renamed { + value: usize, +} + +fn main() { + let _ = as Reflect>::ty(); + let _ = ::ty(); +}