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