1use etdl_parser::ast::{
2 BackoffStrategy, Condition, EtlDocument, EventTree, Node,
3};
4use etdl_parser::asyncapi::AsyncApiRegistry;
5use etdl_parser::ecel;
6use std::collections::BTreeMap;
7
8use crate::validate::Diagnostic;
9
10use super::CodeGenerator;
11
12pub struct RustCodeGenerator {
13 pub version: String,
14}
15
16impl RustCodeGenerator {
17 pub fn new() -> Self {
18 RustCodeGenerator {
19 version: "1.0.0".to_string(),
20 }
21 }
22}
23
24impl CodeGenerator for RustCodeGenerator {
25 fn generate_all(
26 &self,
27 doc: &EtlDocument,
28 fault_tree_probs: &BTreeMap<String, f64>,
29 _registry: &AsyncApiRegistry,
30 _diagnostics: &mut Vec<Diagnostic>,
31 ) -> Result<String, String> {
32 let mut output = String::new();
33
34 output.push_str(&format!(
35 "// AUTOGENERATED BY ETDL COMPILER v{} - DO NOT EDIT DIRECTLY\n\n",
36 self.version
37 ));
38
39 let imports = collect_imports(doc);
40 output.push_str(&imports);
41 output.push('\n');
42
43 let constants = generate_fault_tree_constants(doc, fault_tree_probs);
44 output.push_str(&constants);
45
46 for (_tree_name, tree) in &doc.event_trees {
47 let handler_code = generate_event_tree_handler(doc, tree, fault_tree_probs)?;
48 output.push_str(&handler_code);
49 output.push('\n');
50 }
51
52 Ok(output)
53 }
54}
55
56fn collect_imports(doc: &EtlDocument) -> String {
57 let mut imports = String::from(
58 "use std::time::Duration;\n\
59 use etdl_core::telemetry::BranchMonitor;\n\
60 use etdl_core::retry::{RetryPolicy, BackoffStrategy};\n",
61 );
62
63 let mut alias_set: BTreeMap<String, bool> = BTreeMap::new();
64
65 for (_tree_name, tree) in &doc.event_trees {
66 alias_set.insert(tree.initiating_event.message.alias.clone(), true);
67
68 for (_node_id, node) in &tree.nodes {
69 match node {
70 Node::Operation(op) => {
71 if let Some(ref ext_ref) = op.emits {
72 alias_set.insert(ext_ref.alias.clone(), true);
73 }
74 }
75 Node::Consequence(cons) => {
76 if let Some(ref ext_ref) = cons.channel {
77 alias_set.insert(ext_ref.alias.clone(), true);
78 }
79 if let Some(ref ext_ref) = cons.message {
80 alias_set.insert(ext_ref.alias.clone(), true);
81 }
82 }
83 _ => {}
84 }
85 }
86 }
87
88 for alias in alias_set.keys() {
89 let mod_name = alias.replace('-', "_");
90 imports.push_str(&format!("use {}::messages::*;\n", mod_name));
91 }
92
93 imports
94}
95
96fn generate_fault_tree_constants(
97 doc: &EtlDocument,
98 fault_tree_probs: &BTreeMap<String, f64>,
99) -> String {
100 let mut output = String::new();
101
102 for (_tree_name, tree) in &doc.event_trees {
103 for (node_id, node) in &tree.nodes {
104 if let Node::Operation(op) = node {
105 if op.on_failure_probability_source.is_some() {
106 if let Some((ft_id, prob)) = find_fault_tree_prob(doc, node_id, fault_tree_probs) {
107 output.push_str(&format!(
108 "// Computed from faultTrees.{}.topEvent at build time (Section 5.16)\n",
109 ft_id
110 ));
111 let const_name = to_upper_snake(&format!(
112 "{}_failure_probability",
113 node_id
114 ));
115 output.push_str(&format!(
116 "const {}: f64 = {:.6};\n\n",
117 const_name, prob
118 ));
119 }
120 }
121 }
122 }
123 }
124
125 output
126}
127
128fn find_fault_tree_prob(
129 _doc: &EtlDocument,
130 _node_id: &str,
131 fault_tree_probs: &BTreeMap<String, f64>,
132) -> Option<(String, f64)> {
133 fault_tree_probs.iter().next().map(|(k, &v)| (k.clone(), v))
134}
135
136fn generate_event_tree_handler(
137 doc: &EtlDocument,
138 tree: &EventTree,
139 fault_tree_probs: &BTreeMap<String, f64>,
140) -> Result<String, String> {
141 let mut output = String::new();
142
143 let fn_name = format!("handle_{}", to_snake_case(&tree.initiating_event.id));
144 let message_type = ref_to_rust_type(&tree.initiating_event.message);
145
146 output.push_str(&format!(
147 "pub async fn {}(message: {}) -> Result<(), WorkflowError> {{\n",
148 fn_name, message_type
149 ));
150
151 let first_barrier = find_first_barrier(tree);
152 let monitor_name = first_barrier.map(|id| to_snake_case(id)).unwrap_or_else(|| "monitor".to_string());
153
154 output.push_str(&format!(
155 " let mut {} = BranchMonitor::new(\"{}\");\n\n",
156 monitor_name, first_barrier.unwrap_or("root")
157 ));
158
159 let start_node_id = &tree.initiating_event.next;
160 let body = generate_node_code(
161 doc, tree, start_node_id, 1, fault_tree_probs, &monitor_name,
162 )?;
163 output.push_str(&body);
164
165 output.push_str(" Ok(())\n");
166 output.push_str("}\n");
167
168 Ok(output)
169}
170
171fn find_first_barrier(tree: &EventTree) -> Option<&str> {
172 let mut current = &tree.initiating_event.next;
173 loop {
174 match tree.nodes.get(current.as_str()) {
175 Some(Node::Barrier(_)) => return Some(current.as_str()),
176 Some(Node::Operation(op)) => {
177 current = &op.next;
178 }
179 Some(Node::Consequence(_)) => return None,
180 None => return None,
181 }
182 }
183}
184
185fn generate_node_code(
186 doc: &EtlDocument,
187 tree: &EventTree,
188 node_id: &str,
189 depth: usize,
190 fault_tree_probs: &BTreeMap<String, f64>,
191 monitor_name: &str,
192) -> Result<String, String> {
193 let _indent = " ".repeat(depth);
194
195 let node = tree
196 .nodes
197 .get(node_id)
198 .ok_or_else(|| format!("node '{}' not found", node_id))?;
199
200 match node {
201 Node::Barrier(barrier) => {
202 generate_barrier_code(tree, node_id, barrier, depth, doc, fault_tree_probs, monitor_name)
203 }
204 Node::Operation(op) => {
205 generate_operation_code(tree, node_id, op, depth, doc, fault_tree_probs, monitor_name)
206 }
207 Node::Consequence(cons) => {
208 generate_consequence_code(cons, depth)
209 }
210 }
211}
212
213fn generate_barrier_code(
214 tree: &EventTree,
215 node_id: &str,
216 barrier: &etdl_parser::ast::Barrier,
217 depth: usize,
218 doc: &EtlDocument,
219 fault_tree_probs: &BTreeMap<String, f64>,
220 monitor_name: &str,
221) -> Result<String, String> {
222 let indent = " ".repeat(depth);
223 let mut output = String::new();
224
225 for (i, branch) in barrier.branches.iter().enumerate() {
226 if i == 0 {
227 if branch.condition == Condition::Default {
228 let prob = get_branch_prob(branch, node_id, fault_tree_probs);
229 if let Some(p) = prob {
230 output.push_str(&format!(
231 "{} {}.record_branch(\"{}\", {:.6});\n",
232 indent, monitor_name, branch.outcome, p
233 ));
234 }
235 let body = generate_node_code(
236 doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
237 )?;
238 output.push_str(&body);
239 } else {
240 let cond = condition_to_rust_code(&branch.condition);
241 output.push_str(&format!("{}if {} {{\n", indent, cond));
242
243 let prob = get_branch_prob(branch, node_id, fault_tree_probs);
244 if let Some(p) = prob {
245 output.push_str(&format!(
246 "{} {}.record_branch(\"{}\", {:.6});\n",
247 indent, monitor_name, branch.outcome, p
248 ));
249 }
250
251 let body = generate_node_code(
252 doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
253 )?;
254 output.push_str(&body);
255 output.push_str(&format!("{}}}", indent));
256 }
257 } else if branch.condition == Condition::Default {
258 output.push_str(" else {\n");
259
260 let prob = get_branch_prob(branch, node_id, fault_tree_probs);
261 if let Some(p) = prob {
262 output.push_str(&format!(
263 "{} {}.record_branch(\"{}\", {:.6});\n",
264 indent, monitor_name, branch.outcome, p
265 ));
266 }
267
268 let body = generate_node_code(
269 doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
270 )?;
271 output.push_str(&body);
272 output.push_str(&format!("{}}}\n", indent));
273 } else {
274 let cond = condition_to_rust_code(&branch.condition);
275 output.push_str(&format!(" else if {} {{\n", cond));
276
277 let prob = get_branch_prob(branch, node_id, fault_tree_probs);
278 if let Some(p) = prob {
279 output.push_str(&format!(
280 "{} {}.record_branch(\"{}\", {:.6});\n",
281 indent, monitor_name, branch.outcome, p
282 ));
283 }
284
285 let body = generate_node_code(
286 doc, tree, &branch.next, depth + 1, fault_tree_probs, monitor_name,
287 )?;
288 output.push_str(&body);
289 output.push_str(&format!("{}}}", indent));
290 }
291 }
292
293 Ok(output)
294}
295
296fn generate_operation_code(
297 tree: &EventTree,
298 node_id: &str,
299 op: &etdl_parser::ast::Operation,
300 depth: usize,
301 doc: &EtlDocument,
302 fault_tree_probs: &BTreeMap<String, f64>,
303 monitor_name: &str,
304) -> Result<String, String> {
305 let indent = " ".repeat(depth);
306 let mut output = String::new();
307
308 let handler_name = to_snake_case(&op.handler);
309 let timeout = op.timeout_ms.unwrap_or(5000);
310
311 if let Some(ref retry) = op.retry_policy {
312 let strategy = match retry.backoff_strategy.as_ref().unwrap_or(&BackoffStrategy::Fixed) {
313 BackoffStrategy::Exponential => "BackoffStrategy::Exponential",
314 BackoffStrategy::Fixed => "BackoffStrategy::Fixed",
315 };
316 output.push_str(&format!(
317 "{}let retry = RetryPolicy {{\n\
318 {} max_attempts: {},\n\
319 {} backoff_ms: {},\n\
320 {} strategy: {},\n\
321 {}}};\n",
322 indent, indent, retry.max_attempts, indent, retry.backoff_ms, indent, strategy, indent
323 ));
324 output.push_str(&format!(
325 "{}match retry.execute(|| {}(&message), Duration::from_millis({})).await {{\n",
326 indent, handler_name, timeout
327 ));
328 } else {
329 output.push_str(&format!(
330 "{}match {}(&message).await {{\n",
331 indent, handler_name
332 ));
333 }
334
335 output.push_str(&format!("{} Ok(_result) => {{\n", indent));
336
337 let next_node = tree.nodes.get(&op.next);
338 match next_node {
339 Some(Node::Consequence(cons)) => {
340 match cons.consequence_operation {
341 etdl_parser::ast::ConsequenceOperation::Send => {
342 if let Some(ref channel_ref) = cons.channel {
343 let channel_name = extract_last_segment(channel_ref);
344 output.push_str(&format!(
345 "{} publish_to_channel(\"{}\", _result).await?;\n",
346 indent, channel_name
347 ));
348 }
349 }
350 etdl_parser::ast::ConsequenceOperation::Terminate => {}
351 }
352 }
353 Some(_) => {
354 generate_node_code(
355 doc, tree, &op.next, depth + 2, fault_tree_probs, monitor_name,
356 ).map(|body| output.push_str(&body))?;
357 }
358 None => {}
359 }
360
361 output.push_str(&format!("{} }}\n", indent));
362
363 if let Some(ref _on_failure_id) = op.on_failure {
364 output.push_str(&format!("{} Err(err) => {{\n", indent));
365
366 if let Some(ref _ps) = op.on_failure_probability_source {
367 let const_name = to_upper_snake(&format!(
368 "{}_failure_probability",
369 node_id
370 ));
371 let prob_exists = find_fault_tree_prob(doc, node_id, fault_tree_probs).is_some();
372 if prob_exists {
373 output.push_str(&format!(
374 "{} {}.record_failure(\"{}\", &err, Some({}));\n",
375 indent, monitor_name, node_id, const_name
376 ));
377 } else {
378 output.push_str(&format!(
379 "{} {}.record_failure(\"{}\", &err, None);\n",
380 indent, monitor_name, node_id
381 ));
382 }
383 } else {
384 output.push_str(&format!(
385 "{} {}.record_failure(\"{}\", &err, None);\n",
386 indent, monitor_name, node_id
387 ));
388 }
389
390 let on_failure_id = op.on_failure.as_ref().unwrap();
391 match tree.nodes.get(on_failure_id) {
392 Some(Node::Consequence(cons)) => {
393 match cons.consequence_operation {
394 etdl_parser::ast::ConsequenceOperation::Send => {
395 if let Some(ref channel_ref) = cons.channel {
396 let channel_name = extract_last_segment(channel_ref);
397 output.push_str(&format!(
398 "{} publish_to_channel(\"{}\", message).await?;\n",
399 indent, channel_name
400 ));
401 }
402 }
403 etdl_parser::ast::ConsequenceOperation::Terminate => {}
404 }
405 }
406 _ => {
407 generate_node_code(
408 doc, tree, on_failure_id, depth + 2, fault_tree_probs, monitor_name,
409 ).map(|body| output.push_str(&body))?;
410 }
411 }
412
413 output.push_str(&format!("{} }}\n", indent));
414 } else {
415 output.push_str(&format!(
416 "{} Err(err) => return Err(WorkflowError::new(format!(\"{{}}\", err))),\n",
417 indent
418 ));
419 }
420
421 output.push_str(&format!("{}}}\n", indent));
422
423 Ok(output)
424}
425
426fn generate_consequence_code(
427 cons: &etdl_parser::ast::Consequence,
428 depth: usize,
429) -> Result<String, String> {
430 let indent = " ".repeat(depth);
431 let mut output = String::new();
432
433 match cons.consequence_operation {
434 etdl_parser::ast::ConsequenceOperation::Send => {
435 if let Some(ref channel_ref) = cons.channel {
436 let channel_name = extract_last_segment(channel_ref);
437 output.push_str(&format!(
438 "{}publish_to_channel(\"{}\", message).await?;\n",
439 indent, channel_name
440 ));
441 }
442 }
443 etdl_parser::ast::ConsequenceOperation::Terminate => {}
444 }
445
446 if depth == 1 {
447 output.push_str(&format!("{}Ok(())\n", indent));
448 }
449
450 Ok(output)
451}
452
453fn condition_to_rust_code(condition: &Condition) -> String {
454 match condition {
455 Condition::Default => "true".to_string(),
456 Condition::Comparison(cmp) => {
457 let (left_node, has_wildcard) = build_path_expression(&cmp.left);
458 let right_val = operand_to_val_str(&cmp.right, false);
459
460 if has_wildcard {
461 format!(
462 "{}.iter().all(|item| item{} {} {})",
463 left_node.path_prefix,
464 left_node.remaining_path,
465 comparator_str(&cmp.op),
466 right_val
467 )
468 } else {
469 format!(
470 "{} {} {}",
471 left_node.path_prefix + &left_node.remaining_path,
472 comparator_str(&cmp.op),
473 right_val
474 )
475 }
476 }
477 }
478}
479
480struct PathParts {
481 path_prefix: String,
482 remaining_path: String,
483}
484
485fn build_path_expression(operand: &ecel::Operand) -> (PathParts, bool) {
486 match operand {
487 ecel::Operand::Path(path_expr) => {
488 let segments = &path_expr.segments;
489 let mut pre_wildcard = Vec::new();
490 let mut post_wildcard = Vec::new();
491 let mut has_wildcard = false;
492
493 for (i, seg) in segments.iter().enumerate() {
494 if i == 0 {
495 continue;
496 }
497 if has_wildcard {
498 post_wildcard.push(seg.clone());
499 } else if matches!(seg, ecel::PathSegment::Wildcard) {
500 has_wildcard = true;
501 } else {
502 pre_wildcard.push(seg.clone());
503 }
504 }
505
506 let mut prefix = String::from("message");
507 for seg in &pre_wildcard {
508 match seg {
509 ecel::PathSegment::Field(name) => {
510 prefix.push('.');
511 prefix.push_str(&to_snake_case(name));
512 }
513 ecel::PathSegment::Index(idx) => {
514 prefix.push_str(&format!("[{}]", idx));
515 }
516 ecel::PathSegment::QuotedKey(name) => {
517 prefix.push_str(&format!("[\"{}\"]", name));
518 }
519 _ => {}
520 }
521 }
522
523 let mut suffix = String::new();
524 for seg in &post_wildcard {
525 match seg {
526 ecel::PathSegment::Field(name) => {
527 suffix.push('.');
528 suffix.push_str(&to_snake_case(name));
529 }
530 ecel::PathSegment::Index(idx) => {
531 suffix.push_str(&format!("[{}]", idx));
532 }
533 ecel::PathSegment::QuotedKey(name) => {
534 suffix.push_str(&format!("[\"{}\"]", name));
535 }
536 _ => {}
537 }
538 }
539
540 (PathParts {
541 path_prefix: prefix,
542 remaining_path: suffix,
543 }, has_wildcard)
544 }
545 ecel::Operand::Literal(_) => {
546 (PathParts {
547 path_prefix: String::new(),
548 remaining_path: String::new(),
549 }, false)
550 }
551 }
552}
553
554fn _operand_to_path_str(operand: &ecel::Operand) -> String {
555 match operand {
556 ecel::Operand::Path(path) => {
557 let segments = &path.segments;
558 segments.iter()
559 .skip(1)
560 .map(|seg| match seg {
561 ecel::PathSegment::Field(name) => to_snake_case(name),
562 ecel::PathSegment::Wildcard => "*".to_string(),
563 ecel::PathSegment::Index(idx) => idx.to_string(),
564 ecel::PathSegment::QuotedKey(name) => format!("\"{}\"", name),
565 })
566 .collect::<Vec<_>>()
567 .join(".")
568 }
569 ecel::Operand::Literal(_) => String::new(),
570 }
571}
572
573fn operand_to_val_str(operand: &ecel::Operand, _in_closure: bool) -> String {
574 match operand {
575 ecel::Operand::Path(_) => "".to_string(),
576 ecel::Operand::Literal(lit) => match lit {
577 ecel::Literal::Number(n) => n.to_string(),
578 ecel::Literal::String(s) => format!("\"{}\"", s),
579 ecel::Literal::Bool(b) => b.to_string(),
580 ecel::Literal::Null => "None".to_string(),
581 ecel::Literal::Array(items) => {
582 let inner: Vec<String> = items.iter().map(|item| match item {
583 ecel::Literal::Number(n) => n.to_string(),
584 ecel::Literal::String(s) => format!("\"{}\"", s),
585 ecel::Literal::Bool(b) => b.to_string(),
586 ecel::Literal::Null => "None".to_string(),
587 ecel::Literal::Array(_) => "[]".to_string(),
588 }).collect();
589 format!("vec![{}]", inner.join(", "))
590 }
591 },
592 }
593}
594
595fn comparator_str(op: &ecel::Comparator) -> &str {
596 match op {
597 ecel::Comparator::Eq => "==",
598 ecel::Comparator::Neq => "!=",
599 ecel::Comparator::Gte => ">=",
600 ecel::Comparator::Lte => "<=",
601 ecel::Comparator::Gt => ">",
602 ecel::Comparator::Lt => "<",
603 ecel::Comparator::In => "in",
604 ecel::Comparator::Matches => "matches",
605 }
606}
607
608fn ref_to_rust_type(ext_ref: &etdl_parser::ast::ExternalRef) -> String {
609 extract_last_segment(ext_ref)
610}
611
612fn extract_last_segment(ext_ref: &etdl_parser::ast::ExternalRef) -> String {
613 let pointer = &ext_ref.pointer;
614 let parts: Vec<&str> = pointer.split('/').collect();
615 let last = parts.last().unwrap_or(&"Unknown");
616 to_pascal_case(last)
617}
618
619fn get_branch_prob(
620 branch: &etdl_parser::ast::Branch,
621 _node_id: &str,
622 fault_tree_probs: &BTreeMap<String, f64>,
623) -> Option<f64> {
624 if let Some(ref ps) = branch.probability_source {
625 let ft_id = extract_ft_id(&ps.pointer);
626 return fault_tree_probs.get(&ft_id).copied();
627 }
628 branch.effective_probability()
629}
630
631fn extract_ft_id(pointer: &str) -> String {
632 pointer
633 .trim_start_matches("#/faultTrees/")
634 .trim_end_matches("/topEvent")
635 .to_string()
636}
637
638fn to_snake_case(s: &str) -> String {
639 let mut result = String::new();
640 for (i, c) in s.chars().enumerate() {
641 if c.is_uppercase() {
642 if i > 0 {
643 result.push('_');
644 }
645 result.push(c.to_lowercase().next().unwrap());
646 } else {
647 result.push(c);
648 }
649 }
650 result
651}
652
653fn to_upper_snake(s: &str) -> String {
654 let snake = to_snake_case(s);
655 snake.to_uppercase()
656}
657
658fn to_pascal_case(s: &str) -> String {
659 let mut result = String::new();
660 let mut capitalize = true;
661 for c in s.chars() {
662 if c == '_' || c == '-' || c == ' ' || c == '/' {
663 capitalize = true;
664 } else if capitalize {
665 result.push(c.to_uppercase().next().unwrap());
666 capitalize = false;
667 } else {
668 result.push(c);
669 }
670 }
671 result
672}