telltale_runtime/compiler/

pretty.rs

//! Pretty-printer for the choreography DSL.
//!
//! Emits layout-sensitive syntax for the choreography language.

use crate::ast::{
    Annotations, AuthorityExpr, Branch, CaseBranch, ChoiceGuard, Choreography, Condition,
    MessageType, Protocol, Role, RoleParam,
};
use crate::compiler::parser::parse_choreography_str;

#[derive(Debug, Clone)]
pub struct PrettyConfig {
    pub indent: usize,
}

impl Default for PrettyConfig {
    fn default() -> Self {
        Self { indent: 2 }
    }
}

pub fn format_choreography(choreo: &Choreography) -> String {
    format_choreography_with_config(choreo, &PrettyConfig::default())
}

pub fn format_choreography_with_config(choreo: &Choreography, config: &PrettyConfig) -> String {
    let mut out = String::new();

    if let Some(namespace) = &choreo.namespace {
        out.push_str(&format!(
            "module {} exposing ({})\n\n",
            namespace, choreo.name
        ));
    }

    out.push_str(&format!("protocol {} =\n", choreo.name));
    write_line(
        &mut out,
        config.indent,
        &format!("roles {}", format_role_list(&choreo.roles)),
    );
    format_protocol(&choreo.protocol, config.indent, config, &mut out);

    out
}

pub fn format_choreography_str(input: &str) -> Result<String, crate::compiler::parser::ParseError> {
    let choreo = parse_choreography_str(input)?;
    Ok(format_choreography(&choreo))
}

// RECURSION_SAFE: structural recursion over finite protocol AST depth.
fn format_protocol(protocol: &Protocol, indent: usize, config: &PrettyConfig, out: &mut String) {
    match protocol {
        Protocol::End => {}
        Protocol::Begin {
            operation,
            args,
            progress,
            continuation,
        } => {
            let mut line = if args.is_empty() {
                format!("begin {}", operation)
            } else {
                format!("begin {}({})", operation, args.join(", "))
            };
            if let Some(progress) = progress {
                line.push_str(&format!(" {}", format_progress_attachment(progress)));
            }
            write_line(out, indent, &line);
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Await {
            operation,
            continuation,
        } => {
            write_line(out, indent, &format!("await {}", operation));
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Resolve {
            operation,
            outcome,
            continuation,
        } => {
            write_line(
                out,
                indent,
                &format!(
                    "resolve {} as {}",
                    operation,
                    format_commitment_outcome(outcome)
                ),
            );
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Invalidate {
            operation,
            continuation,
        } => {
            write_line(out, indent, &format!("invalidate {}", operation));
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Send {
            from,
            to,
            message,
            annotations,
            from_annotations,
            continuation,
            ..
        } => format_send_protocol(
            from,
            annotations,
            from_annotations,
            to,
            message,
            continuation,
            indent,
            config,
            out,
        ),
        Protocol::Broadcast {
            from,
            message,
            annotations,
            from_annotations,
            continuation,
            ..
        } => format_broadcast_protocol(
            from,
            annotations,
            from_annotations,
            message,
            continuation,
            indent,
            config,
            out,
        ),
        Protocol::Choice { role, branches, .. } => {
            format_choice_protocol(role, branches, indent, config, out)
        }
        Protocol::Let {
            name,
            expr,
            continuation,
            ..
        } => {
            write_line(
                out,
                indent,
                &format!("let {} = {}", name, format_authority_expr(expr)),
            );
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Case { expr, branches } => {
            write_line(
                out,
                indent,
                &format!("case {} of", format_authority_expr(expr)),
            );
            for branch in branches {
                format_case_branch(branch, indent + config.indent, config, out);
            }
        }
        Protocol::Timeout {
            role,
            duration_ms,
            body,
            on_timeout,
            on_cancel,
        } => {
            write_line(
                out,
                indent,
                &format!("timeout {}ms {} at", duration_ms, format_role_ref(role)),
            );
            format_protocol(body, indent + config.indent, config, out);
            write_line(out, indent, "on timeout =>");
            format_protocol(on_timeout, indent + config.indent, config, out);
            if let Some(on_cancel) = on_cancel.as_deref() {
                write_line(out, indent, "on cancel =>");
                format_protocol(on_cancel, indent + config.indent, config, out);
            }
        }
        Protocol::Loop { condition, body } => {
            format_loop_protocol(condition, body, indent, config, out)
        }
        Protocol::Parallel { protocols } => {
            format_parallel_protocol(protocols, indent, config, out)
        }
        Protocol::Rec { label, body } => format_rec_protocol(label, body, indent, config, out),
        Protocol::Var(label) => {
            write_line(out, indent, &format!("continue {}", label));
        }
        Protocol::Publish {
            event,
            arg,
            continuation,
        } => {
            if let Some(arg) = arg {
                write_line(out, indent, &format!("publish {}{}", event, arg));
            } else {
                write_line(out, indent, &format!("publish {}", event));
            }
            format_protocol(continuation, indent, config, out);
        }
        Protocol::PublishAuthority {
            witness,
            publication_name,
            continuation,
        } => {
            write_line(
                out,
                indent,
                &format!("publish {} as {}", witness, publication_name),
            );
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Materialize {
            proof,
            publication,
            continuation,
        } => {
            write_line(
                out,
                indent,
                &format!("materialize {} from {}", proof, publication),
            );
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Handoff {
            operation,
            target,
            receipt,
            continuation,
        } => {
            write_line(
                out,
                indent,
                &format!(
                    "handoff {} to {} using {}",
                    operation,
                    format_role_ref(target),
                    receipt
                ),
            );
            format_protocol(continuation, indent, config, out);
        }
        Protocol::DependentWork {
            name,
            arg,
            required_for,
            continuation,
        } => {
            let work_head = match arg {
                Some(arg) => format!(
                    "dependent work {}{} required for {}",
                    name, arg, required_for
                ),
                None => format!("dependent work {} required for {}", name, required_for),
            };
            write_line(out, indent, &work_head);
            format_protocol(continuation, indent, config, out);
        }
        Protocol::Extension {
            extension,
            continuation,
            ..
        } => format_extension_protocol(extension.type_name(), continuation, indent, config, out),
    }
}

fn format_progress_attachment(progress: &crate::ast::ProgressAttachment) -> String {
    let mut parts = vec![format!("progress {}", progress.contract_name)];
    if let Some(profile) = &progress.requires_profile {
        parts.push(format!("requires {}", profile));
    }
    if let Some(window) = &progress.within_window {
        parts.push(format!("within {}", window));
    }
    if let Some(timeout) = &progress.on_timeout {
        parts.push(format!("on timeout => {}", timeout));
    }
    if let Some(stall) = &progress.on_stall {
        parts.push(format!("on stall => {}", stall));
    }
    parts.join(" ")
}

fn format_commitment_outcome(outcome: &crate::ast::CommitmentOutcome) -> String {
    match outcome {
        crate::ast::CommitmentOutcome::Success(payload) => payload.as_ref().map_or_else(
            || "Success".to_string(),
            |payload| format!("Success({payload})"),
        ),
        crate::ast::CommitmentOutcome::Failure(payload) => payload.as_ref().map_or_else(
            || "Failure".to_string(),
            |payload| format!("Failure({payload})"),
        ),
        crate::ast::CommitmentOutcome::Timeout(payload) => payload.as_ref().map_or_else(
            || "Timeout".to_string(),
            |payload| format!("Timeout({payload})"),
        ),
        crate::ast::CommitmentOutcome::Cancelled => "Cancelled".to_string(),
    }
}

fn format_send_protocol(
    from: &Role,
    annotations: &Annotations,
    from_annotations: &Annotations,
    to: &Role,
    message: &MessageType,
    continuation: &Protocol,
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    let sender_annotations = merge_sender_annotations(from_annotations, annotations);
    write_line(out, indent, &format_sender_term(from, &sender_annotations));
    write_line(
        out,
        indent + config.indent,
        &format!("-> {} : {}", format_role_ref(to), format_message(message)),
    );
    format_protocol(continuation, indent, config, out);
}

fn format_broadcast_protocol(
    from: &Role,
    annotations: &Annotations,
    from_annotations: &Annotations,
    message: &MessageType,
    continuation: &Protocol,
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    let sender_annotations = merge_sender_annotations(from_annotations, annotations);
    write_line(out, indent, &format_sender_term(from, &sender_annotations));
    write_line(
        out,
        indent + config.indent,
        &format!("->* : {}", format_message(message)),
    );
    format_protocol(continuation, indent, config, out);
}

fn format_choice_protocol(
    role: &Role,
    branches: &[Branch],
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    write_line(out, indent, &format!("choice {} at", format_role_ref(role)));
    for branch in branches {
        format_branch(branch, indent + config.indent, config, out);
    }
}

fn format_loop_protocol(
    condition: &Option<Condition>,
    body: &Protocol,
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    let header = format_loop_header(condition);
    if is_end(body) {
        write_line(out, indent, &format!("{} {{}}", header));
    } else {
        write_line(out, indent, &header);
        format_protocol(body, indent + config.indent, config, out);
    }
}

fn format_parallel_protocol(
    protocols: &[Protocol],
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    write_line(out, indent, "par");
    for branch in protocols {
        if is_end(branch) {
            write_line(out, indent + config.indent, "| {}");
        } else {
            write_line(out, indent + config.indent, "|");
            format_protocol(branch, indent + (2 * config.indent), config, out);
        }
    }
}

fn format_rec_protocol(
    label: &proc_macro2::Ident,
    body: &Protocol,
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    if is_end(body) {
        write_line(out, indent, &format!("rec {} {{}}", label));
    } else {
        write_line(out, indent, &format!("rec {}", label));
        format_protocol(body, indent + config.indent, config, out);
    }
}

fn format_extension_protocol(
    extension_type_name: &str,
    continuation: &Protocol,
    indent: usize,
    config: &PrettyConfig,
    out: &mut String,
) {
    write_line(out, indent, &format!("// extension: {extension_type_name}"));
    format_protocol(continuation, indent, config, out);
}

fn format_branch(branch: &Branch, indent: usize, config: &PrettyConfig, out: &mut String) {
    let mut header = format!("| {}", branch.label);
    if let Some(guard) = &branch.guard {
        header.push_str(&format!(" {}", format_choice_guard(guard)));
    }

    if is_end(&branch.protocol) {
        write_line(out, indent, &format!("{} => {{}}", header));
    } else {
        write_line(out, indent, &format!("{} =>", header));
        format_protocol(&branch.protocol, indent + config.indent, config, out);
    }
}

fn format_case_branch(branch: &CaseBranch, indent: usize, config: &PrettyConfig, out: &mut String) {
    let binders = if branch.pattern.binders.is_empty() {
        String::new()
    } else {
        format!("({})", branch.pattern.binders.join(", "))
    };
    let header = format!("| {}{} =>", branch.pattern.constructor, binders);
    if is_end(&branch.protocol) {
        write_line(out, indent, &format!("{header} {{}}"));
    } else {
        write_line(out, indent, &header);
        format_protocol(&branch.protocol, indent + config.indent, config, out);
    }
}

fn format_choice_guard(guard: &ChoiceGuard) -> String {
    match guard {
        ChoiceGuard::Predicate(tokens) => format!("when ({})", tokens),
        ChoiceGuard::Evidence {
            effect,
            operation,
            args,
            binding,
        } => {
            let args = if args.is_empty() {
                String::new()
            } else {
                args.join(", ")
            };
            format!(
                "when check {}.{}({}) yields {}",
                effect, operation, args, binding
            )
        }
    }
}

fn format_authority_expr(expr: &AuthorityExpr) -> String {
    match expr {
        AuthorityExpr::Var(name) => name.clone(),
        AuthorityExpr::Check {
            effect,
            operation,
            args,
        } => format!("check {}.{}({})", effect, operation, args.join(", ")),
        AuthorityExpr::Observe {
            effect,
            operation,
            args,
        } => format!("observe {}.{}({})", effect, operation, args.join(", ")),
        AuthorityExpr::Transfer { subject, from, to } => {
            format!("transfer {} from {} to {}", subject, from, to)
        }
        AuthorityExpr::Constructor { name, arg } => match arg {
            Some(arg) => format!("{name}({arg})"),
            None => name.clone(),
        },
        AuthorityExpr::Call { name, args } => format!("{name}({})", args.join(", ")),
    }
}

fn format_loop_header(condition: &Option<Condition>) -> String {
    match condition {
        Some(Condition::RoleDecides(role)) => {
            format!("loop decide by {}", format_role_ref(role))
        }
        Some(Condition::Count(count)) => format!("loop repeat {}", count),
        Some(Condition::Custom(tokens)) => format!("loop while {}", tokens),
        Some(Condition::Fuel(count)) => format!("loop while \"fuel:{}\"", count),
        Some(Condition::YieldAfter(count)) => format!("loop while \"yield_after:{}\"", count),
        Some(Condition::YieldWhen(label)) => format!("loop while \"yield_when:{}\"", label),
        None => "loop forever".to_string(),
    }
}

fn format_role_list(roles: &[Role]) -> String {
    roles
        .iter()
        .map(|role| {
            let mut out = role.name().to_string();
            if let Some(param) = role.param() {
                out.push('[');
                out.push_str(&format_role_param(param));
                out.push(']');
            }
            out
        })
        .collect::<Vec<_>>()
        .join(", ")
}

fn format_role_param(param: &RoleParam) -> String {
    param.to_string()
}

fn format_role_ref(role: &Role) -> String {
    let mut out = role.name().to_string();
    if let Some(index) = role.index() {
        out.push('[');
        out.push_str(&index.to_string());
        out.push(']');
    }
    out
}

fn format_sender_term(role: &Role, annotations: &Annotations) -> String {
    let mut out = format_role_ref(role);
    let entries = annotations.dsl_entries();
    if !entries.is_empty() {
        let formatted = entries
            .into_iter()
            .map(|entry| format!("{} : {}", entry.key, entry.value))
            .collect::<Vec<_>>()
            .join(", ");
        out.push_str(" { ");
        out.push_str(&formatted);
        out.push_str(" }");
    }
    out
}

fn merge_sender_annotations(
    sender_annotations: &Annotations,
    statement_annotations: &Annotations,
) -> Annotations {
    let mut merged = sender_annotations.dsl_entries();
    let mut existing_keys = merged
        .iter()
        .map(|entry| entry.key.clone())
        .collect::<std::collections::BTreeSet<_>>();
    for entry in statement_annotations.dsl_entries() {
        if existing_keys.insert(entry.key.clone()) {
            merged.push(entry);
        }
    }
    Annotations::from_dsl_entries(&merged)
}

fn normalize_surface_type_string(s: &str) -> String {
    s.replace(" :: ", ".").replace("::", ".")
}

fn format_message(message: &MessageType) -> String {
    let mut out = message.name.to_string();
    if let Some(payload) = &message.payload {
        let payload_str = payload.to_string();
        if payload_str.starts_with('(') {
            out.push(' ');
            out.push_str(&payload_str);
        } else {
            out.push_str(" of ");
            out.push_str(&normalize_surface_type_string(&payload_str));
        }
    } else if let Some(type_annotation) = &message.type_annotation {
        out.push_str(" of ");
        out.push_str(&normalize_surface_type_string(&type_annotation.to_string()));
    }
    out
}

fn is_end(protocol: &Protocol) -> bool {
    matches!(protocol, Protocol::End)
}

fn write_line(out: &mut String, indent: usize, text: &str) {
    out.push_str(&" ".repeat(indent));
    out.push_str(text);
    out.push('\n');
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pretty_roundtrip_basic() {
        let input = "protocol PingPong =\n  roles Alice, Bob\n  Alice -> Bob : Ping\n  Bob -> Alice : Pong\n";
        let choreo = parse_choreography_str(input).expect("should parse");
        let formatted = format_choreography(&choreo);
        assert!(formatted.contains("Alice\n    -> Bob : Ping"));
        assert!(parse_choreography_str(&formatted).is_ok());
    }

    #[test]
    fn pretty_roundtrip_choice_and_loop() {
        let input = r#"
protocol Demo =
  roles Client, Server
  choice Client at
    | Buy =>
      Client -> Server : Purchase
    | Cancel =>
      Client -> Server : Cancel
  loop repeat 2
    Client -> Server : Ping
    Server -> Client : Pong
"#;
        let choreo = parse_choreography_str(input).expect("should parse");
        let formatted = format_choreography(&choreo);
        assert!(formatted.contains("choice Client at"));
        assert!(formatted.contains("| Buy =>"));
        assert!(parse_choreography_str(&formatted).is_ok());
    }

    #[test]
    fn pretty_emits_aligned_arrows_and_sender_records() {
        let input = r#"
protocol Styled =
  roles A, B, C, D
  A { priority : high } -> B : Request of shop.Order
  par
    | C -> D : Left
    | D -> C : Right
"#;

        let choreo = parse_choreography_str(input).expect("should parse");
        let formatted = format_choreography(&choreo);
        assert!(formatted.contains("A { priority : high }\n    -> B : Request of shop.Order"));
        assert!(formatted.contains("par\n    |\n      C\n        -> D : Left"));
        assert!(parse_choreography_str(&formatted).is_ok());
    }

    #[test]
    fn pretty_is_stable_on_reformat() {
        let input = r#"
protocol Stable =
  roles A, B
  A { priority : high } -> B : Request of shop.Order
"#;

        let first = format_choreography_str(input).expect("first format should succeed");
        let second = format_choreography_str(&first).expect("second format should succeed");
        assert_eq!(first, second);
    }
}