use std::rc::Rc;
use indexmap::{IndexMap, IndexSet};
use oas3::spec::ObjectSchema;
use crate::{
generator::{
ast::{
DiscriminatedVariant, Documentation, EnumMethod, EnumToken, EnumVariantToken, RustType, TypeRef, VariantDef,
},
converter::ConverterContext,
naming::identifiers::{split_pascal_case, strip_parent_prefix, to_rust_type_name},
},
utils::parse_schema_ref_path,
};
#[derive(Debug, Clone)]
pub(crate) struct DiscriminatorConverter {
context: Rc<ConverterContext>,
}
impl DiscriminatorConverter {
pub(crate) fn new(context: Rc<ConverterContext>) -> Self {
Self { context }
}
pub(crate) fn detect_discriminated_parent(&self, schema_name: &str) -> Option<&str> {
self.context.graph().parent(schema_name)
}
pub(crate) fn build_base_discriminated_enum(
&self,
name: &str,
schema: &ObjectSchema,
fallback_type: &str,
) -> anyhow::Result<RustType> {
let tag_field = schema
.discriminator
.as_ref()
.map(|d| &d.property_name)
.ok_or_else(|| anyhow::anyhow!("missing discriminator property for schema '{name}'"))?;
let enum_name = to_rust_type_name(name);
let variants = self.build_variants_from_mapping(&enum_name, schema);
let fallback_name = split_pascal_case(&enum_name)
.last()
.cloned()
.unwrap_or_else(|| enum_name.clone());
let fallback = DiscriminatedVariant::builder()
.variant_name(EnumVariantToken::new(fallback_name))
.type_name(TypeRef::new(fallback_type).with_boxed())
.build();
Ok(
RustType::discriminated_enum()
.name(&EnumToken::from_raw(enum_name))
.docs(Documentation::from_optional(schema.description.as_ref()))
.discriminator_field(tag_field.clone())
.variants(variants)
.maybe_fallback(Some(fallback))
.call(),
)
}
pub(crate) fn build_variants_from_mapping(
&self,
parent_name: &str,
schema: &ObjectSchema,
) -> Vec<DiscriminatedVariant> {
let Some(mapping) = self.context.graph().effective_mapping(schema) else {
return vec![];
};
let is_reachable = |name: &String| {
self
.context
.reachable_schemas
.as_ref()
.is_none_or(|filter| filter.contains(name))
};
mapping
.iter()
.filter_map(|(tag, ref_path)| {
let name = parse_schema_ref_path(ref_path)?;
is_reachable(&name).then_some((tag.clone(), name))
})
.fold(IndexMap::<String, Vec<String>>::new(), |mut acc, (tag, name)| {
acc.entry(name).or_default().push(tag);
acc
})
.into_iter()
.map(|(schema_name, tags)| {
let type_name = to_rust_type_name(&schema_name);
let variant_name = strip_parent_prefix(parent_name, &type_name);
DiscriminatedVariant::builder()
.variant_name(EnumVariantToken::new(variant_name))
.discriminator_values(tags)
.type_name(TypeRef::new(type_name).with_boxed())
.build()
})
.collect()
}
pub(crate) fn try_upgrade_to_discriminated(
&self,
name: &str,
schema: &ObjectSchema,
variants: &[VariantDef],
methods: Vec<EnumMethod>,
) -> Option<RustType> {
let discriminator = schema.discriminator.as_ref()?;
let mapping = &self.context.graph().effective_mapping(schema)?;
if !Self::all_mappings_have_variants(variants, mapping) {
return None;
}
let discriminated_variants = Self::convert_to_discriminated_variants(variants, mapping);
Some(
RustType::discriminated_enum()
.name(&EnumToken::from_raw(name))
.docs(Documentation::from_optional(schema.description.as_ref()))
.discriminator_field(discriminator.property_name.clone())
.variants(discriminated_variants)
.methods(methods)
.call(),
)
}
fn all_mappings_have_variants(variants: &[VariantDef], mapping: &IndexMap<String, String>) -> bool {
if variants.is_empty() || mapping.is_empty() {
return false;
}
let known_types = variants
.iter()
.filter_map(VariantDef::unboxed_type_name)
.collect::<IndexSet<_>>();
mapping.values().all(|ref_path| {
parse_schema_ref_path(ref_path).is_some_and(|name| known_types.contains(&to_rust_type_name(&name)))
})
}
fn convert_to_discriminated_variants(
variants: &[VariantDef],
mapping: &IndexMap<String, String>,
) -> Vec<DiscriminatedVariant> {
mapping
.iter()
.filter_map(|(tag, ref_path)| {
let type_name = parse_schema_ref_path(ref_path).map(|n| to_rust_type_name(&n))?;
Some((type_name, tag.clone()))
})
.fold(IndexMap::<String, Vec<String>>::new(), |mut acc, (type_name, tag)| {
acc.entry(type_name).or_default().push(tag);
acc
})
.into_iter()
.filter_map(|(type_name, tags)| {
let variant = variants
.iter()
.find(|v| v.unboxed_type_name() == Some(type_name.clone()))?;
let inner_type = variant.single_wrapped_type()?;
Some(
DiscriminatedVariant::builder()
.variant_name(variant.name.clone())
.type_name(inner_type.clone())
.discriminator_values(tags)
.build(),
)
})
.collect()
}
}