Skip to main content

telltale_language/ast/
annotation_collection.rs

1//! Typed accessors for protocol annotation collections.
2
3use super::{DslAnnotationEntry, ProtocolAnnotation};
4use std::collections::HashMap;
5use std::time::Duration;
6
7/// A collection of protocol annotations with typed accessors.
8#[derive(Debug, Clone, Default, PartialEq)]
9pub struct Annotations {
10    items: Vec<ProtocolAnnotation>,
11}
12
13impl Annotations {
14    /// Create an empty annotation set.
15    #[must_use]
16    pub fn new() -> Self {
17        Self { items: Vec::new() }
18    }
19
20    /// Create from a single annotation.
21    #[must_use]
22    pub fn single(annotation: ProtocolAnnotation) -> Self {
23        Self {
24            items: vec![annotation],
25        }
26    }
27
28    /// Create from a vector of annotations.
29    #[must_use]
30    pub fn from_vec(items: Vec<ProtocolAnnotation>) -> Self {
31        Self { items }
32    }
33
34    /// Create from ordered raw DSL entries.
35    ///
36    /// Entry order is preserved exactly. Known annotations are lowered to their
37    /// typed variants; unknown entries remain `Custom`.
38    #[must_use]
39    pub fn from_dsl_entries(entries: &[DslAnnotationEntry]) -> Self {
40        let mut items = Vec::new();
41        let mut timed_choice = false;
42        let mut timeout_ms = None;
43
44        for entry in entries {
45            if entry.key == "timed_choice" && entry.value == "true" {
46                timed_choice = true;
47                continue;
48            }
49            if entry.key == "timeout_ms" {
50                timeout_ms = entry.value.parse::<u64>().ok().map(Duration::from_millis);
51                continue;
52            }
53            items.push(ProtocolAnnotation::parse_dsl_entry(
54                &entry.key,
55                &entry.value,
56            ));
57        }
58
59        if timed_choice || timeout_ms.is_some() {
60            items.insert(
61                0,
62                ProtocolAnnotation::TimedChoice {
63                    duration: timeout_ms.unwrap_or_else(|| Duration::from_secs(5)),
64                },
65            );
66        }
67
68        Self { items }
69    }
70
71    pub fn from_dsl_map(map: &HashMap<String, String>) -> Self {
72        let mut entries: Vec<_> = map.iter().collect();
73        entries.sort_by_key(|(key_a, _)| *key_a);
74        let entries = entries
75            .into_iter()
76            .map(|(key, value)| DslAnnotationEntry::new(key.clone(), value.clone()))
77            .collect::<Vec<_>>();
78        Self::from_dsl_entries(&entries)
79    }
80
81    pub fn dsl_map(&self) -> HashMap<String, String> {
82        let mut map = HashMap::new();
83
84        for annotation in &self.items {
85            for (key, value) in annotation.dsl_entries() {
86                map.insert(key, value);
87            }
88        }
89
90        map
91    }
92
93    /// Return raw DSL entries in stable order.
94    #[must_use]
95    pub fn dsl_entries(&self) -> Vec<DslAnnotationEntry> {
96        self.items
97            .iter()
98            .flat_map(|annotation| {
99                annotation
100                    .dsl_entries()
101                    .into_iter()
102                    .map(|(key, value)| DslAnnotationEntry::new(key, value))
103            })
104            .collect()
105    }
106
107    /// Add an annotation.
108    pub fn push(&mut self, annotation: ProtocolAnnotation) {
109        self.items.push(annotation);
110    }
111
112    /// Check if empty.
113    #[must_use]
114    pub fn is_empty(&self) -> bool {
115        self.items.is_empty()
116    }
117
118    /// Get the number of annotations.
119    #[must_use]
120    pub fn len(&self) -> usize {
121        self.items.len()
122    }
123
124    /// Iterate over annotations.
125    pub fn iter(&self) -> impl Iterator<Item = &ProtocolAnnotation> {
126        self.items.iter()
127    }
128
129    /// Iterate mutably over annotations.
130    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut ProtocolAnnotation> {
131        self.items.iter_mut()
132    }
133
134    /// Check if any annotation satisfies a predicate.
135    #[must_use]
136    pub fn any<F>(&self, f: F) -> bool
137    where
138        F: FnMut(&ProtocolAnnotation) -> bool,
139    {
140        self.items.iter().any(f)
141    }
142
143    /// Find the first annotation satisfying a predicate.
144    #[must_use]
145    pub fn find<F>(&self, f: F) -> Option<&ProtocolAnnotation>
146    where
147        F: FnMut(&&ProtocolAnnotation) -> bool,
148    {
149        self.items.iter().find(f)
150    }
151
152    /// Get timed choice annotation if present.
153    #[must_use]
154    pub fn timed_choice(&self) -> Option<Duration> {
155        self.items.iter().find_map(|a| a.timed_choice_duration())
156    }
157
158    /// Check if this has a timed choice annotation.
159    #[must_use]
160    pub fn has_timed_choice(&self) -> bool {
161        self.items.iter().any(|a| a.is_timed_choice())
162    }
163
164    /// Get priority annotation if present.
165    #[must_use]
166    pub fn priority(&self) -> Option<u32> {
167        self.items.iter().find_map(|a| a.priority_value())
168    }
169
170    /// Check if marked as idempotent.
171    #[must_use]
172    pub fn is_idempotent(&self) -> bool {
173        self.items.iter().any(|a| a.is_idempotent())
174    }
175
176    /// Check if has any trace annotation.
177    #[must_use]
178    pub fn has_trace(&self) -> bool {
179        self.items.iter().any(|a| a.is_trace())
180    }
181
182    /// Check if has a heartbeat annotation.
183    #[must_use]
184    pub fn has_heartbeat(&self) -> bool {
185        self.items.iter().any(|a| a.is_heartbeat())
186    }
187
188    /// Get heartbeat parameters if present (interval, on_missing_count).
189    #[must_use]
190    pub fn heartbeat(&self) -> Option<(Duration, u32)> {
191        self.items.iter().find_map(|a| a.heartbeat_params())
192    }
193
194    /// Check if has a runtime timeout annotation.
195    #[must_use]
196    pub fn has_runtime_timeout(&self) -> bool {
197        self.items.iter().any(|a| a.is_runtime_timeout())
198    }
199
200    /// Get runtime timeout duration if present.
201    #[must_use]
202    pub fn runtime_timeout(&self) -> Option<Duration> {
203        self.items.iter().find_map(|a| a.runtime_timeout_duration())
204    }
205
206    /// Check if has a parallel annotation.
207    #[must_use]
208    pub fn has_parallel(&self) -> bool {
209        self.items.iter().any(|a| a.is_parallel())
210    }
211
212    /// Check if has an ordered annotation.
213    #[must_use]
214    pub fn has_ordered(&self) -> bool {
215        self.items.iter().any(|a| a.is_ordered())
216    }
217
218    /// Check if has a min_responses annotation.
219    #[must_use]
220    pub fn has_min_responses(&self) -> bool {
221        self.items.iter().any(|a| a.is_min_responses())
222    }
223
224    /// Get min_responses value if present.
225    #[must_use]
226    pub fn min_responses(&self) -> Option<u32> {
227        self.items.iter().find_map(|a| a.min_responses_value())
228    }
229
230    /// Get a custom annotation value by key.
231    #[must_use]
232    pub fn custom(&self, key: &str) -> Option<&str> {
233        self.items.iter().find_map(|a| a.custom_value(key))
234    }
235
236    /// Get every custom annotation value for a key in source order.
237    #[must_use]
238    pub fn custom_values(&self, key: &str) -> Vec<&str> {
239        self.items
240            .iter()
241            .filter_map(|annotation| annotation.custom_value(key))
242            .collect()
243    }
244
245    /// Merge annotations from another set.
246    pub fn merge(&mut self, other: &Annotations) {
247        self.items.extend(other.items.iter().cloned());
248    }
249
250    /// Clear all annotations.
251    pub fn clear(&mut self) {
252        self.items.clear();
253    }
254}
255
256impl IntoIterator for Annotations {
257    type Item = ProtocolAnnotation;
258    type IntoIter = std::vec::IntoIter<ProtocolAnnotation>;
259
260    fn into_iter(self) -> Self::IntoIter {
261        self.items.into_iter()
262    }
263}
264
265impl<'a> IntoIterator for &'a Annotations {
266    type Item = &'a ProtocolAnnotation;
267    type IntoIter = std::slice::Iter<'a, ProtocolAnnotation>;
268
269    fn into_iter(self) -> Self::IntoIter {
270        self.items.iter()
271    }
272}
273
274impl FromIterator<ProtocolAnnotation> for Annotations {
275    fn from_iter<I: IntoIterator<Item = ProtocolAnnotation>>(iter: I) -> Self {
276        Self {
277            items: iter.into_iter().collect(),
278        }
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_timed_choice_annotation() {
288        let ann = ProtocolAnnotation::timed_choice_ms(5000);
289        assert!(ann.is_timed_choice());
290        assert_eq!(ann.timed_choice_duration(), Some(Duration::from_secs(5)));
291    }
292
293    #[test]
294    fn test_priority_annotation() {
295        let ann = ProtocolAnnotation::priority(10);
296        assert!(ann.is_priority());
297        assert_eq!(ann.priority_value(), Some(10));
298    }
299
300    #[test]
301    fn test_custom_annotation() {
302        let ann = ProtocolAnnotation::custom("my_key", "my_value");
303        assert!(ann.is_custom_key("my_key"));
304        assert_eq!(ann.custom_value("my_key"), Some("my_value"));
305        assert_eq!(ann.custom_value("other"), None);
306    }
307
308    #[test]
309    fn test_parse_dsl_entry() {
310        let ann = ProtocolAnnotation::parse_dsl_entry("priority", "5");
311        assert_eq!(ann, ProtocolAnnotation::Priority(5));
312
313        let ann = ProtocolAnnotation::parse_dsl_entry("unknown", "value");
314        assert!(matches!(ann, ProtocolAnnotation::Custom { .. }));
315    }
316
317    #[test]
318    fn test_annotations_collection() {
319        let mut anns = Annotations::new();
320        anns.push(ProtocolAnnotation::timed_choice_ms(1000));
321        anns.push(ProtocolAnnotation::priority(5));
322
323        assert!(anns.has_timed_choice());
324        assert_eq!(anns.timed_choice(), Some(Duration::from_secs(1)));
325        assert_eq!(anns.priority(), Some(5));
326        assert_eq!(anns.len(), 2);
327    }
328
329    #[test]
330    fn test_map_roundtrip() {
331        let mut original = HashMap::new();
332        original.insert("timed_choice".to_string(), "true".to_string());
333        original.insert("timeout_ms".to_string(), "5000".to_string());
334        original.insert("priority".to_string(), "10".to_string());
335
336        let anns = Annotations::from_dsl_map(&original);
337        assert!(anns.has_timed_choice());
338        assert_eq!(anns.timed_choice(), Some(Duration::from_secs(5)));
339        assert_eq!(anns.priority(), Some(10));
340
341        let restored = anns.dsl_map();
342        assert_eq!(restored.get("timed_choice"), Some(&"true".to_string()));
343        assert_eq!(restored.get("timeout_ms"), Some(&"5000".to_string()));
344    }
345
346    #[test]
347    fn test_dsl_entries_preserve_order() {
348        let entries = vec![
349            DslAnnotationEntry::new("guard_capability", "chat:send"),
350            DslAnnotationEntry::new("flow_cost", "10"),
351            DslAnnotationEntry::new("journal_facts", "sent"),
352        ];
353
354        let annotations = Annotations::from_dsl_entries(&entries);
355        let restored = annotations.dsl_entries();
356
357        assert_eq!(restored, entries);
358    }
359
360    #[test]
361    fn test_custom_values_preserve_duplicates() {
362        let entries = vec![
363            DslAnnotationEntry::new("guard_capability", "chat:send"),
364            DslAnnotationEntry::new("guard_capability", "chat:audit"),
365        ];
366
367        let annotations = Annotations::from_dsl_entries(&entries);
368        assert_eq!(
369            annotations.custom_values("guard_capability"),
370            vec!["chat:send", "chat:audit"]
371        );
372        assert_eq!(annotations.custom("guard_capability"), Some("chat:send"));
373    }
374
375    #[test]
376    fn test_parallel_annotation() {
377        let ann = ProtocolAnnotation::parallel();
378        assert!(ann.is_parallel());
379        assert!(!ann.is_ordered());
380    }
381
382    #[test]
383    fn test_ordered_annotation() {
384        let ann = ProtocolAnnotation::ordered();
385        assert!(ann.is_ordered());
386        assert!(!ann.is_parallel());
387    }
388
389    #[test]
390    fn test_min_responses_annotation() {
391        let ann = ProtocolAnnotation::min_responses(3);
392        assert!(ann.is_min_responses());
393        assert_eq!(ann.min_responses_value(), Some(3));
394    }
395
396    #[test]
397    fn test_annotations_parallel_ordered() {
398        let mut anns = Annotations::new();
399        anns.push(ProtocolAnnotation::parallel());
400        anns.push(ProtocolAnnotation::min_responses(5));
401
402        assert!(anns.has_parallel());
403        assert!(!anns.has_ordered());
404        assert!(anns.has_min_responses());
405        assert_eq!(anns.min_responses(), Some(5));
406    }
407
408    #[test]
409    fn test_parse_dsl_entry_parallel() {
410        let ann = ProtocolAnnotation::parse_dsl_entry("parallel", "true");
411        assert_eq!(ann, ProtocolAnnotation::Parallel);
412
413        let ann = ProtocolAnnotation::parse_dsl_entry("ordered", "true");
414        assert_eq!(ann, ProtocolAnnotation::Ordered);
415
416        let ann = ProtocolAnnotation::parse_dsl_entry("parallel", "");
417        assert_eq!(ann, ProtocolAnnotation::Parallel);
418
419        let ann = ProtocolAnnotation::parse_dsl_entry("ordered", "");
420        assert_eq!(ann, ProtocolAnnotation::Ordered);
421
422        let ann = ProtocolAnnotation::parse_dsl_entry("min_responses", "3");
423        assert_eq!(ann, ProtocolAnnotation::MinResponses(3));
424    }
425
426    #[test]
427    fn test_to_map_new_annotations() {
428        let mut anns = Annotations::new();
429        anns.push(ProtocolAnnotation::Parallel);
430        anns.push(ProtocolAnnotation::Ordered);
431        anns.push(ProtocolAnnotation::MinResponses(5));
432
433        let map = anns.dsl_map();
434        assert_eq!(map.get("parallel"), Some(&"true".to_string()));
435        assert_eq!(map.get("ordered"), Some(&"true".to_string()));
436        assert_eq!(map.get("min_responses"), Some(&"5".to_string()));
437    }
438}