swayfmt/items/item_enum/
mod.rs

1use crate::{
2    comments::rewrite_with_comments,
3    config::user_def::FieldAlignment,
4    formatter::{
5        shape::{ExprKind, LineStyle},
6        *,
7    },
8    utils::{
9        map::byte_span::{ByteSpan, LeafSpans},
10        CurlyBrace,
11    },
12};
13use std::fmt::Write;
14use sway_ast::{
15    keywords::{ColonToken, EnumToken, Keyword, Token},
16    CommaToken, ItemEnum, PubToken,
17};
18use sway_types::{ast::Delimiter, Spanned};
19
20#[cfg(test)]
21mod tests;
22
23impl Format for ItemEnum {
24    fn format(
25        &self,
26        formatted_code: &mut FormattedCode,
27        formatter: &mut Formatter,
28    ) -> Result<(), FormatterError> {
29        formatter.with_shape(
30            formatter
31                .shape
32                .with_code_line_from(LineStyle::Multiline, ExprKind::default()),
33            |formatter| -> Result<(), FormatterError> {
34                // Required for comment formatting
35                let start_len = formatted_code.len();
36                // If there is a visibility token add it to the formatted_code with a ` ` after it.
37                if self.visibility.is_some() {
38                    write!(formatted_code, "{} ", PubToken::AS_STR)?;
39                }
40                // Add enum token and name
41                write!(formatted_code, "{} ", EnumToken::AS_STR)?;
42                self.name.format(formatted_code, formatter)?;
43                // Format `GenericParams`, if any
44                if let Some(generics) = &self.generics {
45                    generics.format(formatted_code, formatter)?;
46                }
47
48                let fields = self.fields.get();
49
50                formatter.shape.code_line.update_expr_new_line(true);
51
52                // Handle opening brace
53                Self::open_curly_brace(formatted_code, formatter)?;
54                // Determine alignment tactic
55                match formatter.config.structures.field_alignment {
56                    FieldAlignment::AlignFields(enum_variant_align_threshold) => {
57                        writeln!(formatted_code)?;
58                        let type_fields = &fields
59                            .value_separator_pairs
60                            .iter()
61                            // TODO: Handle annotations instead of stripping them.
62                            //       See: https://github.com/FuelLabs/sway/issues/6802
63                            .map(|(type_field, _comma_token)| &type_field.value)
64                            .collect::<Vec<_>>();
65                        // In first iteration we are going to be collecting the lengths of the enum variants.
66                        let variants_lengths: Vec<usize> = type_fields
67                            .iter()
68                            .map(|type_field| type_field.name.as_str().len())
69                            .collect();
70
71                        // Find the maximum length that is still smaller than the align threshold.
72                        let mut max_valid_variant_length = 0;
73                        variants_lengths.iter().for_each(|length| {
74                            if *length > max_valid_variant_length
75                                && *length < enum_variant_align_threshold
76                            {
77                                max_valid_variant_length = *length;
78                            }
79                        });
80
81                        for (var_index, type_field) in type_fields.iter().enumerate() {
82                            write!(formatted_code, "{}", formatter.indent_to_str()?)?;
83
84                            // Add name
85                            type_field.name.format(formatted_code, formatter)?;
86                            let current_variant_length = variants_lengths[var_index];
87                            if current_variant_length < max_valid_variant_length {
88                                // We need to add alignment between : and ty
89                                // max_valid_variant_length: the length of the variant that we are taking as a reference to align
90                                // current_variant_length: the length of the current variant that we are trying to format
91                                let mut required_alignment =
92                                    max_valid_variant_length - current_variant_length;
93                                while required_alignment != 0 {
94                                    write!(formatted_code, " ")?;
95                                    required_alignment -= 1;
96                                }
97                            }
98                            // Add `:`, ty & `CommaToken`
99                            write!(formatted_code, " {} ", ColonToken::AS_STR)?;
100                            type_field.ty.format(formatted_code, formatter)?;
101                            writeln!(formatted_code, "{}", CommaToken::AS_STR)?;
102                        }
103                        if let Some(final_value) = &fields.final_value_opt {
104                            final_value.format(formatted_code, formatter)?;
105                            writeln!(formatted_code)?;
106                        }
107                    }
108                    FieldAlignment::Off => fields.format(formatted_code, formatter)?,
109                }
110                // Handle closing brace
111                Self::close_curly_brace(formatted_code, formatter)?;
112
113                rewrite_with_comments::<ItemEnum>(
114                    formatter,
115                    self.span(),
116                    self.leaf_spans(),
117                    formatted_code,
118                    start_len,
119                )?;
120
121                Ok(())
122            },
123        )?;
124
125        Ok(())
126    }
127}
128
129impl CurlyBrace for ItemEnum {
130    fn open_curly_brace(
131        line: &mut String,
132        formatter: &mut Formatter,
133    ) -> Result<(), FormatterError> {
134        let open_brace = Delimiter::Brace.as_open_char();
135        // Add opening brace to the same line
136        write!(line, " {open_brace}")?;
137        formatter.indent();
138
139        Ok(())
140    }
141    fn close_curly_brace(
142        line: &mut String,
143        formatter: &mut Formatter,
144    ) -> Result<(), FormatterError> {
145        // If shape is becoming left-most aligned or - indent just have the default shape
146        formatter.unindent();
147        write!(
148            line,
149            "{}{}",
150            formatter.indent_to_str()?,
151            Delimiter::Brace.as_close_char()
152        )?;
153
154        Ok(())
155    }
156}
157
158impl LeafSpans for ItemEnum {
159    fn leaf_spans(&self) -> Vec<ByteSpan> {
160        let mut collected_spans = Vec::new();
161        if let Some(visibility) = &self.visibility {
162            collected_spans.push(ByteSpan::from(visibility.span()));
163        }
164        collected_spans.push(ByteSpan::from(self.enum_token.span()));
165        collected_spans.push(ByteSpan::from(self.name.span()));
166        if let Some(generics) = &self.generics {
167            collected_spans.push(ByteSpan::from(generics.parameters.span()))
168        }
169        collected_spans.append(&mut self.fields.leaf_spans());
170        collected_spans
171    }
172}