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::ids::IndexVec;
10use crate::pretty::FmtWithCtx;
11use crate::utils::TAB_INCR;
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 include_layouts(&self) -> bool {
28        false
29    }
30    fn include_safety(&self) -> bool {
31        false
32    }
33    fn hide_storage_statements(&self) -> bool {
34        false
35    }
36
37    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a>;
38    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a>;
39    fn set_current_type<'a>(&'a self, type_id: TypeDeclId) -> Self::Reborrow<'a>;
40    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a>;
41    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a>;
42    fn push_bound_regions<'a>(
43        &'a self,
44        regions: &'a IndexVec<RegionId, RegionParam>,
45    ) -> Self::Reborrow<'a> {
46        self.push_binder(Cow::Owned(GenericParams {
47            regions: regions.clone(),
48            ..Default::default()
49        }))
50    }
51    /// Return the depth of binders we're under.
52    fn binder_depth(&self) -> usize;
53
54    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a>;
55    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a>;
56    fn indent(&self) -> String;
57
58    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result;
59    fn format_bound_var<Id: Idx + Display, T>(
60        &self,
61        f: &mut fmt::Formatter<'_>,
62        var: DeBruijnVar<Id>,
63        var_prefix: &str,
64        fmt_var: impl Fn(&T) -> Option<String>,
65    ) -> fmt::Result
66    where
67        GenericParams: HasIdxVecOf<Id, Output = T>;
68
69    fn format_method_name(
70        &self,
71        f: &mut fmt::Formatter<'_>,
72        trait_id: TraitDeclId,
73        method_id: TraitMethodId,
74    ) -> fmt::Result {
75        if let Some(translated) = self.get_crate()
76            && let Some(names) = translated.assoc_item_names.get(trait_id)
77            && let Some(name) = names.methods.get(method_id).copied()
78        {
79            write!(f, "{name}")
80        } else {
81            write!(f, "{}", method_id.to_pretty_string())
82        }
83    }
84    fn format_assoc_type_name(
85        &self,
86        f: &mut fmt::Formatter<'_>,
87        trait_id: TraitDeclId,
88        type_id: AssocTypeId,
89    ) -> fmt::Result {
90        if let Some(translated) = self.get_crate()
91            && let Some(names) = translated.assoc_item_names.get(trait_id)
92            && let Some(name) = names.types.get(type_id).copied()
93        {
94            write!(f, "{name}")
95        } else {
96            write!(f, "{}", type_id.to_pretty_string())
97        }
98    }
99    fn format_assoc_const_name(
100        &self,
101        f: &mut fmt::Formatter<'_>,
102        trait_id: TraitDeclId,
103        const_id: AssocConstId,
104    ) -> fmt::Result {
105        if let Some(translated) = self.get_crate()
106            && let Some(names) = translated.assoc_item_names.get(trait_id)
107            && let Some(name) = names.consts.get(const_id).copied()
108        {
109            write!(f, "{name}")
110        } else {
111            write!(f, "{}", const_id.to_pretty_string())
112        }
113    }
114    fn format_assoc_item_name(
115        &self,
116        f: &mut fmt::Formatter<'_>,
117        trait_id: TraitDeclId,
118        item_id: AssocItemId,
119    ) -> fmt::Result {
120        match item_id {
121            AssocItemId::Type(id) => self.format_assoc_type_name(f, trait_id, id),
122            AssocItemId::Method(id) => self.format_method_name(f, trait_id, id),
123            AssocItemId::Const(id) => self.format_assoc_const_name(f, trait_id, id),
124        }
125    }
126
127    fn format_enum_variant_name(
128        &self,
129        f: &mut fmt::Formatter<'_>,
130        type_id: TypeDeclId,
131        variant_id: VariantId,
132    ) -> fmt::Result {
133        let variant = if let Some(translated) = self.get_crate()
134            && let Some(def) = translated.type_decls.get(type_id)
135        {
136            match &def.kind {
137                TypeDeclKind::Enum(variants) => &variants.get(variant_id).unwrap().name,
138                TypeDeclKind::Struct(..) | TypeDeclKind::Union(..) => "_",
139                TypeDeclKind::Opaque | TypeDeclKind::Alias(..) | TypeDeclKind::Error(..) => {
140                    &variant_id.to_pretty_string()
141                }
142            }
143        } else {
144            &variant_id.to_pretty_string()
145        };
146        write!(f, "{variant}")
147    }
148    fn format_current_variant_name(
149        &self,
150        f: &mut fmt::Formatter<'_>,
151        variant_id: VariantId,
152    ) -> fmt::Result;
153    fn format_enum_variant(
154        &self,
155        f: &mut fmt::Formatter<'_>,
156        type_id: TypeDeclId,
157        variant_id: VariantId,
158    ) -> fmt::Result {
159        write!(f, "{}::", type_id.with_ctx(self))?;
160        self.format_enum_variant_name(f, type_id, variant_id)?;
161        Ok(())
162    }
163
164    fn format_field_name(
165        &self,
166        f: &mut fmt::Formatter<'_>,
167        type_id: TypeDeclId,
168        opt_variant_id: Option<VariantId>,
169        field_id: FieldId,
170    ) -> fmt::Result {
171        let field_name = if let Some(translated) = self.get_crate()
172            && let Some(def) = translated.type_decls.get(type_id)
173        {
174            match (&def.kind, opt_variant_id) {
175                (TypeDeclKind::Enum(variants), Some(variant_id)) => {
176                    Some(&variants[variant_id].fields[field_id].name)
177                }
178                (TypeDeclKind::Struct(fields) | TypeDeclKind::Union(fields), _) => {
179                    Some(&fields[field_id].name)
180                }
181                _ => None,
182            }
183        } else {
184            None
185        };
186        if let Some(field_name) = field_name {
187            write!(f, "{field_name}")
188        } else {
189            write!(f, "{field_id}")
190        }
191    }
192    fn format_current_field_name(
193        &self,
194        f: &mut fmt::Formatter<'_>,
195        _variant_id: VariantId,
196        field_id: FieldId,
197    ) -> fmt::Result;
198}
199
200/// Context for formatting.
201#[derive(Default)]
202pub struct FmtCtx<'a> {
203    pub translated: Option<&'a TranslatedCrate>,
204    pub include_layouts: bool,
205    pub include_safety: bool,
206    pub hide_storage_statements: bool,
207    pub current_type: Option<TypeDeclId>,
208    /// Generics form a stack, where each binder introduces a new level. For DeBruijn indices to
209    /// work, we keep the innermost parameters at the start of the vector.
210    pub generics: BindingStack<Cow<'a, GenericParams>>,
211    pub local_names: Option<IndexVec<LocalId, String>>,
212    pub indent_level: usize,
213}
214
215impl<'c> AstFormatter for FmtCtx<'c> {
216    type Reborrow<'a>
217        = FmtCtx<'a>
218    where
219        Self: 'a;
220
221    fn get_crate(&self) -> Option<&TranslatedCrate> {
222        self.translated
223    }
224
225    fn include_layouts(&self) -> bool {
226        self.include_layouts
227    }
228    fn include_safety(&self) -> bool {
229        self.include_safety
230    }
231    fn hide_storage_statements(&self) -> bool {
232        self.hide_storage_statements
233    }
234
235    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a> {
236        FmtCtx {
237            generics: BindingStack::empty(),
238            ..self.reborrow()
239        }
240    }
241    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a> {
242        FmtCtx {
243            generics: BindingStack::new(Cow::Borrowed(generics)),
244            ..self.reborrow()
245        }
246    }
247    fn set_current_type<'a>(&'a self, type_id: TypeDeclId) -> Self::Reborrow<'a> {
248        FmtCtx {
249            current_type: Some(type_id),
250            ..self.reborrow()
251        }
252    }
253    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a> {
254        FmtCtx {
255            local_names: Some(compute_local_names(locals)),
256            ..self.reborrow()
257        }
258    }
259    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a> {
260        let mut ret = self.reborrow();
261        ret.generics.push(new_params);
262        ret
263    }
264    fn binder_depth(&self) -> usize {
265        self.generics.len()
266    }
267
268    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a> {
269        FmtCtx {
270            indent_level: self.indent_level + 1,
271            ..self.reborrow()
272        }
273    }
274    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a> {
275        FmtCtx {
276            indent_level: 0,
277            ..self.reborrow()
278        }
279    }
280    fn indent(&self) -> String {
281        TAB_INCR.repeat(self.indent_level)
282    }
283
284    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result {
285        if let Some(local_names) = &self.local_names {
286            write!(f, "{}", local_names[id])
287        } else {
288            write!(f, "_{id}")
289        }
290    }
291
292    fn format_current_variant_name(
293        &self,
294        f: &mut fmt::Formatter<'_>,
295        variant_id: VariantId,
296    ) -> fmt::Result {
297        match self.current_type {
298            Some(type_id) => self.format_enum_variant_name(f, type_id, variant_id),
299            None => write!(f, "{variant_id}"),
300        }
301    }
302
303    fn format_current_field_name(
304        &self,
305        f: &mut fmt::Formatter<'_>,
306        variant_id: VariantId,
307        field_id: FieldId,
308    ) -> fmt::Result {
309        match self.current_type {
310            Some(type_id) => self.format_field_name(f, type_id, Some(variant_id), field_id),
311            None => write!(f, "{field_id}"),
312        }
313    }
314
315    fn format_bound_var<Id: Idx + Display, T>(
316        &self,
317        f: &mut fmt::Formatter<'_>,
318        var: DeBruijnVar<Id>,
319        var_prefix: &str,
320        fmt_var: impl Fn(&T) -> Option<String>,
321    ) -> fmt::Result
322    where
323        GenericParams: HasIdxVecOf<Id, Output = T>,
324    {
325        if self.generics.is_empty() {
326            return write!(f, "{var_prefix}{var}");
327        }
328        match self.generics.get_var::<_, GenericParams>(var) {
329            None => write!(f, "missing({var_prefix}{var})"),
330            Some(v) => match fmt_var(v) {
331                Some(name) => write!(f, "{name}"),
332                None => {
333                    write!(f, "{var_prefix}")?;
334                    let (dbid, varid) = self.generics.as_bound_var(var);
335                    let depth = self.generics.depth().index - dbid.index;
336                    if depth == 0 {
337                        write!(f, "{varid}")
338                    } else {
339                        write!(f, "{varid}_{depth}")
340                    }
341                }
342            },
343        }
344    }
345}
346
347impl<'a> FmtCtx<'a> {
348    pub fn new() -> Self {
349        FmtCtx::default()
350    }
351
352    pub fn get_item(&self, id: ItemId) -> Result<ItemRef<'_>, Option<&Name>> {
353        let Some(translated) = &self.translated else {
354            return Err(None);
355        };
356        translated
357            .get_item(id)
358            .ok_or_else(|| Some(translated.item_short_name(id)))
359    }
360
361    /// Print the whole definition.
362    pub fn format_decl_id(&self, id: impl Into<ItemId>) -> String {
363        let id = id.into();
364        match self.get_item(id) {
365            Ok(d) => d.to_string_with_ctx(self),
366            Err(opt_name) => {
367                let opt_name = opt_name
368                    .map(|n| format!(" ({})", n.with_ctx(self)))
369                    .unwrap_or_default();
370                format!("Missing decl: {id:?}{opt_name}")
371            }
372        }
373    }
374
375    fn reborrow<'b>(&'b self) -> FmtCtx<'b> {
376        FmtCtx {
377            translated: self.translated,
378            include_layouts: self.include_layouts,
379            include_safety: self.include_safety,
380            hide_storage_statements: self.hide_storage_statements,
381            current_type: self.current_type,
382            generics: self.generics.clone(),
383            local_names: self.local_names.clone(),
384            indent_level: self.indent_level,
385        }
386    }
387}
388
389/// Compute a unique name for each local.
390pub fn compute_local_names(locals: &Locals) -> IndexVec<LocalId, String> {
391    let mut local_names = locals.locals.map_ref(|local| {
392        format!(
393            "{}_{}",
394            local.name.as_deref().unwrap_or_default(),
395            local.index
396        )
397    });
398
399    let mut name_counts = HashMap::<String, usize>::new();
400    for local in &locals.locals {
401        *name_counts
402            .entry(local_names[local.index].clone())
403            .or_default() += 1;
404        if let Some(name) = &local.name {
405            *name_counts.entry(name.clone()).or_default() += 1;
406        }
407    }
408
409    for (id, local) in locals.locals.iter_enumerated() {
410        if let Some(name) = &local.name
411            && !name.is_empty()
412            && name_counts[name] == 1
413        {
414            local_names[id] = name.clone();
415        }
416    }
417    local_names
418}