use std::collections::{HashMap, HashSet};
use crate::core::ir::{EnumDef, TypeDef};
use crate::e2e::codegen::call_ir::named_type;
use super::parse::{parse_path, segment_name};
use super::types::PathSegment;
use super::types::{IrEnumMap, TaggedEnumWire};
pub(super) fn build_ir_enum_map(type_defs: &[TypeDef], enums: &[EnumDef]) -> IrEnumMap {
let enum_names: HashSet<&str> = enums.iter().map(|e| e.name.as_str()).collect();
let struct_names: HashSet<&str> = type_defs.iter().map(|t| t.name.as_str()).collect();
let mut field_types: HashMap<String, HashMap<String, String>> = HashMap::new();
let mut enum_fields: HashMap<String, HashSet<String>> = HashMap::new();
let mut enum_field_types: HashMap<String, HashMap<String, String>> = HashMap::new();
for type_def in type_defs {
for field in &type_def.fields {
let Some(named) = named_type(&field.ty) else {
continue;
};
if enum_names.contains(named) {
enum_fields
.entry(type_def.name.clone())
.or_default()
.insert(field.name.clone());
enum_field_types
.entry(type_def.name.clone())
.or_default()
.insert(field.name.clone(), named.to_string());
} else if struct_names.contains(named) {
field_types
.entry(type_def.name.clone())
.or_default()
.insert(field.name.clone(), named.to_string());
}
}
}
let (variant_payload_types, variant_payload_is_collection) = build_variant_payload_types(enums);
IrEnumMap {
field_types,
enum_fields,
enum_field_types,
variant_payload_types,
variant_payload_is_collection,
tagged_enum_wire: build_tagged_enum_wire(enums),
data_carrying_enum_names: data_carrying_enum_names(enums),
enum_wire_variants: build_enum_wire_variants(enums),
root_type: None,
}
}
fn data_carrying_enum_names(enums: &[EnumDef]) -> HashSet<String> {
enums
.iter()
.filter(|enum_def| enum_def.variants.iter().any(|variant| !variant.fields.is_empty()))
.map(|enum_def| enum_def.name.clone())
.collect()
}
fn build_enum_wire_variants(enums: &[EnumDef]) -> HashMap<String, HashMap<String, String>> {
let mut per_enum = HashMap::new();
for enum_def in enums {
let identifiers: HashSet<&str> = enum_def.variants.iter().map(|v| v.name.as_str()).collect();
let mut by_wire: HashMap<String, String> = HashMap::new();
let mut ambiguous: HashSet<String> = HashSet::new();
for variant in &enum_def.variants {
let wire = crate::codegen::naming::wire_variant_value(
&variant.name,
variant.serde_rename.as_deref(),
enum_def.serde_rename_all.as_deref(),
);
if wire == variant.name || identifiers.contains(wire.as_str()) {
continue;
}
if by_wire.insert(wire.clone(), variant.name.clone()).is_some() {
ambiguous.insert(wire);
}
}
for wire in &ambiguous {
by_wire.remove(wire);
}
if !by_wire.is_empty() {
per_enum.insert(enum_def.name.clone(), by_wire);
}
}
per_enum
}
fn build_tagged_enum_wire(enums: &[EnumDef]) -> HashMap<String, TaggedEnumWire> {
enums
.iter()
.filter_map(|enum_def| {
let tag = enum_def.serde_tag.clone()?;
let variants = enum_def
.variants
.iter()
.map(|variant| {
let wire = crate::codegen::naming::wire_variant_value(
&variant.name,
variant.serde_rename.as_deref(),
enum_def.serde_rename_all.as_deref(),
);
(variant.name.clone(), wire)
})
.collect();
Some((enum_def.name.clone(), TaggedEnumWire { tag, variants }))
})
.collect()
}
type VariantPayloadTypeMap = HashMap<String, HashMap<String, (String, String)>>;
type VariantPayloadCollectionMap = HashMap<String, HashSet<String>>;
fn build_variant_payload_types(enums: &[EnumDef]) -> (VariantPayloadTypeMap, VariantPayloadCollectionMap) {
let mut variant_payload_types: HashMap<String, HashMap<String, (String, String)>> = HashMap::new();
let mut variant_payload_is_collection: HashMap<String, HashSet<String>> = HashMap::new();
for enum_def in enums {
for variant in &enum_def.variants {
let [only_field] = variant.fields.as_slice() else {
continue;
};
let Some(named) = named_type(&only_field.ty) else {
continue;
};
variant_payload_types
.entry(enum_def.name.clone())
.or_default()
.insert(variant.name.clone(), (only_field.name.clone(), named.to_string()));
if super::ir_collection::is_vec_type(&only_field.ty) {
variant_payload_is_collection
.entry(enum_def.name.clone())
.or_default()
.insert(variant.name.clone());
}
}
}
(variant_payload_types, variant_payload_is_collection)
}
fn resolve_owner<'a>(map: &'a IrEnumMap, root: &'a str, prefix: &[PathSegment]) -> Option<&'a str> {
let mut owner = root;
for segment in prefix {
let name = segment_name(segment)?;
let next = map.field_types.get(owner).and_then(|fields| fields.get(name))?;
owner = next.as_str();
}
Some(owner)
}
pub(super) fn is_enum_path(map: &IrEnumMap, path: &str) -> bool {
let Some(root) = map.root_type.as_deref() else {
return false;
};
let segments = parse_path(path);
let Some((last, prefix)) = segments.split_last() else {
return false;
};
let Some(owner) = resolve_owner(map, root, prefix) else {
return false;
};
let Some(name) = segment_name(last) else {
return false;
};
map.enum_fields.get(owner).is_some_and(|fields| fields.contains(name))
}
pub(super) fn enum_type_at_path(map: &IrEnumMap, path: &str) -> Option<String> {
let root = map.root_type.as_deref()?;
enum_type_at_path_from(map, root, path)
}
pub(super) fn enum_type_at_path_from(map: &IrEnumMap, owner: &str, path: &str) -> Option<String> {
let segments = parse_path(path);
let (last, prefix) = segments.split_last()?;
let owner = resolve_owner(map, owner, prefix)?;
let name = segment_name(last)?;
map.enum_field_types
.get(owner)
.and_then(|fields| fields.get(name))
.cloned()
}