use etdl_parser::ast::{
BackoffStrategy, Condition, EtlDocument, EventTree, Node,
};
use etdl_parser::asyncapi::AsyncApiRegistry;
use etdl_parser::ecel;
use std::collections::BTreeMap;
use crate::validate::Diagnostic;
use super::CodeGenerator;
pub struct RustCodeGenerator {
pub version: String,
}
impl RustCodeGenerator {
pub fn new() -> Self {
RustCodeGenerator {
version: "1.0.0".to_string(),
}
}
}
impl CodeGenerator for RustCodeGenerator {
fn generate_all(
&self,
doc: &EtlDocument,
fault_tree_probs: &BTreeMap<String, f64>,
_registry: &AsyncApiRegistry,
_diagnostics: &mut Vec<Diagnostic>,
) -> Result<String, String> {
let mut output = String::new();
output.push_str(&format!(
"// AUTOGENERATED BY ETDL COMPILER v{} - DO NOT EDIT DIRECTLY\n\n",
self.version
));
let imports = collect_imports(doc);
output.push_str(&imports);
output.push('\n');
let constants = generate_fault_tree_constants(doc, fault_tree_probs);
output.push_str(&constants);
for (_tree_name, tree) in &doc.event_trees {
let handler_code = generate_event_tree_handler(doc, tree, fault_tree_probs)?;
output.push_str(&handler_code);
output.push('\n');
}
Ok(output)
}
}
fn collect_imports(doc: &EtlDocument) -> String {
let mut imports = String::from(
"use std::time::Duration;\n\
use etdl_core::telemetry::BranchMonitor;\n\
use etdl_core::retry::{RetryPolicy, BackoffStrategy};\n",
);
let mut alias_set: BTreeMap<String, bool> = BTreeMap::new();
for (_tree_name, tree) in &doc.event_trees {
alias_set.insert(tree.initiating_event.message.alias.clone(), true);
for (_node_id, node) in &tree.nodes {
match node {
Node::Operation(op) => {
if let Some(ref ext_ref) = op.emits {
alias_set.insert(ext_ref.alias.clone(), true);
}
}
Node::Consequence(cons) => {
if let Some(ref ext_ref) = cons.channel {
alias_set.insert(ext_ref.alias.clone(), true);
}
if let Some(ref ext_ref) = cons.message {
alias_set.insert(ext_ref.alias.clone(), true);
}
}
_ => {}
}
}
}
for alias in alias_set.keys() {
let mod_name = alias.replace('-', "_");
imports.push_str(&format!("use {}::messages::*;\n", mod_name));
}
imports
}
fn generate_fault_tree_constants(
doc: &EtlDocument,
fault_tree_probs: &BTreeMap<String, f64>,
) -> String {
let mut output = String::new();
for (_tree_name, tree) in &doc.event_trees {
for (node_id, node) in &tree.nodes {
if let Node::Operation(op) = node {
if op.on_failure_probability_source.is_some() {
if let Some((ft_id, prob)) = find_fault_tree_prob(doc, node_id, fault_tree_probs) {
output.push_str(&format!(
"// Computed from faultTrees.{}.topEvent at build time (Section 5.16)\n",
ft_id
));
let const_name = to_upper_snake(&format!(
"{}_failure_probability",
node_id
));
output.push_str(&format!(
"const {}: f64 = {:.6};\n\n",
const_name, prob
));
}
}
}
}
}
output
}
fn find_fault_tree_prob(
_doc: &EtlDocument,
_node_id: &str,
fault_tree_probs: &BTreeMap<String, f64>,
) -> Option<(String, f64)> {
fault_tree_probs.iter().next().map(|(k, &v)| (k.clone(), v))
}
fn generate_event_tree_handler(
doc: &EtlDocument,
tree: &EventTree,
fault_tree_probs: &BTreeMap<String, f64>,
) -> Result<String, String> {
let mut output = String::new();
let fn_name = format!("handle_{}", to_snake_case(&tree.initiating_event.id));
let message_type = ref_to_rust_type(&tree.initiating_event.message);
output.push_str(&format!(
"pub async fn {}(message: {}) -> Result<(), WorkflowError> {{\n",
fn_name, message_type
));
let first_barrier = find_first_barrier(tree);
let monitor_name = first_barrier.map(|id| to_snake_case(id)).unwrap_or_else(|| "monitor".to_string());
output.push_str(&format!(
" let mut {} = BranchMonitor::new(\"{}\");\n\n",
monitor_name, first_barrier.unwrap_or("root")
));
let start_node_id = &tree.initiating_event.next;
let body = generate_node_code(
doc, tree, start_node_id, 1, fault_tree_probs, &monitor_name,
)?;
output.push_str(&body);
output.push_str(" Ok(())\n");
output.push_str("}\n");
Ok(output)
}
fn find_first_barrier(tree: &EventTree) -> Option<&str> {
let mut current = &tree.initiating_event.next;
loop {
match tree.nodes.get(current.as_str()) {
Some(Node::Barrier(_)) => return Some(current.as_str()),
Some(Node::Operation(op)) => {
current = &op.next;
}
Some(Node::Consequence(_)) => return None,
None => return None,
}
}
}
fn generate_node_code(
doc: &EtlDocument,
tree: &EventTree,
node_id: &str,
depth: usize,
fault_tree_probs: &BTreeMap<String, f64>,
monitor_name: &str,
) -> Result<String, String> {
let _indent = " ".repeat(depth);
let node = tree
.nodes
.get(node_id)
.ok_or_else(|| format!("node '{}' not found", node_id))?;
match node {
Node::Barrier(barrier) => {
generate_barrier_code(tree, node_id, barrier, depth, doc, fault_tree_probs, monitor_name)
}
Node::Operation(op) => {
generate_operation_code(tree, node_id, op, depth, doc, fault_tree_probs, monitor_name)
}
Node::Consequence(cons) => {
generate_consequence_code(cons, depth)
}
}
}
fn generate_barrier_code(
tree: &EventTree,
node_id: &str,
barrier: &etdl_parser::ast::Barrier,
depth: usize,
doc: &EtlDocument,
fault_tree_probs: &BTreeMap<String, f64>,
monitor_name: &str,
) -> Result<String, String> {
let indent = " ".repeat(depth);
let mut output = String::new();
for (i, branch) in barrier.branches.iter().enumerate() {
if i == 0 {
if branch.condition == Condition::Default {
let prob = get_branch_prob(branch, node_id, fault_tree_probs);
if let Some(p) = prob {
output.push_str(&format!(
"{} {}.record_branch(\"{}\", {:.6});\n",
indent, monitor_name, branch.outcome, p
));
}
let body = generate_node_code(
doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
)?;
output.push_str(&body);
} else {
let cond = condition_to_rust_code(&branch.condition);
output.push_str(&format!("{}if {} {{\n", indent, cond));
let prob = get_branch_prob(branch, node_id, fault_tree_probs);
if let Some(p) = prob {
output.push_str(&format!(
"{} {}.record_branch(\"{}\", {:.6});\n",
indent, monitor_name, branch.outcome, p
));
}
let body = generate_node_code(
doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
)?;
output.push_str(&body);
output.push_str(&format!("{}}}", indent));
}
} else if branch.condition == Condition::Default {
output.push_str(" else {\n");
let prob = get_branch_prob(branch, node_id, fault_tree_probs);
if let Some(p) = prob {
output.push_str(&format!(
"{} {}.record_branch(\"{}\", {:.6});\n",
indent, monitor_name, branch.outcome, p
));
}
let body = generate_node_code(
doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
)?;
output.push_str(&body);
output.push_str(&format!("{}}}\n", indent));
} else {
let cond = condition_to_rust_code(&branch.condition);
output.push_str(&format!(" else if {} {{\n", cond));
let prob = get_branch_prob(branch, node_id, fault_tree_probs);
if let Some(p) = prob {
output.push_str(&format!(
"{} {}.record_branch(\"{}\", {:.6});\n",
indent, monitor_name, branch.outcome, p
));
}
let body = generate_node_code(
doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
)?;
output.push_str(&body);
output.push_str(&format!("{}}}", indent));
}
}
Ok(output)
}
fn generate_operation_code(
tree: &EventTree,
node_id: &str,
op: &etdl_parser::ast::Operation,
depth: usize,
doc: &EtlDocument,
fault_tree_probs: &BTreeMap<String, f64>,
monitor_name: &str,
) -> Result<String, String> {
let indent = " ".repeat(depth);
let mut output = String::new();
let handler_name = to_snake_case(&op.handler);
let timeout = op.timeout_ms.unwrap_or(5000);
if let Some(ref retry) = op.retry_policy {
let strategy = match retry.backoff_strategy.as_ref().unwrap_or(&BackoffStrategy::Fixed) {
BackoffStrategy::Exponential => "BackoffStrategy::Exponential",
BackoffStrategy::Fixed => "BackoffStrategy::Fixed",
};
output.push_str(&format!(
"{}let retry = RetryPolicy {{\n\
{} max_attempts: {},\n\
{} backoff_ms: {},\n\
{} strategy: {},\n\
{}}};\n",
indent, indent, retry.max_attempts, indent, retry.backoff_ms, indent, strategy, indent
));
output.push_str(&format!(
"{}match retry.execute(|| {}(&message), Duration::from_millis({})).await {{\n",
indent, handler_name, timeout
));
} else {
output.push_str(&format!(
"{}match {}(&message).await {{\n",
indent, handler_name
));
}
output.push_str(&format!("{} Ok(_result) => {{\n", indent));
let next_node = tree.nodes.get(&op.next);
match next_node {
Some(Node::Consequence(cons)) => {
match cons.consequence_operation {
etdl_parser::ast::ConsequenceOperation::Send => {
if let Some(ref channel_ref) = cons.channel {
let channel_name = extract_last_segment(channel_ref);
output.push_str(&format!(
"{} publish_to_channel(\"{}\", _result).await?;\n",
indent, channel_name
));
}
}
etdl_parser::ast::ConsequenceOperation::Terminate => {}
}
}
Some(_) => {
generate_node_code(
doc, tree, &op.next, depth + 2, fault_tree_probs, monitor_name,
).map(|body| output.push_str(&body))?;
}
None => {}
}
output.push_str(&format!("{} }}\n", indent));
if let Some(ref _on_failure_id) = op.on_failure {
output.push_str(&format!("{} Err(err) => {{\n", indent));
if let Some(ref _ps) = op.on_failure_probability_source {
let const_name = to_upper_snake(&format!(
"{}_failure_probability",
node_id
));
let prob_exists = find_fault_tree_prob(doc, node_id, fault_tree_probs).is_some();
if prob_exists {
output.push_str(&format!(
"{} {}.record_failure(\"{}\", &err, Some({}));\n",
indent, monitor_name, node_id, const_name
));
} else {
output.push_str(&format!(
"{} {}.record_failure(\"{}\", &err, None);\n",
indent, monitor_name, node_id
));
}
} else {
output.push_str(&format!(
"{} {}.record_failure(\"{}\", &err, None);\n",
indent, monitor_name, node_id
));
}
let on_failure_id = op.on_failure.as_ref().unwrap();
match tree.nodes.get(on_failure_id) {
Some(Node::Consequence(cons)) => {
match cons.consequence_operation {
etdl_parser::ast::ConsequenceOperation::Send => {
if let Some(ref channel_ref) = cons.channel {
let channel_name = extract_last_segment(channel_ref);
output.push_str(&format!(
"{} publish_to_channel(\"{}\", message).await?;\n",
indent, channel_name
));
}
}
etdl_parser::ast::ConsequenceOperation::Terminate => {}
}
}
_ => {
generate_node_code(
doc, tree, on_failure_id, depth + 2, fault_tree_probs, monitor_name,
).map(|body| output.push_str(&body))?;
}
}
output.push_str(&format!("{} }}\n", indent));
} else {
output.push_str(&format!(
"{} Err(err) => return Err(WorkflowError::new(format!(\"{{}}\", err))),\n",
indent
));
}
output.push_str(&format!("{}}}\n", indent));
Ok(output)
}
fn generate_consequence_code(
cons: &etdl_parser::ast::Consequence,
depth: usize,
) -> Result<String, String> {
let indent = " ".repeat(depth);
let mut output = String::new();
match cons.consequence_operation {
etdl_parser::ast::ConsequenceOperation::Send => {
if let Some(ref channel_ref) = cons.channel {
let channel_name = extract_last_segment(channel_ref);
output.push_str(&format!(
"{}publish_to_channel(\"{}\", message).await?;\n",
indent, channel_name
));
}
}
etdl_parser::ast::ConsequenceOperation::Terminate => {}
}
if depth == 1 {
output.push_str(&format!("{}Ok(())\n", indent));
}
Ok(output)
}
fn condition_to_rust_code(condition: &Condition) -> String {
match condition {
Condition::Default => "true".to_string(),
Condition::Comparison(cmp) => {
let (left_node, has_wildcard) = build_path_expression(&cmp.left);
let right_val = operand_to_val_str(&cmp.right, false);
if has_wildcard {
format!(
"{}.iter().all(|item| item{} {} {})",
left_node.path_prefix,
left_node.remaining_path,
comparator_str(&cmp.op),
right_val
)
} else {
format!(
"{} {} {}",
left_node.path_prefix + &left_node.remaining_path,
comparator_str(&cmp.op),
right_val
)
}
}
}
}
struct PathParts {
path_prefix: String,
remaining_path: String,
}
fn build_path_expression(operand: &ecel::Operand) -> (PathParts, bool) {
match operand {
ecel::Operand::Path(path_expr) => {
let segments = &path_expr.segments;
let mut pre_wildcard = Vec::new();
let mut post_wildcard = Vec::new();
let mut has_wildcard = false;
for (i, seg) in segments.iter().enumerate() {
if i == 0 {
continue;
}
if has_wildcard {
post_wildcard.push(seg.clone());
} else if matches!(seg, ecel::PathSegment::Wildcard) {
has_wildcard = true;
} else {
pre_wildcard.push(seg.clone());
}
}
let mut prefix = String::from("message");
for seg in &pre_wildcard {
match seg {
ecel::PathSegment::Field(name) => {
prefix.push('.');
prefix.push_str(&to_snake_case(name));
}
ecel::PathSegment::Index(idx) => {
prefix.push_str(&format!("[{}]", idx));
}
ecel::PathSegment::QuotedKey(name) => {
prefix.push_str(&format!("[\"{}\"]", name));
}
_ => {}
}
}
let mut suffix = String::new();
for seg in &post_wildcard {
match seg {
ecel::PathSegment::Field(name) => {
suffix.push('.');
suffix.push_str(&to_snake_case(name));
}
ecel::PathSegment::Index(idx) => {
suffix.push_str(&format!("[{}]", idx));
}
ecel::PathSegment::QuotedKey(name) => {
suffix.push_str(&format!("[\"{}\"]", name));
}
_ => {}
}
}
(PathParts {
path_prefix: prefix,
remaining_path: suffix,
}, has_wildcard)
}
ecel::Operand::Literal(_) => {
(PathParts {
path_prefix: String::new(),
remaining_path: String::new(),
}, false)
}
}
}
fn _operand_to_path_str(operand: &ecel::Operand) -> String {
match operand {
ecel::Operand::Path(path) => {
let segments = &path.segments;
segments.iter()
.skip(1)
.map(|seg| match seg {
ecel::PathSegment::Field(name) => to_snake_case(name),
ecel::PathSegment::Wildcard => "*".to_string(),
ecel::PathSegment::Index(idx) => idx.to_string(),
ecel::PathSegment::QuotedKey(name) => format!("\"{}\"", name),
})
.collect::<Vec<_>>()
.join(".")
}
ecel::Operand::Literal(_) => String::new(),
}
}
fn operand_to_val_str(operand: &ecel::Operand, _in_closure: bool) -> String {
match operand {
ecel::Operand::Path(_) => "".to_string(),
ecel::Operand::Literal(lit) => match lit {
ecel::Literal::Number(n) => n.to_string(),
ecel::Literal::String(s) => format!("\"{}\"", s),
ecel::Literal::Bool(b) => b.to_string(),
ecel::Literal::Null => "None".to_string(),
ecel::Literal::Array(items) => {
let inner: Vec<String> = items.iter().map(|item| match item {
ecel::Literal::Number(n) => n.to_string(),
ecel::Literal::String(s) => format!("\"{}\"", s),
ecel::Literal::Bool(b) => b.to_string(),
ecel::Literal::Null => "None".to_string(),
ecel::Literal::Array(_) => "[]".to_string(),
}).collect();
format!("vec![{}]", inner.join(", "))
}
},
}
}
fn comparator_str(op: &ecel::Comparator) -> &str {
match op {
ecel::Comparator::Eq => "==",
ecel::Comparator::Neq => "!=",
ecel::Comparator::Gte => ">=",
ecel::Comparator::Lte => "<=",
ecel::Comparator::Gt => ">",
ecel::Comparator::Lt => "<",
ecel::Comparator::In => "in",
ecel::Comparator::Matches => "matches",
}
}
fn ref_to_rust_type(ext_ref: &etdl_parser::ast::ExternalRef) -> String {
extract_last_segment(ext_ref)
}
fn extract_last_segment(ext_ref: &etdl_parser::ast::ExternalRef) -> String {
let pointer = &ext_ref.pointer;
let parts: Vec<&str> = pointer.split('/').collect();
let last = parts.last().unwrap_or(&"Unknown");
to_pascal_case(last)
}
fn get_branch_prob(
branch: &etdl_parser::ast::Branch,
_node_id: &str,
fault_tree_probs: &BTreeMap<String, f64>,
) -> Option<f64> {
if let Some(ref ps) = branch.probability_source {
let ft_id = extract_ft_id(&ps.pointer);
return fault_tree_probs.get(&ft_id).copied();
}
branch.effective_probability()
}
fn extract_ft_id(pointer: &str) -> String {
pointer
.trim_start_matches("#/faultTrees/")
.trim_end_matches("/topEvent")
.to_string()
}
fn to_snake_case(s: &str) -> String {
let mut result = String::new();
for (i, c) in s.chars().enumerate() {
if c.is_uppercase() {
if i > 0 {
result.push('_');
}
result.push(c.to_lowercase().next().unwrap());
} else {
result.push(c);
}
}
result
}
fn to_upper_snake(s: &str) -> String {
let snake = to_snake_case(s);
snake.to_uppercase()
}
fn to_pascal_case(s: &str) -> String {
let mut result = String::new();
let mut capitalize = true;
for c in s.chars() {
if c == '_' || c == '-' || c == ' ' || c == '/' {
capitalize = true;
} else if capitalize {
result.push(c.to_uppercase().next().unwrap());
capitalize = false;
} else {
result.push(c);
}
}
result
}