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!(f, "{}@[{}, {}]", a.pretty(self.arena), interval.start, interval.end)
185            }
186        }
187    }
188}
189
190impl<'a> fmt::Display for Pretty<'a, TransformStmt<'a>> {
191    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
192        if let Some(var) = self.inner.var {
193            write!(f, "let {var} = ")?;
194        }
195        write!(f, "{}", self.inner.app.pretty(self.arena))
196    }
197}
198
199impl<'a> fmt::Display for Pretty<'a, Clause<'a>> {
200    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
201        write!(f, "{}", self.inner.head.pretty(self.arena))?;
202        if let Some(interval) = &self.inner.head_time {
203            write!(f, "@[{}, {}]", interval.start, interval.end)?;
204        }
205        if !self.inner.premises.is_empty() || !self.inner.transform.is_empty() {
206            write!(f, " :- ")?;
207            let mut first = true;
208            for premise in self.inner.premises {
209                if !first {
210                    write!(f, ", ")?;
211                }
212                write!(f, "{}", premise.pretty(self.arena))?;
213                first = false;
214            }
215            for transform in self.inner.transform {
216                if !first {
217                    write!(f, ", ")?;
218                }
219                write!(f, "{}", transform.pretty(self.arena))?;
220                first = false;
221            }
222        }
223        write!(f, ".")
224    }
225}
226
227impl<'a> fmt::Display for Pretty<'a, Constraints<'a>> {
228    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
229        if !self.inner.consequences.is_empty() {
230            write!(f, " |> ")?;
231            for (i, c) in self.inner.consequences.iter().enumerate() {
232                if i > 0 {
233                    write!(f, ", ")?;
234                }
235                write!(f, "{}", c.pretty(self.arena))?;
236            }
237        }
238        if !self.inner.alternatives.is_empty() {
239            for alt in self.inner.alternatives.iter() {
240                write!(f, " | ")?;
241                for (i, c) in alt.iter().enumerate() {
242                    if i > 0 {
243                        write!(f, ", ")?;
244                    }
245                    write!(f, "{}", c.pretty(self.arena))?;
246                }
247            }
248        }
249        Ok(())
250    }
251}
252
253impl<'a> fmt::Display for Pretty<'a, BoundDecl<'a>> {
254    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
255        for (i, b) in self.inner.base_terms.iter().enumerate() {
256            if i > 0 {
257                write!(f, ", ")?;
258            }
259            write!(f, "{}", b.pretty(self.arena))?;
260        }
261        Ok(())
262    }
263}
264
265impl<'a> fmt::Display for Pretty<'a, Decl<'a>> {
266    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
267        write!(f, "{}", self.inner.atom.pretty(self.arena))?;
268        if self.inner.is_temporal {
269            write!(f, " temporal")?;
270        }
271        if !self.inner.descr.is_empty() {
272            write!(f, " [")?;
273            for (i, d) in self.inner.descr.iter().enumerate() {
274                if i > 0 {
275                    write!(f, ", ")?;
276                }
277                write!(f, "{}", d.pretty(self.arena))?;
278            }
279            write!(f, "]")?;
280        }
281        if let Some(bounds) = self.inner.bounds
282            && !bounds.is_empty()
283        {
284            write!(f, " bound ")?;
285            for (i, b) in bounds.iter().enumerate() {
286                if i > 0 {
287                    write!(f, " | ")?;
288                }
289                write!(f, "{}", b.pretty(self.arena))?;
290            }
291        }
292        if let Some(constraints) = &self.inner.constraints {
293            write!(f, "{}", constraints.pretty(self.arena))?;
294        }
295        write!(f, ".")
296    }
297}
298
299impl<'a> fmt::Display for Pretty<'a, Unit<'a>> {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        for decl in self.inner.decls {
302            writeln!(f, "{}", decl.pretty(self.arena))?;
303        }
304        for clause in self.inner.clauses {
305            writeln!(f, "{}", clause.pretty(self.arena))?;
306        }
307        Ok(())
308    }
309}