Skip to main content

mangle_ast/
pretty.rs

1// Copyright 2025 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::{
16    Arena, Atom, BaseTerm, BoundDecl, Clause, Const, Constraints, Decl, FunctionIndex,
17    PredicateIndex, Term, TransformStmt, Unit, VariableIndex,
18};
19use std::fmt;
20
21/// Provides pretty-printing, including name lookup from arena.
22/// Usage:
23/// ```
24/// # fn print(arena: &mangle_ast::Arena, clause: mangle_ast::Clause) -> String {
25/// use mangle_ast::PrettyPrint;
26/// clause.pretty(&arena).to_string()
27/// # }
28/// ```
29pub struct Pretty<'a, T: ?Sized> {
30    arena: &'a Arena,
31    inner: &'a T,
32}
33
34pub trait PrettyPrint {
35    fn pretty<'a>(&'a self, arena: &'a Arena) -> Pretty<'a, Self>
36    where
37        Self: Sized,
38    {
39        Pretty { arena, inner: self }
40    }
41}
42
43impl PrettyPrint for VariableIndex {}
44impl PrettyPrint for PredicateIndex {}
45impl PrettyPrint for FunctionIndex {}
46impl<'a> PrettyPrint for Const<'a> {}
47impl<'a> PrettyPrint for BaseTerm<'a> {}
48impl<'a> PrettyPrint for Atom<'a> {}
49impl<'a> PrettyPrint for Term<'a> {}
50impl<'a> PrettyPrint for TransformStmt<'a> {}
51impl<'a> PrettyPrint for Clause<'a> {}
52impl<'a> PrettyPrint for Constraints<'a> {}
53impl<'a> PrettyPrint for BoundDecl<'a> {}
54impl<'a> PrettyPrint for Decl<'a> {}
55impl<'a> PrettyPrint for Unit<'a> {}
56
57impl<'a> fmt::Display for Pretty<'a, VariableIndex> {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        if self.inner.0 == 0 {
60            write!(f, "_")
61        } else {
62            match self.arena.lookup_name(self.inner.0) {
63                Some(name) => write!(f, "{name}"),
64                None => write!(f, "v${}", self.inner.0),
65            }
66        }
67    }
68}
69
70impl<'a> fmt::Display for Pretty<'a, PredicateIndex> {
71    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72        match self.arena.predicate_name(*self.inner) {
73            Some(name) => write!(f, "{name}"),
74            None => write!(f, "p${}", self.inner.0),
75        }
76    }
77}
78
79impl<'a> fmt::Display for Pretty<'a, FunctionIndex> {
80    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
81        match self.arena.function_name(*self.inner) {
82            Some(name) => write!(f, "{name}"),
83            None => write!(f, "f${}", self.inner.0),
84        }
85    }
86}
87
88impl<'a> fmt::Display for Pretty<'a, Const<'a>> {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        match self.inner {
91            Const::Name(n) => match self.arena.lookup_name(*n) {
92                Some(name) => write!(f, "{name}"),
93                None => write!(f, "n${n}"),
94            },
95            Const::Bool(b) => write!(f, "{b}"),
96            Const::Number(n) => write!(f, "{n}"),
97            Const::Float(fl) => write!(f, "{fl}"),
98            Const::String(s) => write!(f, "{s:?}"), // Use Debug for quoting
99            Const::Bytes(b) => write!(f, "{b:?}"),
100            Const::Time(t) => write!(f, "t#{t}"),
101            Const::Duration(d) => write!(f, "d#{d}"),
102            Const::List(l) => {
103                write!(f, "[")?;
104                for (i, c) in l.iter().enumerate() {
105                    if i > 0 {
106                        write!(f, ", ")?;
107                    }
108                    write!(f, "{}", c.pretty(self.arena))?;
109                }
110                write!(f, "]")
111            }
112            Const::Map { keys, values } => {
113                if keys.is_empty() {
114                    write!(f, "fn:map()")
115                } else {
116                    write!(f, "[")?;
117                    for (i, (k, v)) in keys.iter().zip(values.iter()).enumerate() {
118                        if i > 0 {
119                            write!(f, ", ")?;
120                        }
121                        write!(f, "{}: {}", k.pretty(self.arena), v.pretty(self.arena))?;
122                    }
123                    write!(f, "]")
124                }
125            }
126            Const::Struct { fields, values } => {
127                write!(f, "{{")?;
128                for (i, (field, val)) in fields.iter().zip(values.iter()).enumerate() {
129                    if i > 0 {
130                        write!(f, ", ")?;
131                    }
132                    write!(f, "{field}: {}", val.pretty(self.arena))?;
133                }
134                write!(f, "}}")
135            }
136        }
137    }
138}
139
140impl<'a> fmt::Display for Pretty<'a, BaseTerm<'a>> {
141    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142        match self.inner {
143            BaseTerm::Const(c) => write!(f, "{}", c.pretty(self.arena)),
144            BaseTerm::Variable(v) => write!(f, "{}", v.pretty(self.arena)),
145            BaseTerm::ApplyFn(fun, args) => {
146                write!(f, "{}(", fun.pretty(self.arena))?;
147                for (i, arg) in args.iter().enumerate() {
148                    if i > 0 {
149                        write!(f, ", ")?;
150                    }
151                    write!(f, "{}", arg.pretty(self.arena))?;
152                }
153                write!(f, ")")
154            }
155        }
156    }
157}
158
159impl<'a> fmt::Display for Pretty<'a, Atom<'a>> {
160    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161        write!(f, "{}(", self.inner.sym.pretty(self.arena))?;
162        for (i, arg) in self.inner.args.iter().enumerate() {
163            if i > 0 {
164                write!(f, ", ")?;
165            }
166            write!(f, "{}", arg.pretty(self.arena))?;
167        }
168        write!(f, ")")
169    }
170}
171
172impl<'a> fmt::Display for Pretty<'a, Term<'a>> {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        match self.inner {
175            Term::Atom(a) => write!(f, "{}", a.pretty(self.arena)),
176            Term::NegAtom(a) => write!(f, "!{}", a.pretty(self.arena)),
177            Term::Eq(l, r) => {
178                write!(f, "{} = {}", l.pretty(self.arena), r.pretty(self.arena))
179            }
180            Term::Ineq(l, r) => {
181                write!(f, "{} != {}", l.pretty(self.arena), r.pretty(self.arena))
182            }
183            Term::TemporalAtom(a, interval) => {
184                write!(
185                    f,
186                    "{}@[{}, {}]",
187                    a.pretty(self.arena),
188                    interval.start,
189                    interval.end
190                )
191            }
192        }
193    }
194}
195
196impl<'a> fmt::Display for Pretty<'a, TransformStmt<'a>> {
197    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
198        if let Some(var) = self.inner.var {
199            write!(f, "let {var} = ")?;
200        }
201        write!(f, "{}", self.inner.app.pretty(self.arena))
202    }
203}
204
205impl<'a> fmt::Display for Pretty<'a, Clause<'a>> {
206    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
207        write!(f, "{}", self.inner.head.pretty(self.arena))?;
208        if let Some(interval) = &self.inner.head_time {
209            write!(f, "@[{}, {}]", interval.start, interval.end)?;
210        }
211        if !self.inner.premises.is_empty() || !self.inner.transform.is_empty() {
212            write!(f, " :- ")?;
213            let mut first = true;
214            for premise in self.inner.premises {
215                if !first {
216                    write!(f, ", ")?;
217                }
218                write!(f, "{}", premise.pretty(self.arena))?;
219                first = false;
220            }
221            for transform in self.inner.transform {
222                if !first {
223                    write!(f, ", ")?;
224                }
225                write!(f, "{}", transform.pretty(self.arena))?;
226                first = false;
227            }
228        }
229        write!(f, ".")
230    }
231}
232
233impl<'a> fmt::Display for Pretty<'a, Constraints<'a>> {
234    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235        if !self.inner.consequences.is_empty() {
236            write!(f, " |> ")?;
237            for (i, c) in self.inner.consequences.iter().enumerate() {
238                if i > 0 {
239                    write!(f, ", ")?;
240                }
241                write!(f, "{}", c.pretty(self.arena))?;
242            }
243        }
244        if !self.inner.alternatives.is_empty() {
245            for alt in self.inner.alternatives.iter() {
246                write!(f, " | ")?;
247                for (i, c) in alt.iter().enumerate() {
248                    if i > 0 {
249                        write!(f, ", ")?;
250                    }
251                    write!(f, "{}", c.pretty(self.arena))?;
252                }
253            }
254        }
255        Ok(())
256    }
257}
258
259impl<'a> fmt::Display for Pretty<'a, BoundDecl<'a>> {
260    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261        for (i, b) in self.inner.base_terms.iter().enumerate() {
262            if i > 0 {
263                write!(f, ", ")?;
264            }
265            write!(f, "{}", b.pretty(self.arena))?;
266        }
267        Ok(())
268    }
269}
270
271impl<'a> fmt::Display for Pretty<'a, Decl<'a>> {
272    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
273        write!(f, "{}", self.inner.atom.pretty(self.arena))?;
274        if self.inner.is_temporal {
275            write!(f, " temporal")?;
276        }
277        if !self.inner.descr.is_empty() {
278            write!(f, " [")?;
279            for (i, d) in self.inner.descr.iter().enumerate() {
280                if i > 0 {
281                    write!(f, ", ")?;
282                }
283                write!(f, "{}", d.pretty(self.arena))?;
284            }
285            write!(f, "]")?;
286        }
287        if let Some(bounds) = self.inner.bounds
288            && !bounds.is_empty()
289        {
290            write!(f, " bound ")?;
291            for (i, b) in bounds.iter().enumerate() {
292                if i > 0 {
293                    write!(f, " | ")?;
294                }
295                write!(f, "{}", b.pretty(self.arena))?;
296            }
297        }
298        if let Some(constraints) = &self.inner.constraints {
299            write!(f, "{}", constraints.pretty(self.arena))?;
300        }
301        write!(f, ".")
302    }
303}
304
305impl<'a> fmt::Display for Pretty<'a, Unit<'a>> {
306    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
307        for decl in self.inner.decls {
308            writeln!(f, "{}", decl.pretty(self.arena))?;
309        }
310        for clause in self.inner.clauses {
311            writeln!(f, "{}", clause.pretty(self.arena))?;
312        }
313        Ok(())
314    }
315}