Skip to main content

conjure_cp_core/ast/
pretty.rs

1//! Functions for pretty printing Conjure models.
2//!
3//! Most things can be pretty printed using `Display`; however some, notably collections
4//! can not, for example, Vec<Expression>
5
6use std::fmt::Display;
7
8use super::{Atom, Expression, Name, SatEncodingDecision, SymbolTable};
9use crate::ast::domains::HasDomain;
10use itertools::Itertools;
11
12/// Pretty prints a `Vec<Expression>` as if it were a top level constraint list in a `such that`.
13///
14/// Each expression is printed on a new line, and expressions are delimited by commas.
15///
16/// For some input expressions A,B,C:
17/// ```text
18/// A,
19/// B,
20/// C
21/// ```
22///
23/// Each `Expression` is printed using its underlying `Display` implementation.
24pub fn pretty_expressions_as_top_level(expressions: &[Expression]) -> String {
25    expressions.iter().map(|x| format!("{x}")).join(",\n")
26}
27
28/// Pretty-print semantic SAT encoding decisions.
29pub fn pretty_sat_decisions(clauses: &[SatEncodingDecision]) -> String {
30    clauses.iter().map(|clause| format!("{clause}")).join(",\n")
31}
32
33/// Pretty prints a `Vec<Expression>` as if it were a conjunction.
34///
35/// For some input expressions A,B,C:
36///
37/// ```text
38/// (A /\ B /\ C)
39/// ```
40///
41/// Each `Expression` is printed using its underlying `Display` implementation.
42pub fn pretty_expressions_as_conjunction(expressions: &[Expression]) -> String {
43    let mut str = expressions.iter().map(|x| format!("{x}")).join(" /\\ ");
44
45    str.insert(0, '(');
46    str.push(')');
47
48    str
49}
50
51/// Pretty prints a `Vec<T>` in a vector like syntax.
52///
53/// For some input values A,B,C:
54///
55/// ```text
56/// [A,B,C]
57/// ````
58///
59/// Each element is printed using its underlying `Display` implementation.
60pub fn pretty_vec<T: Display>(elems: &[T]) -> String {
61    let mut str = elems.iter().map(|x| format!("{x}")).join(", ");
62    str.insert(0, '[');
63    str.push(']');
64
65    str
66}
67
68/// Pretty prints an expression with an Essence-style type annotation.
69///
70/// `::` is treated as an expression operator for parenthesisation purposes.
71pub fn pretty_expression_type_annotation(expr: &Expression, ty: impl Display) -> String {
72    pretty_expression_annotation(expr, "::", ty)
73}
74
75/// Pretty prints an expression with an Essence-style domain annotation.
76///
77/// `:` is treated as an expression operator for parenthesisation purposes.
78pub fn pretty_expression_domain_annotation(expr: &Expression, domain: impl Display) -> String {
79    pretty_expression_annotation(expr, ":", domain)
80}
81
82fn pretty_expression_annotation(
83    expr: &Expression,
84    operator: &str,
85    annotation: impl Display,
86) -> String {
87    let expr = parenthesise_if_needed(expr, Precedence::ANNOTATION);
88    format!("{expr} {operator} {annotation}")
89}
90
91fn parenthesise_if_needed(expr: &Expression, parent_precedence: Precedence) -> String {
92    let rendered = expr.to_string();
93    if expression_precedence(expr).binds_weaker_than(parent_precedence) {
94        format!("({rendered})")
95    } else {
96        rendered
97    }
98}
99
100/// Expression precedence used for minimal parenthesisation.
101///
102/// These levels follow the relative operator ordering in
103/// `crates/tree-sitter-essence/grammar.js`, with `:` and `::` inserted as annotation operators
104/// that bind tighter than binary arithmetic/comparison expressions and looser than unary/postfix
105/// expressions. The numeric values are local to the pretty-printer.
106#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
107struct Precedence(i8);
108
109impl Precedence {
110    const LOWEST: Self = Self(-100);
111    const IMPLICATION: Self = Self(-4);
112    const OR: Self = Self(-2);
113    const AND: Self = Self(-1);
114    const COMPARISON: Self = Self(5);
115    const ADDITIVE: Self = Self(10);
116    const MULTIPLICATIVE: Self = Self(20);
117    const ANNOTATION: Self = Self(25);
118    const UNARY: Self = Self(30);
119    const POSTFIX: Self = Self(40);
120    const ATOM: Self = Self(100);
121
122    fn binds_weaker_than(self, parent: Self) -> bool {
123        self < parent
124    }
125}
126
127fn expression_precedence(expr: &Expression) -> Precedence {
128    match expr {
129        Expression::Imply(_, _, _) | Expression::Iff(_, _, _) => Precedence::IMPLICATION,
130        Expression::Or(_, _) => Precedence::OR,
131        Expression::And(_, _) => Precedence::AND,
132        Expression::Eq(_, _, _)
133        | Expression::Neq(_, _, _)
134        | Expression::Geq(_, _, _)
135        | Expression::Leq(_, _, _)
136        | Expression::Gt(_, _, _)
137        | Expression::Lt(_, _, _)
138        | Expression::In(_, _, _)
139        | Expression::Supset(_, _, _)
140        | Expression::SupsetEq(_, _, _)
141        | Expression::Subset(_, _, _)
142        | Expression::SubsetEq(_, _, _)
143        | Expression::LexLt(_, _, _)
144        | Expression::LexLeq(_, _, _)
145        | Expression::LexGt(_, _, _)
146        | Expression::LexGeq(_, _, _) => Precedence::COMPARISON,
147        Expression::Sum(_, _) | Expression::Minus(_, _, _) | Expression::PairwiseSum(_, _, _) => {
148            Precedence::ADDITIVE
149        }
150        Expression::Product(_, _)
151        | Expression::UnsafeDiv(_, _, _)
152        | Expression::SafeDiv(_, _, _)
153        | Expression::UnsafeMod(_, _, _)
154        | Expression::SafeMod(_, _, _)
155        | Expression::PairwiseProduct(_, _, _) => Precedence::MULTIPLICATIVE,
156        Expression::Not(_, _)
157        | Expression::Neg(_, _)
158        | Expression::Abs(_, _)
159        | Expression::Card(_, _)
160        | Expression::ToInt(_, _) => Precedence::UNARY,
161        Expression::Factorial(_, _)
162        | Expression::UnsafePow(_, _, _)
163        | Expression::SafePow(_, _, _)
164        | Expression::UnsafeIndex(_, _, _)
165        | Expression::SafeIndex(_, _, _)
166        | Expression::UnsafeSlice(_, _, _)
167        | Expression::SafeSlice(_, _, _) => Precedence::POSTFIX,
168        Expression::Union(_, _, _) | Expression::Intersect(_, _, _) => Precedence::LOWEST,
169        Expression::TypeAnnotation(_, _, _) | Expression::DomainAnnotation(_, _, _) => {
170            Precedence::ANNOTATION
171        }
172        Expression::Atomic(_, Atom::Reference(_))
173        | Expression::Atomic(_, Atom::Literal(_))
174        | Expression::AbstractLiteral(_, _)
175        | Expression::Comprehension(_, _)
176        | Expression::Metavar(_, _)
177        | Expression::FromSolution(_, _) => Precedence::ATOM,
178        _ => Precedence::ATOM,
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use crate::ast::{Atom, DeclarationPtr, Domain, Metadata, Moo, Name, Range};
186
187    fn atom(name: &str) -> Expression {
188        Expression::Atomic(
189            Metadata::new(),
190            Atom::new_ref(DeclarationPtr::new_find(
191                Name::user(name),
192                Domain::int(vec![Range::Bounded(0, 10)]),
193            )),
194        )
195    }
196
197    #[test]
198    fn domain_annotation_parenthesises_comparison_lhs() {
199        let expr = Expression::Eq(
200            Metadata::new(),
201            Moo::new(atom("y")),
202            Moo::new(Expression::Card(Metadata::new(), Moo::new(atom("x")))),
203        );
204
205        assert_eq!(
206            pretty_expression_domain_annotation(&expr, "bool"),
207            "(y = |x|) : bool"
208        );
209    }
210
211    #[test]
212    fn domain_annotation_does_not_parenthesise_unary_or_atom_lhs() {
213        let x = atom("x");
214        let card = Expression::Card(Metadata::new(), Moo::new(x.clone()));
215
216        assert_eq!(pretty_expression_domain_annotation(&x, "int"), "x : int");
217        assert_eq!(
218            pretty_expression_domain_annotation(&card, "int"),
219            "|x| : int"
220        );
221    }
222
223    #[test]
224    fn type_annotation_uses_same_expression_precedence_as_domain_annotation() {
225        let expr = Expression::Eq(Metadata::new(), Moo::new(atom("x")), Moo::new(atom("y")));
226
227        assert_eq!(
228            pretty_expression_type_annotation(&expr, "bool"),
229            "(x = y) :: bool"
230        );
231    }
232}
233
234/// Pretty prints, in essence syntax, the variable declaration for the given symbol.
235///
236/// E.g.
237///
238/// ```text
239/// find a: int(1..5)
240/// ```
241///
242/// When a representation has been selected for an abstract domain, the domain is printed with
243/// that representation's short name, e.g.
244/// `find x: set (representation occurrence, maxSize 3) of int(1..4)`.
245///
246/// Returns None if the symbol is not in the symbol table, or if it is not a variable.
247pub fn pretty_variable_declaration(symbol_table: &SymbolTable, var_name: &Name) -> Option<String> {
248    let decl = symbol_table.lookup(var_name)?;
249    let var = decl.as_find()?;
250    let domain = var.domain_of();
251    let domain_str = format_domain_with_selected_representation(&decl, &domain);
252    let keyword = if decl.is_find_auxiliary() {
253        "findAux"
254    } else {
255        "find"
256    };
257    Some(format!("{keyword} {var_name}: {domain_str}"))
258}
259
260/// Format a domain for display, overlaying a selected representation short name when present.
261fn format_domain_with_selected_representation(
262    decl: &crate::ast::DeclarationPtr,
263    domain: &crate::ast::DomainPtr,
264) -> String {
265    let selected = decl
266        .reprs()
267        .iter()
268        .map(|(_, state)| state.rule().short_name())
269        .next();
270
271    // Only sets and multisets show their selected representation in a declaration; rule traces
272    // and snapshots have always printed every other domain as written.
273    let shows_selection = domain.as_set().is_some() || domain.as_mset().is_some();
274
275    match selected {
276        Some(short_name)
277            if shows_selection && domain.representation_preference() != Some(short_name) =>
278        {
279            format_domain_with_representation(domain, short_name)
280        }
281        _ => domain.to_string(),
282    }
283}
284
285/// Pretty-print a find/findAux declaration with `repr` annotated on its domain.
286///
287/// Example: `find x: set (representation packed, maxSize 3) of int(1..4)`.
288pub fn pretty_find_with_representation(
289    decl: &crate::ast::DeclarationPtr,
290    repr: &str,
291) -> Option<String> {
292    let domain = decl.domain()?;
293    let domain_str = format_domain_with_representation(&domain, repr);
294    let keyword = if decl.is_find_auxiliary() {
295        "findAux"
296    } else {
297        "find"
298    };
299    Some(format!("{keyword} {}: {domain_str}", decl.name()))
300}
301
302/// Print `domain` as usual, but with `repr` as its top-level representation preference.
303pub fn format_domain_with_representation(domain: &crate::ast::DomainPtr, repr: &str) -> String {
304    domain.with_representation_preference(repr).to_string()
305}
306
307/// Pretty prints, in essence syntax, the declaration for the given value letting.
308///
309/// E.g.
310///
311/// ```text
312/// letting A be 1+2+3
313/// ```
314///
315/// Returns None if the symbol is not in the symbol table, or if it is not a value letting.
316pub fn pretty_value_letting_declaration(symbol_table: &SymbolTable, name: &Name) -> Option<String> {
317    let decl = symbol_table.lookup(name)?;
318    let letting = decl.as_value_letting()?;
319    Some(format!("letting {name} be {letting}"))
320}
321
322/// Pretty prints, in essence syntax, the declaration for the given domain letting.
323///
324/// E.g.
325///
326/// ```text
327/// letting A be domain bool
328/// ```
329///
330/// Returns None if the symbol is not in the symbol table, or if it is not a domain letting.
331pub fn pretty_domain_letting_declaration(
332    symbol_table: &SymbolTable,
333    name: &Name,
334) -> Option<String> {
335    let decl = symbol_table.lookup(name)?;
336    let letting = decl.as_domain_letting()?;
337    Some(format!("letting {name} be domain {letting}"))
338}