use std::{
collections::{BTreeMap, BTreeSet},
sync::Arc,
};
use anyhow::Context;
use oas3::spec::ObjectSchema;
use super::{
CodegenConfig, cache::SharedSchemaCache, common::SchemaExt, field_optionality::FieldOptionalityPolicy, metadata,
string_enum_optimizer::StringEnumOptimizer, structs::StructConverter, type_resolver::TypeResolver,
};
use crate::generator::{
ast::{
DiscriminatedEnumDef, DiscriminatedVariant, EnumDef, EnumMethod, EnumMethodKind, EnumToken, EnumVariantToken,
RustType, SerdeAttribute, StructDef, TypeRef, VariantContent, VariantDef,
},
naming::{
identifiers::{ensure_unique, to_rust_type_name},
inference::{
VariantNameNormalizer, derive_method_names, extract_enum_values, infer_variant_name, strip_common_affixes,
},
},
schema_registry::{ReferenceExtractor, SchemaRegistry},
};
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum UnionKind {
OneOf,
AnyOf,
}
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum CollisionStrategy {
Preserve,
Deduplicate,
}
#[derive(Clone)]
pub(crate) struct EnumConverter {
graph: Arc<SchemaRegistry>,
type_resolver: TypeResolver,
struct_converter: StructConverter,
preserve_case_variants: bool,
case_insensitive_enums: bool,
pub(crate) no_helpers: bool,
}
impl EnumConverter {
pub(crate) fn new(graph: &Arc<SchemaRegistry>, type_resolver: TypeResolver, config: CodegenConfig) -> Self {
let struct_converter = StructConverter::new(graph, config, None, FieldOptionalityPolicy::standard());
Self {
graph: graph.clone(),
type_resolver,
struct_converter,
preserve_case_variants: config.preserve_case_variants,
case_insensitive_enums: config.case_insensitive_enums,
no_helpers: config.no_helpers,
}
}
pub(crate) fn convert_simple_enum(
&self,
name: &str,
schema: &ObjectSchema,
cache: Option<&mut SharedSchemaCache>,
) -> Option<RustType> {
let mut enum_values: Vec<String> = schema
.enum_values
.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect();
enum_values.sort();
if cache.as_ref().is_some_and(|c| c.is_enum_generated(&enum_values)) {
return None;
}
let strategy = if self.preserve_case_variants {
CollisionStrategy::Preserve
} else {
CollisionStrategy::Deduplicate
};
let enum_def = self.build_simple_enum(name, schema, strategy);
if let (Some(c), RustType::Enum(e)) = (cache, &enum_def) {
c.register_enum(enum_values, e.name.to_string());
c.mark_name_used(e.name.to_string());
}
Some(enum_def)
}
pub(crate) fn convert_union_enum(
&self,
name: &str,
schema: &ObjectSchema,
kind: UnionKind,
mut cache: Option<&mut SharedSchemaCache>,
) -> anyhow::Result<Vec<RustType>> {
if kind == UnionKind::AnyOf {
let optimizer = StringEnumOptimizer::new(&self.graph, self.case_insensitive_enums);
if let Some(result) = optimizer.try_convert(name, schema, cache.as_deref_mut()) {
return Ok(result);
}
}
let result = self.process_union(name, schema, kind, cache.as_deref_mut())?;
if let Some(c) = cache
&& let Some(values) = extract_enum_values(schema)
&& let Some(RustType::Enum(e)) = result.last()
{
c.register_enum(values, e.name.to_string());
}
Ok(result)
}
fn build_simple_enum(&self, name: &str, schema: &ObjectSchema, strategy: CollisionStrategy) -> RustType {
let mut variants: Vec<VariantDef> = vec![];
let mut seen_names: BTreeMap<String, usize> = BTreeMap::new();
for (i, value) in schema.enum_values.iter().enumerate() {
let Some(normalized) = VariantNameNormalizer::normalize(value) else {
continue;
};
match seen_names.get(&normalized.name) {
Some(&existing_idx) if strategy == CollisionStrategy::Deduplicate => {
variants[existing_idx]
.serde_attrs
.push(SerdeAttribute::Alias(normalized.rename_value));
}
Some(_) => {
let unique_name = format!("{}{i}", normalized.name);
let idx = variants.len();
seen_names.insert(unique_name.clone(), idx);
Self::push_variant(&mut variants, unique_name, &normalized.rename_value);
}
None => {
let idx = variants.len();
seen_names.insert(normalized.name.clone(), idx);
Self::push_variant(&mut variants, normalized.name, &normalized.rename_value);
}
}
}
RustType::Enum(EnumDef {
name: EnumToken::from_raw(name),
docs: metadata::extract_docs(schema.description.as_ref()),
variants,
case_insensitive: self.case_insensitive_enums,
..Default::default()
})
}
fn push_variant(variants: &mut Vec<VariantDef>, name: impl Into<EnumVariantToken>, rename: &str) {
variants.push(VariantDef {
name: name.into(),
content: VariantContent::Unit,
serde_attrs: vec![SerdeAttribute::Rename(rename.to_string())],
deprecated: false,
..Default::default()
});
}
fn process_union(
&self,
name: &str,
schema: &ObjectSchema,
kind: UnionKind,
mut cache: Option<&mut SharedSchemaCache>,
) -> anyhow::Result<Vec<RustType>> {
let variants_src = match kind {
UnionKind::OneOf => &schema.one_of,
UnionKind::AnyOf => &schema.any_of,
};
let mut inline_types = vec![];
let mut variants = vec![];
let mut seen_names = BTreeSet::new();
for (i, variant_ref) in variants_src.iter().enumerate() {
let resolved = variant_ref
.resolve(self.graph.spec())
.with_context(|| format!("Schema resolution failed for union variant {i}"))?;
if resolved.is_null() {
continue;
}
let ref_name_opt = ReferenceExtractor::extract_ref_name_from_obj_ref(variant_ref).or_else(|| {
if resolved.all_of.len() == 1 {
ReferenceExtractor::extract_ref_name_from_obj_ref(&resolved.all_of[0])
} else {
None
}
});
let (variant, mut generated) = if let Some(schema_name) = ref_name_opt {
self.create_ref_variant(&schema_name, &resolved, &mut seen_names)
} else {
self.create_inline_variant(i, &resolved, name, &mut seen_names, cache.as_deref_mut())?
};
variants.push(variant);
inline_types.append(&mut generated);
}
strip_common_affixes(&mut variants);
let methods = if self.no_helpers {
vec![]
} else {
self.generate_methods(&variants, &inline_types, name, cache)
};
let main_enum = Self::build_union_enum_def(name, schema, kind, variants, methods);
inline_types.push(main_enum);
Ok(inline_types)
}
fn generate_methods(
&self,
variants: &[VariantDef],
inline_types: &[RustType],
enum_name: &str,
mut cache: Option<&mut SharedSchemaCache>,
) -> Vec<EnumMethod> {
let enum_name = to_rust_type_name(enum_name);
let struct_map: BTreeMap<_, _> = inline_types
.iter()
.filter_map(|t| match t {
RustType::Struct(s) => Some((s.name.to_string(), s)),
_ => None,
})
.collect();
let eligible: Vec<_> = variants
.iter()
.filter_map(|v| {
let type_ref = v.single_wrapped_type()?;
let method_kind = self.get_method_kind_for_type(type_ref, &v.name, &struct_map, cache.as_deref_mut())?;
Some((v.name.clone(), method_kind))
})
.collect();
if eligible.is_empty() {
return vec![];
}
let variant_names: Vec<String> = eligible.iter().map(|(name, _)| name.to_string()).collect();
let method_names = derive_method_names(&enum_name, &variant_names);
let mut seen = BTreeSet::new();
eligible
.into_iter()
.zip(method_names)
.map(|((variant_name, kind), base_name)| {
let method_name = ensure_unique(&base_name, &seen);
seen.insert(method_name.clone());
EnumMethod::new(method_name, &variant_name, kind)
})
.collect()
}
fn get_method_kind_for_type(
&self,
type_ref: &TypeRef,
variant_name: &EnumVariantToken,
struct_map: &BTreeMap<String, &StructDef>,
cache: Option<&mut SharedSchemaCache>,
) -> Option<EnumMethodKind> {
let base_name = type_ref.unboxed_base_type_name();
let struct_def = if let Some(&s) = struct_map.get(&base_name) {
Some(s.clone())
} else {
self.lookup_struct_def(type_ref, cache)
};
let s = struct_def.as_ref()?;
if !s.has_default() || type_ref.is_array {
return None;
}
let required: Vec<_> = s.required_fields().collect();
match required.len() {
0 => Some(EnumMethodKind::SimpleConstructor {
variant_name: variant_name.clone(),
wrapped_type: type_ref.clone(),
}),
1 => Some(EnumMethodKind::ParameterizedConstructor {
variant_name: variant_name.clone(),
wrapped_type: type_ref.clone(),
param_name: required[0].name.to_string(),
param_type: required[0].rust_type.clone(),
}),
_ => None,
}
}
fn lookup_struct_def(&self, type_ref: &TypeRef, cache: Option<&mut SharedSchemaCache>) -> Option<StructDef> {
let schema_name = type_ref.unboxed_base_type_name();
let schema = self.graph.get_schema(&schema_name)?;
if !schema.is_object() && schema.properties.is_empty() {
return None;
}
let struct_result = self
.struct_converter
.convert_struct(&schema_name, schema, None, cache)
.ok()?;
match struct_result.result {
RustType::Struct(s) => Some(s),
_ => None,
}
}
fn create_ref_variant(
&self,
schema_name: &str,
resolved_schema: &ObjectSchema,
seen_names: &mut BTreeSet<String>,
) -> (VariantDef, Vec<RustType>) {
let rust_type_name = to_rust_type_name(schema_name);
let mut type_ref = TypeRef::new(&rust_type_name);
if self.graph.is_cyclic(schema_name) {
type_ref = type_ref.with_boxed();
}
let variant_name = ensure_unique(&rust_type_name, seen_names);
let variant = VariantDef {
name: EnumVariantToken::from(variant_name),
docs: metadata::extract_docs(resolved_schema.description.as_ref()),
content: VariantContent::Tuple(vec![type_ref]),
deprecated: resolved_schema.deprecated.unwrap_or(false),
..Default::default()
};
(variant, vec![])
}
fn create_inline_variant(
&self,
index: usize,
resolved_schema: &ObjectSchema,
enum_name: &str,
seen_names: &mut BTreeSet<String>,
mut cache: Option<&mut SharedSchemaCache>,
) -> anyhow::Result<(VariantDef, Vec<RustType>)> {
let mut resolved_schema_merged = resolved_schema.clone();
if !resolved_schema.all_of.is_empty() {
resolved_schema_merged = self.type_resolver.merge_all_of_schema(resolved_schema)?;
}
let resolved_schema = &resolved_schema_merged;
if let Some(const_value) = &resolved_schema.const_value {
let normalized = VariantNameNormalizer::normalize(const_value)
.ok_or_else(|| anyhow::anyhow!("Unsupported const value type: {const_value}"))?;
let variant_name = ensure_unique(&normalized.name, seen_names);
let variant = VariantDef {
name: EnumVariantToken::from(variant_name),
docs: metadata::extract_docs(resolved_schema.description.as_ref()),
content: VariantContent::Unit,
serde_attrs: vec![SerdeAttribute::Rename(normalized.rename_value)],
deprecated: resolved_schema.deprecated.unwrap_or(false),
};
return Ok((variant, vec![]));
}
let base_name = resolved_schema
.title
.as_ref()
.map_or_else(|| infer_variant_name(resolved_schema, index), |t| to_rust_type_name(t));
let variant_name = ensure_unique(&base_name, seen_names);
let (content, generated_types) = if resolved_schema.properties.is_empty() {
let mut array_conversion = None;
if resolved_schema.is_array() {
array_conversion =
self
.type_resolver
.resolve_nullable_array_union(&variant_name, resolved_schema, cache.as_deref_mut())?;
}
if let Some(conversion) = array_conversion {
(VariantContent::Tuple(vec![conversion.result]), conversion.inline_types)
} else if !resolved_schema.one_of.is_empty() || !resolved_schema.any_of.is_empty() {
let uses_one_of = !resolved_schema.one_of.is_empty();
let result = self.type_resolver.resolve_inline_union_type(
enum_name,
&variant_name,
resolved_schema,
uses_one_of,
cache,
)?;
(VariantContent::Tuple(vec![result.result]), result.inline_types)
} else {
let type_ref = self.type_resolver.resolve_type(resolved_schema)?;
(VariantContent::Tuple(vec![type_ref]), vec![])
}
} else {
let struct_name_prefix = format!("{enum_name}{variant_name}");
let result = self
.struct_converter
.convert_struct(&struct_name_prefix, resolved_schema, None, cache)?;
let (struct_def, mut inline_types) = (result.result, result.inline_types);
let struct_name = match &struct_def {
RustType::Struct(s) => s.name.clone(),
_ => unreachable!("convert_struct must return a Struct"),
};
inline_types.push(struct_def);
(VariantContent::Tuple(vec![TypeRef::new(struct_name)]), inline_types)
};
let variant = VariantDef {
name: EnumVariantToken::from(variant_name),
docs: metadata::extract_docs(resolved_schema.description.as_ref()),
content,
serde_attrs: vec![],
deprecated: resolved_schema.deprecated.unwrap_or(false),
};
Ok((variant, generated_types))
}
fn build_union_enum_def(
name: &str,
schema: &ObjectSchema,
_kind: UnionKind,
variants: Vec<VariantDef>,
methods: Vec<EnumMethod>,
) -> RustType {
if let Some(discriminator) = &schema.discriminator
&& let Some(mapping) = &discriminator.mapping
&& Self::all_variants_are_refs(&variants, mapping)
{
let disc_variants = Self::build_discriminated_variants(&variants, mapping);
return RustType::DiscriminatedEnum(DiscriminatedEnumDef {
name: EnumToken::from_raw(name),
docs: metadata::extract_docs(schema.description.as_ref()),
discriminator_field: discriminator.property_name.clone(),
variants: disc_variants,
..Default::default()
});
}
RustType::Enum(EnumDef {
name: EnumToken::from_raw(name),
docs: metadata::extract_docs(schema.description.as_ref()),
variants,
serde_attrs: vec![SerdeAttribute::Untagged],
case_insensitive: false,
methods,
..Default::default()
})
}
fn all_variants_are_refs(variants: &[VariantDef], mapping: &BTreeMap<String, String>) -> bool {
if variants.is_empty() || mapping.is_empty() {
return false;
}
let variant_types: BTreeSet<String> = variants.iter().filter_map(VariantDef::unboxed_type_name).collect();
mapping
.values()
.filter_map(|ref_path| Self::ref_path_to_type_name(ref_path))
.all(|type_name| variant_types.contains(&type_name))
}
fn build_discriminated_variants(
variants: &[VariantDef],
mapping: &BTreeMap<String, String>,
) -> Vec<DiscriminatedVariant> {
mapping
.iter()
.filter_map(|(disc_value, ref_path)| {
let expected_type = Self::ref_path_to_type_name(ref_path)?;
let variant = Self::find_variant_by_type(variants, &expected_type)?;
let type_ref = variant.single_wrapped_type()?;
Some(DiscriminatedVariant {
discriminator_value: disc_value.clone(),
variant_name: variant.name.to_string(),
type_name: type_ref.clone(),
})
})
.collect()
}
fn ref_path_to_type_name(ref_path: &str) -> Option<String> {
SchemaRegistry::extract_ref_name(ref_path).map(|name| to_rust_type_name(&name))
}
fn find_variant_by_type<'a>(variants: &'a [VariantDef], type_name: &str) -> Option<&'a VariantDef> {
variants
.iter()
.find(|v| v.unboxed_type_name().is_some_and(|name| name == type_name))
}
}