unitscale_macros 0.2.0

// Copyright 2025 GooseGrid Technologies
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![doc = include_str!("../README.md")]

use proc_macro::TokenStream;
use proc_macro2::{Literal, Span};
use quote::quote;
use syn::parse::{Parse, ParseStream};
use syn::{Ident, ItemStruct, LitFloat, LitStr, Result, Token, parse_macro_input};

/// Detect trait type and return `Ident` of `UnitScaleF32` or `UnitScaleF64`.
fn determine_trait_type(with_string: Option<String>) -> Result<Ident> {
    match &with_string {
        None => syn::parse_str::<Ident>("UnitScaleF32"),
        Some(s) if s == "f64" => syn::parse_str::<Ident>("UnitScaleF64"),
        Some(s) if s == "f32" => syn::parse_str::<Ident>("UnitScaleF32"),
        Some(_other) => Err(syn::Error::new_spanned(
            &with_string,
            "Not valid trait type; expected \"f32\" or \"f64\" (requires feature `f64`)",
        )),
    }
}

/// Detect `SCALE` type and return `Ident` of `f32` or `f64`.
fn determine_scalar_type(with_lit_str: Option<LitStr>) -> Result<Ident> {
    if with_lit_str.is_none() {
        syn::parse_str::<Ident>("f32")
    } else if let Some(scalar_type_lit) = &with_lit_str
        && let lit = scalar_type_lit.value()
        && (lit == "f64" || lit == "f32")
    {
        syn::parse_str::<Ident>(&lit)
    } else {
        Err(syn::Error::new_spanned(
            &with_lit_str,
            "Not valid scalar type; expected \"f32\" or \"f64\" (requires feature `f64`)",
        ))
    }
}

/// Unit type selects the `f32` or `f64` (feature: "f64") implementation trait and associated scalar.
trait UnitTypes {
    fn trait_type(&self) -> Result<Ident>;
    fn scalar_type(&self) -> Result<Ident>;
}

/// Unit type macro arguments
struct UnitTypeArgs {
    // Scalar type: "f32" or "f64" (requires feature `f64`).
    with: Option<LitStr>,
}

impl UnitTypeArgs {
    fn with_as_string(&self) -> Option<String> {
        self.with.as_ref().map(LitStr::value)
    }
}

impl UnitTypes for UnitTypeArgs {
    fn trait_type(&self) -> Result<Ident> {
        let with = self.with_as_string();
        determine_trait_type(with)
    }

    fn scalar_type(&self) -> Result<Ident> {
        let with = self.with.clone();
        determine_scalar_type(with)
    }
}

impl Parse for UnitTypeArgs {
    fn parse(input: ParseStream) -> Result<Self> {
        let mut with: Option<LitStr> = if !input.is_empty() {
            None
        } else {
            return Ok(Self { with: None });
        };

        while !input.is_empty() {
            let ident: Ident = input.parse()?;

            // Expect `=` after identifier.
            input.parse::<Token![=]>()?;

            match ident.to_string().as_str() {
                "with" => {
                    with = input.parse().ok();
                }
                "" => {}
                _other => {
                    return Err(input.error("Expects an optional `with`"));
                }
            }

            if input.peek(Token![,]) {
                input.parse::<Token![,]>()?;
            } else {
                break;
            }
        }

        Ok(Self { with })
    }
}

/// `unit_type` macro to create a valid scaled value bounds of `U`
///
/// Trys to create bounded value from f32 that is scaled by
/// `unitscale_core::UnitScaleF32`. If using `with = "f64"` you
/// will create a `unitscale_core::UnitScaleF64`. Typically this
/// would look like,
///
/// ```
/// use unitscale_core::{Scaled, UnitScaleError};
///
/// #[unitscale_macros::unit_type]
/// struct Volts;
///
/// struct Scale0_1;
///
/// impl unitscale_core::UnitScaleF32 for Scale0_1 {
///     const SCALE: f32 = 0.1;
/// }
///
/// fn func() {
///     let data: Volts<Scale0_1, u16> = Volts::try_from(3.11).expect("Out of bounds");
///     println!("Value {} should be 31", data.scaled_value());
/// }
/// ```
///
/// You gain functions to `TryFrom<f32>` to convert within bounds of a given
/// type that supports `FromPrimitive`. This allows you to encapsulate the
/// unit, scale and type in one shot easily. Or `TryFrom<f64>` if using
/// `with = "f64"`. For example,
///
///
/// Where the value of `data` will be 31 and trims off the rest, but means 3.1
/// encoded in `u16`. Therefore, name the `unitscale_core::UnitScale` to
/// something meaningful for readability.
///
/// **NB** If you want your `derive` macros to work as intended append them after the
/// `unit_type` macro, *e.g.*,
///
/// ```ignore
/// /// My custom docs for my type
/// #[unitscale_macros::unit_type]
/// #[derive(Copy, Clone)]
/// pub struct Datums;
/// ```
#[proc_macro_attribute]
pub fn unit_type(attrs: TokenStream, items: TokenStream) -> TokenStream {
    let input_struct = parse_macro_input!(items as ItemStruct);
    let args = parse_macro_input!(attrs as UnitTypeArgs);

    create_unit_type(&input_struct, &args)
}

/// Redefines the struct to include the correct parameters
fn create_unit_type(ast: &ItemStruct, args: &UnitTypeArgs) -> TokenStream {
    let name = &ast.ident;
    let vis = &ast.vis;

    let trait_name = match args.trait_type() {
        Ok(tt) => tt,
        Err(e) => {
            return e.to_compile_error().into();
        }
    };
    let (scalar_type, fn_to_scalar_type, fn_from_scalar_type) = match args.scalar_type() {
        Ok(st) => {
            let scalar_type_str = st.to_string();
            let fn_to_scalar_type =
                Ident::new(&format!("to_{}", scalar_type_str), Span::call_site());
            let fn_from_scalar_type =
                Ident::new(&format!("from_{}", scalar_type_str), Span::call_site());
            (st, fn_to_scalar_type, fn_from_scalar_type)
        }
        Err(e) => {
            return e.to_compile_error().into();
        }
    };

    let try_from_doc = {
        let txt = format!(
            "Try to scale from `{}` to a type `U` within bounds.",
            scalar_type,
        );

        let lit = Literal::string(&txt);
        quote! { #[doc = #lit] }
    };

    let docs: Vec<_> = ast
        .attrs
        .iter()
        .filter(|attr| attr.path().is_ident("doc") || attr.path().is_ident("derive"))
        .collect();

    let generated = quote! {
        #(#docs)*
        ///
        /// Scaled value of the desired type `U`. Typically an unsigned integer.
        ///
        /// `S` is a `unitscale_core::#trait_name` trait
        #vis struct #name<S, U> {
            value: U,
            _scale: core::marker::PhantomData<S>,
        }

        /// Return the encapsulated value.
        impl<S, U> unitscale_core::Scaled<U> for #name<S, U>
        where
            S: unitscale_core::#trait_name,
            U: Copy
        {
            /// Retrieve scaled value `x / SCALE`.
            fn scaled_value(&self) -> U {
                self.value
            }
        }

        impl<S, U> unitscale_core::ScaledPrimitiveByteSize for #name<S, U>
        where
            S: unitscale_core::#trait_name,
            U: Copy + num_traits::ToPrimitive,
        {
            /// Get the primitive byte size of `U`. Useful for converting into your protocol implementation.
            fn primitive_byte_size() -> usize {
                core::mem::size_of::<U>()
            }
        }

        impl<S, U> #name<S, U>
        where
            S: unitscale_core::#trait_name,
            U: Copy + num_traits::ToPrimitive,
        {
            /// Create new value from `U` *as is*. The value is not **scaled**.
            #vis fn from_scaled_value(value: U) -> Self {
                Self {
                    value,
                    _scale: core::marker::PhantomData,
                }
            }

            /// Revert scaled value back to intended user data `scaled_value * SCALE`.
            #vis fn #fn_to_scalar_type(&self) -> Option<#scalar_type> {
                self.value.#fn_to_scalar_type().map(|v| v * S::SCALE)
            }
        }

        /// Allows us to convert and scale value from any numeric type
        /// that goes up to `u64, i64, f64` to the desired type. This is
        /// typical for `u8`, `u16` or `u32`.
        impl<S, U> core::convert::TryFrom<#scalar_type> for #name<S, U>
        where
            S: unitscale_core::#trait_name,
            U: num_traits::FromPrimitive,
        {
            type Error = unitscale_core::UnitScaleError;

            #try_from_doc
            fn try_from(value: #scalar_type) -> Result<Self, Self::Error> {
                let scaled_value = value / S::SCALE;

                Ok(Self {
                    value: U::#fn_from_scalar_type(scaled_value).ok_or(UnitScaleError::Conversion(
                        format!(
                            "Scaled {} is out of {} bounds",
                            scaled_value,
                            core::any::type_name::<U>(),
                        )
                    ))?,
                    _scale: core::marker::PhantomData,
                })
            }
        }
    };

    TokenStream::from(generated)
}

/// Scale macro arguments
struct UnitScaleArgs {
    /// The scaling factor.
    to: LitFloat,
    /// Scalar datatype: `f32` or `f64`.
    with: Option<LitStr>,
}

impl UnitScaleArgs {
    fn with_as_string(&self) -> Option<String> {
        self.with.as_ref().map(LitStr::value)
    }
}

impl UnitTypes for UnitScaleArgs {
    fn trait_type(&self) -> Result<Ident> {
        let with = self.with_as_string();
        determine_trait_type(with)
    }

    fn scalar_type(&self) -> Result<Ident> {
        let with = self.with.clone();
        determine_scalar_type(with)
    }
}

impl Parse for UnitScaleArgs {
    fn parse(input: ParseStream) -> Result<Self> {
        let mut to: Option<LitFloat> = None;
        let mut with: Option<LitStr> = None;

        while !input.is_empty() {
            let ident: Ident = input.parse()?;

            // Expect `=` after identifier.
            input.parse::<Token![=]>()?;

            match ident.to_string().as_str() {
                "to" => {
                    to = input.parse().ok();
                }
                "with" => {
                    with = input.parse().ok();
                }
                _ => {
                    return Err(input.error("Expects `to` and optionally `with`"));
                }
            }

            if input.peek(Token![,]) {
                input.parse::<Token![,]>()?;
            } else {
                break;
            }
        }

        let to = to.ok_or_else(|| input.error("Missing required `to` arguement"))?;

        Ok(Self { to, with })
    }
}

/// `unit_scale` macro to create a valid scaler
///
/// This is used in conjuction with `unitscale_core::Scaled` to
/// couple the scaling and type `U` to the struct. It is for
/// convenience to create many scalars of arbritary scale. It takes
/// the following options
///
/// -  `to`: Scalar float
/// - **optional** `with`: Specify `f32` or `f64` (feature "f64")
///
/// ```
/// use unitscale_core::UnitScaleF32;
///
/// // Simple example showing 0.1, with an aptly named struct
/// #[unitscale_macros::unit_scale(to = 0.1)]
/// struct Scale0_1;
///
/// fn func() {
///     println!("Scalar is {}", Scale0_1::SCALE);
/// }
/// ```
/// This is analogous to,
///
/// ```
/// struct Scale0_1;
///
/// impl unitscale_core::UnitScaleF32 for Scale0_1 {
///     const SCALE: f32 = 0.1;
/// }
/// ```
///
/// See `unitscale_macros::datatype` for an example
#[proc_macro_attribute]
pub fn unit_scale(attrs: TokenStream, items: TokenStream) -> TokenStream {
    let input_struct = parse_macro_input!(items as ItemStruct);
    let args = parse_macro_input!(attrs as UnitScaleArgs);

    create_unit_scale(&input_struct, &args)
}

/// Creates the impl for the intended struct
fn create_unit_scale(ast: &ItemStruct, args: &UnitScaleArgs) -> TokenStream {
    let name = &ast.ident;
    let scalar = &args.to;

    let trait_name = match args.trait_type() {
        Ok(tt) => tt,
        Err(e) => {
            return e.to_compile_error().into();
        }
    };
    let scalar_type = match args.scalar_type() {
        Ok(st) => st,
        Err(e) => {
            return e.to_compile_error().into();
        }
    };

    let expanded = quote! {
        #ast

        impl unitscale_core::#trait_name for #name {
            const SCALE: #scalar_type = #scalar;
        }
    };

    TokenStream::from(expanded)
}