telltale_types/

local.rs

//! Local Session Types for Multiparty Protocols
//!
//! This module defines local types that describe protocols from a single participant's
//! perspective. Local types are obtained by projecting global types onto specific roles.
//!
//! Based on: "A Very Gentle Introduction to Multiparty Session Types" (Yoshida & Gheri)
//!
//! # Lean Correspondence
//!
//! The core `LocalTypeR` enum mirrors `lean/SessionTypes/LocalTypeR/Core/Base.lean`:
//! - `LocalTypeR::End` ↔ Lean's `LocalTypeR.end`
//! - `LocalTypeR::Send` ↔ Lean's `LocalTypeR.send`
//! - `LocalTypeR::Recv` ↔ Lean's `LocalTypeR.recv`
//! - `LocalTypeR::Mu` ↔ Lean's `LocalTypeR.mu`
//! - `LocalTypeR::Var` ↔ Lean's `LocalTypeR.var`

use crate::{Label, ValType};
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;

/// Core local type matching Lean's `LocalTypeR`.
///
/// This is the minimal type used for validation and correspondence proofs.
/// For code generation, see the extended `LocalType` in the DSL crate.
///
/// # Examples
///
/// ```
/// use telltale_types::{LocalTypeR, Label};
///
/// // Simple send: !B{msg.end}
/// let lt = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
/// assert!(lt.well_formed());
///
/// // Recursive type: μt. !B{msg.t}
/// let rec = LocalTypeR::mu(
///     "t",
///     LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("t")),
/// );
/// assert!(rec.well_formed());
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum LocalTypeR {
    /// Protocol termination
    End,
    /// Internal choice: send to partner with choice of continuations
    Send {
        partner: String,
        branches: Vec<(Label, Option<ValType>, LocalTypeR)>,
    },
    /// External choice: receive from partner with offered continuations
    Recv {
        partner: String,
        branches: Vec<(Label, Option<ValType>, LocalTypeR)>,
    },
    /// Recursive type: μt.T binds variable t in body T
    Mu { var: String, body: Box<LocalTypeR> },
    /// Type variable: reference to enclosing μ-binder
    Var(String),
}

impl LocalTypeR {
    fn collect_all_var_names(&self, names: &mut BTreeSet<String>) {
        match self {
            LocalTypeR::End => {}
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                for (_, _, cont) in branches {
                    cont.collect_all_var_names(names);
                }
            }
            LocalTypeR::Mu { var, body } => {
                names.insert(var.clone());
                body.collect_all_var_names(names);
            }
            LocalTypeR::Var(t) => {
                names.insert(t.clone());
            }
        }
    }

    fn all_var_names(&self) -> BTreeSet<String> {
        let mut names = BTreeSet::new();
        self.collect_all_var_names(&mut names);
        names
    }

    fn fresh_var(base: &str, avoid: &BTreeSet<String>) -> String {
        let mut idx = 0usize;
        loop {
            // bounded: idx increments until fresh name found (finite avoid set)
            let candidate = format!("{base}_{idx}");
            if !avoid.contains(&candidate) {
                return candidate;
            }
            idx += 1;
        }
    }

    /// Create a simple send with one label
    #[must_use]
    pub fn send(partner: impl Into<String>, label: Label, cont: LocalTypeR) -> Self {
        LocalTypeR::Send {
            partner: partner.into(),
            branches: vec![(label, None, cont)],
        }
    }

    /// Create a send with multiple branches
    #[must_use]
    pub fn send_choice(
        partner: impl Into<String>,
        branches: Vec<(Label, Option<ValType>, LocalTypeR)>,
    ) -> Self {
        LocalTypeR::Send {
            partner: partner.into(),
            branches,
        }
    }

    /// Create a simple recv with one label
    #[must_use]
    pub fn recv(partner: impl Into<String>, label: Label, cont: LocalTypeR) -> Self {
        LocalTypeR::Recv {
            partner: partner.into(),
            branches: vec![(label, None, cont)],
        }
    }

    /// Create a recv with multiple branches
    #[must_use]
    pub fn recv_choice(
        partner: impl Into<String>,
        branches: Vec<(Label, Option<ValType>, LocalTypeR)>,
    ) -> Self {
        LocalTypeR::Recv {
            partner: partner.into(),
            branches,
        }
    }

    /// Create a recursive type
    #[must_use]
    pub fn mu(var: impl Into<String>, body: LocalTypeR) -> Self {
        LocalTypeR::Mu {
            var: var.into(),
            body: Box::new(body),
        }
    }

    /// Create a type variable
    #[must_use]
    pub fn var(name: impl Into<String>) -> Self {
        LocalTypeR::Var(name.into())
    }

    /// Extract free type variables from a local type.
    ///
    /// Corresponds to Lean's `LocalTypeR.freeVars`.
    #[must_use]
    pub fn free_vars(&self) -> Vec<String> {
        let mut result = Vec::new();
        let mut bound = BTreeSet::new();
        self.collect_free_vars(&mut result, &mut bound);
        result
    }

    fn collect_free_vars(&self, free: &mut Vec<String>, bound: &mut BTreeSet<String>) {
        match self {
            LocalTypeR::End => {}
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                for (_, _, cont) in branches {
                    cont.collect_free_vars(free, bound);
                }
            }
            LocalTypeR::Mu { var, body } => {
                bound.insert(var.clone());
                body.collect_free_vars(free, bound);
                bound.remove(var);
            }
            LocalTypeR::Var(t) => {
                if !bound.contains(t) && !free.contains(t) {
                    free.push(t.clone());
                }
            }
        }
    }

    /// Substitute a local type for a variable.
    ///
    /// Corresponds to Lean's `LocalTypeR.substitute`.
    #[must_use]
    pub fn substitute(&self, var_name: &str, replacement: &LocalTypeR) -> LocalTypeR {
        match self {
            LocalTypeR::End => LocalTypeR::End,
            LocalTypeR::Send { partner, branches } => {
                Self::substitute_branching(partner, branches, var_name, replacement, true)
            }
            LocalTypeR::Recv { partner, branches } => {
                Self::substitute_branching(partner, branches, var_name, replacement, false)
            }
            LocalTypeR::Mu { var, body } => Self::substitute_mu(var, body, var_name, replacement),
            LocalTypeR::Var(t) => {
                if t == var_name {
                    replacement.clone()
                } else {
                    LocalTypeR::Var(t.clone())
                }
            }
        }
    }

    fn substitute_branching(
        partner: &str,
        branches: &[(Label, Option<ValType>, LocalTypeR)],
        var_name: &str,
        replacement: &LocalTypeR,
        is_send: bool,
    ) -> LocalTypeR {
        let substituted = branches
            .iter()
            .map(|(l, vt, cont)| {
                (
                    l.clone(),
                    vt.clone(),
                    cont.substitute(var_name, replacement),
                )
            })
            .collect();
        if is_send {
            LocalTypeR::Send {
                partner: partner.to_string(),
                branches: substituted,
            }
        } else {
            LocalTypeR::Recv {
                partner: partner.to_string(),
                branches: substituted,
            }
        }
    }

    fn substitute_mu(
        var: &str,
        body: &LocalTypeR,
        var_name: &str,
        replacement: &LocalTypeR,
    ) -> LocalTypeR {
        if var == var_name {
            return LocalTypeR::Mu {
                var: var.to_string(),
                body: Box::new(body.clone()),
            };
        }
        let replacement_free = replacement.free_vars();
        if replacement_free.iter().any(|v| v == var) {
            // Alpha-rename binder to avoid capture before descending.
            let mut avoid = body.all_var_names();
            avoid.extend(replacement.all_var_names());
            avoid.insert(var_name.to_string());
            let fresh = Self::fresh_var(var, &avoid);
            let renamed_body = body.substitute(var, &LocalTypeR::Var(fresh.clone()));
            return LocalTypeR::Mu {
                var: fresh,
                body: Box::new(renamed_body.substitute(var_name, replacement)),
            };
        }
        LocalTypeR::Mu {
            var: var.to_string(),
            body: Box::new(body.substitute(var_name, replacement)),
        }
    }

    /// Unfold one level of recursion: μt.T ↦ T[μt.T/t]
    ///
    /// Corresponds to Lean's `LocalTypeR.unfold`.
    #[must_use]
    pub fn unfold(&self) -> LocalTypeR {
        match self {
            LocalTypeR::Mu { var, body } => body.substitute(var, self),
            _ => self.clone(),
        }
    }

    /// Count nested recursion binders at the current head position.
    ///
    /// Corresponds to Lean's `LocalTypeR.muHeight`.
    #[must_use]
    pub fn mu_height(&self) -> usize {
        match self {
            LocalTypeR::Mu { body, .. } => 1 + body.mu_height(),
            _ => 0,
        }
    }

    /// Fully unfold leading recursion until a non-`Mu` head is exposed.
    ///
    /// Corresponds to Lean's `LocalTypeR.fullUnfold`.
    #[must_use]
    pub fn full_unfold(&self) -> LocalTypeR {
        let mut current = self.clone();
        let fuel = self.mu_height();
        for _ in 0..fuel {
            current = current.unfold();
        }
        current
    }

    /// Compute the dual of a local type (swap send/recv).
    ///
    /// The dual of role A's view is role B's view when A and B are the only participants.
    /// Corresponds to Lean's `LocalTypeR.dual`.
    #[must_use]
    pub fn dual(&self) -> LocalTypeR {
        match self {
            LocalTypeR::End => LocalTypeR::End,
            LocalTypeR::Send { partner, branches } => LocalTypeR::Recv {
                partner: partner.clone(),
                branches: branches
                    .iter()
                    .map(|(l, vt, cont)| (l.clone(), vt.clone(), cont.dual()))
                    .collect(),
            },
            LocalTypeR::Recv { partner, branches } => LocalTypeR::Send {
                partner: partner.clone(),
                branches: branches
                    .iter()
                    .map(|(l, vt, cont)| (l.clone(), vt.clone(), cont.dual()))
                    .collect(),
            },
            LocalTypeR::Mu { var, body } => LocalTypeR::Mu {
                var: var.clone(),
                body: Box::new(body.dual()),
            },
            LocalTypeR::Var(t) => LocalTypeR::Var(t.clone()),
        }
    }

    /// Check if all recursion variables are bound.
    ///
    /// Corresponds to Lean's `LocalTypeR.allVarsBound`.
    #[must_use]
    pub fn all_vars_bound(&self) -> bool {
        self.check_vars_bound(&BTreeSet::new())
    }

    fn check_vars_bound(&self, bound: &BTreeSet<String>) -> bool {
        match self {
            LocalTypeR::End => true,
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => branches
                .iter()
                .all(|(_, _, cont)| cont.check_vars_bound(bound)),
            LocalTypeR::Mu { var, body } => {
                let mut new_bound = bound.clone();
                new_bound.insert(var.clone());
                body.check_vars_bound(&new_bound)
            }
            LocalTypeR::Var(t) => bound.contains(t),
        }
    }

    /// Check if each choice has at least one branch.
    ///
    /// Corresponds to Lean's `LocalTypeR.allChoicesNonEmpty`.
    #[must_use]
    pub fn all_choices_non_empty(&self) -> bool {
        match self {
            LocalTypeR::End => true,
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                !branches.is_empty()
                    && branches
                        .iter()
                        .all(|(_, _, cont)| cont.all_choices_non_empty())
            }
            LocalTypeR::Mu { body, .. } => body.all_choices_non_empty(),
            LocalTypeR::Var(_) => true,
        }
    }

    fn branches_have_unique_names(branches: &[(Label, Option<ValType>, LocalTypeR)]) -> bool {
        let mut seen = BTreeSet::new();
        branches
            .iter()
            .all(|(label, _, _)| seen.insert(label.name.clone()))
    }

    /// Check if all send/recv nodes have unique branch label names.
    #[must_use]
    pub fn unique_branch_labels(&self) -> bool {
        match self {
            LocalTypeR::End | LocalTypeR::Var(_) => true,
            LocalTypeR::Mu { body, .. } => body.unique_branch_labels(),
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                Self::branches_have_unique_names(branches)
                    && branches
                        .iter()
                        .all(|(_, _, cont)| cont.unique_branch_labels())
            }
        }
    }

    /// Well-formedness predicate for local types.
    ///
    /// Corresponds to Lean's `LocalTypeR.wellFormed`.
    /// A local type is well-formed if:
    /// 1. All recursion variables are bound
    /// 2. Each choice has at least one branch
    /// 3. All recursion is guarded (no immediate recursion without communication)
    #[must_use]
    pub fn well_formed(&self) -> bool {
        self.all_vars_bound()
            && self.all_choices_non_empty()
            && self.unique_branch_labels()
            && self.is_guarded()
    }

    /// Check whether full unfolding exposes a communication head.
    ///
    /// Corresponds to Lean's `LocalTypeR.reachesCommunication`.
    #[must_use]
    pub fn reaches_communication(&self) -> bool {
        match self.full_unfold() {
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                !branches.is_empty()
            }
            LocalTypeR::End | LocalTypeR::Var(_) | LocalTypeR::Mu { .. } => false,
        }
    }

    /// Mechanically characterized practical fragment that needs no manual `ReachesComm`.
    ///
    /// Corresponds to Lean's `LocalTypeR.regularPracticalFragment`.
    #[must_use]
    pub fn regular_practical_fragment(&self) -> bool {
        self.well_formed() && self.reaches_communication()
    }

    /// Count the depth of a local type (for termination proofs).
    ///
    /// Corresponds to Lean's `LocalTypeR.depth`.
    #[must_use]
    pub fn depth(&self) -> usize {
        match self {
            LocalTypeR::End => 0,
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                1 + branches
                    .iter()
                    .map(|(_, _, t)| t.depth())
                    .max()
                    .unwrap_or(0)
            }
            LocalTypeR::Mu { body, .. } => 1 + body.depth(),
            LocalTypeR::Var(_) => 0,
        }
    }

    /// Check if a local type is guarded (no immediate recursion).
    ///
    /// Corresponds to Lean's `LocalTypeR.isGuarded`.
    #[must_use]
    pub fn is_guarded(&self) -> bool {
        match self {
            LocalTypeR::End => true,
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                branches.iter().all(|(_, _, cont)| cont.is_guarded())
            }
            LocalTypeR::Mu { body, .. } => match body.as_ref() {
                LocalTypeR::Var(_) | LocalTypeR::Mu { .. } => false,
                _ => body.is_guarded(),
            },
            LocalTypeR::Var(_) => true,
        }
    }

    /// Check if a specific recursion variable is guarded in this local type.
    ///
    /// A variable `v` is guarded if it does not appear "bare" (i.e., every
    /// occurrence is under a send or recv). Corresponds to Lean's
    /// `LocalTypeR.isGuarded v`.
    #[must_use]
    pub fn is_var_guarded(&self, var: &str) -> bool {
        match self {
            LocalTypeR::End => true,
            LocalTypeR::Var(w) => w != var,
            LocalTypeR::Send { .. } | LocalTypeR::Recv { .. } => true,
            LocalTypeR::Mu { var: t, body, .. } => {
                if var == t {
                    true
                } else {
                    body.is_var_guarded(var)
                }
            }
        }
    }

    /// Extract labels from the outermost visible choice node.
    ///
    /// For recursive terms (`Mu`), this inspects the wrapped body.
    #[must_use]
    pub fn head_labels(&self) -> Vec<String> {
        match self {
            LocalTypeR::End | LocalTypeR::Var(_) => vec![],
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                branches.iter().map(|(l, _, _)| l.name.clone()).collect()
            }
            LocalTypeR::Mu { body, .. } => body.head_labels(),
        }
    }

    /// Backward-compatible alias for `head_labels`.
    #[must_use]
    pub fn labels(&self) -> Vec<String> {
        self.head_labels()
    }

    /// Extract all labels that appear anywhere in the local type.
    #[must_use]
    pub fn all_labels(&self) -> Vec<String> {
        let mut labels = BTreeSet::new();
        self.collect_all_labels(&mut labels);
        labels.into_iter().collect()
    }

    fn collect_all_labels(&self, labels: &mut BTreeSet<String>) {
        match self {
            LocalTypeR::End | LocalTypeR::Var(_) => {}
            LocalTypeR::Mu { body, .. } => body.collect_all_labels(labels),
            LocalTypeR::Send { branches, .. } | LocalTypeR::Recv { branches, .. } => {
                for (label, _, cont) in branches {
                    labels.insert(label.name.clone());
                    cont.collect_all_labels(labels);
                }
            }
        }
    }

    /// Extract all partners from a local type.
    ///
    /// Corresponds to Lean's `LocalTypeR.partners`.
    #[must_use]
    pub fn partners(&self) -> Vec<String> {
        let mut result = BTreeSet::new();
        self.collect_partners(&mut result);
        result.into_iter().collect()
    }

    fn collect_partners(&self, partners: &mut BTreeSet<String>) {
        match self {
            LocalTypeR::End | LocalTypeR::Var(_) => {}
            LocalTypeR::Send { partner, branches } | LocalTypeR::Recv { partner, branches } => {
                partners.insert(partner.clone());
                for (_, _, cont) in branches {
                    cont.collect_partners(partners);
                }
            }
            LocalTypeR::Mu { body, .. } => body.collect_partners(partners),
        }
    }

    /// Check if a local type mentions a specific partner.
    #[must_use]
    pub fn mentions_partner(&self, role: &str) -> bool {
        self.mentions_partner_inner(role)
    }

    fn mentions_partner_inner(&self, role: &str) -> bool {
        match self {
            LocalTypeR::End | LocalTypeR::Var(_) => false,
            LocalTypeR::Mu { body, .. } => body.mentions_partner_inner(role),
            LocalTypeR::Send { partner, branches } | LocalTypeR::Recv { partner, branches } => {
                partner == role
                    || branches
                        .iter()
                        .any(|(_, _, cont)| cont.mentions_partner_inner(role))
            }
        }
    }

    /// Check if this is an internal choice (send)
    #[must_use]
    pub fn is_send(&self) -> bool {
        matches!(self, LocalTypeR::Send { .. })
    }

    /// Check if this is an external choice (recv)
    #[must_use]
    pub fn is_recv(&self) -> bool {
        matches!(self, LocalTypeR::Recv { .. })
    }

    /// Check if this is a terminated type
    #[must_use]
    pub fn is_end(&self) -> bool {
        matches!(self, LocalTypeR::End)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::PayloadSort;
    use assert_matches::assert_matches;
    use proptest::prelude::*;

    #[test]
    fn test_simple_local_type() {
        // !B{msg.end}
        let lt = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
        assert!(lt.well_formed());
        assert_eq!(lt.partners().len(), 1);
        assert!(lt.mentions_partner("B"));
    }

    #[test]
    fn test_recursive_local_type() {
        // μt. !B{msg.t}
        let lt = LocalTypeR::mu(
            "t",
            LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("t")),
        );
        assert!(lt.well_formed());
        assert!(lt.all_vars_bound());
    }

    #[test]
    fn test_dual() {
        // !B{msg.end} dual is ?B{msg.end}
        let send = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
        let recv = send.dual();

        assert_matches!(recv, LocalTypeR::Recv { partner, branches } => {
            assert_eq!(partner, "B");
            assert_eq!(branches.len(), 1);
            assert_eq!((branches[0].0).name, "msg");
        });
    }

    #[test]
    fn test_unfold() {
        // μt. !B{msg.t} unfolds to !B{msg.(μt. !B{msg.t})}
        let lt = LocalTypeR::mu(
            "t",
            LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("t")),
        );
        let unfolded = lt.unfold();

        assert_matches!(unfolded, LocalTypeR::Send { partner, branches } => {
            assert_eq!(partner, "B");
            assert_matches!(branches[0].2, LocalTypeR::Mu { .. });
        });
    }

    #[test]
    fn test_substitute() {
        let lt = LocalTypeR::var("t");
        let replacement = LocalTypeR::End;
        assert_eq!(lt.substitute("t", &replacement), LocalTypeR::End);
        assert_eq!(lt.substitute("s", &replacement), LocalTypeR::var("t"));
    }

    #[test]
    fn test_substitute_avoids_capture() {
        let lt = LocalTypeR::mu("y", LocalTypeR::var("x"));
        let result = lt.substitute("x", &LocalTypeR::var("y"));

        assert_matches!(result, LocalTypeR::Mu { var, body } => {
            assert_ne!(var, "y");
            assert_matches!(*body, LocalTypeR::Var(ref name) => {
                assert_eq!(name, "y");
            });
        });
    }

    #[test]
    fn test_substitute_avoids_capture_with_nested_binders() {
        let lt = LocalTypeR::mu("y", LocalTypeR::mu("y_0", LocalTypeR::var("x")));
        let result = lt.substitute("x", &LocalTypeR::var("y"));
        assert!(result.free_vars().contains(&"y".to_string()));
    }

    #[test]
    fn test_unbound_variable() {
        // !B{msg.t} where t is unbound
        let lt = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("t"));
        assert!(!lt.all_vars_bound());
        assert!(!lt.well_formed());
    }

    #[test]
    fn test_guarded() {
        // μt. t is not guarded
        let unguarded = LocalTypeR::mu("t", LocalTypeR::var("t"));
        assert!(!unguarded.is_guarded());
        assert!(!unguarded.well_formed()); // Unguarded recursion should fail well_formed()

        // μt. !B{msg.t} is guarded
        let guarded = LocalTypeR::mu(
            "t",
            LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("t")),
        );
        assert!(guarded.is_guarded());
        assert!(guarded.well_formed()); // Guarded recursion should pass well_formed()
    }

    #[test]
    fn test_mu_height_and_full_unfold() {
        let lt = LocalTypeR::mu(
            "outer",
            LocalTypeR::mu(
                "inner",
                LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("inner")),
            ),
        );
        assert_eq!(lt.mu_height(), 2);
        assert_matches!(lt.full_unfold(), LocalTypeR::Send { partner, branches } => {
            assert_eq!(partner, "B");
            assert_eq!(branches.len(), 1);
        });
    }

    #[test]
    fn test_reaches_communication_matches_practical_fragment_intent() {
        assert!(!LocalTypeR::End.reaches_communication());
        assert!(!LocalTypeR::End.regular_practical_fragment());

        let send = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
        assert!(send.reaches_communication());
        assert!(send.regular_practical_fragment());

        let guarded = LocalTypeR::mu(
            "t",
            LocalTypeR::recv("B", Label::new("msg"), LocalTypeR::var("t")),
        );
        assert!(guarded.reaches_communication());
        assert!(guarded.regular_practical_fragment());

        let unguarded = LocalTypeR::mu("t", LocalTypeR::var("t"));
        assert!(!unguarded.reaches_communication());
        assert!(!unguarded.regular_practical_fragment());
    }

    #[test]
    fn test_free_vars() {
        // μt. !B{msg.s} has free var s
        let lt = LocalTypeR::mu(
            "t",
            LocalTypeR::send("B", Label::new("msg"), LocalTypeR::var("s")),
        );
        let free = lt.free_vars();
        assert_eq!(free, vec!["s"]);
    }

    #[test]
    fn test_choice_with_payload() {
        let branches = vec![
            (
                Label::with_sort("accept", PayloadSort::Bool),
                None,
                LocalTypeR::End,
            ),
            (
                Label::with_sort("data", PayloadSort::Nat),
                None,
                LocalTypeR::End,
            ),
        ];
        let lt = LocalTypeR::recv_choice("A", branches);
        assert!(lt.well_formed());
        assert_eq!(lt.head_labels(), vec!["accept", "data"]);
        assert_eq!(lt.labels(), vec!["accept", "data"]);
    }

    #[test]
    fn test_depth() {
        let lt = LocalTypeR::send(
            "B",
            Label::new("outer"),
            LocalTypeR::send("C", Label::new("inner"), LocalTypeR::End),
        );
        assert_eq!(lt.depth(), 2);
    }

    #[test]
    fn test_is_send_recv() {
        let send = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
        let recv = LocalTypeR::recv("B", Label::new("msg"), LocalTypeR::End);

        assert!(send.is_send());
        assert!(!send.is_recv());
        assert!(recv.is_recv());
        assert!(!recv.is_send());
    }

    #[test]
    fn test_duplicate_branch_labels_not_well_formed() {
        let lt = LocalTypeR::recv_choice(
            "A",
            vec![
                (Label::new("msg"), None, LocalTypeR::End),
                (Label::new("msg"), None, LocalTypeR::End),
            ],
        );
        assert!(!lt.unique_branch_labels());
        assert!(!lt.well_formed());
    }

    #[test]
    fn test_partners_are_sorted() {
        let lt = LocalTypeR::send(
            "Zed",
            Label::new("outer"),
            LocalTypeR::recv(
                "Alice",
                Label::new("mid"),
                LocalTypeR::send("Bob", Label::new("inner"), LocalTypeR::End),
            ),
        );

        assert_eq!(lt.partners(), vec!["Alice", "Bob", "Zed"]);
    }

    #[test]
    fn test_mentions_partner_direct_traversal_matches_expectation() {
        let lt = LocalTypeR::mu(
            "t",
            LocalTypeR::send(
                "A",
                Label::new("x"),
                LocalTypeR::recv("B", Label::new("y"), LocalTypeR::var("t")),
            ),
        );
        assert!(lt.mentions_partner("A"));
        assert!(lt.mentions_partner("B"));
        assert!(!lt.mentions_partner("C"));
    }

    #[test]
    fn test_all_labels_collects_nested_labels() {
        let lt = LocalTypeR::send(
            "B",
            Label::new("outer"),
            LocalTypeR::recv(
                "A",
                Label::new("inner"),
                LocalTypeR::send("C", Label::new("leaf"), LocalTypeR::End),
            ),
        );
        assert_eq!(lt.head_labels(), vec!["outer"]);
        assert_eq!(lt.all_labels(), vec!["inner", "leaf", "outer"]);
    }

    fn arb_var_name() -> impl Strategy<Value = String> {
        prop_oneof![
            Just("x".to_string()),
            Just("y".to_string()),
            Just("z".to_string()),
            Just("t".to_string()),
            Just("u".to_string()),
        ]
    }

    proptest! {
        #[test]
        fn prop_substitute_identity_when_var_absent(var in arb_var_name()) {
            let lt = LocalTypeR::send("B", Label::new("msg"), LocalTypeR::End);
            let replacement = LocalTypeR::var("r");
            prop_assert_eq!(lt.substitute(&var, &replacement), lt);
        }

        #[test]
        fn prop_substitute_avoids_capture_simple(
            binder in arb_var_name(),
            target in arb_var_name(),
        ) {
            prop_assume!(binder != target);
            let lt = LocalTypeR::mu(&binder, LocalTypeR::var(&target));
            let replacement = LocalTypeR::var(&binder);
            let result = lt.substitute(&target, &replacement);
            prop_assert!(result.free_vars().contains(&binder));
        }
    }
}