Skip to main content

zen_engine/model/
decision_content.rs

1use crate::nodes::function::v2::strip::TypeStripper;
2use crate::policy::PolicyDocument;
3use ahash::{HashMap, HashMapExt};
4use serde::{Deserialize, Deserializer, Serialize};
5use std::sync::Arc;
6use zen_expression::{ExpressionKind, Isolate, OpcodeCache};
7use zen_types::decision::{DecisionEdge, DecisionNode, DecisionNodeKind, FunctionNodeContent};
8
9#[derive(Clone, Debug, Serialize)]
10#[serde(untagged)]
11pub enum DecisionContent {
12    Graph(GraphContent),
13    Policy(PolicyContent),
14}
15
16impl<'de> Deserialize<'de> for DecisionContent {
17    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
18    where
19        D: Deserializer<'de>,
20    {
21        let value = serde_json::Value::deserialize(deserializer)?;
22        let is_policy = value
23            .as_object()
24            .is_some_and(|object| object.contains_key("blocks"));
25
26        let content = if is_policy {
27            serde_path_to_error::deserialize::<_, PolicyContent>(value).map(Self::Policy)
28        } else {
29            serde_path_to_error::deserialize::<_, GraphContent>(value).map(Self::Graph)
30        };
31
32        content.map_err(serde::de::Error::custom)
33    }
34}
35
36impl Default for DecisionContent {
37    fn default() -> Self {
38        Self::Graph(GraphContent::default())
39    }
40}
41
42impl DecisionContent {
43    pub fn as_graph(&self) -> Option<&GraphContent> {
44        match self {
45            Self::Graph(g) => Some(g),
46            Self::Policy(_) => None,
47        }
48    }
49
50    pub fn as_policy(&self) -> Option<&PolicyContent> {
51        match self {
52            Self::Policy(p) => Some(p),
53            Self::Graph(_) => None,
54        }
55    }
56
57    pub fn kind(&self) -> &'static str {
58        match self {
59            Self::Graph(_) => "graph",
60            Self::Policy(_) => "policy",
61        }
62    }
63
64    pub fn into_graph_arc(self: Arc<Self>) -> Option<Arc<GraphContent>> {
65        match Arc::try_unwrap(self) {
66            Ok(Self::Graph(g)) => Some(Arc::new(g)),
67            Ok(Self::Policy(_)) => None,
68            Err(arc) => match arc.as_ref() {
69                Self::Graph(g) => Some(Arc::new(g.clone())),
70                Self::Policy(_) => None,
71            },
72        }
73    }
74}
75
76impl From<GraphContent> for DecisionContent {
77    fn from(value: GraphContent) -> Self {
78        Self::Graph(value)
79    }
80}
81
82impl From<PolicyContent> for DecisionContent {
83    fn from(value: PolicyContent) -> Self {
84        Self::Policy(value)
85    }
86}
87
88impl From<Arc<PolicyDocument>> for PolicyContent {
89    fn from(value: Arc<PolicyDocument>) -> Self {
90        Self(value)
91    }
92}
93
94#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Default)]
95#[serde(rename_all = "camelCase")]
96pub struct GraphContent {
97    pub nodes: Vec<Arc<DecisionNode>>,
98    pub edges: Vec<Arc<DecisionEdge>>,
99
100    #[serde(default, skip_serializing_if = "Vec::is_empty")]
101    pub imports: Vec<Arc<str>>,
102
103    #[serde(skip)]
104    pub compiled_cache: Option<Arc<OpcodeCache>>,
105
106    #[serde(skip)]
107    pub stripped_functions: Option<Arc<HashMap<Arc<str>, Arc<str>>>>,
108}
109
110#[derive(Clone, Debug, Deserialize, Serialize)]
111#[serde(transparent)]
112pub struct PolicyContent(pub Arc<PolicyDocument>);
113
114impl GraphContent {
115    pub fn compile(&mut self) {
116        self.compile_functions();
117        if self.compiled_cache.is_some() {
118            return;
119        }
120
121        let mut sources: Vec<(Arc<str>, ExpressionKind)> = Vec::new();
122        for node in &self.nodes {
123            match &node.kind {
124                DecisionNodeKind::ExpressionNode { content } => {
125                    for expr in content.expressions.iter() {
126                        if !expr.key.is_empty() && !expr.value.is_empty() {
127                            sources.push((expr.value.clone(), ExpressionKind::Standard));
128                        }
129                    }
130                }
131                DecisionNodeKind::DecisionTableNode { content } => {
132                    for rule in content.rules.iter() {
133                        for input in content.inputs.iter() {
134                            let Some(rule_value) = rule.get(&input.id) else {
135                                continue;
136                            };
137
138                            let kind = if input.field.is_some() {
139                                ExpressionKind::Unary
140                            } else {
141                                ExpressionKind::Standard
142                            };
143
144                            sources.push((rule_value.clone(), kind));
145                        }
146
147                        for output in content.outputs.iter() {
148                            let Some(rule_value) = rule.get(&output.id) else {
149                                continue;
150                            };
151
152                            sources.push((rule_value.clone(), ExpressionKind::Standard));
153                        }
154                    }
155                }
156                _ => {}
157            }
158        }
159
160        let mut cache: OpcodeCache = OpcodeCache::new();
161        let mut isolate = Isolate::new();
162
163        for (source, kind) in &sources {
164            let map = match kind {
165                ExpressionKind::Standard => &mut cache.standard,
166                ExpressionKind::Unary => &mut cache.unary,
167            };
168            if map.contains_key(source) {
169                continue;
170            }
171
172            let result = match kind {
173                ExpressionKind::Standard => isolate
174                    .compile_standard(source)
175                    .map(|e| e.bytecode().to_vec()),
176                ExpressionKind::Unary => {
177                    isolate.compile_unary(source).map(|e| e.bytecode().to_vec())
178                }
179            };
180            if let Ok(bytecode) = result {
181                map.insert(source.clone(), Arc::from(bytecode));
182            }
183        }
184
185        self.compiled_cache.replace(Arc::new(cache));
186    }
187
188    fn compile_functions(&mut self) {
189        if self.stripped_functions.is_some() {
190            return;
191        }
192        let mut stripped: HashMap<Arc<str>, Arc<str>> = HashMap::new();
193        for node in &self.nodes {
194            let DecisionNodeKind::FunctionNode {
195                content: FunctionNodeContent::Version2(function),
196            } = &node.kind
197            else {
198                continue;
199            };
200            stripped
201                .entry(function.source.clone())
202                .or_insert_with(|| TypeStripper::strip(&function.source));
203        }
204        self.stripped_functions.replace(Arc::new(stripped));
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn malformed_graph_error_mentions_field() {
214        let error = serde_json::from_str::<DecisionContent>(r#"{"nodes":[],"edges":"bad"}"#)
215            .unwrap_err()
216            .to_string();
217        assert!(error.contains("edges"), "{error}");
218    }
219
220    #[test]
221    fn malformed_policy_error_mentions_inner_field() {
222        let json = r#"{"blocks":[{"type":"assertion","id":"b1","props":{"data":{}}}]}"#;
223        let error = serde_json::from_str::<DecisionContent>(json)
224            .unwrap_err()
225            .to_string();
226        assert!(error.contains("output"), "{error}");
227    }
228
229    #[test]
230    fn compile_strips_function_sources() {
231        let json = r#"{"nodes":[{"id":"f","name":"f","type":"functionNode","content":{"source":"export const handler = (input: { age: number }) => ({ total: input.age });"}}],"edges":[]}"#;
232        let mut content: GraphContent = serde_json::from_str(json).unwrap();
233        content.compile();
234        let stripped = content.stripped_functions.as_ref().unwrap();
235        assert_eq!(stripped.len(), 1);
236        let value = stripped.values().next().unwrap();
237        assert!(!value.contains(": number"), "{value}");
238    }
239
240    #[test]
241    fn valid_graph_routes_to_graph_variant() {
242        let content: DecisionContent = serde_json::from_str(r#"{"nodes":[],"edges":[]}"#).unwrap();
243        assert!(content.as_graph().is_some());
244    }
245
246    #[test]
247    fn valid_policy_routes_to_policy_variant() {
248        let content: DecisionContent = serde_json::from_str(r#"{"blocks":[]}"#).unwrap();
249        assert!(content.as_policy().is_some());
250    }
251}