1use crate::error::Result;
7use crate::types::{Document, Value};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub enum ComputedExpression {
14 Concat(Vec<String>),
16 Sum(Vec<String>),
18 Product(Vec<String>),
20 Average(Vec<String>),
22 Custom(String),
24}
25
26impl ComputedExpression {
27 pub fn evaluate(&self, doc: &Document) -> Option<Value> {
29 match self {
30 ComputedExpression::Concat(fields) => {
31 let mut result = String::new();
32 for field in fields {
33 if let Some(value) = doc.data.get(field) {
34 if let Some(s) = value.as_str() {
35 if !result.is_empty() {
36 result.push(' ');
37 }
38 result.push_str(s);
39 }
40 }
41 }
42 Some(Value::String(result))
43 }
44
45 ComputedExpression::Sum(fields) => {
46 let mut sum = 0i64;
47 for field in fields {
48 if let Some(value) = doc.data.get(field) {
49 if let Some(i) = value.as_i64() {
50 sum += i;
51 }
52 }
53 }
54 Some(Value::Int(sum))
55 }
56
57 ComputedExpression::Product(fields) => {
58 let mut product = 1i64;
59 for field in fields {
60 if let Some(value) = doc.data.get(field) {
61 if let Some(i) = value.as_i64() {
62 product *= i;
63 }
64 }
65 }
66 Some(Value::Int(product))
67 }
68
69 ComputedExpression::Average(fields) => {
70 let mut sum = 0.0;
71 let mut count = 0;
72 for field in fields {
73 if let Some(value) = doc.data.get(field) {
74 if let Some(f) = value.as_f64() {
75 sum += f;
76 count += 1;
77 } else if let Some(i) = value.as_i64() {
78 sum += i as f64;
79 count += 1;
80 }
81 }
82 }
83 if count > 0 {
84 Some(Value::Float(sum / count as f64))
85 } else {
86 None
87 }
88 }
89
90 ComputedExpression::Custom(_expr) => {
91 None
94 }
95 }
96 }
97}
98
99pub struct ComputedFields {
101 fields: HashMap<String, HashMap<String, ComputedExpression>>,
103}
104
105impl ComputedFields {
106 pub fn new() -> Self {
107 Self {
108 fields: HashMap::new(),
109 }
110 }
111
112 pub fn register(
114 &mut self,
115 collection: impl Into<String>,
116 field: impl Into<String>,
117 expression: ComputedExpression,
118 ) {
119 let collection = collection.into();
120 self.fields
121 .entry(collection)
122 .or_insert_with(HashMap::new)
123 .insert(field.into(), expression);
124 }
125
126 pub fn apply(&self, collection: &str, doc: &mut Document) -> Result<()> {
128 if let Some(computed) = self.fields.get(collection) {
129 for (field_name, expression) in computed {
130 if let Some(value) = expression.evaluate(doc) {
131 doc.data.insert(field_name.clone(), value);
132 }
133 }
134 }
135 Ok(())
136 }
137
138 pub fn get_fields(&self, collection: &str) -> Option<&HashMap<String, ComputedExpression>> {
140 self.fields.get(collection)
141 }
142}
143
144impl Default for ComputedFields {
145 fn default() -> Self {
146 Self::new()
147 }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153
154 #[test]
155 fn test_concat_expression() {
156 let expr =
157 ComputedExpression::Concat(vec!["first_name".to_string(), "last_name".to_string()]);
158
159 let mut doc = Document::new();
160 doc.data
161 .insert("first_name".to_string(), Value::String("John".to_string()));
162 doc.data
163 .insert("last_name".to_string(), Value::String("Doe".to_string()));
164
165 let result = expr.evaluate(&doc);
166 assert_eq!(result, Some(Value::String("John Doe".to_string())));
167 }
168
169 #[test]
170 fn test_sum_expression() {
171 let expr = ComputedExpression::Sum(vec!["a".to_string(), "b".to_string(), "c".to_string()]);
172
173 let mut doc = Document::new();
174 doc.data.insert("a".to_string(), Value::Int(10));
175 doc.data.insert("b".to_string(), Value::Int(20));
176 doc.data.insert("c".to_string(), Value::Int(30));
177
178 let result = expr.evaluate(&doc);
179 assert_eq!(result, Some(Value::Int(60)));
180 }
181
182 #[test]
183 fn test_average_expression() {
184 let expr = ComputedExpression::Average(vec!["score1".to_string(), "score2".to_string()]);
185
186 let mut doc = Document::new();
187 doc.data.insert("score1".to_string(), Value::Float(85.5));
188 doc.data.insert("score2".to_string(), Value::Float(92.5));
189
190 let result = expr.evaluate(&doc);
191 assert_eq!(result, Some(Value::Float(89.0)));
192 }
193
194 #[test]
195 fn test_computed_fields_registry() {
196 let mut registry = ComputedFields::new();
197
198 registry.register(
199 "users",
200 "full_name",
201 ComputedExpression::Concat(vec!["first_name".to_string(), "last_name".to_string()]),
202 );
203
204 let mut doc = Document::new();
205 doc.data
206 .insert("first_name".to_string(), Value::String("Jane".to_string()));
207 doc.data
208 .insert("last_name".to_string(), Value::String("Smith".to_string()));
209
210 registry.apply("users", &mut doc).unwrap();
211
212 assert_eq!(
213 doc.data.get("full_name"),
214 Some(&Value::String("Jane Smith".to_string()))
215 );
216 }
217}