Skip to main content

telltale_language/ast/
annotation_collection.rs

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