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 = entry
197                            .dot(row.key.as_ref())
198                            .and_then(|slot| slot.dot("result"))
199                            .unwrap_or(Variable::Null);
200                        let element = iteration_element(node_trace, &paths, loop_mode, index);
201                        let dollar = expression_dollar_scope(&content.expressions, entry);
202                        state.executions.push(BlockExecution {
203                            block_id: row_block_id.clone(),
204                            policy_path: None,
205                            instance_path: instance_for(
206                                node,
207                                &paths,
208                                loop_mode,
209                                iterations.len(),
210                                index,
211                                inherited_instance,
212                            ),
213                            trace: BlockTrace::Expression {
214                                property: property.clone(),
215                                value,
216                            },
217                            operand_values: operand_values(
218                                &local_reads,
219                                &paths,
220                                &node_trace.input,
221                                element.as_ref(),
222                                Some(&dollar),
223                            ),
224                            writes: Vec::new(),
225                            reads: reads.clone(),
226                        });
227                    }
228                }
229            }
230            DecisionNodeKind::DecisionTableNode { content } => {
231                let loop_mode = matches!(
232                    content.transform_attributes.execution_mode,
233                    TransformExecutionMode::Loop
234                );
235                let collect = matches!(content.hit_policy, DecisionTableHitPolicy::Collect);
236                let iterations = trace_entries(node_trace.trace_data.as_ref(), loop_mode);
237                let reads = state.db.node_global_reads(node, &paths, None);
238                let environment_root = state.dt_environment(content, node_trace, trace);
239                let outputs_root = match &paths.output_path {
240                    Some(path) => node_trace.output.dot(path).unwrap_or(Variable::Null),
241                    None => node_trace.output.clone(),
242                };
243                for (index, entry) in iterations.iter().enumerate() {
244                    let row_traces: Vec<Variable> = if collect {
245                        entry
246                            .as_array()
247                            .map(|rows| rows.borrow().iter().cloned().collect())
248                            .unwrap_or_default()
249                    } else {
250                        vec![entry.clone()]
251                    };
252                    let matched_rows: Vec<u32> = row_traces
253                        .iter()
254                        .filter_map(|row| row.dot("index"))
255                        .filter_map(|value| match value {
256                            Variable::Number(number) => number.to_u32(),
257                            _ => None,
258                        })
259                        .collect();
260
261                    let iter_result = if loop_mode {
262                        element_at(&outputs_root, index).unwrap_or(Variable::Null)
263                    } else {
264                        outputs_root.clone()
265                    };
266                    let mut evaluations: Vec<HashMap<Arc<str>, Variable>> = Vec::new();
267                    if collect {
268                        for (row_index, _) in row_traces.iter().enumerate() {
269                            let Some(row) = element_at(&iter_result, row_index) else {
270                                continue;
271                            };
272                            let mut evaluation: HashMap<Arc<str>, Variable> = HashMap::new();
273                            for column in content.outputs.iter() {
274                                if let Some(value) = row.dot(column.field.as_ref()) {
275                                    evaluation.insert(
276                                        output_prefixed(&paths, &column.field),
277                                        value.deep_clone(),
278                                    );
279                                }
280                            }
281                            if !evaluation.is_empty() {
282                                evaluations.push(evaluation);
283                            }
284                        }
285                    } else {
286                        let mut evaluation: HashMap<Arc<str>, Variable> = HashMap::new();
287                        for column in content.outputs.iter() {
288                            if let Some(value) = iter_result.dot(column.field.as_ref()) {
289                                evaluation.insert(
290                                    output_prefixed(&paths, &column.field),
291                                    value.deep_clone(),
292                                );
293                            }
294                        }
295                        evaluations.push(evaluation);
296                    }
297
298                    let element = iteration_element(node_trace, &paths, loop_mode, index);
299                    let mut operands: HashMap<Arc<str>, Variable> = HashMap::new();
300                    for row in &row_traces {
301                        let Some(reference_map) =
302                            row.dot("reference_map").and_then(|value| value.as_object())
303                        else {
304                            continue;
305                        };
306                        for (field, value) in reference_map.borrow().iter() {
307                            operands.insert(Arc::from(field.as_ref()), value.deep_clone());
308                        }
309                    }
310                    if operands.is_empty() {
311                        let local_reads = state.db.node_local_reads(node, None);
312                        operands = operand_values(
313                            &local_reads,
314                            &paths,
315                            &node_trace.input,
316                            element.as_ref(),
317                            None,
318                        );
319                    }
320
321                    let environment = if loop_mode {
322                        environment_root
323                            .as_ref()
324                            .and_then(|env| element_at(env, index))
325                    } else {
326                        environment_root.clone()
327                    };
328                    let extras = environment.map(|env| state.dt_extras(content, env));
329                    state.executions.push(BlockExecution {
330                        block_id: block_id.clone(),
331                        policy_path: None,
332                        instance_path: instance_for(
333                            node,
334                            &paths,
335                            loop_mode,
336                            iterations.len(),
337                            index,
338                            inherited_instance,
339                        ),
340                        trace: BlockTrace::DecisionTable {
341                            matched_rows,
342                            evaluations,
343                            extras,
344                        },
345                        operand_values: operands,
346                        writes: Vec::new(),
347                        reads: reads.clone(),
348                    });
349                }
350            }
351            DecisionNodeKind::SwitchNode { content } => {
352                let taken: HashSet<Arc<str>> = node_trace
353                    .trace_data
354                    .as_ref()
355                    .and_then(|data| data.dot("statements"))
356                    .and_then(|statements| statements.as_array())
357                    .map(|statements| {
358                        statements
359                            .borrow()
360                            .iter()
361                            .filter_map(|statement| statement.dot("id"))
362                            .filter_map(|id| id.as_str().map(Arc::from))
363                            .collect()
364                    })
365                    .unwrap_or_default();
366                let arms: Vec<ConditionTrace> = content
367                    .statements
368                    .iter()
369                    .map(|statement| ConditionTrace {
370                        id: statement.id.clone(),
371                        result: taken.contains(&statement.id),
372                    })
373                    .collect();
374                let matched_arm = content
375                    .statements
376                    .iter()
377                    .find(|statement| taken.contains(&statement.id))
378                    .map(|statement| statement.id.clone());
379                let reads = state.db.node_global_reads(node, &paths, None);
380                let local_reads = state.db.node_local_reads(node, None);
381                state.executions.push(BlockExecution {
382                    block_id: block_id.clone(),
383                    policy_path: None,
384                    instance_path: inherited_instance.cloned(),
385                    trace: BlockTrace::Match {
386                        matched_arm,
387                        value: Variable::Null,
388                        arms,
389                    },
390                    operand_values: operand_values(
391                        &local_reads,
392                        &paths,
393                        &node_trace.input,
394                        None,
395                        None,
396                    ),
397                    writes: Vec::new(),
398                    reads,
399                });
400            }
401            DecisionNodeKind::FunctionNode { content } => {
402                let source = function_source(content);
403                let local_reads: Vec<Arc<str>> = Db::function_input_reads(&source)
404                    .into_iter()
405                    .map(Arc::from)
406                    .collect();
407                let reads: Vec<Arc<str>> = local_reads
408                    .iter()
409                    .filter_map(|read| map_read(&paths, read))
410                    .collect();
411                state.executions.push(BlockExecution {
412                    block_id: block_id.clone(),
413                    policy_path: None,
414                    instance_path: inherited_instance.cloned(),
415                    trace: BlockTrace::Expression {
416                        property: Arc::from(""),
417                        value: node_trace.output.clone(),
418                    },
419                    operand_values: operand_values(
420                        &local_reads,
421                        &paths,
422                        &node_trace.input,
423                        None,
424                        None,
425                    ),
426                    writes: shallow_writes(&node_trace.input, &node_trace.output),
427                    reads,
428                });
429            }
430            DecisionNodeKind::DecisionNode { content } => {
431                let key = content.key.clone();
432                let sub_content = (!state.visiting.contains(&key))
433                    .then(|| {
434                        state
435                            .snapshot
436                            .graphs
437                            .get(&key)
438                            .and_then(|content| content.as_graph())
439                            .cloned()
440                    })
441                    .flatten();
442                let sub_traces = sub_content
443                    .as_ref()
444                    .map(|_| sub_trace_maps(node_trace.trace_data.as_ref()))
445                    .unwrap_or_default();
446
447                if let (Some(sub_content), false) = (sub_content, sub_traces.is_empty()) {
448                    let group_index = state.executions.len();
449                    state.executions.push(BlockExecution {
450                        block_id: block_id.clone(),
451                        policy_path: None,
452                        instance_path: inherited_instance.cloned(),
453                        trace: BlockTrace::Expression {
454                            property: Arc::from(""),
455                            value: node_trace.output.clone(),
456                        },
457                        operand_values: HashMap::new(),
458                        writes: shallow_writes(&node_trace.input, &node_trace.output),
459                        reads: Vec::new(),
460                    });
461                    let child_start = state.executions.len();
462                    state.visiting.insert(key.clone());
463                    let looped = sub_traces.len() > 1;
464                    for (index, sub_trace) in sub_traces.iter().enumerate() {
465                        let prefix = if looped {
466                            format!("{block_id}[{index}]/")
467                        } else {
468                            format!("{block_id}/")
469                        };
470                        let instance: Option<Arc<str>> = if looped {
471                            Some(Arc::from(format!("{}.{index}", loop_label(node, &paths))))
472                        } else {
473                            inherited_instance.cloned()
474                        };
475                        walk_graph(state, &sub_content, sub_trace, &prefix, instance.as_ref());
476                    }
477                    state.visiting.remove(&key);
478                    let free_reads = subtree_free_reads(&state.executions[child_start..], &paths);
479                    state.executions[group_index].reads = free_reads;
480                } else {
481                    push_opaque(state, &block_id, node_trace, inherited_instance);
482                }
483            }
484            _ => {
485                push_opaque(state, &block_id, node_trace, inherited_instance);
486            }
487        }
488    }
489}
490
491fn push_opaque(
492    state: &mut EnhanceState,
493    block_id: &Arc<str>,
494    node_trace: &DecisionGraphTrace,
495    inherited_instance: Option<&Arc<str>>,
496) {
497    state.executions.push(BlockExecution {
498        block_id: block_id.clone(),
499        policy_path: None,
500        instance_path: inherited_instance.cloned(),
501        trace: BlockTrace::Expression {
502            property: Arc::from(""),
503            value: node_trace.output.clone(),
504        },
505        operand_values: HashMap::new(),
506        writes: shallow_writes(&node_trace.input, &node_trace.output),
507        reads: Vec::new(),
508    });
509}
510
511fn prefixed(prefix: &str, id: &str) -> Arc<str> {
512    if prefix.is_empty() {
513        Arc::from(id)
514    } else {
515        Arc::from(format!("{prefix}{id}"))
516    }
517}
518
519fn output_prefixed(paths: &NodePaths, key: &str) -> Arc<str> {
520    if paths.output_prefix.is_empty() {
521        Arc::from(key)
522    } else {
523        Arc::from(format!("{}.{key}", paths.output_prefix.join(".")))
524    }
525}
526
527fn map_read(paths: &NodePaths, path: &str) -> Option<Arc<str>> {
528    match &paths.read_base {
529        ReadBase::NodeInput => Some(Arc::from(path)),
530        ReadBase::Opaque => None,
531        ReadBase::Prefixed(prefix) => Some(Arc::from(format!("{}.{path}", prefix.join(".")))),
532    }
533}
534
535fn loop_label(node: &DecisionNode, paths: &NodePaths) -> String {
536    match &paths.read_base {
537        ReadBase::Prefixed(segments) => segments.join("."),
538        _ => node.name.to_string(),
539    }
540}
541
542fn instance_for(
543    node: &DecisionNode,
544    paths: &NodePaths,
545    loop_mode: bool,
546    total: usize,
547    index: usize,
548    inherited: Option<&Arc<str>>,
549) -> Option<Arc<str>> {
550    if loop_mode && total > 1 {
551        Some(Arc::from(format!("{}.{index}", loop_label(node, paths))))
552    } else {
553        inherited.cloned()
554    }
555}
556
557fn trace_entries(trace_data: Option<&Variable>, loop_mode: bool) -> Vec<Variable> {
558    match trace_data {
559        Some(Variable::Array(items)) if loop_mode => items.borrow().iter().cloned().collect(),
560        Some(data) => vec![data.clone()],
561        None => vec![Variable::Null],
562    }
563}
564
565fn element_at(value: &Variable, index: usize) -> Option<Variable> {
566    value
567        .as_array()
568        .and_then(|items| items.borrow().get(index).cloned())
569}
570
571fn iteration_element(
572    node_trace: &DecisionGraphTrace,
573    paths: &NodePaths,
574    loop_mode: bool,
575    index: usize,
576) -> Option<Variable> {
577    if !loop_mode {
578        return None;
579    }
580    let ReadBase::Prefixed(prefix) = &paths.read_base else {
581        return None;
582    };
583    node_trace
584        .input
585        .dot(&prefix.join("."))
586        .and_then(|collection| element_at(&collection, index))
587}
588
589fn operand_values(
590    local_reads: &[Arc<str>],
591    paths: &NodePaths,
592    node_input: &Variable,
593    element: Option<&Variable>,
594    dollar: Option<&Variable>,
595) -> HashMap<Arc<str>, Variable> {
596    let mut out: HashMap<Arc<str>, Variable> = HashMap::new();
597    for read in local_reads {
598        let value = if let Some(rest) = read.strip_prefix("$.") {
599            dollar.and_then(|scope| scope.dot(rest))
600        } else if let Some(element) = element {
601            element.dot(read)
602        } else {
603            match &paths.read_base {
604                ReadBase::NodeInput => node_input.dot(read),
605                ReadBase::Prefixed(prefix) => {
606                    node_input.dot(&format!("{}.{read}", prefix.join(".")))
607                }
608                ReadBase::Opaque => None,
609            }
610        };
611        if let Some(value) = value {
612            out.insert(read.clone(), value.deep_clone());
613        }
614    }
615    out
616}
617
618fn expression_dollar_scope(rows: &[zen_types::decision::Expression], entry: &Variable) -> Variable {
619    let scope = Variable::empty_object();
620    for row in rows {
621        if row.key.is_empty() {
622            continue;
623        }
624        let Some(value) = entry
625            .dot(row.key.as_ref())
626            .and_then(|slot| slot.dot("result"))
627        else {
628            continue;
629        };
630        scope.dot_insert(row.key.as_ref(), value);
631    }
632    scope
633}
634
635fn shallow_writes(input: &Variable, output: &Variable) -> Vec<WriteTrace> {
636    let Some(entries) = output.as_object() else {
637        return Vec::new();
638    };
639    let mut writes: Vec<WriteTrace> = Vec::new();
640    for (key, value) in entries.borrow().iter() {
641        if key.starts_with('$') {
642            continue;
643        }
644        if input.dot(key).is_some_and(|previous| previous == *value) {
645            continue;
646        }
647        writes.push(WriteTrace {
648            path: Arc::from(key.as_ref()),
649            value: value.deep_clone(),
650        });
651    }
652    writes.sort_by(|a, b| a.path.cmp(&b.path));
653    writes
654}
655
656fn sub_trace_maps(trace_data: Option<&Variable>) -> Vec<GraphTraceMap> {
657    match trace_data {
658        Some(Variable::Array(items)) => items.borrow().iter().filter_map(as_trace_map).collect(),
659        Some(data) => as_trace_map(data).map(|map| vec![map]).unwrap_or_default(),
660        None => Vec::new(),
661    }
662}
663
664fn as_trace_map(value: &Variable) -> Option<GraphTraceMap> {
665    if !matches!(value, Variable::Object(_)) {
666        return None;
667    }
668    let json = serde_json::to_value(value).ok()?;
669    let map: GraphTraceMap = serde_json::from_value(json).ok()?;
670    (!map.is_empty()).then_some(map)
671}
672
673fn subtree_free_reads(children: &[BlockExecution], paths: &NodePaths) -> Vec<Arc<str>> {
674    let mut written: Vec<Arc<str>> = Vec::new();
675    for child in children {
676        for write in &child.writes {
677            written.push(write.path.clone());
678        }
679        match &child.trace {
680            BlockTrace::Expression { property, .. } if !property.is_empty() => {
681                written.push(property.clone());
682            }
683            BlockTrace::DecisionTable { evaluations, .. } => {
684                for evaluation in evaluations {
685                    written.extend(evaluation.keys().cloned());
686                }
687            }
688            _ => {}
689        }
690    }
691    let covered = |read: &str| {
692        written
693            .iter()
694            .any(|path| read == path.as_ref() || read.starts_with(&format!("{path}.")))
695    };
696    let mut out: Vec<Arc<str>> = Vec::new();
697    for child in children {
698        for read in &child.reads {
699            if covered(read) {
700                continue;
701            }
702            if let Some(mapped) = map_read(paths, read) {
703                out.push(mapped);
704            }
705        }
706    }
707    out.sort();
708    out.dedup();
709    out
710}