Skip to main content

charon_lib/pretty/
formatter.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3use std::fmt;
4use std::fmt::Display;
5
6use index_vec::Idx;
7
8use crate::ast::*;
9use crate::common::TAB_INCR;
10use crate::ids::IndexVec;
11use crate::pretty::FmtWithCtx;
12
13pub trait IntoFormatter {
14    type C: AstFormatter;
15    fn into_fmt(self) -> Self::C;
16}
17
18/// An [`AstFormatter`] contains the context required to pretty-print the ast. An ast type can then
19/// be pretty-printed using the [`FmtWithCtx`] trait.
20pub trait AstFormatter: Sized {
21    type Reborrow<'a>: AstFormatter + 'a
22    where
23        Self: 'a;
24
25    fn get_crate(&self) -> Option<&TranslatedCrate>;
26
27    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a>;
28    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a>;
29    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a>;
30    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a>;
31    fn push_bound_regions<'a>(
32        &'a self,
33        regions: &'a IndexVec<RegionId, RegionParam>,
34    ) -> Self::Reborrow<'a> {
35        self.push_binder(Cow::Owned(GenericParams {
36            regions: regions.clone(),
37            ..Default::default()
38        }))
39    }
40    /// Return the depth of binders we're under.
41    fn binder_depth(&self) -> usize;
42
43    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a>;
44    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a>;
45    fn indent(&self) -> String;
46
47    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result;
48    fn format_bound_var<Id: Idx + Display, T>(
49        &self,
50        f: &mut fmt::Formatter<'_>,
51        var: DeBruijnVar<Id>,
52        var_prefix: &str,
53        fmt_var: impl Fn(&T) -> Option<String>,
54    ) -> fmt::Result
55    where
56        GenericParams: HasIdxVecOf<Id, Output = T>;
57
58    fn format_method_name(
59        &self,
60        f: &mut fmt::Formatter<'_>,
61        trait_id: TraitDeclId,
62        method_id: TraitMethodId,
63    ) -> fmt::Result {
64        if let Some(translated) = self.get_crate()
65            && let Some(names) = translated.assoc_item_names.get(trait_id)
66            && let Some(name) = names.methods.get(method_id).copied()
67        {
68            write!(f, "{name}")
69        } else {
70            write!(f, "{}", method_id.to_pretty_string())
71        }
72    }
73    fn format_assoc_type_name(
74        &self,
75        f: &mut fmt::Formatter<'_>,
76        trait_id: TraitDeclId,
77        type_id: AssocTypeId,
78    ) -> fmt::Result {
79        if let Some(translated) = self.get_crate()
80            && let Some(names) = translated.assoc_item_names.get(trait_id)
81            && let Some(name) = names.types.get(type_id).copied()
82        {
83            write!(f, "{name}")
84        } else {
85            write!(f, "{}", type_id.to_pretty_string())
86        }
87    }
88    fn format_assoc_const_name(
89        &self,
90        f: &mut fmt::Formatter<'_>,
91        trait_id: TraitDeclId,
92        const_id: AssocConstId,
93    ) -> fmt::Result {
94        if let Some(translated) = self.get_crate()
95            && let Some(names) = translated.assoc_item_names.get(trait_id)
96            && let Some(name) = names.consts.get(const_id).copied()
97        {
98            write!(f, "{name}")
99        } else {
100            write!(f, "{}", const_id.to_pretty_string())
101        }
102    }
103    fn format_enum_variant_name(
104        &self,
105        f: &mut fmt::Formatter<'_>,
106        type_id: TypeDeclId,
107        variant_id: VariantId,
108    ) -> fmt::Result {
109        let variant = if let Some(translated) = self.get_crate()
110            && let Some(def) = translated.type_decls.get(type_id)
111            && let Some(variants) = def.kind.as_enum()
112        {
113            &variants.get(variant_id).unwrap().name
114        } else {
115            &variant_id.to_pretty_string()
116        };
117        write!(f, "{variant}")
118    }
119    fn format_enum_variant(
120        &self,
121        f: &mut fmt::Formatter<'_>,
122        type_id: TypeDeclId,
123        variant_id: VariantId,
124    ) -> fmt::Result {
125        write!(f, "{}::", type_id.with_ctx(self))?;
126        self.format_enum_variant_name(f, type_id, variant_id)?;
127        Ok(())
128    }
129
130    fn format_field_name(
131        &self,
132        f: &mut fmt::Formatter<'_>,
133        type_id: TypeDeclId,
134        opt_variant_id: Option<VariantId>,
135        field_id: FieldId,
136    ) -> fmt::Result {
137        let field_name = if let Some(translated) = self.get_crate()
138            && let Some(def) = translated.type_decls.get(type_id)
139        {
140            match (&def.kind, opt_variant_id) {
141                (TypeDeclKind::Enum(variants), Some(variant_id)) => {
142                    variants[variant_id].fields[field_id].name.as_ref()
143                }
144                (TypeDeclKind::Struct(fields) | TypeDeclKind::Union(fields), None) => {
145                    fields[field_id].name.as_ref()
146                }
147                _ => None,
148            }
149        } else {
150            None
151        };
152        if let Some(field_name) = field_name {
153            write!(f, "{field_name}")
154        } else {
155            write!(f, "{field_id}")
156        }
157    }
158}
159
160/// Context for formatting.
161#[derive(Default)]
162pub struct FmtCtx<'a> {
163    pub translated: Option<&'a TranslatedCrate>,
164    /// Generics form a stack, where each binder introduces a new level. For DeBruijn indices to
165    /// work, we keep the innermost parameters at the start of the vector.
166    pub generics: BindingStack<Cow<'a, GenericParams>>,
167    pub local_names: Option<IndexVec<LocalId, String>>,
168    pub indent_level: usize,
169}
170
171impl<'c> AstFormatter for FmtCtx<'c> {
172    type Reborrow<'a>
173        = FmtCtx<'a>
174    where
175        Self: 'a;
176
177    fn get_crate(&self) -> Option<&TranslatedCrate> {
178        self.translated
179    }
180
181    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a> {
182        FmtCtx {
183            generics: BindingStack::empty(),
184            ..self.reborrow()
185        }
186    }
187    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a> {
188        FmtCtx {
189            generics: BindingStack::new(Cow::Borrowed(generics)),
190            ..self.reborrow()
191        }
192    }
193    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a> {
194        FmtCtx {
195            local_names: Some(compute_local_names(locals)),
196            ..self.reborrow()
197        }
198    }
199    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a> {
200        let mut ret = self.reborrow();
201        ret.generics.push(new_params);
202        ret
203    }
204    fn binder_depth(&self) -> usize {
205        self.generics.len()
206    }
207
208    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a> {
209        FmtCtx {
210            indent_level: self.indent_level + 1,
211            ..self.reborrow()
212        }
213    }
214    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a> {
215        FmtCtx {
216            indent_level: 0,
217            ..self.reborrow()
218        }
219    }
220    fn indent(&self) -> String {
221        TAB_INCR.repeat(self.indent_level)
222    }
223
224    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result {
225        if let Some(local_names) = &self.local_names {
226            write!(f, "{}", local_names[id])
227        } else {
228            write!(f, "_{id}")
229        }
230    }
231
232    fn format_bound_var<Id: Idx + Display, T>(
233        &self,
234        f: &mut fmt::Formatter<'_>,
235        var: DeBruijnVar<Id>,
236        var_prefix: &str,
237        fmt_var: impl Fn(&T) -> Option<String>,
238    ) -> fmt::Result
239    where
240        GenericParams: HasIdxVecOf<Id, Output = T>,
241    {
242        if self.generics.is_empty() {
243            return write!(f, "{var_prefix}{var}");
244        }
245        match self.generics.get_var::<_, GenericParams>(var) {
246            None => write!(f, "missing({var_prefix}{var})"),
247            Some(v) => match fmt_var(v) {
248                Some(name) => write!(f, "{name}"),
249                None => {
250                    write!(f, "{var_prefix}")?;
251                    let (dbid, varid) = self.generics.as_bound_var(var);
252                    let depth = self.generics.depth().index - dbid.index;
253                    if depth == 0 {
254                        write!(f, "{varid}")
255                    } else {
256                        write!(f, "{varid}_{depth}")
257                    }
258                }
259            },
260        }
261    }
262}
263
264impl<'a> FmtCtx<'a> {
265    pub fn new() -> Self {
266        FmtCtx::default()
267    }
268
269    pub fn get_item(&self, id: ItemId) -> Result<ItemRef<'_>, Option<&Name>> {
270        let Some(translated) = &self.translated else {
271            return Err(None);
272        };
273        translated
274            .get_item(id)
275            .ok_or_else(|| Some(translated.item_short_name(id)))
276    }
277
278    /// Print the whole definition.
279    pub fn format_decl_id(&self, id: impl Into<ItemId>) -> String {
280        let id = id.into();
281        match self.get_item(id) {
282            Ok(d) => d.to_string_with_ctx(self),
283            Err(opt_name) => {
284                let opt_name = opt_name
285                    .map(|n| format!(" ({})", n.with_ctx(self)))
286                    .unwrap_or_default();
287                format!("Missing decl: {id:?}{opt_name}")
288            }
289        }
290    }
291
292    fn reborrow<'b>(&'b self) -> FmtCtx<'b> {
293        FmtCtx {
294            translated: self.translated,
295            generics: self.generics.clone(),
296            local_names: self.local_names.clone(),
297            indent_level: self.indent_level,
298        }
299    }
300}
301
302/// Compute a unique name for each local.
303pub fn compute_local_names(locals: &Locals) -> IndexVec<LocalId, String> {
304    let mut local_names = locals.locals.map_ref(|local| {
305        format!(
306            "{}_{}",
307            local.name.as_deref().unwrap_or_default(),
308            local.index
309        )
310    });
311
312    let mut name_counts = HashMap::<String, usize>::new();
313    for local in &locals.locals {
314        *name_counts
315            .entry(local_names[local.index].clone())
316            .or_default() += 1;
317        if let Some(name) = &local.name {
318            *name_counts.entry(name.clone()).or_default() += 1;
319        }
320    }
321
322    for (id, local) in locals.locals.iter_enumerated() {
323        if let Some(name) = &local.name
324            && !name.is_empty()
325            && name_counts[name] == 1
326        {
327            local_names[id] = name.clone();
328        }
329    }
330    local_names
331}