Skip to main content

telltale_language/ast/
protocol.rs

1// Protocol AST definitions
2
3use super::annotation::Annotations;
4use super::{MessageType, NonEmptyVec, ProgressAttachment, Role, ValidationError};
5use proc_macro2::{Ident, TokenStream};
6
7#[path = "protocol_validation.rs"]
8mod validation;
9use validation::{ensure_declared_role, validate_choice_branches};
10
11fn annotation_has_key(annotation: &super::annotation::ProtocolAnnotation, key: &str) -> bool {
12    annotation
13        .dsl_entries()
14        .iter()
15        .any(|(entry_key, _)| entry_key == key)
16}
17
18fn annotation_has_value(
19    annotation: &super::annotation::ProtocolAnnotation,
20    key: &str,
21    value: &str,
22) -> bool {
23    annotation
24        .dsl_entries()
25        .iter()
26        .any(|(entry_key, entry_value)| entry_key == key && entry_value.eq_ignore_ascii_case(value))
27}
28
29/// Protocol specification using choreographic constructs
30#[derive(Debug)]
31pub enum Protocol {
32    /// Begin one explicit semantic operation instance.
33    Begin {
34        operation: String,
35        args: Vec<String>,
36        progress: Option<ProgressAttachment>,
37        continuation: Box<Protocol>,
38    },
39
40    /// Await one previously begun semantic operation instance.
41    Await {
42        operation: String,
43        continuation: Box<Protocol>,
44    },
45
46    /// Resolve one previously begun semantic operation instance.
47    Resolve {
48        operation: String,
49        outcome: CommitmentOutcome,
50        continuation: Box<Protocol>,
51    },
52
53    /// Invalidate one previously begun semantic operation instance.
54    Invalidate {
55        operation: String,
56        continuation: Box<Protocol>,
57    },
58
59    /// Message send: A -> B: Message
60    Send {
61        from: Role,
62        to: Role,
63        message: MessageType,
64        continuation: Box<Protocol>,
65        /// Statement-level annotations
66        annotations: Annotations,
67        /// From role annotations
68        from_annotations: Annotations,
69        /// To role annotations
70        to_annotations: Annotations,
71    },
72
73    /// Broadcast: A -> *: Message
74    Broadcast {
75        from: Role,
76        to_all: NonEmptyVec<Role>,
77        message: MessageType,
78        continuation: Box<Protocol>,
79        /// Statement-level annotations
80        annotations: Annotations,
81        /// From role annotations
82        from_annotations: Annotations,
83    },
84
85    /// Choice made by a role
86    Choice {
87        role: Role,
88        branches: NonEmptyVec<Branch>,
89        /// Statement-level annotations
90        annotations: Annotations,
91    },
92
93    /// Local authority/evidence binding.
94    Let {
95        /// Bound variable name.
96        name: String,
97        /// Whether the binding is authoritative, observational, or plain.
98        mode: AuthorityBindingMode,
99        /// Bound expression.
100        expr: AuthorityExpr,
101        /// Whether the binding is linear/single-use.
102        linear: bool,
103        /// Continuation after the binding.
104        continuation: Box<Protocol>,
105    },
106
107    /// Local authority/result match.
108    Case {
109        /// Scrutinee expression.
110        expr: AuthorityExpr,
111        /// Match branches.
112        branches: NonEmptyVec<CaseBranch>,
113    },
114
115    /// Explicit timeout/cancel surface syntax prior to projection lowering.
116    Timeout {
117        /// Role that owns the timeout decision.
118        role: Role,
119        /// Timeout duration in milliseconds.
120        duration_ms: u64,
121        /// Main body before timeout fires.
122        body: Box<Protocol>,
123        /// Timeout branch.
124        on_timeout: Box<Protocol>,
125        /// Optional explicit cancellation branch.
126        on_cancel: Option<Box<Protocol>>,
127    },
128
129    /// Loop construct
130    Loop {
131        condition: Option<Condition>,
132        body: Box<Protocol>,
133    },
134
135    /// Parallel composition
136    Parallel { protocols: NonEmptyVec<Protocol> },
137
138    /// Recursive protocol with label
139    Rec { label: Ident, body: Box<Protocol> },
140
141    /// Reference to recursive label
142    Var(Ident),
143
144    /// Canonical semantic publication surface.
145    Publish {
146        event: String,
147        arg: Option<String>,
148        continuation: Box<Protocol>,
149    },
150
151    /// Canonical publication that lifts an authoritative witness into a named publication.
152    PublishAuthority {
153        witness: String,
154        publication_name: String,
155        continuation: Box<Protocol>,
156    },
157
158    /// Canonical materialization from one named publication.
159    Materialize {
160        proof: String,
161        publication: String,
162        continuation: Box<Protocol>,
163    },
164
165    /// Explicit semantic owner handoff.
166    Handoff {
167        operation: String,
168        target: Role,
169        receipt: String,
170        continuation: Box<Protocol>,
171    },
172
173    /// Declared semantically required dependent work.
174    DependentWork {
175        name: String,
176        arg: Option<String>,
177        required_for: String,
178        continuation: Box<Protocol>,
179    },
180
181    /// Protocol extension point for custom behaviors
182    Extension {
183        /// The extension implementation
184        extension: Box<dyn crate::extensions::ProtocolExtension>,
185        /// Continuation after this extension
186        continuation: Box<Protocol>,
187        /// Statement-level annotations
188        annotations: Annotations,
189    },
190
191    /// Protocol termination
192    End,
193}
194
195/// A branch in a choice
196#[derive(Debug, Clone)]
197pub struct Branch {
198    pub label: Ident,
199    pub guard: Option<ChoiceGuard>,
200    pub protocol: Protocol,
201}
202
203/// Match branch in a `case/of` expression.
204#[derive(Debug, Clone)]
205pub struct CaseBranch {
206    pub pattern: CasePattern,
207    pub protocol: Protocol,
208}
209
210/// Pattern accepted by DSL `case/of`.
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub struct CasePattern {
213    pub constructor: String,
214    pub binders: Vec<String>,
215}
216
217/// Explicit authority/observation mode for one local binding.
218#[derive(Debug, Clone, Copy, PartialEq, Eq)]
219pub enum AuthorityBindingMode {
220    Plain,
221    Authoritative,
222    Observe,
223}
224
225impl Clone for Protocol {
226    fn clone(&self) -> Self {
227        match self {
228            Protocol::Begin {
229                operation,
230                args,
231                progress,
232                continuation,
233            } => Protocol::Begin {
234                operation: operation.clone(),
235                args: args.clone(),
236                progress: progress.clone(),
237                continuation: continuation.clone(),
238            },
239            Protocol::Await {
240                operation,
241                continuation,
242            } => Protocol::Await {
243                operation: operation.clone(),
244                continuation: continuation.clone(),
245            },
246            Protocol::Resolve {
247                operation,
248                outcome,
249                continuation,
250            } => Protocol::Resolve {
251                operation: operation.clone(),
252                outcome: outcome.clone(),
253                continuation: continuation.clone(),
254            },
255            Protocol::Invalidate {
256                operation,
257                continuation,
258            } => Protocol::Invalidate {
259                operation: operation.clone(),
260                continuation: continuation.clone(),
261            },
262            Protocol::Send {
263                from,
264                to,
265                message,
266                continuation,
267                annotations,
268                from_annotations,
269                to_annotations,
270            } => Protocol::Send {
271                from: from.clone(),
272                to: to.clone(),
273                message: message.clone(),
274                continuation: continuation.clone(),
275                annotations: annotations.clone(),
276                from_annotations: from_annotations.clone(),
277                to_annotations: to_annotations.clone(),
278            },
279            Protocol::Broadcast {
280                from,
281                to_all,
282                message,
283                continuation,
284                annotations,
285                from_annotations,
286            } => Protocol::Broadcast {
287                from: from.clone(),
288                to_all: to_all.clone(),
289                message: message.clone(),
290                continuation: continuation.clone(),
291                annotations: annotations.clone(),
292                from_annotations: from_annotations.clone(),
293            },
294            Protocol::Choice {
295                role,
296                branches,
297                annotations,
298            } => Protocol::Choice {
299                role: role.clone(),
300                branches: branches.clone(),
301                annotations: annotations.clone(),
302            },
303            Protocol::Let {
304                name,
305                mode,
306                expr,
307                linear,
308                continuation,
309            } => Protocol::Let {
310                name: name.clone(),
311                mode: *mode,
312                expr: expr.clone(),
313                linear: *linear,
314                continuation: continuation.clone(),
315            },
316            Protocol::Case { expr, branches } => Protocol::Case {
317                expr: expr.clone(),
318                branches: branches.clone(),
319            },
320            Protocol::Timeout {
321                role,
322                duration_ms,
323                body,
324                on_timeout,
325                on_cancel,
326            } => Protocol::Timeout {
327                role: role.clone(),
328                duration_ms: *duration_ms,
329                body: body.clone(),
330                on_timeout: on_timeout.clone(),
331                on_cancel: on_cancel.clone(),
332            },
333            Protocol::Loop { condition, body } => Protocol::Loop {
334                condition: condition.clone(),
335                body: body.clone(),
336            },
337            Protocol::Parallel { protocols } => Protocol::Parallel {
338                protocols: protocols.clone(),
339            },
340            Protocol::Rec { label, body } => Protocol::Rec {
341                label: label.clone(),
342                body: body.clone(),
343            },
344            Protocol::Var(label) => Protocol::Var(label.clone()),
345            Protocol::Publish {
346                event,
347                arg,
348                continuation,
349            } => Protocol::Publish {
350                event: event.clone(),
351                arg: arg.clone(),
352                continuation: continuation.clone(),
353            },
354            Protocol::PublishAuthority {
355                witness,
356                publication_name,
357                continuation,
358            } => Protocol::PublishAuthority {
359                witness: witness.clone(),
360                publication_name: publication_name.clone(),
361                continuation: continuation.clone(),
362            },
363            Protocol::Materialize {
364                proof,
365                publication,
366                continuation,
367            } => Protocol::Materialize {
368                proof: proof.clone(),
369                publication: publication.clone(),
370                continuation: continuation.clone(),
371            },
372            Protocol::Handoff {
373                operation,
374                target,
375                receipt,
376                continuation,
377            } => Protocol::Handoff {
378                operation: operation.clone(),
379                target: target.clone(),
380                receipt: receipt.clone(),
381                continuation: continuation.clone(),
382            },
383            Protocol::DependentWork {
384                name,
385                arg,
386                required_for,
387                continuation,
388            } => Protocol::DependentWork {
389                name: name.clone(),
390                arg: arg.clone(),
391                required_for: required_for.clone(),
392                continuation: continuation.clone(),
393            },
394            Protocol::Extension {
395                extension,
396                continuation,
397                annotations,
398            } => Protocol::Extension {
399                extension: extension.clone(),
400                continuation: continuation.clone(),
401                annotations: annotations.clone(),
402            },
403            Protocol::End => Protocol::End,
404        }
405    }
406}
407
408/// Authority- or evidence-oriented expression surface syntax.
409#[derive(Debug, Clone, PartialEq, Eq)]
410pub enum AuthorityExpr {
411    Var(String),
412    Check {
413        effect: String,
414        operation: String,
415        args: Vec<String>,
416    },
417    Observe {
418        effect: String,
419        operation: String,
420        args: Vec<String>,
421    },
422    Transfer {
423        subject: String,
424        from: String,
425        to: String,
426    },
427    Constructor {
428        name: String,
429        arg: Option<String>,
430    },
431    Call {
432        name: String,
433        args: Vec<String>,
434    },
435}
436
437/// Explicit outcome used to resolve a begun semantic operation.
438#[derive(Debug, Clone, PartialEq, Eq)]
439pub enum CommitmentOutcome {
440    Success(Option<String>),
441    Failure(Option<String>),
442    Timeout(Option<String>),
443    Cancelled,
444}
445
446/// Guard surface syntax attached to one choice branch.
447#[derive(Debug, Clone)]
448pub enum ChoiceGuard {
449    Predicate(TokenStream),
450    Evidence {
451        effect: String,
452        operation: String,
453        args: Vec<String>,
454        binding: String,
455    },
456}
457
458/// Loop condition
459#[derive(Debug, Clone)]
460pub enum Condition {
461    /// Loop while a role decides
462    RoleDecides(Role),
463    /// Fixed iteration count
464    Count(usize),
465    /// Custom condition
466    Custom(TokenStream),
467    /// Fuel-based bounding: maximum number of loop iterations
468    Fuel(usize),
469    /// Yield control after N communication steps
470    YieldAfter(usize),
471    /// Yield when a specific label/condition is encountered
472    YieldWhen(String),
473}
474
475impl Protocol {
476    #[must_use]
477    pub fn mentions_role(&self, role: &Role) -> bool {
478        match self {
479            Protocol::Begin { continuation, .. }
480            | Protocol::Await { continuation, .. }
481            | Protocol::Resolve { continuation, .. }
482            | Protocol::Invalidate { continuation, .. } => continuation.mentions_role(role),
483            Protocol::Send {
484                from,
485                to,
486                continuation,
487                ..
488            } => {
489                from.matches_family(role)
490                    || to.matches_family(role)
491                    || continuation.mentions_role(role)
492            }
493            Protocol::Broadcast {
494                from,
495                to_all,
496                continuation,
497                ..
498            } => {
499                from.matches_family(role)
500                    || to_all.iter().any(|r| r.matches_family(role))
501                    || continuation.mentions_role(role)
502            }
503            Protocol::Choice {
504                role: r, branches, ..
505            } => r.matches_family(role) || branches.iter().any(|b| b.protocol.mentions_role(role)),
506            Protocol::Let { continuation, .. } => continuation.mentions_role(role),
507            Protocol::Case { branches, .. } => {
508                branches.iter().any(|b| b.protocol.mentions_role(role))
509            }
510            Protocol::Timeout {
511                role: timeout_role,
512                body,
513                on_timeout,
514                on_cancel,
515                ..
516            } => {
517                timeout_role.matches_family(role)
518                    || body.mentions_role(role)
519                    || on_timeout.mentions_role(role)
520                    || on_cancel
521                        .as_deref()
522                        .is_some_and(|branch| branch.mentions_role(role))
523            }
524            Protocol::Loop { body, .. } => body.mentions_role(role),
525            Protocol::Parallel { protocols } => protocols.iter().any(|p| p.mentions_role(role)),
526            Protocol::Rec { body, .. } => body.mentions_role(role),
527            Protocol::Publish { continuation, .. }
528            | Protocol::PublishAuthority { continuation, .. }
529            | Protocol::Materialize { continuation, .. }
530            | Protocol::DependentWork { continuation, .. } => continuation.mentions_role(role),
531            Protocol::Handoff {
532                target,
533                continuation,
534                ..
535            } => target.matches_family(role) || continuation.mentions_role(role),
536            Protocol::Extension {
537                extension,
538                continuation,
539                ..
540            } => extension.mentions_role(role) || continuation.mentions_role(role),
541            Protocol::Var(_) | Protocol::End => false,
542        }
543    }
544
545    pub(crate) fn validate(&self, roles: &[Role]) -> Result<(), ValidationError> {
546        match self {
547            Protocol::Begin { continuation, .. }
548            | Protocol::Await { continuation, .. }
549            | Protocol::Resolve { continuation, .. }
550            | Protocol::Invalidate { continuation, .. } => continuation.validate(roles),
551            Protocol::Send {
552                from,
553                to,
554                continuation,
555                ..
556            } => {
557                ensure_declared_role(roles, from)?;
558                ensure_declared_role(roles, to)?;
559                continuation.validate(roles)
560            }
561            Protocol::Broadcast {
562                from,
563                to_all,
564                continuation,
565                ..
566            } => {
567                ensure_declared_role(roles, from)?;
568                for to in to_all {
569                    ensure_declared_role(roles, to)?;
570                }
571                continuation.validate(roles)
572            }
573            Protocol::Choice { role, branches, .. } => {
574                ensure_declared_role(roles, role)?;
575                validate_choice_branches(role, branches)?;
576                Ok(())
577            }
578            Protocol::Let { continuation, .. } => continuation.validate(roles),
579            Protocol::Case { branches, .. } => {
580                for branch in branches {
581                    branch.protocol.validate(roles)?;
582                }
583                Ok(())
584            }
585            Protocol::Timeout {
586                role,
587                body,
588                on_timeout,
589                on_cancel,
590                ..
591            } => {
592                ensure_declared_role(roles, role)?;
593                body.validate(roles)?;
594                on_timeout.validate(roles)?;
595                if let Some(on_cancel) = on_cancel.as_deref() {
596                    on_cancel.validate(roles)?;
597                }
598                Ok(())
599            }
600            Protocol::Loop { body, .. } => body.validate(roles),
601            Protocol::Parallel { protocols } => {
602                for p in protocols {
603                    p.validate(roles)?;
604                }
605                Ok(())
606            }
607            Protocol::Rec { body, .. } => body.validate(roles),
608            Protocol::Publish { continuation, .. }
609            | Protocol::PublishAuthority { continuation, .. }
610            | Protocol::Materialize { continuation, .. }
611            | Protocol::DependentWork { continuation, .. } => continuation.validate(roles),
612            Protocol::Handoff {
613                target,
614                continuation,
615                ..
616            } => {
617                ensure_declared_role(roles, target)?;
618                continuation.validate(roles)
619            }
620            Protocol::Extension {
621                extension,
622                continuation,
623                ..
624            } => {
625                // Validate the extension with the extension system's validation
626                extension.validate(roles).map_err(|e| {
627                    ValidationError::ExtensionError(format!("Extension validation failed: {}", e))
628                })?;
629                continuation.validate(roles)
630            }
631            Protocol::Var(_) | Protocol::End => Ok(()),
632        }
633    }
634
635    /// Get statement-level annotations for this protocol node
636    pub fn get_annotations(&self) -> &Annotations {
637        match self {
638            Protocol::Send { annotations, .. } => annotations,
639            Protocol::Broadcast { annotations, .. } => annotations,
640            Protocol::Choice { annotations, .. } => annotations,
641            Protocol::Begin { .. }
642            | Protocol::Await { .. }
643            | Protocol::Resolve { .. }
644            | Protocol::Invalidate { .. }
645            | Protocol::Let { .. }
646            | Protocol::Case { .. }
647            | Protocol::Timeout { .. }
648            | Protocol::Publish { .. }
649            | Protocol::PublishAuthority { .. }
650            | Protocol::Materialize { .. }
651            | Protocol::Handoff { .. }
652            | Protocol::DependentWork { .. } => {
653                static EMPTY: std::sync::OnceLock<Annotations> = std::sync::OnceLock::new();
654                EMPTY.get_or_init(Annotations::new)
655            }
656            Protocol::Extension { annotations, .. } => annotations,
657            Protocol::Loop { .. }
658            | Protocol::Parallel { .. }
659            | Protocol::Rec { .. }
660            | Protocol::Var(_)
661            | Protocol::End => {
662                // Return empty annotations for protocol nodes that don't have annotations yet
663                static EMPTY: std::sync::OnceLock<Annotations> = std::sync::OnceLock::new();
664                EMPTY.get_or_init(Annotations::new)
665            }
666        }
667    }
668
669    /// Get from-role annotations for Send/Broadcast statements
670    pub fn get_from_annotations(&self) -> Option<&Annotations> {
671        match self {
672            Protocol::Send {
673                from_annotations, ..
674            } => Some(from_annotations),
675            Protocol::Broadcast {
676                from_annotations, ..
677            } => Some(from_annotations),
678            _ => None,
679        }
680    }
681
682    /// Get to-role annotations for Send statements
683    pub fn get_to_annotations(&self) -> Option<&Annotations> {
684        match self {
685            Protocol::Send { to_annotations, .. } => Some(to_annotations),
686            _ => None,
687        }
688    }
689
690    /// Get mutable reference to annotations for modification
691    pub fn get_annotations_mut(&mut self) -> Option<&mut Annotations> {
692        match self {
693            Protocol::Send { annotations, .. } => Some(annotations),
694            Protocol::Broadcast { annotations, .. } => Some(annotations),
695            Protocol::Choice { annotations, .. } => Some(annotations),
696            Protocol::Begin { .. }
697            | Protocol::Await { .. }
698            | Protocol::Resolve { .. }
699            | Protocol::Invalidate { .. }
700            | Protocol::Let { .. }
701            | Protocol::Case { .. }
702            | Protocol::Timeout { .. }
703            | Protocol::Publish { .. }
704            | Protocol::PublishAuthority { .. }
705            | Protocol::Materialize { .. }
706            | Protocol::Handoff { .. }
707            | Protocol::DependentWork { .. } => None,
708            Protocol::Extension { annotations, .. } => Some(annotations),
709            Protocol::Loop { .. }
710            | Protocol::Parallel { .. }
711            | Protocol::Rec { .. }
712            | Protocol::Var(_)
713            | Protocol::End => None,
714        }
715    }
716
717    /// Get mutable reference to from-role annotations
718    pub fn get_from_annotations_mut(&mut self) -> Option<&mut Annotations> {
719        match self {
720            Protocol::Send {
721                from_annotations, ..
722            } => Some(from_annotations),
723            Protocol::Broadcast {
724                from_annotations, ..
725            } => Some(from_annotations),
726            _ => None,
727        }
728    }
729
730    /// Get mutable reference to to-role annotations
731    pub fn get_to_annotations_mut(&mut self) -> Option<&mut Annotations> {
732        match self {
733            Protocol::Send { to_annotations, .. } => Some(to_annotations),
734            _ => None,
735        }
736    }
737
738    /// Add a typed annotation
739    pub fn add_annotation(&mut self, annotation: super::annotation::ProtocolAnnotation) -> bool {
740        if let Some(annotations) = self.get_annotations_mut() {
741            annotations.push(annotation);
742            true
743        } else {
744            false
745        }
746    }
747
748    /// Clear all annotations on this protocol node
749    pub fn clear_annotations(&mut self) {
750        if let Some(annotations) = self.get_annotations_mut() {
751            annotations.clear();
752        }
753        if let Some(from_annotations) = self.get_from_annotations_mut() {
754            from_annotations.clear();
755        }
756        if let Some(to_annotations) = self.get_to_annotations_mut() {
757            to_annotations.clear();
758        }
759    }
760
761    /// Check if any annotations are present
762    pub fn has_any_annotations(&self) -> bool {
763        !self.get_annotations().is_empty()
764            || self.get_from_annotations().is_some_and(|a| !a.is_empty())
765            || self.get_to_annotations().is_some_and(|a| !a.is_empty())
766    }
767
768    /// Count total number of annotations (statement + role annotations)
769    pub fn annotation_count(&self) -> usize {
770        self.get_annotations().len()
771            + self.get_from_annotations().map_or(0, |a| a.len())
772            + self.get_to_annotations().map_or(0, |a| a.len())
773    }
774
775    /// Merge annotations from another protocol node
776    pub fn merge_annotations_from(&mut self, other: &Protocol) {
777        // Merge statement annotations
778        if let Some(self_annotations) = self.get_annotations_mut() {
779            self_annotations.merge(other.get_annotations());
780        }
781
782        // Merge from-role annotations
783        if let (Some(self_from), Some(other_from)) = (
784            self.get_from_annotations_mut(),
785            other.get_from_annotations(),
786        ) {
787            self_from.merge(other_from);
788        }
789
790        // Merge to-role annotations
791        if let (Some(self_to), Some(other_to)) =
792            (self.get_to_annotations_mut(), other.get_to_annotations())
793        {
794            self_to.merge(other_to);
795        }
796    }
797
798    /// Validate that required annotations are present
799    pub fn validate_required_annotations(&self, required_keys: &[&str]) -> Result<(), Vec<String>> {
800        let missing: Vec<String> = required_keys
801            .iter()
802            .filter(|&key| {
803                !self
804                    .get_annotations()
805                    .iter()
806                    .any(|annotation| annotation_has_key(annotation, key))
807            })
808            .map(|&key| key.to_string())
809            .collect();
810
811        if missing.is_empty() {
812            Ok(())
813        } else {
814            Err(missing)
815        }
816    }
817
818    /// Collect all protocol nodes that have a specific annotation (recursive traversal)
819    pub fn collect_nodes_with_annotation<'a>(&'a self, key: &str, nodes: &mut Vec<&'a Protocol>) {
820        if self
821            .get_annotations()
822            .iter()
823            .any(|annotation| annotation_has_key(annotation, key))
824        {
825            nodes.push(self);
826        }
827
828        // Recursively traverse protocol structure
829        match self {
830            Protocol::Begin { continuation, .. }
831            | Protocol::Await { continuation, .. }
832            | Protocol::Resolve { continuation, .. }
833            | Protocol::Invalidate { continuation, .. } => {
834                continuation.collect_nodes_with_annotation(key, nodes);
835            }
836            Protocol::Send { continuation, .. } => {
837                continuation.collect_nodes_with_annotation(key, nodes);
838            }
839            Protocol::Broadcast { continuation, .. } => {
840                continuation.collect_nodes_with_annotation(key, nodes);
841            }
842            Protocol::Let { continuation, .. } => {
843                continuation.collect_nodes_with_annotation(key, nodes);
844            }
845            Protocol::Choice { branches, .. } => {
846                for branch in branches {
847                    branch.protocol.collect_nodes_with_annotation(key, nodes);
848                }
849            }
850            Protocol::Case { branches, .. } => {
851                for branch in branches {
852                    branch.protocol.collect_nodes_with_annotation(key, nodes);
853                }
854            }
855            Protocol::Timeout {
856                body,
857                on_timeout,
858                on_cancel,
859                ..
860            } => {
861                body.collect_nodes_with_annotation(key, nodes);
862                on_timeout.collect_nodes_with_annotation(key, nodes);
863                if let Some(on_cancel) = on_cancel.as_deref() {
864                    on_cancel.collect_nodes_with_annotation(key, nodes);
865                }
866            }
867            Protocol::Loop { body, .. } => {
868                body.collect_nodes_with_annotation(key, nodes);
869            }
870            Protocol::Parallel { protocols } => {
871                for protocol in protocols {
872                    protocol.collect_nodes_with_annotation(key, nodes);
873                }
874            }
875            Protocol::Rec { body, .. } => {
876                body.collect_nodes_with_annotation(key, nodes);
877            }
878            Protocol::Publish { continuation, .. }
879            | Protocol::PublishAuthority { continuation, .. }
880            | Protocol::Materialize { continuation, .. }
881            | Protocol::Handoff { continuation, .. }
882            | Protocol::DependentWork { continuation, .. } => {
883                continuation.collect_nodes_with_annotation(key, nodes);
884            }
885            Protocol::Extension { continuation, .. } => {
886                continuation.collect_nodes_with_annotation(key, nodes);
887            }
888            Protocol::Var(_) | Protocol::End => {
889                // Terminal nodes - no further traversal needed
890            }
891        }
892    }
893
894    /// Collect all protocol nodes that have a specific annotation with a specific value
895    pub fn collect_nodes_with_annotation_value<'a>(
896        &'a self,
897        key: &str,
898        value: &str,
899        nodes: &mut Vec<&'a Protocol>,
900    ) {
901        if self
902            .get_annotations()
903            .iter()
904            .any(|annotation| annotation_has_value(annotation, key, value))
905        {
906            nodes.push(self);
907        }
908
909        // Recursively traverse protocol structure
910        match self {
911            Protocol::Begin { continuation, .. }
912            | Protocol::Await { continuation, .. }
913            | Protocol::Resolve { continuation, .. }
914            | Protocol::Invalidate { continuation, .. } => {
915                continuation.collect_nodes_with_annotation_value(key, value, nodes);
916            }
917            Protocol::Send { continuation, .. } => {
918                continuation.collect_nodes_with_annotation_value(key, value, nodes);
919            }
920            Protocol::Broadcast { continuation, .. } => {
921                continuation.collect_nodes_with_annotation_value(key, value, nodes);
922            }
923            Protocol::Let { continuation, .. } => {
924                continuation.collect_nodes_with_annotation_value(key, value, nodes);
925            }
926            Protocol::Choice { branches, .. } => {
927                for branch in branches {
928                    branch
929                        .protocol
930                        .collect_nodes_with_annotation_value(key, value, nodes);
931                }
932            }
933            Protocol::Case { branches, .. } => {
934                for branch in branches {
935                    branch
936                        .protocol
937                        .collect_nodes_with_annotation_value(key, value, nodes);
938                }
939            }
940            Protocol::Timeout {
941                body,
942                on_timeout,
943                on_cancel,
944                ..
945            } => {
946                body.collect_nodes_with_annotation_value(key, value, nodes);
947                on_timeout.collect_nodes_with_annotation_value(key, value, nodes);
948                if let Some(on_cancel) = on_cancel.as_deref() {
949                    on_cancel.collect_nodes_with_annotation_value(key, value, nodes);
950                }
951            }
952            Protocol::Loop { body, .. } => {
953                body.collect_nodes_with_annotation_value(key, value, nodes);
954            }
955            Protocol::Parallel { protocols } => {
956                for protocol in protocols {
957                    protocol.collect_nodes_with_annotation_value(key, value, nodes);
958                }
959            }
960            Protocol::Rec { body, .. } => {
961                body.collect_nodes_with_annotation_value(key, value, nodes);
962            }
963            Protocol::Publish { continuation, .. }
964            | Protocol::PublishAuthority { continuation, .. }
965            | Protocol::Materialize { continuation, .. }
966            | Protocol::Handoff { continuation, .. }
967            | Protocol::DependentWork { continuation, .. } => {
968                continuation.collect_nodes_with_annotation_value(key, value, nodes);
969            }
970            Protocol::Extension { continuation, .. } => {
971                continuation.collect_nodes_with_annotation_value(key, value, nodes);
972            }
973            Protocol::Var(_) | Protocol::End => {
974                // Terminal nodes - no further traversal needed
975            }
976        }
977    }
978
979    /// Count total annotations throughout the protocol tree (recursive)
980    pub fn deep_annotation_count(&self) -> usize {
981        let mut count = self.annotation_count();
982
983        match self {
984            Protocol::Begin { continuation, .. }
985            | Protocol::Await { continuation, .. }
986            | Protocol::Resolve { continuation, .. }
987            | Protocol::Invalidate { continuation, .. } => {
988                count += continuation.deep_annotation_count();
989            }
990            Protocol::Send { continuation, .. } => {
991                count += continuation.deep_annotation_count();
992            }
993            Protocol::Broadcast { continuation, .. } => {
994                count += continuation.deep_annotation_count();
995            }
996            Protocol::Let { continuation, .. } => {
997                count += continuation.deep_annotation_count();
998            }
999            Protocol::Choice { branches, .. } => {
1000                for branch in branches {
1001                    count += branch.protocol.deep_annotation_count();
1002                }
1003            }
1004            Protocol::Case { branches, .. } => {
1005                for branch in branches {
1006                    count += branch.protocol.deep_annotation_count();
1007                }
1008            }
1009            Protocol::Timeout {
1010                body,
1011                on_timeout,
1012                on_cancel,
1013                ..
1014            } => {
1015                count += body.deep_annotation_count();
1016                count += on_timeout.deep_annotation_count();
1017                if let Some(on_cancel) = on_cancel.as_deref() {
1018                    count += on_cancel.deep_annotation_count();
1019                }
1020            }
1021            Protocol::Loop { body, .. } => {
1022                count += body.deep_annotation_count();
1023            }
1024            Protocol::Parallel { protocols } => {
1025                for protocol in protocols {
1026                    count += protocol.deep_annotation_count();
1027                }
1028            }
1029            Protocol::Rec { body, .. } => {
1030                count += body.deep_annotation_count();
1031            }
1032            Protocol::Publish { continuation, .. }
1033            | Protocol::PublishAuthority { continuation, .. }
1034            | Protocol::Materialize { continuation, .. }
1035            | Protocol::Handoff { continuation, .. }
1036            | Protocol::DependentWork { continuation, .. } => {
1037                count += continuation.deep_annotation_count();
1038            }
1039            Protocol::Extension { continuation, .. } => {
1040                count += continuation.deep_annotation_count();
1041            }
1042            Protocol::Var(_) | Protocol::End => {
1043                // Terminal nodes - no additional annotations
1044            }
1045        }
1046
1047        count
1048    }
1049
1050    /// Apply a function to all protocol nodes that have annotations (visitor pattern)
1051    pub fn visit_annotated_nodes<F>(&self, f: &mut F)
1052    where
1053        F: FnMut(&Protocol),
1054    {
1055        if self.has_any_annotations() {
1056            f(self);
1057        }
1058
1059        match self {
1060            Protocol::Begin { continuation, .. }
1061            | Protocol::Await { continuation, .. }
1062            | Protocol::Resolve { continuation, .. }
1063            | Protocol::Invalidate { continuation, .. } => {
1064                continuation.visit_annotated_nodes(f);
1065            }
1066            Protocol::Send { continuation, .. } => {
1067                continuation.visit_annotated_nodes(f);
1068            }
1069            Protocol::Broadcast { continuation, .. } => {
1070                continuation.visit_annotated_nodes(f);
1071            }
1072            Protocol::Let { continuation, .. } => {
1073                continuation.visit_annotated_nodes(f);
1074            }
1075            Protocol::Choice { branches, .. } => {
1076                for branch in branches {
1077                    branch.protocol.visit_annotated_nodes(f);
1078                }
1079            }
1080            Protocol::Case { branches, .. } => {
1081                for branch in branches {
1082                    branch.protocol.visit_annotated_nodes(f);
1083                }
1084            }
1085            Protocol::Timeout {
1086                body,
1087                on_timeout,
1088                on_cancel,
1089                ..
1090            } => {
1091                body.visit_annotated_nodes(f);
1092                on_timeout.visit_annotated_nodes(f);
1093                if let Some(on_cancel) = on_cancel.as_deref() {
1094                    on_cancel.visit_annotated_nodes(f);
1095                }
1096            }
1097            Protocol::Loop { body, .. } => {
1098                body.visit_annotated_nodes(f);
1099            }
1100            Protocol::Parallel { protocols } => {
1101                for protocol in protocols {
1102                    protocol.visit_annotated_nodes(f);
1103                }
1104            }
1105            Protocol::Rec { body, .. } => {
1106                body.visit_annotated_nodes(f);
1107            }
1108            Protocol::Publish { continuation, .. }
1109            | Protocol::PublishAuthority { continuation, .. }
1110            | Protocol::Materialize { continuation, .. }
1111            | Protocol::Handoff { continuation, .. }
1112            | Protocol::DependentWork { continuation, .. } => {
1113                continuation.visit_annotated_nodes(f);
1114            }
1115            Protocol::Extension { continuation, .. } => {
1116                continuation.visit_annotated_nodes(f);
1117            }
1118            Protocol::Var(_) | Protocol::End => {
1119                // Terminal nodes
1120            }
1121        }
1122    }
1123
1124    /// Apply a mutable function to all protocol nodes that have annotations
1125    pub fn visit_annotated_nodes_mut<F>(&mut self, f: &mut F)
1126    where
1127        F: FnMut(&mut Protocol),
1128    {
1129        if self.has_any_annotations() {
1130            f(self);
1131        }
1132
1133        match self {
1134            Protocol::Begin { continuation, .. }
1135            | Protocol::Await { continuation, .. }
1136            | Protocol::Resolve { continuation, .. }
1137            | Protocol::Invalidate { continuation, .. } => {
1138                continuation.visit_annotated_nodes_mut(f);
1139            }
1140            Protocol::Send { continuation, .. } => {
1141                continuation.visit_annotated_nodes_mut(f);
1142            }
1143            Protocol::Broadcast { continuation, .. } => {
1144                continuation.visit_annotated_nodes_mut(f);
1145            }
1146            Protocol::Let { continuation, .. } => {
1147                continuation.visit_annotated_nodes_mut(f);
1148            }
1149            Protocol::Choice { branches, .. } => {
1150                for branch in branches {
1151                    branch.protocol.visit_annotated_nodes_mut(f);
1152                }
1153            }
1154            Protocol::Case { branches, .. } => {
1155                for branch in branches {
1156                    branch.protocol.visit_annotated_nodes_mut(f);
1157                }
1158            }
1159            Protocol::Timeout {
1160                body,
1161                on_timeout,
1162                on_cancel,
1163                ..
1164            } => {
1165                body.visit_annotated_nodes_mut(f);
1166                on_timeout.visit_annotated_nodes_mut(f);
1167                if let Some(on_cancel) = on_cancel.as_deref_mut() {
1168                    on_cancel.visit_annotated_nodes_mut(f);
1169                }
1170            }
1171            Protocol::Loop { body, .. } => {
1172                body.visit_annotated_nodes_mut(f);
1173            }
1174            Protocol::Parallel { protocols } => {
1175                for protocol in protocols {
1176                    protocol.visit_annotated_nodes_mut(f);
1177                }
1178            }
1179            Protocol::Rec { body, .. } => {
1180                body.visit_annotated_nodes_mut(f);
1181            }
1182            Protocol::Publish { continuation, .. }
1183            | Protocol::PublishAuthority { continuation, .. }
1184            | Protocol::Materialize { continuation, .. }
1185            | Protocol::Handoff { continuation, .. }
1186            | Protocol::DependentWork { continuation, .. } => {
1187                continuation.visit_annotated_nodes_mut(f);
1188            }
1189            Protocol::Extension { continuation, .. } => {
1190                continuation.visit_annotated_nodes_mut(f);
1191            }
1192            Protocol::Var(_) | Protocol::End => {
1193                // Terminal nodes
1194            }
1195        }
1196    }
1197}