kaptan/firmware/libs/kaptan-config-derive/src/expand.rs

use proc_macro2::TokenStream;
use quote::quote;
use syn::Result;

use crate::parse::{Check, Model, ModelKind};

pub(crate) fn expand(model: Model) -> Result<TokenStream> {
    let Model {
        ident,
        generics,
        crate_path,
        post_validate,
        kind,
        fields,
    } = model;
    match kind {
        ModelKind::Scalars => expand_scalars(ident, generics, crate_path, post_validate, fields),
        ModelKind::Sections => expand_children(ident, generics, crate_path, post_validate, fields),
    }
}

fn expand_scalars(
    ident: syn::Ident,
    generics: syn::Generics,
    crate_path: syn::Path,
    post_validate: Option<syn::Path>,
    fields: Vec<crate::parse::FieldModel>,
) -> Result<TokenStream> {
    let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
    let mut checks = Vec::new();
    let mut descriptors = Vec::new();
    let mut reads = Vec::new();
    let mut writes = Vec::new();
    let mut snapshot_widths = Vec::new();
    for field in &fields {
        let FieldMetadata {
            descriptor,
            checks: field_checks,
        } = expand_field(&crate_path, field)?;
        descriptors.push(descriptor);
        checks.extend(field_checks.iter().cloned());
        let id = field.id.expect("scalar configuration fields have IDs");
        let ty = &field.ty;
        let field_ident = &field.ident;
        let (kind, read, decode, out_of_domain_rule) = if !field.enum_backed {
            snapshot_widths.push(quote! { <#ty as #crate_path::ConfigScalar>::KIND.width() });
            (
                quote! { <#ty as #crate_path::ConfigScalar>::KIND },
                quote! { #crate_path::ConfigScalar::into_config_value(self.#field_ident) },
                quote! { <#ty as #crate_path::ConfigScalar>::try_from_config_value(value) },
                quote! { #crate_path::ValidationRule::AboveMaximum },
            )
        } else {
            snapshot_widths.push(quote! { 1usize });
            (
                quote! { #crate_path::ValueKind::U8 },
                quote! { #crate_path::ConfigValue::U8(self.#field_ident.wire_value()) },
                quote! {
                    match value {
                        #crate_path::ConfigValue::U8(value) => <#ty>::try_from_wire(value)
                            .ok_or(#crate_path::ScalarValueError::OutOfDomain),
                        value => Err(#crate_path::ScalarValueError::WrongKind(value.kind())),
                    }
                },
                quote! { #crate_path::ValidationRule::InvalidEnumValue },
            )
        };
        reads.push(quote! {
            #id => Some(#read),
        });
        writes.push(quote! {
            #id => {
                let next = match #decode {
                    Ok(next) => next,
                    Err(#crate_path::ScalarValueError::WrongKind(actual)) => {
                        return Err(#crate_path::FieldWriteError::WrongType {
                            field: #crate_path::FieldId::new(#id),
                            expected: #kind,
                            actual,
                        });
                    }
                    Err(#crate_path::ScalarValueError::OutOfDomain) => {
                        return Err(#crate_path::FieldWriteError::Violation(
                            #crate_path::ConfigViolation::field(
                                #crate_path::FieldId::new(#id),
                                #out_of_domain_rule,
                            ),
                        ));
                    }
                };
                let previous = self.#field_ident;
                self.#field_ident = next;
                let validation = (|| -> Result<(), #crate_path::ConfigViolation> {
                    #(#field_checks)*
                    Ok(())
                })();
                if let Err(violation) = validation {
                    self.#field_ident = previous;
                    return Err(#crate_path::FieldWriteError::Violation(violation));
                }
                Ok(())
            }
        });
    }
    let post_validate = post_validate.map(|hook| quote! { #hook(self)?; });
    Ok(quote! {
        impl #impl_generics #ident #type_generics #where_clause {
            /// Field metadata generated from the same declarations as validation.
            pub const DESCRIPTORS: &'static [#crate_path::FieldDescriptor] = &[
                #(#descriptors),*
            ];

            fn __validate_local_configuration(
                &self,
            ) -> Result<(), #crate_path::ConfigViolation> {
                if !#crate_path::validation::__valid_enum_options(Self::DESCRIPTORS) {
                    return Err(#crate_path::ConfigViolation::global(
                        #crate_path::ValidationRule::Inconsistent,
                    ));
                }
                #(#checks)*
                Ok(())
            }
        }

        impl #impl_generics #crate_path::Config for #ident #type_generics #where_clause {
            const FIELD_COUNT: usize = Self::DESCRIPTORS.len();
            const SNAPSHOT_VALUE_BYTES: usize = 0 #(+ #snapshot_widths)*;
            const SCHEMA_ID: [u8; 16] =
                #crate_path::validation::__schema_id(Self::DESCRIPTORS);

            fn descriptor(index: usize) -> Option<&'static #crate_path::FieldDescriptor> {
                Self::DESCRIPTORS.get(index)
            }

            fn field_value(&self, id: #crate_path::FieldId) -> Option<#crate_path::ConfigValue> {
                match id.get() {
                    #(#reads)*
                    _ => None,
                }
            }

            fn set_field(
                &mut self,
                id: #crate_path::FieldId,
                value: #crate_path::ConfigValue,
            ) -> Result<(), #crate_path::FieldWriteError> {
                match id.get() {
                    #(#writes,)*
                    _ => Err(#crate_path::FieldWriteError::UnknownField(id)),
                }
            }

            fn validate(&self) -> Result<(), #crate_path::ConfigViolation> {
                self.__validate_local_configuration()?;
                #post_validate
                Ok(())
            }
        }
    })
}

fn expand_children(
    ident: syn::Ident,
    generics: syn::Generics,
    crate_path: syn::Path,
    post_validate: Option<syn::Path>,
    fields: Vec<crate::parse::FieldModel>,
) -> Result<TokenStream> {
    let (impl_generics, type_generics, where_clause) = generics.split_for_impl();
    let mut validations = Vec::new();
    let mut reads = Vec::new();
    let mut writes = Vec::new();
    let mut counts = Vec::new();
    let mut snapshot_widths = Vec::new();
    let mut schema_ids = Vec::new();
    let mut descriptor_lookups = Vec::new();
    for field in &fields {
        let field_ident = &field.ident;
        let ty = &field.ty;
        validations.push(quote! {
            <#ty as #crate_path::Config>::validate(&self.#field_ident)?;
        });
        reads.push(quote! {
            if let Some(value) = <#ty as #crate_path::Config>::field_value(
                &self.#field_ident,
                id,
            ) {
                return Some(value);
            }
        });
        writes.push(quote! {
            if <#ty as #crate_path::Config>::field_value(&self.#field_ident, id).is_some() {
                return <#ty as #crate_path::Config>::set_field(
                    &mut self.#field_ident,
                    id,
                    value,
                );
            }
        });
        counts.push(quote! { <#ty as #crate_path::Config>::FIELD_COUNT });
        snapshot_widths.push(quote! { <#ty as #crate_path::Config>::SNAPSHOT_VALUE_BYTES });
        schema_ids.push(quote! { <#ty as #crate_path::Config>::SCHEMA_ID });
        descriptor_lookups.push(quote! {
            if index < <#ty as #crate_path::Config>::FIELD_COUNT {
                return <#ty as #crate_path::Config>::descriptor(index);
            }
            index -= <#ty as #crate_path::Config>::FIELD_COUNT;
        });
    }
    let post_validate = post_validate.map(|hook| quote! { #hook(self)?; });

    Ok(quote! {
        impl #impl_generics #crate_path::Config for #ident #type_generics #where_clause {
            const FIELD_COUNT: usize = 0 #(+ #counts)*;
            const SNAPSHOT_VALUE_BYTES: usize = 0 #(+ #snapshot_widths)*;
            const SCHEMA_ID: [u8; 16] =
                #crate_path::validation::__combine_schema_ids(&[#(#schema_ids),*]);

            fn descriptor(mut index: usize) -> Option<&'static #crate_path::FieldDescriptor> {
                #(#descriptor_lookups)*
                None
            }

            fn field_value(&self, id: #crate_path::FieldId) -> Option<#crate_path::ConfigValue> {
                #(#reads)*
                None
            }

            fn set_field(
                &mut self,
                id: #crate_path::FieldId,
                value: #crate_path::ConfigValue,
            ) -> Result<(), #crate_path::FieldWriteError> {
                #(#writes)*
                Err(#crate_path::FieldWriteError::UnknownField(id))
            }

            fn validate(&self) -> Result<(), #crate_path::ConfigViolation> {
                if !#crate_path::validation::__unique_field_ids::<Self>() {
                    return Err(#crate_path::ConfigViolation::global(
                        #crate_path::ValidationRule::Inconsistent,
                    ));
                }
                if !#crate_path::validation::__unique_display_keys::<Self>() {
                    return Err(#crate_path::ConfigViolation::global(
                        #crate_path::ValidationRule::Inconsistent,
                    ));
                }
                #(#validations)*
                #post_validate
                Ok(())
            }
        }

        impl #impl_generics #ident #type_generics #where_clause {
            /// Number of fields contributed by the complete generated tree.
            pub const FIELD_COUNT: usize = <Self as #crate_path::Config>::FIELD_COUNT;
        }
    })
}

struct FieldMetadata {
    descriptor: TokenStream,
    checks: Vec<TokenStream>,
}

fn expand_field(crate_path: &syn::Path, field: &crate::parse::FieldModel) -> Result<FieldMetadata> {
    let mut checks = Vec::new();
    let id = field.id.expect("scalar configuration fields have IDs");
    let minimum = field.checks.iter().find_map(|check| match check {
        Check::Minimum(value) => Some(value),
        _ => None,
    });
    let maximum = field.checks.iter().find_map(|check| match check {
        Check::Maximum(value) => Some(value),
        _ => None,
    });
    let finite = field
        .checks
        .iter()
        .any(|check| matches!(check, Check::Finite));
    let nonzero = field
        .checks
        .iter()
        .any(|check| matches!(check, Check::Nonzero));
    let float = matches!(&field.ty, syn::Type::Path(path)
        if path.path.is_ident("f32") || path.path.is_ident("f64"));
    if float && !finite && (minimum.is_some() || maximum.is_some()) {
        let field_ident = &field.ident;
        checks.push(quote! {
            if self.#field_ident.is_nan() {
                return Err(#crate_path::ConfigViolation::field(
                    #crate_path::FieldId::new(#id),
                    #crate_path::ValidationRule::Unordered,
                ));
            }
        });
    }
    for check in &field.checks {
        checks.push(expand_check(
            crate_path,
            &field.ident,
            id,
            check.clone(),
            float,
        ));
    }
    let key = field.display_key.as_ref().expect("id fields have a key");
    let unit = field.unit.as_ref().expect("id fields have a unit");
    let ty = &field.ty;
    let (kind, enum_options) = if !field.enum_backed {
        (
            quote! { <#ty as #crate_path::ConfigScalar>::KIND },
            quote! { &[] },
        )
    } else {
        (
            quote! { #crate_path::ValueKind::U8 },
            quote! { <#ty>::CONFIG_ENUM_OPTIONS },
        )
    };
    let value = config_value_constructor_tokens(crate_path, &field.ty)?;
    let minimum = option_value_tokens(value.clone(), minimum);
    let maximum = option_value_tokens(value.clone(), maximum);
    Ok(FieldMetadata {
        descriptor: quote! {
            #crate_path::FieldDescriptor {
                id: #crate_path::FieldId::new(#id),
                display_key: #key,
                kind: #kind,
                minimum: #minimum,
                maximum: #maximum,
                enum_options: #enum_options,
                unit: #unit,
                finite: #finite,
                nonzero: #nonzero,
            }
        },
        checks,
    })
}

fn option_value_tokens(value: TokenStream, expression: Option<&syn::Expr>) -> TokenStream {
    match expression {
        Some(expression) => quote! { Some(#value(#expression)) },
        None => quote! { None },
    }
}

fn config_value_constructor_tokens(crate_path: &syn::Path, ty: &syn::Type) -> Result<TokenStream> {
    let syn::Type::Path(path) = ty else {
        return Err(syn::Error::new_spanned(
            ty,
            "configuration metadata requires a scalar type",
        ));
    };
    let Some(segment) = path.path.segments.last() else {
        return Err(syn::Error::new_spanned(
            ty,
            "configuration metadata requires a scalar type",
        ));
    };
    let value = match segment.ident.to_string().as_str() {
        "bool" => quote! { #crate_path::ConfigValue::Bool },
        "i32" => quote! { #crate_path::ConfigValue::I32 },
        "u8" => quote! { #crate_path::ConfigValue::U8 },
        "u16" => quote! { #crate_path::ConfigValue::U16 },
        "u32" => quote! { #crate_path::ConfigValue::U32 },
        "u64" => quote! { #crate_path::ConfigValue::U64 },
        "f32" => quote! { #crate_path::ConfigValue::F32 },
        "f64" => quote! { #crate_path::ConfigValue::F64 },
        _ => {
            return Ok(quote! { #crate_path::ConfigValue::U32 });
        }
    };
    Ok(value)
}

fn expand_check(
    crate_path: &syn::Path,
    field: &syn::Ident,
    id: u16,
    check: Check,
    float: bool,
) -> TokenStream {
    let zero = if float {
        quote! { 0.0 }
    } else {
        quote! { 0 }
    };
    match check {
        Check::Finite => quote! {
            if !self.#field.is_finite() {
                return Err(#crate_path::ConfigViolation::field(
                    #crate_path::FieldId::new(#id),
                    #crate_path::ValidationRule::NonFinite,
                ));
            }
        },
        Check::Nonzero => quote! {
            if self.#field == #zero {
                return Err(#crate_path::ConfigViolation::field(
                    #crate_path::FieldId::new(#id),
                    #crate_path::ValidationRule::Zero,
                ));
            }
        },
        Check::Minimum(minimum) => quote! {
            if self.#field < (#minimum) {
                return Err(#crate_path::ConfigViolation::field(
                    #crate_path::FieldId::new(#id),
                    #crate_path::ValidationRule::BelowMinimum,
                ));
            }
        },
        Check::Maximum(maximum) => quote! {
            if self.#field > (#maximum) {
                return Err(#crate_path::ConfigViolation::field(
                    #crate_path::FieldId::new(#id),
                    #crate_path::ValidationRule::AboveMaximum,
                ));
            }
        },
    }
}