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 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_assoc_item_name(
104        &self,
105        f: &mut fmt::Formatter<'_>,
106        trait_id: TraitDeclId,
107        item_id: AssocItemId,
108    ) -> fmt::Result {
109        match item_id {
110            AssocItemId::Type(id) => self.format_assoc_type_name(f, trait_id, id),
111            AssocItemId::Method(id) => self.format_method_name(f, trait_id, id),
112            AssocItemId::Const(id) => self.format_assoc_const_name(f, trait_id, id),
113        }
114    }
115
116    fn format_enum_variant_name(
117        &self,
118        f: &mut fmt::Formatter<'_>,
119        type_id: TypeDeclId,
120        variant_id: VariantId,
121    ) -> fmt::Result {
122        let variant = if let Some(translated) = self.get_crate()
123            && let Some(def) = translated.type_decls.get(type_id)
124            && let Some(variants) = def.kind.as_enum()
125        {
126            &variants.get(variant_id).unwrap().name
127        } else {
128            &variant_id.to_pretty_string()
129        };
130        write!(f, "{variant}")
131    }
132    fn format_enum_variant(
133        &self,
134        f: &mut fmt::Formatter<'_>,
135        type_id: TypeDeclId,
136        variant_id: VariantId,
137    ) -> fmt::Result {
138        write!(f, "{}::", type_id.with_ctx(self))?;
139        self.format_enum_variant_name(f, type_id, variant_id)?;
140        Ok(())
141    }
142
143    fn format_field_name(
144        &self,
145        f: &mut fmt::Formatter<'_>,
146        type_id: TypeDeclId,
147        opt_variant_id: Option<VariantId>,
148        field_id: FieldId,
149    ) -> fmt::Result {
150        let field_name = if let Some(translated) = self.get_crate()
151            && let Some(def) = translated.type_decls.get(type_id)
152        {
153            match (&def.kind, opt_variant_id) {
154                (TypeDeclKind::Enum(variants), Some(variant_id)) => {
155                    Some(&variants[variant_id].fields[field_id].name)
156                }
157                (TypeDeclKind::Struct(fields) | TypeDeclKind::Union(fields), None) => {
158                    Some(&fields[field_id].name)
159                }
160                _ => None,
161            }
162        } else {
163            None
164        };
165        if let Some(field_name) = field_name {
166            write!(f, "{field_name}")
167        } else {
168            write!(f, "{field_id}")
169        }
170    }
171}
172
173/// Context for formatting.
174#[derive(Default)]
175pub struct FmtCtx<'a> {
176    pub translated: Option<&'a TranslatedCrate>,
177    /// Generics form a stack, where each binder introduces a new level. For DeBruijn indices to
178    /// work, we keep the innermost parameters at the start of the vector.
179    pub generics: BindingStack<Cow<'a, GenericParams>>,
180    pub local_names: Option<IndexVec<LocalId, String>>,
181    pub indent_level: usize,
182}
183
184impl<'c> AstFormatter for FmtCtx<'c> {
185    type Reborrow<'a>
186        = FmtCtx<'a>
187    where
188        Self: 'a;
189
190    fn get_crate(&self) -> Option<&TranslatedCrate> {
191        self.translated
192    }
193
194    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a> {
195        FmtCtx {
196            generics: BindingStack::empty(),
197            ..self.reborrow()
198        }
199    }
200    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a> {
201        FmtCtx {
202            generics: BindingStack::new(Cow::Borrowed(generics)),
203            ..self.reborrow()
204        }
205    }
206    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a> {
207        FmtCtx {
208            local_names: Some(compute_local_names(locals)),
209            ..self.reborrow()
210        }
211    }
212    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a> {
213        let mut ret = self.reborrow();
214        ret.generics.push(new_params);
215        ret
216    }
217    fn binder_depth(&self) -> usize {
218        self.generics.len()
219    }
220
221    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a> {
222        FmtCtx {
223            indent_level: self.indent_level + 1,
224            ..self.reborrow()
225        }
226    }
227    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a> {
228        FmtCtx {
229            indent_level: 0,
230            ..self.reborrow()
231        }
232    }
233    fn indent(&self) -> String {
234        TAB_INCR.repeat(self.indent_level)
235    }
236
237    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result {
238        if let Some(local_names) = &self.local_names {
239            write!(f, "{}", local_names[id])
240        } else {
241            write!(f, "_{id}")
242        }
243    }
244
245    fn format_bound_var<Id: Idx + Display, T>(
246        &self,
247        f: &mut fmt::Formatter<'_>,
248        var: DeBruijnVar<Id>,
249        var_prefix: &str,
250        fmt_var: impl Fn(&T) -> Option<String>,
251    ) -> fmt::Result
252    where
253        GenericParams: HasIdxVecOf<Id, Output = T>,
254    {
255        if self.generics.is_empty() {
256            return write!(f, "{var_prefix}{var}");
257        }
258        match self.generics.get_var::<_, GenericParams>(var) {
259            None => write!(f, "missing({var_prefix}{var})"),
260            Some(v) => match fmt_var(v) {
261                Some(name) => write!(f, "{name}"),
262                None => {
263                    write!(f, "{var_prefix}")?;
264                    let (dbid, varid) = self.generics.as_bound_var(var);
265                    let depth = self.generics.depth().index - dbid.index;
266                    if depth == 0 {
267                        write!(f, "{varid}")
268                    } else {
269                        write!(f, "{varid}_{depth}")
270                    }
271                }
272            },
273        }
274    }
275}
276
277impl<'a> FmtCtx<'a> {
278    pub fn new() -> Self {
279        FmtCtx::default()
280    }
281
282    pub fn get_item(&self, id: ItemId) -> Result<ItemRef<'_>, Option<&Name>> {
283        let Some(translated) = &self.translated else {
284            return Err(None);
285        };
286        translated
287            .get_item(id)
288            .ok_or_else(|| Some(translated.item_short_name(id)))
289    }
290
291    /// Print the whole definition.
292    pub fn format_decl_id(&self, id: impl Into<ItemId>) -> String {
293        let id = id.into();
294        match self.get_item(id) {
295            Ok(d) => d.to_string_with_ctx(self),
296            Err(opt_name) => {
297                let opt_name = opt_name
298                    .map(|n| format!(" ({})", n.with_ctx(self)))
299                    .unwrap_or_default();
300                format!("Missing decl: {id:?}{opt_name}")
301            }
302        }
303    }
304
305    fn reborrow<'b>(&'b self) -> FmtCtx<'b> {
306        FmtCtx {
307            translated: self.translated,
308            generics: self.generics.clone(),
309            local_names: self.local_names.clone(),
310            indent_level: self.indent_level,
311        }
312    }
313}
314
315/// Compute a unique name for each local.
316pub fn compute_local_names(locals: &Locals) -> IndexVec<LocalId, String> {
317    let mut local_names = locals.locals.map_ref(|local| {
318        format!(
319            "{}_{}",
320            local.name.as_deref().unwrap_or_default(),
321            local.index
322        )
323    });
324
325    let mut name_counts = HashMap::<String, usize>::new();
326    for local in &locals.locals {
327        *name_counts
328            .entry(local_names[local.index].clone())
329            .or_default() += 1;
330        if let Some(name) = &local.name {
331            *name_counts.entry(name.clone()).or_default() += 1;
332        }
333    }
334
335    for (id, local) in locals.locals.iter_enumerated() {
336        if let Some(name) = &local.name
337            && !name.is_empty()
338            && name_counts[name] == 1
339        {
340            local_names[id] = name.clone();
341        }
342    }
343    local_names
344}