Skip to main content

zen_engine/model/
decision_content.rs

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