Skip to main content

zen_engine/workspace/graph/
enhance.rs

1use std::rc::Rc;
2use std::sync::Arc;
3
4use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
5use base64::Engine as _;
6use rust_decimal::prelude::ToPrimitive;
7use zen_expression::variable::Variable;
8use zen_expression::Isolate;
9use zen_types::decision::{
10    DecisionNode, DecisionNodeKind, DecisionTableContent, DecisionTableHitPolicy,
11    TransformExecutionMode,
12};
13
14use crate::model::GraphContent;
15use crate::nodes::decision_table::DecisionTableNodeHandler;
16use crate::workspace::db::{Db, Snapshot};
17use crate::workspace::graph::editor::{NodePaths, ReadBase};
18use crate::workspace::graph::function_source;
19use crate::workspace::types::{
20    BlockExecution, BlockTrace, ConditionTrace, DecisionTableExtras, EvaluationError, Trace,
21    WriteTrace,
22};
23use crate::DecisionGraphTrace;
24
25pub type GraphTraceMap = HashMap<Arc<str>, DecisionGraphTrace>;
26
27struct EnhanceState<'a> {
28    db: &'a Db,
29    snapshot: Arc<Snapshot>,
30    executions: Vec<BlockExecution>,
31    visiting: HashSet<Arc<str>>,
32}
33
34impl EnhanceState<'_> {
35    fn dt_environment(
36        &self,
37        content: &DecisionTableContent,
38        node_trace: &DecisionGraphTrace,
39        trace: &GraphTraceMap,
40    ) -> Option<Variable> {
41        let nodes = Variable::from_object(
42            trace
43                .values()
44                .filter(|entry| entry.order < node_trace.order)
45                .map(|entry| (Rc::from(entry.name.as_ref()), entry.output.clone()))
46                .collect(),
47        );
48        let base = node_trace.input.depth_clone(1);
49        base.dot_insert("$nodes", nodes.clone());
50        let Some(input_field) = &content.transform_attributes.input_field else {
51            return Some(base);
52        };
53        let mut isolate = Isolate::with_environment(base);
54        let calculated = isolate.run_standard(input_field.as_ref()).ok()?;
55        match &calculated {
56            Variable::Array(items) => {
57                let items = items
58                    .borrow()
59                    .iter()
60                    .map(|item| {
61                        let item = item.depth_clone(1);
62                        item.dot_insert("$nodes", nodes.clone());
63                        item
64                    })
65                    .collect();
66                Some(Variable::from_array(items))
67            }
68            _ => {
69                let calculated = calculated.depth_clone(1);
70                calculated.dot_insert("$nodes", nodes);
71                Some(calculated)
72            }
73        }
74    }
75
76    fn dt_extras(
77        &self,
78        content: &DecisionTableContent,
79        environment: Variable,
80    ) -> DecisionTableExtras {
81        let mut isolate = Isolate::with_environment(environment.depth_clone(1));
82        let bytes_per_row = content.inputs.len().div_ceil(8);
83        let mut bits = vec![0u8; bytes_per_row * content.rules.len()];
84        for (row, rule) in content.rules.iter().enumerate() {
85            for (col, input) in content.inputs.iter().enumerate() {
86                if DecisionTableNodeHandler::cell_passes(rule, input, &mut isolate) {
87                    bits[row * bytes_per_row + (col >> 3)] |= 1 << (col & 7);
88                }
89            }
90        }
91        DecisionTableExtras {
92            input_pass: base64::engine::general_purpose::STANDARD.encode(&bits),
93        }
94    }
95}
96
97impl Db {
98    pub fn enhance_graph_trace(
99        &self,
100        document: &Arc<str>,
101        trace: &GraphTraceMap,
102    ) -> Result<Trace, EvaluationError> {
103        let snapshot = self.snapshot();
104        let Some(content) = snapshot
105            .graphs
106            .get(document)
107            .and_then(|content| content.as_graph())
108            .cloned()
109        else {
110            return Err(EvaluationError::PolicyNotFound(document.clone()));
111        };
112
113        let mut state = EnhanceState {
114            db: self,
115            snapshot: snapshot.clone(),
116            executions: Vec::new(),
117            visiting: HashSet::new(),
118        };
119        state.visiting.insert(document.clone());
120        walk_graph(&mut state, &content, trace, "", None);
121
122        let mut properties: HashMap<Arc<str>, Variable> = HashMap::new();
123        for execution in &state.executions {
124            for write in &execution.writes {
125                properties.insert(write.path.clone(), write.value.clone());
126            }
127            match &execution.trace {
128                BlockTrace::Expression { property, value } if !property.is_empty() => {
129                    properties.insert(property.clone(), value.clone());
130                }
131                BlockTrace::DecisionTable { evaluations, .. } => {
132                    for evaluation in evaluations {
133                        for (path, value) in evaluation {
134                            properties.insert(path.clone(), value.clone());
135                        }
136                    }
137                }
138                _ => {}
139            }
140        }
141
142        Ok(Trace {
143            engine_version: Arc::from(crate::ENGINE_VERSION),
144            properties,
145            executions: state.executions,
146        })
147    }
148}
149
150fn walk_graph(
151    state: &mut EnhanceState,
152    content: &GraphContent,
153    trace: &GraphTraceMap,
154    id_prefix: &str,
155    inherited_instance: Option<&Arc<str>>,
156) {
157    let mut executed: Vec<(&DecisionNode, &DecisionGraphTrace)> = content
158        .nodes
159        .iter()
160        .filter_map(|node| trace.get(node.id.as_ref()).map(|t| (node.as_ref(), t)))
161        .collect();
162    executed.sort_by_key(|(_, node_trace)| node_trace.order);
163
164    for (node, node_trace) in executed {
165        if matches!(
166            node.kind,
167            DecisionNodeKind::InputNode { .. } | DecisionNodeKind::OutputNode { .. }
168        ) {
169            continue;
170        }
171        let block_id = prefixed(id_prefix, &node.id);
172        let paths = NodePaths::new(node);
173
174        match &node.kind {
175            DecisionNodeKind::ExpressionNode { content } => {
176                let loop_mode = matches!(
177                    content.transform_attributes.execution_mode,
178                    TransformExecutionMode::Loop
179                );
180                let iterations = trace_entries(node_trace.trace_data.as_ref(), loop_mode);
181                for row in content.expressions.iter() {
182                    if row.key.is_empty() || row.value.is_empty() {
183                        continue;
184                    }
185                    let row_block_id: Arc<str> = Arc::from(format!("{block_id}:{}", row.id));
186                    let reads = state.db.node_global_reads(
187                        node,
188                        &paths,
189                        Some(std::slice::from_ref(&row.id)),
190                    );
191                    let local_reads = state
192                        .db
193                        .node_local_reads(node, Some(std::slice::from_ref(&row.id)));
194                    let property = output_prefixed(&paths, &row.key);
195                    for (index, entry) in iterations.iter().enumerate() {
196                        let value = expression_trace_result(entry, row.key.as_ref())
197                            .unwrap_or(Variable::Null);
198                        let element = iteration_element(node_trace, &paths, loop_mode, index);
199                        let dollar = expression_dollar_scope(&content.expressions, entry);
200                        state.executions.push(BlockExecution {
201                            block_id: row_block_id.clone(),
202                            policy_path: None,
203                            instance_path: instance_for(
204                                node,
205                                &paths,
206                                loop_mode,
207                                iterations.len(),
208                                index,
209                                inherited_instance,
210                            ),
211                            trace: BlockTrace::Expression {
212                                property: property.clone(),
213                                value,
214                            },
215                            operand_values: operand_values(
216                                &local_reads,
217                                &paths,
218                                &node_trace.input,
219                                element.as_ref(),
220                                Some(&dollar),
221                            ),
222                            writes: Vec::new(),
223                            reads: reads.clone(),
224                        });
225                    }
226                }
227            }
228            DecisionNodeKind::DecisionTableNode { content } => {
229                let loop_mode = matches!(
230                    content.transform_attributes.execution_mode,
231                    TransformExecutionMode::Loop
232                );
233                let collect = matches!(content.hit_policy, DecisionTableHitPolicy::Collect);
234                let iterations = trace_entries(node_trace.trace_data.as_ref(), loop_mode);
235                let reads = state.db.node_global_reads(node, &paths, None);
236                let environment_root = state.dt_environment(content, node_trace, trace);
237                let outputs_root = match &paths.output_path {
238                    Some(path) => node_trace.output.dot(path).unwrap_or(Variable::Null),
239                    None => node_trace.output.clone(),
240                };
241                for (index, entry) in iterations.iter().enumerate() {
242                    let row_traces: Vec<Variable> = if collect {
243                        entry
244                            .as_array()
245                            .map(|rows| rows.borrow().iter().cloned().collect())
246                            .unwrap_or_default()
247                    } else {
248                        vec![entry.clone()]
249                    };
250                    let matched_rows: Vec<u32> = row_traces
251                        .iter()
252                        .filter_map(|row| row.dot("index"))
253                        .filter_map(|value| match value {
254                            Variable::Number(number) => number.to_u32(),
255                            _ => None,
256                        })
257                        .collect();
258
259                    let iter_result = if loop_mode {
260                        element_at(&outputs_root, index).unwrap_or(Variable::Null)
261                    } else {
262                        outputs_root.clone()
263                    };
264                    let mut evaluations: Vec<HashMap<Arc<str>, Variable>> = Vec::new();
265                    if collect {
266                        for (row_index, _) in row_traces.iter().enumerate() {
267                            let Some(row) = element_at(&iter_result, row_index) else {
268                                continue;
269                            };
270                            let mut evaluation: HashMap<Arc<str>, Variable> = HashMap::new();
271                            for column in content.outputs.iter() {
272                                if let Some(value) = row.dot(column.field.as_ref()) {
273                                    evaluation.insert(
274                                        output_prefixed(&paths, &column.field),
275                                        value.deep_clone(),
276                                    );
277                                }
278                            }
279                            if !evaluation.is_empty() {
280                                evaluations.push(evaluation);
281                            }
282                        }
283                    } else {
284                        let mut evaluation: HashMap<Arc<str>, Variable> = HashMap::new();
285                        for column in content.outputs.iter() {
286                            if let Some(value) = iter_result.dot(column.field.as_ref()) {
287                                evaluation.insert(
288                                    output_prefixed(&paths, &column.field),
289                                    value.deep_clone(),
290                                );
291                            }
292                        }
293                        evaluations.push(evaluation);
294                    }
295
296                    let element = iteration_element(node_trace, &paths, loop_mode, index);
297                    let mut operands: HashMap<Arc<str>, Variable> = HashMap::new();
298                    for row in &row_traces {
299                        let Some(reference_map) =
300                            row.dot("reference_map").and_then(|value| value.as_object())
301                        else {
302                            continue;
303                        };
304                        for (field, value) in reference_map.borrow().iter() {
305                            operands.insert(Arc::from(field.as_ref()), value.deep_clone());
306                        }
307                    }
308                    if operands.is_empty() {
309                        let local_reads = state.db.node_local_reads(node, None);
310                        operands = operand_values(
311                            &local_reads,
312                            &paths,
313                            &node_trace.input,
314                            element.as_ref(),
315                            None,
316                        );
317                    }
318
319                    let environment = if loop_mode {
320                        environment_root
321                            .as_ref()
322                            .and_then(|env| element_at(env, index))
323                    } else {
324                        environment_root.clone()
325                    };
326                    let extras = environment.map(|env| state.dt_extras(content, env));
327                    state.executions.push(BlockExecution {
328                        block_id: block_id.clone(),
329                        policy_path: None,
330                        instance_path: instance_for(
331                            node,
332                            &paths,
333                            loop_mode,
334                            iterations.len(),
335                            index,
336                            inherited_instance,
337                        ),
338                        trace: BlockTrace::DecisionTable {
339                            matched_rows,
340                            evaluations,
341                            extras,
342                        },
343                        operand_values: operands,
344                        writes: Vec::new(),
345                        reads: reads.clone(),
346                    });
347                }
348            }
349            DecisionNodeKind::SwitchNode { content } => {
350                let taken: HashSet<Arc<str>> = node_trace
351                    .trace_data
352                    .as_ref()
353                    .and_then(|data| data.dot("statements"))
354                    .and_then(|statements| statements.as_array())
355                    .map(|statements| {
356                        statements
357                            .borrow()
358                            .iter()
359                            .filter_map(|statement| statement.dot("id"))
360                            .filter_map(|id| id.as_str().map(Arc::from))
361                            .collect()
362                    })
363                    .unwrap_or_default();
364                let arms: Vec<ConditionTrace> = content
365                    .statements
366                    .iter()
367                    .map(|statement| ConditionTrace {
368                        id: statement.id.clone(),
369                        result: taken.contains(&statement.id),
370                    })
371                    .collect();
372                let matched_arm = content
373                    .statements
374                    .iter()
375                    .find(|statement| taken.contains(&statement.id))
376                    .map(|statement| statement.id.clone());
377                let reads = state.db.node_global_reads(node, &paths, None);
378                let local_reads = state.db.node_local_reads(node, None);
379                state.executions.push(BlockExecution {
380                    block_id: block_id.clone(),
381                    policy_path: None,
382                    instance_path: inherited_instance.cloned(),
383                    trace: BlockTrace::Match {
384                        matched_arm,
385                        value: Variable::Null,
386                        arms,
387                    },
388                    operand_values: operand_values(
389                        &local_reads,
390                        &paths,
391                        &node_trace.input,
392                        None,
393                        None,
394                    ),
395                    writes: Vec::new(),
396                    reads,
397                });
398            }
399            DecisionNodeKind::FunctionNode { content } => {
400                let source = function_source(content);
401                let local_reads: Vec<Arc<str>> = Db::function_input_reads(&source)
402                    .into_iter()
403                    .map(Arc::from)
404                    .collect();
405                let reads: Vec<Arc<str>> = local_reads
406                    .iter()
407                    .filter_map(|read| map_read(&paths, read))
408                    .collect();
409                state.executions.push(BlockExecution {
410                    block_id: block_id.clone(),
411                    policy_path: None,
412                    instance_path: inherited_instance.cloned(),
413                    trace: BlockTrace::Expression {
414                        property: Arc::from(""),
415                        value: node_trace.output.clone(),
416                    },
417                    operand_values: operand_values(
418                        &local_reads,
419                        &paths,
420                        &node_trace.input,
421                        None,
422                        None,
423                    ),
424                    writes: shallow_writes(&node_trace.input, &node_trace.output),
425                    reads,
426                });
427            }
428            DecisionNodeKind::DecisionNode { content } => {
429                let key = content.key.clone();
430                let sub_content = (!state.visiting.contains(&key))
431                    .then(|| {
432                        state
433                            .snapshot
434                            .graphs
435                            .get(&key)
436                            .and_then(|content| content.as_graph())
437                            .cloned()
438                    })
439                    .flatten();
440                let sub_traces = sub_content
441                    .as_ref()
442                    .map(|_| sub_trace_maps(node_trace.trace_data.as_ref()))
443                    .unwrap_or_default();
444
445                if let (Some(sub_content), false) = (sub_content, sub_traces.is_empty()) {
446                    let group_index = state.executions.len();
447                    state.executions.push(BlockExecution {
448                        block_id: block_id.clone(),
449                        policy_path: None,
450                        instance_path: inherited_instance.cloned(),
451                        trace: BlockTrace::Expression {
452                            property: Arc::from(""),
453                            value: node_trace.output.clone(),
454                        },
455                        operand_values: HashMap::new(),
456                        writes: shallow_writes(&node_trace.input, &node_trace.output),
457                        reads: Vec::new(),
458                    });
459                    let child_start = state.executions.len();
460                    state.visiting.insert(key.clone());
461                    let looped = sub_traces.len() > 1;
462                    for (index, sub_trace) in sub_traces.iter().enumerate() {
463                        let prefix = if looped {
464                            format!("{block_id}[{index}]/")
465                        } else {
466                            format!("{block_id}/")
467                        };
468                        let instance: Option<Arc<str>> = if looped {
469                            Some(Arc::from(format!("{}.{index}", loop_label(node, &paths))))
470                        } else {
471                            inherited_instance.cloned()
472                        };
473                        walk_graph(state, &sub_content, sub_trace, &prefix, instance.as_ref());
474                    }
475                    state.visiting.remove(&key);
476                    let free_reads = subtree_free_reads(&state.executions[child_start..], &paths);
477                    state.executions[group_index].reads = free_reads;
478                } else {
479                    push_opaque(state, &block_id, node_trace, inherited_instance);
480                }
481            }
482            _ => {
483                push_opaque(state, &block_id, node_trace, inherited_instance);
484            }
485        }
486    }
487}
488
489fn push_opaque(
490    state: &mut EnhanceState,
491    block_id: &Arc<str>,
492    node_trace: &DecisionGraphTrace,
493    inherited_instance: Option<&Arc<str>>,
494) {
495    state.executions.push(BlockExecution {
496        block_id: block_id.clone(),
497        policy_path: None,
498        instance_path: inherited_instance.cloned(),
499        trace: BlockTrace::Expression {
500            property: Arc::from(""),
501            value: node_trace.output.clone(),
502        },
503        operand_values: HashMap::new(),
504        writes: shallow_writes(&node_trace.input, &node_trace.output),
505        reads: Vec::new(),
506    });
507}
508
509fn prefixed(prefix: &str, id: &str) -> Arc<str> {
510    if prefix.is_empty() {
511        Arc::from(id)
512    } else {
513        Arc::from(format!("{prefix}{id}"))
514    }
515}
516
517fn output_prefixed(paths: &NodePaths, key: &str) -> Arc<str> {
518    if paths.output_prefix.is_empty() {
519        Arc::from(key)
520    } else {
521        Arc::from(format!("{}.{key}", paths.output_prefix.join(".")))
522    }
523}
524
525fn map_read(paths: &NodePaths, path: &str) -> Option<Arc<str>> {
526    match &paths.read_base {
527        ReadBase::NodeInput => Some(Arc::from(path)),
528        ReadBase::Opaque => None,
529        ReadBase::Prefixed(prefix) => Some(Arc::from(format!("{}.{path}", prefix.join(".")))),
530    }
531}
532
533fn loop_label(node: &DecisionNode, paths: &NodePaths) -> String {
534    match &paths.read_base {
535        ReadBase::Prefixed(segments) => segments.join("."),
536        _ => node.name.to_string(),
537    }
538}
539
540fn instance_for(
541    node: &DecisionNode,
542    paths: &NodePaths,
543    loop_mode: bool,
544    total: usize,
545    index: usize,
546    inherited: Option<&Arc<str>>,
547) -> Option<Arc<str>> {
548    if loop_mode && total > 1 {
549        Some(Arc::from(format!("{}.{index}", loop_label(node, paths))))
550    } else {
551        inherited.cloned()
552    }
553}
554
555fn trace_entries(trace_data: Option<&Variable>, loop_mode: bool) -> Vec<Variable> {
556    match trace_data {
557        Some(Variable::Array(items)) if loop_mode => items.borrow().iter().cloned().collect(),
558        Some(data) => vec![data.clone()],
559        None => vec![Variable::Null],
560    }
561}
562
563fn element_at(value: &Variable, index: usize) -> Option<Variable> {
564    value
565        .as_array()
566        .and_then(|items| items.borrow().get(index).cloned())
567}
568
569fn iteration_element(
570    node_trace: &DecisionGraphTrace,
571    paths: &NodePaths,
572    loop_mode: bool,
573    index: usize,
574) -> Option<Variable> {
575    if !loop_mode {
576        return None;
577    }
578    let ReadBase::Prefixed(prefix) = &paths.read_base else {
579        return None;
580    };
581    node_trace
582        .input
583        .dot(&prefix.join("."))
584        .and_then(|collection| element_at(&collection, index))
585}
586
587fn operand_values(
588    local_reads: &[Arc<str>],
589    paths: &NodePaths,
590    node_input: &Variable,
591    element: Option<&Variable>,
592    dollar: Option<&Variable>,
593) -> HashMap<Arc<str>, Variable> {
594    let mut out: HashMap<Arc<str>, Variable> = HashMap::new();
595    for read in local_reads {
596        let value = if let Some(rest) = read.strip_prefix("$.") {
597            dollar.and_then(|scope| scope.dot(rest))
598        } else if let Some(element) = element {
599            element.dot(read)
600        } else {
601            match &paths.read_base {
602                ReadBase::NodeInput => node_input.dot(read),
603                ReadBase::Prefixed(prefix) => {
604                    node_input.dot(&format!("{}.{read}", prefix.join(".")))
605                }
606                ReadBase::Opaque => None,
607            }
608        };
609        if let Some(value) = value {
610            out.insert(read.clone(), value.deep_clone());
611        }
612    }
613    out
614}
615
616/// Expression-node trace maps are keyed by the literal row key, which may itself
617/// contain dots (e.g. "quote.annualPremium"), so a `dot()` path lookup would miss
618/// those entries — read the map key directly.
619fn expression_trace_result(entry: &Variable, key: &str) -> Option<Variable> {
620    let obj = entry.as_object()?;
621    let slot = obj.borrow().get(key)?.shallow_clone();
622    slot.dot("result")
623}
624
625fn expression_dollar_scope(rows: &[zen_types::decision::Expression], entry: &Variable) -> Variable {
626    let scope = Variable::empty_object();
627    for row in rows {
628        if row.key.is_empty() {
629            continue;
630        }
631        let Some(value) = expression_trace_result(entry, row.key.as_ref()) else {
632            continue;
633        };
634        scope.dot_insert(row.key.as_ref(), value);
635    }
636    scope
637}
638
639fn shallow_writes(input: &Variable, output: &Variable) -> Vec<WriteTrace> {
640    let Some(entries) = output.as_object() else {
641        return Vec::new();
642    };
643    let mut writes: Vec<WriteTrace> = Vec::new();
644    for (key, value) in entries.borrow().iter() {
645        if key.starts_with('$') {
646            continue;
647        }
648        if input.dot(key).is_some_and(|previous| previous == *value) {
649            continue;
650        }
651        writes.push(WriteTrace {
652            path: Arc::from(key.as_ref()),
653            value: value.deep_clone(),
654        });
655    }
656    writes.sort_by(|a, b| a.path.cmp(&b.path));
657    writes
658}
659
660fn sub_trace_maps(trace_data: Option<&Variable>) -> Vec<GraphTraceMap> {
661    match trace_data {
662        Some(Variable::Array(items)) => items.borrow().iter().filter_map(as_trace_map).collect(),
663        Some(data) => as_trace_map(data).map(|map| vec![map]).unwrap_or_default(),
664        None => Vec::new(),
665    }
666}
667
668fn as_trace_map(value: &Variable) -> Option<GraphTraceMap> {
669    if !matches!(value, Variable::Object(_)) {
670        return None;
671    }
672    let json = serde_json::to_value(value).ok()?;
673    let map: GraphTraceMap = serde_json::from_value(json).ok()?;
674    (!map.is_empty()).then_some(map)
675}
676
677fn subtree_free_reads(children: &[BlockExecution], paths: &NodePaths) -> Vec<Arc<str>> {
678    let mut written: Vec<Arc<str>> = Vec::new();
679    for child in children {
680        for write in &child.writes {
681            written.push(write.path.clone());
682        }
683        match &child.trace {
684            BlockTrace::Expression { property, .. } if !property.is_empty() => {
685                written.push(property.clone());
686            }
687            BlockTrace::DecisionTable { evaluations, .. } => {
688                for evaluation in evaluations {
689                    written.extend(evaluation.keys().cloned());
690                }
691            }
692            _ => {}
693        }
694    }
695    let covered = |read: &str| {
696        written
697            .iter()
698            .any(|path| read == path.as_ref() || read.starts_with(&format!("{path}.")))
699    };
700    let mut out: Vec<Arc<str>> = Vec::new();
701    for child in children {
702        for read in &child.reads {
703            if covered(read) {
704                continue;
705            }
706            if let Some(mapped) = map_read(paths, read) {
707                out.push(mapped);
708            }
709        }
710    }
711    out.sort();
712    out.dedup();
713    out
714}