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