solverforge_core/constraints/
collectors.rs

1use crate::constraints::WasmFunction;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
5#[serde(tag = "name")]
6pub enum Collector {
7    #[serde(rename = "count")]
8    Count {
9        #[serde(skip_serializing_if = "Option::is_none")]
10        distinct: Option<bool>,
11        #[serde(skip_serializing_if = "Option::is_none")]
12        map: Option<WasmFunction>,
13    },
14    #[serde(rename = "sum")]
15    Sum { map: WasmFunction },
16    #[serde(rename = "average")]
17    Average { map: WasmFunction },
18    #[serde(rename = "min")]
19    Min {
20        map: WasmFunction,
21        comparator: WasmFunction,
22    },
23    #[serde(rename = "max")]
24    Max {
25        map: WasmFunction,
26        comparator: WasmFunction,
27    },
28    #[serde(rename = "toList")]
29    ToList {
30        #[serde(skip_serializing_if = "Option::is_none")]
31        map: Option<WasmFunction>,
32    },
33    #[serde(rename = "toSet")]
34    ToSet {
35        #[serde(skip_serializing_if = "Option::is_none")]
36        map: Option<WasmFunction>,
37    },
38    #[serde(rename = "compose")]
39    Compose {
40        collectors: Vec<Collector>,
41        combiner: WasmFunction,
42    },
43    #[serde(rename = "conditionally")]
44    Conditionally {
45        predicate: WasmFunction,
46        collector: Box<Collector>,
47    },
48    #[serde(rename = "collectAndThen")]
49    CollectAndThen {
50        collector: Box<Collector>,
51        mapper: WasmFunction,
52    },
53    #[serde(rename = "loadBalance")]
54    LoadBalance {
55        map: WasmFunction,
56        #[serde(skip_serializing_if = "Option::is_none")]
57        load: Option<WasmFunction>,
58    },
59    #[serde(rename = "toSortedSet")]
60    ToSortedSet {
61        #[serde(skip_serializing_if = "Option::is_none")]
62        map: Option<WasmFunction>,
63        #[serde(skip_serializing_if = "Option::is_none")]
64        comparator: Option<WasmFunction>,
65    },
66    #[serde(rename = "toMap")]
67    ToMap {
68        key_mapper: WasmFunction,
69        value_mapper: WasmFunction,
70        #[serde(skip_serializing_if = "Option::is_none")]
71        merge_function: Option<WasmFunction>,
72    },
73    #[serde(rename = "toSortedMap")]
74    ToSortedMap {
75        key_mapper: WasmFunction,
76        value_mapper: WasmFunction,
77        #[serde(skip_serializing_if = "Option::is_none")]
78        merge_function: Option<WasmFunction>,
79    },
80}
81
82impl Collector {
83    pub fn count() -> Self {
84        Collector::Count {
85            distinct: None,
86            map: None,
87        }
88    }
89
90    pub fn count_distinct() -> Self {
91        Collector::Count {
92            distinct: Some(true),
93            map: None,
94        }
95    }
96
97    pub fn count_with_map(map: WasmFunction) -> Self {
98        Collector::Count {
99            distinct: None,
100            map: Some(map),
101        }
102    }
103
104    pub fn sum(map: WasmFunction) -> Self {
105        Collector::Sum { map }
106    }
107
108    pub fn average(map: WasmFunction) -> Self {
109        Collector::Average { map }
110    }
111
112    pub fn min(map: WasmFunction, comparator: WasmFunction) -> Self {
113        Collector::Min { map, comparator }
114    }
115
116    pub fn max(map: WasmFunction, comparator: WasmFunction) -> Self {
117        Collector::Max { map, comparator }
118    }
119
120    pub fn to_list() -> Self {
121        Collector::ToList { map: None }
122    }
123
124    pub fn to_list_with_map(map: WasmFunction) -> Self {
125        Collector::ToList { map: Some(map) }
126    }
127
128    pub fn to_set() -> Self {
129        Collector::ToSet { map: None }
130    }
131
132    pub fn to_set_with_map(map: WasmFunction) -> Self {
133        Collector::ToSet { map: Some(map) }
134    }
135
136    pub fn compose(collectors: Vec<Collector>, combiner: WasmFunction) -> Self {
137        Collector::Compose {
138            collectors,
139            combiner,
140        }
141    }
142
143    pub fn conditionally(predicate: WasmFunction, collector: Collector) -> Self {
144        Collector::Conditionally {
145            predicate,
146            collector: Box::new(collector),
147        }
148    }
149
150    pub fn collect_and_then(collector: Collector, mapper: WasmFunction) -> Self {
151        Collector::CollectAndThen {
152            collector: Box::new(collector),
153            mapper,
154        }
155    }
156
157    pub fn load_balance(map: WasmFunction) -> Self {
158        Collector::LoadBalance { map, load: None }
159    }
160
161    pub fn load_balance_with_load(map: WasmFunction, load: WasmFunction) -> Self {
162        Collector::LoadBalance {
163            map,
164            load: Some(load),
165        }
166    }
167
168    /// Create a toSortedSet collector with natural ordering.
169    pub fn to_sorted_set() -> Self {
170        Collector::ToSortedSet {
171            map: None,
172            comparator: None,
173        }
174    }
175
176    /// Create a toSortedSet collector with a mapper function.
177    pub fn to_sorted_set_with_map(map: WasmFunction) -> Self {
178        Collector::ToSortedSet {
179            map: Some(map),
180            comparator: None,
181        }
182    }
183
184    /// Create a toSortedSet collector with a custom comparator.
185    pub fn to_sorted_set_with_comparator(comparator: WasmFunction) -> Self {
186        Collector::ToSortedSet {
187            map: None,
188            comparator: Some(comparator),
189        }
190    }
191
192    /// Create a toSortedSet collector with a mapper and comparator.
193    pub fn to_sorted_set_with_map_and_comparator(
194        map: WasmFunction,
195        comparator: WasmFunction,
196    ) -> Self {
197        Collector::ToSortedSet {
198            map: Some(map),
199            comparator: Some(comparator),
200        }
201    }
202
203    /// Create a toMap collector that groups values into sets by key.
204    pub fn to_map(key_mapper: WasmFunction, value_mapper: WasmFunction) -> Self {
205        Collector::ToMap {
206            key_mapper,
207            value_mapper,
208            merge_function: None,
209        }
210    }
211
212    /// Create a toMap collector with a merge function for duplicate keys.
213    pub fn to_map_with_merge(
214        key_mapper: WasmFunction,
215        value_mapper: WasmFunction,
216        merge_function: WasmFunction,
217    ) -> Self {
218        Collector::ToMap {
219            key_mapper,
220            value_mapper,
221            merge_function: Some(merge_function),
222        }
223    }
224
225    /// Create a toSortedMap collector that groups values by key with sorted keys.
226    pub fn to_sorted_map(key_mapper: WasmFunction, value_mapper: WasmFunction) -> Self {
227        Collector::ToSortedMap {
228            key_mapper,
229            value_mapper,
230            merge_function: None,
231        }
232    }
233
234    /// Create a toSortedMap collector with a merge function for duplicate keys.
235    pub fn to_sorted_map_with_merge(
236        key_mapper: WasmFunction,
237        value_mapper: WasmFunction,
238        merge_function: WasmFunction,
239    ) -> Self {
240        Collector::ToSortedMap {
241            key_mapper,
242            value_mapper,
243            merge_function: Some(merge_function),
244        }
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn test_count() {
254        let collector = Collector::count();
255        match collector {
256            Collector::Count { distinct, map } => {
257                assert!(distinct.is_none());
258                assert!(map.is_none());
259            }
260            _ => panic!("Expected Count collector"),
261        }
262    }
263
264    #[test]
265    fn test_count_distinct() {
266        let collector = Collector::count_distinct();
267        match collector {
268            Collector::Count { distinct, .. } => {
269                assert_eq!(distinct, Some(true));
270            }
271            _ => panic!("Expected Count collector"),
272        }
273    }
274
275    #[test]
276    fn test_count_with_map() {
277        let collector = Collector::count_with_map(WasmFunction::new("get_id"));
278        match collector {
279            Collector::Count { map, .. } => {
280                assert!(map.is_some());
281                assert_eq!(map.unwrap().name(), "get_id");
282            }
283            _ => panic!("Expected Count collector"),
284        }
285    }
286
287    #[test]
288    fn test_sum() {
289        let collector = Collector::sum(WasmFunction::new("get_value"));
290        match collector {
291            Collector::Sum { map } => {
292                assert_eq!(map.name(), "get_value");
293            }
294            _ => panic!("Expected Sum collector"),
295        }
296    }
297
298    #[test]
299    fn test_average() {
300        let collector = Collector::average(WasmFunction::new("get_score"));
301        match collector {
302            Collector::Average { map } => {
303                assert_eq!(map.name(), "get_score");
304            }
305            _ => panic!("Expected Average collector"),
306        }
307    }
308
309    #[test]
310    fn test_min() {
311        let collector = Collector::min(
312            WasmFunction::new("get_time"),
313            WasmFunction::new("compare_time"),
314        );
315        match collector {
316            Collector::Min { map, comparator } => {
317                assert_eq!(map.name(), "get_time");
318                assert_eq!(comparator.name(), "compare_time");
319            }
320            _ => panic!("Expected Min collector"),
321        }
322    }
323
324    #[test]
325    fn test_max() {
326        let collector = Collector::max(
327            WasmFunction::new("get_priority"),
328            WasmFunction::new("compare_priority"),
329        );
330        match collector {
331            Collector::Max { map, comparator } => {
332                assert_eq!(map.name(), "get_priority");
333                assert_eq!(comparator.name(), "compare_priority");
334            }
335            _ => panic!("Expected Max collector"),
336        }
337    }
338
339    #[test]
340    fn test_to_list() {
341        let collector = Collector::to_list();
342        match collector {
343            Collector::ToList { map } => {
344                assert!(map.is_none());
345            }
346            _ => panic!("Expected ToList collector"),
347        }
348    }
349
350    #[test]
351    fn test_to_list_with_map() {
352        let collector = Collector::to_list_with_map(WasmFunction::new("get_name"));
353        match collector {
354            Collector::ToList { map } => {
355                assert!(map.is_some());
356            }
357            _ => panic!("Expected ToList collector"),
358        }
359    }
360
361    #[test]
362    fn test_to_set() {
363        let collector = Collector::to_set();
364        match collector {
365            Collector::ToSet { map } => {
366                assert!(map.is_none());
367            }
368            _ => panic!("Expected ToSet collector"),
369        }
370    }
371
372    #[test]
373    fn test_compose() {
374        let collector = Collector::compose(
375            vec![
376                Collector::count(),
377                Collector::sum(WasmFunction::new("get_value")),
378            ],
379            WasmFunction::new("combine"),
380        );
381        match collector {
382            Collector::Compose {
383                collectors,
384                combiner,
385            } => {
386                assert_eq!(collectors.len(), 2);
387                assert_eq!(combiner.name(), "combine");
388            }
389            _ => panic!("Expected Compose collector"),
390        }
391    }
392
393    #[test]
394    fn test_conditionally() {
395        let collector = Collector::conditionally(WasmFunction::new("is_valid"), Collector::count());
396        match collector {
397            Collector::Conditionally {
398                predicate,
399                collector,
400            } => {
401                assert_eq!(predicate.name(), "is_valid");
402                matches!(*collector, Collector::Count { .. });
403            }
404            _ => panic!("Expected Conditionally collector"),
405        }
406    }
407
408    #[test]
409    fn test_collect_and_then() {
410        let collector =
411            Collector::collect_and_then(Collector::count(), WasmFunction::new("to_string"));
412        match collector {
413            Collector::CollectAndThen { collector, mapper } => {
414                matches!(*collector, Collector::Count { .. });
415                assert_eq!(mapper.name(), "to_string");
416            }
417            _ => panic!("Expected CollectAndThen collector"),
418        }
419    }
420
421    #[test]
422    fn test_load_balance() {
423        let collector = Collector::load_balance(WasmFunction::new("get_employee"));
424        match collector {
425            Collector::LoadBalance { map, load } => {
426                assert_eq!(map.name(), "get_employee");
427                assert!(load.is_none());
428            }
429            _ => panic!("Expected LoadBalance collector"),
430        }
431    }
432
433    #[test]
434    fn test_load_balance_with_load() {
435        let collector = Collector::load_balance_with_load(
436            WasmFunction::new("get_employee"),
437            WasmFunction::new("get_load"),
438        );
439        match collector {
440            Collector::LoadBalance { map, load } => {
441                assert_eq!(map.name(), "get_employee");
442                assert!(load.is_some());
443            }
444            _ => panic!("Expected LoadBalance collector"),
445        }
446    }
447
448    #[test]
449    fn test_count_json_serialization() {
450        let collector = Collector::count();
451        let json = serde_json::to_string(&collector).unwrap();
452        assert!(json.contains("\"name\":\"count\""));
453
454        let parsed: Collector = serde_json::from_str(&json).unwrap();
455        assert_eq!(parsed, collector);
456    }
457
458    #[test]
459    fn test_sum_json_serialization() {
460        let collector = Collector::sum(WasmFunction::new("get_value"));
461        let json = serde_json::to_string(&collector).unwrap();
462        assert!(json.contains("\"name\":\"sum\""));
463        assert!(json.contains("\"map\":\"get_value\""));
464
465        let parsed: Collector = serde_json::from_str(&json).unwrap();
466        assert_eq!(parsed, collector);
467    }
468
469    #[test]
470    fn test_compose_json_serialization() {
471        let collector = Collector::compose(vec![Collector::count()], WasmFunction::new("wrap"));
472        let json = serde_json::to_string(&collector).unwrap();
473        assert!(json.contains("\"name\":\"compose\""));
474        assert!(json.contains("\"collectors\""));
475        assert!(json.contains("\"combiner\":\"wrap\""));
476
477        let parsed: Collector = serde_json::from_str(&json).unwrap();
478        assert_eq!(parsed, collector);
479    }
480
481    #[test]
482    fn test_conditionally_json_serialization() {
483        let collector = Collector::conditionally(WasmFunction::new("pred"), Collector::count());
484        let json = serde_json::to_string(&collector).unwrap();
485        assert!(json.contains("\"name\":\"conditionally\""));
486        assert!(json.contains("\"predicate\":\"pred\""));
487        assert!(json.contains("\"collector\""));
488
489        let parsed: Collector = serde_json::from_str(&json).unwrap();
490        assert_eq!(parsed, collector);
491    }
492
493    #[test]
494    fn test_collect_and_then_json_serialization() {
495        let collector =
496            Collector::collect_and_then(Collector::count(), WasmFunction::new("transform"));
497        let json = serde_json::to_string(&collector).unwrap();
498        assert!(json.contains("\"name\":\"collectAndThen\""));
499        assert!(json.contains("\"mapper\":\"transform\""));
500
501        let parsed: Collector = serde_json::from_str(&json).unwrap();
502        assert_eq!(parsed, collector);
503    }
504
505    #[test]
506    fn test_load_balance_json_serialization() {
507        let collector = Collector::load_balance(WasmFunction::new("get_item"));
508        let json = serde_json::to_string(&collector).unwrap();
509        assert!(json.contains("\"name\":\"loadBalance\""));
510
511        let parsed: Collector = serde_json::from_str(&json).unwrap();
512        assert_eq!(parsed, collector);
513    }
514
515    #[test]
516    fn test_nested_collectors_json() {
517        let collector = Collector::compose(
518            vec![
519                Collector::conditionally(
520                    WasmFunction::new("is_valid"),
521                    Collector::sum(WasmFunction::new("get_value")),
522                ),
523                Collector::collect_and_then(Collector::count(), WasmFunction::new("double")),
524            ],
525            WasmFunction::new("combine_results"),
526        );
527        let json = serde_json::to_string(&collector).unwrap();
528        let parsed: Collector = serde_json::from_str(&json).unwrap();
529        assert_eq!(parsed, collector);
530    }
531
532    #[test]
533    fn test_collector_clone() {
534        let collector = Collector::sum(WasmFunction::new("get_value"));
535        let cloned = collector.clone();
536        assert_eq!(collector, cloned);
537    }
538
539    #[test]
540    fn test_collector_debug() {
541        let collector = Collector::count();
542        let debug = format!("{:?}", collector);
543        assert!(debug.contains("Count"));
544    }
545
546    // Tests for toSortedSet
547    #[test]
548    fn test_to_sorted_set() {
549        let collector = Collector::to_sorted_set();
550        match collector {
551            Collector::ToSortedSet { map, comparator } => {
552                assert!(map.is_none());
553                assert!(comparator.is_none());
554            }
555            _ => panic!("Expected ToSortedSet collector"),
556        }
557    }
558
559    #[test]
560    fn test_to_sorted_set_with_map() {
561        let collector = Collector::to_sorted_set_with_map(WasmFunction::new("get_name"));
562        match collector {
563            Collector::ToSortedSet { map, comparator } => {
564                assert!(map.is_some());
565                assert_eq!(map.unwrap().name(), "get_name");
566                assert!(comparator.is_none());
567            }
568            _ => panic!("Expected ToSortedSet collector"),
569        }
570    }
571
572    #[test]
573    fn test_to_sorted_set_with_comparator() {
574        let collector = Collector::to_sorted_set_with_comparator(WasmFunction::new("compare"));
575        match collector {
576            Collector::ToSortedSet { map, comparator } => {
577                assert!(map.is_none());
578                assert!(comparator.is_some());
579                assert_eq!(comparator.unwrap().name(), "compare");
580            }
581            _ => panic!("Expected ToSortedSet collector"),
582        }
583    }
584
585    #[test]
586    fn test_to_sorted_set_with_map_and_comparator() {
587        let collector = Collector::to_sorted_set_with_map_and_comparator(
588            WasmFunction::new("get_key"),
589            WasmFunction::new("compare_keys"),
590        );
591        match collector {
592            Collector::ToSortedSet { map, comparator } => {
593                assert!(map.is_some());
594                assert!(comparator.is_some());
595                assert_eq!(map.unwrap().name(), "get_key");
596                assert_eq!(comparator.unwrap().name(), "compare_keys");
597            }
598            _ => panic!("Expected ToSortedSet collector"),
599        }
600    }
601
602    #[test]
603    fn test_to_sorted_set_json_serialization() {
604        let collector = Collector::to_sorted_set();
605        let json = serde_json::to_string(&collector).unwrap();
606        assert!(json.contains("\"name\":\"toSortedSet\""));
607
608        let parsed: Collector = serde_json::from_str(&json).unwrap();
609        assert_eq!(parsed, collector);
610    }
611
612    #[test]
613    fn test_to_sorted_set_with_comparator_json() {
614        let collector = Collector::to_sorted_set_with_comparator(WasmFunction::new("cmp"));
615        let json = serde_json::to_string(&collector).unwrap();
616        assert!(json.contains("\"comparator\":\"cmp\""));
617
618        let parsed: Collector = serde_json::from_str(&json).unwrap();
619        assert_eq!(parsed, collector);
620    }
621
622    // Tests for toMap
623    #[test]
624    fn test_to_map() {
625        let collector =
626            Collector::to_map(WasmFunction::new("get_key"), WasmFunction::new("get_value"));
627        match collector {
628            Collector::ToMap {
629                key_mapper,
630                value_mapper,
631                merge_function,
632            } => {
633                assert_eq!(key_mapper.name(), "get_key");
634                assert_eq!(value_mapper.name(), "get_value");
635                assert!(merge_function.is_none());
636            }
637            _ => panic!("Expected ToMap collector"),
638        }
639    }
640
641    #[test]
642    fn test_to_map_with_merge() {
643        let collector = Collector::to_map_with_merge(
644            WasmFunction::new("get_key"),
645            WasmFunction::new("get_value"),
646            WasmFunction::new("merge_values"),
647        );
648        match collector {
649            Collector::ToMap {
650                key_mapper,
651                value_mapper,
652                merge_function,
653            } => {
654                assert_eq!(key_mapper.name(), "get_key");
655                assert_eq!(value_mapper.name(), "get_value");
656                assert!(merge_function.is_some());
657                assert_eq!(merge_function.unwrap().name(), "merge_values");
658            }
659            _ => panic!("Expected ToMap collector"),
660        }
661    }
662
663    #[test]
664    fn test_to_map_json_serialization() {
665        let collector = Collector::to_map(WasmFunction::new("get_k"), WasmFunction::new("get_v"));
666        let json = serde_json::to_string(&collector).unwrap();
667        assert!(json.contains("\"name\":\"toMap\""));
668        assert!(json.contains("\"key_mapper\":\"get_k\""));
669        assert!(json.contains("\"value_mapper\":\"get_v\""));
670
671        let parsed: Collector = serde_json::from_str(&json).unwrap();
672        assert_eq!(parsed, collector);
673    }
674
675    #[test]
676    fn test_to_map_with_merge_json_serialization() {
677        let collector = Collector::to_map_with_merge(
678            WasmFunction::new("k"),
679            WasmFunction::new("v"),
680            WasmFunction::new("merge"),
681        );
682        let json = serde_json::to_string(&collector).unwrap();
683        assert!(json.contains("\"merge_function\":\"merge\""));
684
685        let parsed: Collector = serde_json::from_str(&json).unwrap();
686        assert_eq!(parsed, collector);
687    }
688
689    // Tests for toSortedMap
690    #[test]
691    fn test_to_sorted_map() {
692        let collector =
693            Collector::to_sorted_map(WasmFunction::new("get_key"), WasmFunction::new("get_value"));
694        match collector {
695            Collector::ToSortedMap {
696                key_mapper,
697                value_mapper,
698                merge_function,
699            } => {
700                assert_eq!(key_mapper.name(), "get_key");
701                assert_eq!(value_mapper.name(), "get_value");
702                assert!(merge_function.is_none());
703            }
704            _ => panic!("Expected ToSortedMap collector"),
705        }
706    }
707
708    #[test]
709    fn test_to_sorted_map_with_merge() {
710        let collector = Collector::to_sorted_map_with_merge(
711            WasmFunction::new("get_key"),
712            WasmFunction::new("get_value"),
713            WasmFunction::new("merge"),
714        );
715        match collector {
716            Collector::ToSortedMap {
717                key_mapper,
718                value_mapper,
719                merge_function,
720            } => {
721                assert_eq!(key_mapper.name(), "get_key");
722                assert_eq!(value_mapper.name(), "get_value");
723                assert!(merge_function.is_some());
724                assert_eq!(merge_function.unwrap().name(), "merge");
725            }
726            _ => panic!("Expected ToSortedMap collector"),
727        }
728    }
729
730    #[test]
731    fn test_to_sorted_map_json_serialization() {
732        let collector = Collector::to_sorted_map(WasmFunction::new("k"), WasmFunction::new("v"));
733        let json = serde_json::to_string(&collector).unwrap();
734        assert!(json.contains("\"name\":\"toSortedMap\""));
735
736        let parsed: Collector = serde_json::from_str(&json).unwrap();
737        assert_eq!(parsed, collector);
738    }
739}
solverforge_core/constraints/collectors.rs

solverforge_core/constraints/
collectors.rs
