Skip to main content

rh_codegen/generators/
trait_generator.rs

1//! Trait generation functionality
2//!
3//! This module handles the generation of Rust traits from FHIR StructureDefinitions.
4
5use crate::fhir_types::StructureDefinition;
6use crate::generators::accessor_trait_generator::AccessorTraitGenerator;
7use crate::generators::existence_trait_generator::ExistenceTraitGenerator;
8use crate::generators::mutator_trait_generator::MutatorTraitGenerator;
9use crate::generators::DocumentationGenerator;
10use crate::rust_types::RustTrait;
11use crate::{CodegenError, CodegenResult};
12
13/// Trait generator for FHIR StructureDefinitions
14pub struct TraitGenerator {
15    accessor_generator: AccessorTraitGenerator,
16    mutator_generator: MutatorTraitGenerator,
17    existence_generator: ExistenceTraitGenerator,
18}
19
20impl Default for TraitGenerator {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl TraitGenerator {
27    /// Create a new trait generator
28    pub fn new() -> Self {
29        Self {
30            accessor_generator: AccessorTraitGenerator::new(),
31            mutator_generator: MutatorTraitGenerator::new(),
32            existence_generator: ExistenceTraitGenerator::new(),
33        }
34    }
35
36    /// Create a new trait generator with a specific crate name for generated doc examples
37    pub fn new_with_crate_name(crate_name: impl Into<String>) -> Self {
38        Self {
39            accessor_generator: AccessorTraitGenerator::new(),
40            mutator_generator: MutatorTraitGenerator::with_crate_name(crate_name),
41            existence_generator: ExistenceTraitGenerator::new(),
42        }
43    }
44
45    /// Generate a Rust trait from a FHIR StructureDefinition
46    pub fn generate_trait(
47        &mut self,
48        structure_def: &StructureDefinition,
49        category: &str,
50    ) -> CodegenResult<RustTrait> {
51        if structure_def.kind == "logical" {
52            return Err(CodegenError::Generation {
53                message: format!(
54                    "Skipping LogicalModel '{}' - logical models are not generated as traits",
55                    structure_def.name
56                ),
57            });
58        }
59
60        let trait_name = format!(
61            "{}{}",
62            crate::naming::Naming::struct_name(structure_def),
63            category
64        );
65        let mut rust_trait = RustTrait::new(trait_name);
66        rust_trait.doc_comment =
67            DocumentationGenerator::generate_trait_documentation(structure_def);
68
69        self.add_inheritance_relationship(&mut rust_trait, structure_def, category)?;
70
71        match category {
72            "Accessors" => {
73                self.accessor_generator
74                    .add_accessor_methods(&mut rust_trait, structure_def)?;
75            }
76            "Mutators" => {
77                self.mutator_generator
78                    .add_mutator_methods(&mut rust_trait, structure_def)?;
79            }
80            "Existence" => {
81                // For existence traits, we need to generate it differently
82                return self
83                    .existence_generator
84                    .generate_existence_trait(structure_def);
85            }
86            _ => {
87                return Err(CodegenError::Generation {
88                    message: format!("Unknown trait category: {category}"),
89                });
90            }
91        }
92
93        Ok(rust_trait)
94    }
95
96    /// Add inheritance relationship if the StructureDefinition has a base definition
97    fn add_inheritance_relationship(
98        &mut self,
99        rust_trait: &mut RustTrait,
100        structure_def: &StructureDefinition,
101        category: &str,
102    ) -> CodegenResult<()> {
103        // Special case: Resource itself should not have super traits to avoid circular dependency
104        if structure_def.id == "Resource" {
105            return Ok(());
106        }
107
108        if let Some(base_def) = &structure_def.base_definition {
109            if let Some(parent_trait_name) = self.extract_trait_name_from_url(base_def) {
110                if self.is_valid_fhir_base_type(&parent_trait_name) {
111                    rust_trait
112                        .super_traits
113                        .push(format!("{parent_trait_name}{category}"));
114                }
115            }
116        }
117        Ok(())
118    }
119
120    /// Extract trait name from FHIR StructureDefinition URL
121    fn extract_trait_name_from_url(&self, url: &str) -> Option<String> {
122        url.split('/').next_back().map(|s| s.to_string())
123    }
124
125    /// Check if the given type name is a valid FHIR base type for inheritance
126    fn is_valid_fhir_base_type(&self, type_name: &str) -> bool {
127        matches!(
128            type_name,
129            "Resource" | "DomainResource" | "Element" | "BackboneElement"
130        )
131    }
132}