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
18pub 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 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#[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 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 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
381pub 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}