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::List(l) => {
101                write!(f, "[")?;
102                for (i, c) in l.iter().enumerate() {
103                    if i > 0 {
104                        write!(f, ", ")?;
105                    }
106                    write!(f, "{}", c.pretty(self.arena))?;
107                }
108                write!(f, "]")
109            }
110            Const::Map { keys, values } => {
111                if keys.is_empty() {
112                    write!(f, "fn:map()")
113                } else {
114                    write!(f, "[")?;
115                    for (i, (k, v)) in keys.iter().zip(values.iter()).enumerate() {
116                        if i > 0 {
117                            write!(f, ", ")?;
118                        }
119                        write!(f, "{}: {}", k.pretty(self.arena), v.pretty(self.arena))?;
120                    }
121                    write!(f, "]")
122                }
123            }
124            Const::Struct { fields, values } => {
125                write!(f, "{{")?;
126                for (i, (field, val)) in fields.iter().zip(values.iter()).enumerate() {
127                    if i > 0 {
128                        write!(f, ", ")?;
129                    }
130                    write!(f, "{field}: {}", val.pretty(self.arena))?;
131                }
132                write!(f, "}}")
133            }
134        }
135    }
136}
137
138impl<'a> fmt::Display for Pretty<'a, BaseTerm<'a>> {
139    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140        match self.inner {
141            BaseTerm::Const(c) => write!(f, "{}", c.pretty(self.arena)),
142            BaseTerm::Variable(v) => write!(f, "{}", v.pretty(self.arena)),
143            BaseTerm::ApplyFn(fun, args) => {
144                write!(f, "{}(", fun.pretty(self.arena))?;
145                for (i, arg) in args.iter().enumerate() {
146                    if i > 0 {
147                        write!(f, ", ")?;
148                    }
149                    write!(f, "{}", arg.pretty(self.arena))?;
150                }
151                write!(f, ")")
152            }
153        }
154    }
155}
156
157impl<'a> fmt::Display for Pretty<'a, Atom<'a>> {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        write!(f, "{}(", self.inner.sym.pretty(self.arena))?;
160        for (i, arg) in self.inner.args.iter().enumerate() {
161            if i > 0 {
162                write!(f, ", ")?;
163            }
164            write!(f, "{}", arg.pretty(self.arena))?;
165        }
166        write!(f, ")")
167    }
168}
169
170impl<'a> fmt::Display for Pretty<'a, Term<'a>> {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        match self.inner {
173            Term::Atom(a) => write!(f, "{}", a.pretty(self.arena)),
174            Term::NegAtom(a) => write!(f, "!{}", a.pretty(self.arena)),
175            Term::Eq(l, r) => {
176                write!(f, "{} = {}", l.pretty(self.arena), r.pretty(self.arena))
177            }
178            Term::Ineq(l, r) => {
179                write!(f, "{} != {}", l.pretty(self.arena), r.pretty(self.arena))
180            }
181        }
182    }
183}
184
185impl<'a> fmt::Display for Pretty<'a, TransformStmt<'a>> {
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        if let Some(var) = self.inner.var {
188            write!(f, "let {var} = ")?;
189        }
190        write!(f, "{}", self.inner.app.pretty(self.arena))
191    }
192}
193
194impl<'a> fmt::Display for Pretty<'a, Clause<'a>> {
195    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196        write!(f, "{}", self.inner.head.pretty(self.arena))?;
197        if !self.inner.premises.is_empty() || !self.inner.transform.is_empty() {
198            write!(f, " :- ")?;
199            let mut first = true;
200            for premise in self.inner.premises {
201                if !first {
202                    write!(f, ", ")?;
203                }
204                write!(f, "{}", premise.pretty(self.arena))?;
205                first = false;
206            }
207            for transform in self.inner.transform {
208                if !first {
209                    write!(f, ", ")?;
210                }
211                write!(f, "{}", transform.pretty(self.arena))?;
212                first = false;
213            }
214        }
215        write!(f, ".")
216    }
217}
218
219impl<'a> fmt::Display for Pretty<'a, Constraints<'a>> {
220    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221        if !self.inner.consequences.is_empty() {
222            write!(f, " |> ")?;
223            for (i, c) in self.inner.consequences.iter().enumerate() {
224                if i > 0 {
225                    write!(f, ", ")?;
226                }
227                write!(f, "{}", c.pretty(self.arena))?;
228            }
229        }
230        if !self.inner.alternatives.is_empty() {
231            for alt in self.inner.alternatives.iter() {
232                write!(f, " | ")?;
233                for (i, c) in alt.iter().enumerate() {
234                    if i > 0 {
235                        write!(f, ", ")?;
236                    }
237                    write!(f, "{}", c.pretty(self.arena))?;
238                }
239            }
240        }
241        Ok(())
242    }
243}
244
245impl<'a> fmt::Display for Pretty<'a, BoundDecl<'a>> {
246    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247        for (i, b) in self.inner.base_terms.iter().enumerate() {
248            if i > 0 {
249                write!(f, ", ")?;
250            }
251            write!(f, "{}", b.pretty(self.arena))?;
252        }
253        Ok(())
254    }
255}
256
257impl<'a> fmt::Display for Pretty<'a, Decl<'a>> {
258    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
259        write!(f, "{}", self.inner.atom.pretty(self.arena))?;
260        if !self.inner.descr.is_empty() {
261            write!(f, " [")?;
262            for (i, d) in self.inner.descr.iter().enumerate() {
263                if i > 0 {
264                    write!(f, ", ")?;
265                }
266                write!(f, "{}", d.pretty(self.arena))?;
267            }
268            write!(f, "]")?;
269        }
270        if let Some(bounds) = self.inner.bounds
271            && !bounds.is_empty()
272        {
273            write!(f, " bound ")?;
274            for (i, b) in bounds.iter().enumerate() {
275                if i > 0 {
276                    write!(f, " | ")?;
277                }
278                write!(f, "{}", b.pretty(self.arena))?;
279            }
280        }
281        if let Some(constraints) = &self.inner.constraints {
282            write!(f, "{}", constraints.pretty(self.arena))?;
283        }
284        write!(f, ".")
285    }
286}
287
288impl<'a> fmt::Display for Pretty<'a, Unit<'a>> {
289    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
290        for decl in self.inner.decls {
291            writeln!(f, "{}", decl.pretty(self.arena))?;
292        }
293        for clause in self.inner.clauses {
294            writeln!(f, "{}", clause.pretty(self.arena))?;
295        }
296        Ok(())
297    }
298}
mangle_ast/pretty.rs

mangle_ast/
pretty.rs
