editoast: use uom to type units

editoast/Cargo.toml

@@ -78,6 +78,7 @@ tracing = { version = "0.1.41", default-features = false, features = [
   "attributes",
   "log",
 ] }
+uom = { version = "0.36.0", default-features = false, features = ["f64", "si"] }
 url = { version = "2.5.4", features = ["serde"] }
 urlencoding = "2.1.3"
 utoipa = { version = "4.2.3", features = ["chrono", "uuid"] }
@@ -186,6 +187,7 @@ tracing-opentelemetry = { version = "0.28.0", default-features = false, features
   "tracing-log",
 ] }
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
+uom.workspace = true
 url.workspace = true
 utoipa.workspace = true
 uuid.workspace = true

editoast/editoast_common/Cargo.toml

@@ -11,6 +11,7 @@ rangemap.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
 thiserror.workspace = true
+uom.workspace = true
 utoipa.workspace = true
 
 [lints]

editoast/editoast_common/src/lib.rs

@@ -2,6 +2,7 @@ pub mod geometry;
 mod hash_rounded_float;
 pub mod rangemap_utils;
 pub mod schemas;
+pub mod units;
 
 pub use hash_rounded_float::hash_float;
 pub use hash_rounded_float::hash_float_slice;

editoast/editoast_common/src/units.rs

@@ -0,0 +1,162 @@
+//! Module to allow the use of serde with uom quantities
+//!
+//! The serde feature of uom doesn’t allow to specify in which unit the value will be serialized.
+//!
+//! Two helpers are provided for convenience:
+//! * `unit::new` (e.g. `meter::new(32)`) to build a new quantity form a f64 in the given unit
+//! * `unit::from` (e.g. `millimeter::from(length)`) to have the quantity as f64 in the given unit
+//!
+//! ## Usage
+//!
+//! ```ignore
+//! use editoast_model::units::*;
+//! #[derive(Debug, Serialize, Derivative)]
+//! struct Train {
+//!     // This means that serde with read and write the velocity in meters per second
+//!     #[serde(with="meter_per_second")]
+//!     max_speed: Velocity,
+//!     // When using optional values, we must add `default` and use ::option unit
+//!     // See https://stackoverflow.com/a/44303505
+//!     #[serde(default, with="meter::option")]
+//!     length: Option<Length>,
+//! }
+//!
+//! impl Train {
+//!     fn from_meter_per_seconds(mps: f64) -> Self {
+//!         Self {
+//!             max_speed: meter_per_second::new(mps),
+//!         }
+//!     }
+//!
+//!     fn print(&self) {
+//!         println!("The max speed is: {} km/h", kilometer_per_hour::from(self.max_speed));
+//!     }
+//! }
+//! ```
+
+/// Re-export the Quantities that are used in OSRD
+pub use uom::si::f64::{
+    Acceleration, Force, Frequency, Length, LinearMassDensity, LinearNumberDensity, Mass, MassRate,
+    Time, Velocity,
+};
+
+macro_rules! quantity_to_path {
+    (Length, $unit:ident) => {
+        uom::si::length::$unit
+    };
+    (Velocity, $unit:ident) => {
+        uom::si::velocity::$unit
+    };
+    (Acceleration, $unit:ident) => {
+        uom::si::acceleration::$unit
+    };
+    (Mass, $unit:ident) => {
+        uom::si::mass::$unit
+    };
+    (Force, $unit:ident) => {
+        uom::si::force::$unit
+    };
+    (MassRate, $unit:ident) => {
+        uom::si::mass_rate::$unit
+    };
+    (Frequency, $unit:ident) => {
+        uom::si::frequency::$unit
+    };
+    (LinearMassDensity, $unit:ident) => {
+        uom::si::linear_mass_density::$unit
+    };
+    (LinearNumberDensity, $unit:ident) => {
+        uom::si::linear_number_density::$unit
+    };
+    (Time, $unit:ident) => {
+        uom::si::time::$unit
+    };
+}
+
+macro_rules! define_unit {
+    ($unit:ident, $quantity:ident) => {
+        pub mod $unit {
+            use serde::{Deserialize, Deserializer, Serialize, Serializer};
+            use uom::si::f64::*;
+            type Unit = quantity_to_path!($quantity, $unit);
+            pub type ReprType = f64;
+
+            pub fn serialize<S>(value: &$quantity, serializer: S) -> Result<S::Ok, S::Error>
+            where
+                S: Serializer,
+            {
+                value.get::<Unit>().serialize(serializer)
+            }
+
+            pub fn deserialize<'de, D>(deserializer: D) -> Result<$quantity, D::Error>
+            where
+                D: Deserializer<'de>,
+            {
+                let value = ReprType::deserialize(deserializer)?;
+                Ok($quantity::new::<Unit>(value))
+            }
+
+            pub fn new(value: ReprType) -> $quantity {
+                $quantity::new::<Unit>(value)
+            }
+
+            pub fn from(qty: $quantity) -> ReprType {
+                qty.get::<Unit>()
+            }
+
+            pub fn hash<H: std::hash::Hasher>(value: &$quantity, state: &mut H) {
+                crate::hash_float::<5, H>(&from(*value), state);
+            }
+
+            pub mod option {
+                use super::*;
+                pub type ReprType = Option<super::ReprType>;
+
+                pub fn serialize<S>(
+                    value: &Option<$quantity>,
+                    serializer: S,
+                ) -> Result<S::Ok, S::Error>
+                where
+                    S: Serializer,
+                {
+                    value.map(|value| value.get::<Unit>()).serialize(serializer)
+                }
+
+                pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<$quantity>, D::Error>
+                where
+                    D: Deserializer<'de>,
+                {
+                    let value = Option::deserialize(deserializer)?;
+                    Ok(value.map(|value| $quantity::new::<Unit>(value)))
+                }
+
+                pub fn new(value: ReprType) -> Option<$quantity> {
+                    value.map(|v| $quantity::new::<Unit>(v))
+                }
+
+                pub fn from(qty: Option<$quantity>) -> ReprType {
+                    qty.map(|q| q.get::<Unit>())
+                }
+
+                pub fn hash<H: std::hash::Hasher>(value: &Option<$quantity>, state: &mut H) {
+                    super::hash(&value.unwrap_or_default(), state);
+                }
+            }
+        }
+    };
+}
+
+// Any new value here must also be added in editoast_derive/src/annotate_units.rs
+define_unit!(meter, Length);
+define_unit!(millimeter, Length);
+define_unit!(meter_per_second, Velocity);
+define_unit!(kilometer_per_hour, Velocity);
+define_unit!(meter_per_second_squared, Acceleration);
+define_unit!(kilogram, Mass);
+define_unit!(newton, Force);
+define_unit!(kilogram_per_second, MassRate);
+define_unit!(hertz, Frequency);
+define_unit!(kilogram_per_meter, LinearMassDensity);
+define_unit!(per_meter, LinearNumberDensity);
+define_unit!(second, Time);
+define_unit!(millisecond, Time);

editoast/editoast_derive/src/annotate_units.rs

@@ -0,0 +1,54 @@
+use darling::Result;
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{parse_quote, DeriveInput, LitStr};
+
+pub fn get_abbreviation(value: String) -> Option<&'static str> {
+    // Any new value here must also be added in editoast_common/src/units.rs
+    match value.replace("::option", "").as_str() {
+        "second" => Some("Duration in s"),
+        "millisecond" => Some("Duration in ms"),
+        "meter" => Some("Length in m"),
+        "millimeter" => Some("Length in mm"),
+        "meter_per_second" => Some("Velocity in m·s⁻¹"),
+        "kilometer_per_hour" => Some("Velocity in km·h⁻¹"),
+        "meter_per_second_squared" => Some("Acceleration in m·s⁻²"),
+        "kilogram" => Some("Mass in kg"),
+        "newton" => Some("Force in N"),
+        "hertz" => Some("in s⁻¹"),
+        "kilogram_per_meter" => Some("in kg·m⁻¹"),
+        "kilogram_per_second" => Some("in kg·s⁻¹"),
+        "per_meter" => Some("Aerodynamic drag per kg in m⁻¹"),
+        _ => None,
+    }
+}
+
+pub fn annotate_units(input: &mut DeriveInput) -> Result<TokenStream> {
+    // We look for fields that have #[sered(with="meter")] attributes
+    // and we push two new attributes to it to improve the OpenAPI
+    if let syn::Data::Struct(s) = &mut input.data {
+        for f in s.fields.iter_mut() {
+            for attr in f.attrs.clone() {
+                if attr.path().is_ident("serde") {
+                    let _ = attr.parse_nested_meta(|meta| {
+                        if meta.path.is_ident("with") {
+                            let value = meta.value()?;
+                            let s: LitStr = value.parse()?;
+                            if let Some(abbreviation) = get_abbreviation(s.value()) {
+                                if s.value().ends_with("::option") {
+                                    f.attrs
+                                        .push(parse_quote! {#[schema(value_type = Option<f64>)]});
+                                } else {
+                                    f.attrs.push(parse_quote! {#[schema(value_type = f64)]});
+                                }
+                                f.attrs.push(parse_quote! {#[doc = #abbreviation]});
+                            }
+                        }
+                        Ok(())
+                    });
+                }
+            }
+        }
+    }
+    Ok(quote! {#input})
+}

editoast/editoast_derive/src/lib.rs

@@ -1,5 +1,6 @@
 extern crate proc_macro;
 
+mod annotate_units;
 mod error;
 mod model;
 mod search;
@@ -217,6 +218,7 @@ pub fn search_config_store(input: proc_macro::TokenStream) -> proc_macro::TokenS
 /// * `#[model(to_string)]`: calls `to_string()` before writing the field to the database and calls `String::from` after reading (diesel column type: String)
 /// * `#[model(to_enum)]`: is converted as `u8` before writing the field to the database and calls `FromRepr::from_repr` after reading (diesel column type: TinyInt)
 /// * `#[model(remote = "T")]`: calls `Into::<T>::into` before writing the field to the database and calls `T::from` after reading (diesel column type: T)
+/// * `#[model(uom_unit = "T")]`: the value is the path to an unit defined in editoast_common, e.g. `"editoast_common::units::meter"`
 /// * `#[model(geo)]` **TODO**: TBD
 ///
 /// #### A note on identifiers
@@ -269,3 +271,20 @@ mod test_utils {
 
 #[cfg(test)]
 use test_utils::assert_macro_expansion;
+/// Annotates fields of a structs with documentation and value_type for a better utoipa schema
+///
+/// It must be used on structs that use #[derive(ToSchema)]
+///
+/// On every field that has an attribute such as #[serde(with="millimeter")]
+/// It will add:
+/// * #[schema(value_type = f64)]
+/// * /// Length in mm
+#[proc_macro_attribute]
+pub fn annotate_units(_attr: TokenStream, input: TokenStream) -> TokenStream {
+    // We are using a macro attribute to modify in place the attributes of fields to annotate
+    // This requires to mutate the input
+    let mut input = parse_macro_input!(input as DeriveInput);
+    annotate_units::annotate_units(&mut input)
+        .unwrap_or_else(darling::Error::write_errors)
+        .into()
+}

editoast/editoast_derive/src/model/args.rs

@@ -79,6 +79,8 @@ pub(super) struct ModelFieldArgs {
     pub(super) to_enum: bool,
     #[darling(default)]
     pub(super) remote: Option<syn::Type>,
+    #[darling(default)]
+    pub(super) uom_unit: Option<syn::Path>,
 }
 
 impl GeneratedTypeArgs {

editoast/editoast_derive/src/model/config.rs

@@ -45,6 +45,7 @@ pub(crate) enum FieldTransformation {
     Geo,
     ToString,
     ToEnum(syn::Type),
+    UomUnit(syn::Path),
 }
 
 #[derive(Debug, PartialEq)]
@@ -132,6 +133,9 @@ impl ModelField {
             Some(FieldTransformation::ToEnum(_)) => {
                 parse_quote! { #expr as i16 }
             }
+            Some(FieldTransformation::UomUnit(ref unit)) => {
+                parse_quote! { #unit::from(#expr) }
+            }
             None => parse_quote! { #expr },
         }
     }
@@ -146,6 +150,9 @@ impl ModelField {
             Some(FieldTransformation::ToEnum(ref ty)) => {
                 parse_quote! { #ty::from_repr(#expr as usize).expect("Invalid variant repr") }
             }
+            Some(FieldTransformation::UomUnit(ref unit)) => {
+                parse_quote! { #unit::new(#expr) }
+            }
             None => parse_quote! { #expr },
         }
     }
@@ -158,6 +165,7 @@ impl ModelField {
             Some(FieldTransformation::Geo) => unimplemented!("to be designed"),
             Some(FieldTransformation::ToString) => parse_quote! { String },
             Some(FieldTransformation::ToEnum(_)) => parse_quote! { i16 },
+            Some(FieldTransformation::UomUnit(ref unit)) => parse_quote! { #unit::ReprType },
             None => ty.clone(),
         }
     }