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))
}
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);
}
}