kaptan/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_sections(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 defaults = Vec::new();
    let mut reads = Vec::new();
    let mut writes = Vec::new();
    for field in &fields {
        let FieldMetadata {
            descriptor,
            checks: field_checks,
        } = expand_field(&crate_path, field)?;
        if let Some(descriptor) = descriptor {
            descriptors.push(descriptor);
        }
        checks.extend(field_checks);
        let id = field.id.expect("ConfigValidate requires configured fields");
        let default = field
            .default
            .as_ref()
            .expect("ConfigValidate requires defaults");
        let (kind, value) = config_value_tokens(&crate_path, &field.ty)?;
        let field_ident = &field.ident;
        defaults.push(quote! { #field_ident: #default });
        reads.push(quote! {
            #id => Some(#value(self.#field_ident)),
        });
        writes.push(quote! {
            #id => {
                let next = match value {
                    #value(next) => next,
                    actual => {
                        return Err(#crate_path::FieldWriteError::WrongType {
                            field: #crate_path::FieldId::new(#id),
                            expected: #kind,
                            actual: actual.kind(),
                        });
                    }
                };
                let previous = self.#field_ident;
                self.#field_ident = next;
                if let Err(violation) =
                    <Self as #crate_path::ConfigSection>::validate_local_fields(self)
                {
                    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 #crate_path::ValidateFields for #ident #type_generics #where_clause {
            fn validate_fields(&self) -> Result<(), #crate_path::ConfigViolation> {
                <Self as #crate_path::ConfigSection>::validate_local_fields(self)?;
                #post_validate
                Ok(())
            }
        }

        impl #impl_generics #ident #type_generics #where_clause {
            /// Field metadata generated from the same attributes as validation.
            pub const DESCRIPTORS: &'static [#crate_path::FieldDescriptor] = &[
                #(#descriptors),*
            ];
        }

        impl #impl_generics #crate_path::ConfigSection for #ident #type_generics #where_clause {
            const FIELD_COUNT: usize = Self::DESCRIPTORS.len();

            fn default_section() -> Result<Self, #crate_path::ConfigViolation> {
                let section = Self {
                    #(#defaults),*
                };
                <Self as #crate_path::ValidateFields>::validate_fields(&section)?;
                Ok(section)
            }

            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_local_fields(&self) -> Result<(), #crate_path::ConfigViolation> {
                #(#checks)*
                Ok(())
            }
        }
    })
}

fn expand_sections(
    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 validations = fields.iter().map(|field| {
        let field_ident = &field.ident;
        let ty = &field.ty;
        quote! {
            <#ty as #crate_path::ValidateFields>::validate_fields(&self.#field_ident)?;
        }
    });
    let defaults = fields.iter().map(|field| {
        let field_ident = &field.ident;
        let ty = &field.ty;
        quote! {
            #field_ident: <#ty as #crate_path::ConfigSection>::default_section()?
        }
    });
    let reads = fields.iter().map(|field| {
        let field_ident = &field.ident;
        let ty = &field.ty;
        quote! {
            if let Some(value) = <#ty as #crate_path::ConfigSection>::field_value(
                &self.#field_ident,
                id,
            ) {
                return Some(value);
            }
        }
    });
    let writes = fields.iter().map(|field| {
        let field_ident = &field.ident;
        let ty = &field.ty;
        quote! {
            if <#ty as #crate_path::ConfigSection>::field_value(&self.#field_ident, id).is_some() {
                return <#ty as #crate_path::ConfigSection>::set_field(
                    &mut self.#field_ident,
                    id,
                    value,
                );
            }
        }
    });
    let counts = fields.iter().map(|field| {
        let ty = &field.ty;
        quote! { <#ty as #crate_path::ConfigSection>::FIELD_COUNT }
    });
    let post_validate = post_validate.map(|hook| quote! { #hook(self)?; });
    Ok(quote! {
        impl #impl_generics #crate_path::ValidateFields for #ident #type_generics #where_clause {
            fn validate_fields(&self) -> Result<(), #crate_path::ConfigViolation> {
                #(#validations)*
                #post_validate
                Ok(())
            }
        }

        impl #impl_generics #ident #type_generics #where_clause {
            /// Number of fields aggregated from derived nested sections.
            pub const FIELD_COUNT: usize = 0 #(+ #counts)*;
        }

        impl #impl_generics #crate_path::ConfigAggregate for #ident #type_generics #where_clause {
            fn default_aggregate() -> Result<Self, #crate_path::ConfigViolation> {
                let aggregate = Self {
                    #(#defaults),*
                };
                <Self as #crate_path::ValidateFields>::validate_fields(&aggregate)?;
                Ok(aggregate)
            }

            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))
            }
        }
    })
}

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

fn expand_field(crate_path: &syn::Path, field: &crate::parse::FieldModel) -> Result<FieldMetadata> {
    let mut checks = Vec::new();
    let Some(id) = field.id else {
        return Ok(FieldMetadata {
            descriptor: None,
            checks,
        });
    };
    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,
    });
    for check in &field.checks {
        checks.push(expand_check(crate_path, &field.ident, id, check.clone()));
    }
    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 default = field.default.as_ref().expect("id fields have a default");
    let (kind, value) = config_value_tokens(crate_path, &field.ty)?;
    let minimum = option_value_tokens(value.clone(), minimum);
    let maximum = option_value_tokens(value.clone(), maximum);
    Ok(FieldMetadata {
        descriptor: Some(quote! {
            #crate_path::FieldDescriptor {
                id: #crate_path::FieldId::new(#id),
                display_key: #key,
                kind: #kind,
                minimum: #minimum,
                maximum: #maximum,
                default: #value(#default),
                unit: #unit,
                apply_policy: #crate_path::ApplyPolicy::RestartRequired,
            }
        }),
        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_tokens(
    crate_path: &syn::Path,
    ty: &syn::Type,
) -> Result<(TokenStream, 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 (kind, value) = match segment.ident.to_string().as_str() {
        "bool" => (
            quote! { #crate_path::ValueKind::Bool },
            quote! { #crate_path::ConfigValue::Bool },
        ),
        "i32" => (
            quote! { #crate_path::ValueKind::I32 },
            quote! { #crate_path::ConfigValue::I32 },
        ),
        "u32" => (
            quote! { #crate_path::ValueKind::U32 },
            quote! { #crate_path::ConfigValue::U32 },
        ),
        "u64" => (
            quote! { #crate_path::ValueKind::U64 },
            quote! { #crate_path::ConfigValue::U64 },
        ),
        "f32" => (
            quote! { #crate_path::ValueKind::F32 },
            quote! { #crate_path::ConfigValue::F32 },
        ),
        "f64" => (
            quote! { #crate_path::ValueKind::F64 },
            quote! { #crate_path::ConfigValue::F64 },
        ),
        _ => {
            return Err(syn::Error::new_spanned(
                ty,
                "configuration metadata supports bool, i32, u32, u64, f32, and f64",
            ));
        }
    };
    Ok((kind, value))
}

fn expand_check(crate_path: &syn::Path, field: &syn::Ident, id: u16, check: Check) -> TokenStream {
    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 == 0 {
                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,
                ));
            }
        },
    }
}