Skip to main content

zen_engine/workspace/graph/
enhance.rs

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