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