sway_fmt/
traversal.rs

1use sway_core::{
2    AstNode, AstNodeContent, Declaration, Expression, ExpressionKind, IfExpression, ParseTree,
3    ReturnStatement,
4};
5
6use sway_types::span::Span;
7
8use crate::traversal_helper::{
9    format_data_types, format_delineated_path, format_include_statement, format_use_statement,
10};
11
12/// Change contains the formatted change itself.
13/// `start` and `end` denote the start and end of that change,
14/// which are used to caluclate the position for inserting this change in the existing file.
15#[derive(Debug)]
16pub struct Change {
17    pub text: String,
18    pub start: usize,
19    pub end: usize,
20}
21
22impl Change {
23    fn new(span: &Span, change_type: ChangeType) -> Self {
24        let text = match change_type {
25            ChangeType::Struct => format_data_types(span.as_str()),
26            ChangeType::Enum => format_data_types(span.as_str()),
27            ChangeType::IncludeStatement => format_include_statement(span.as_str()),
28            ChangeType::UseStatement => format_use_statement(span.as_str()),
29            ChangeType::DelineatedPath => format_delineated_path(span.as_str()),
30        };
31
32        Self {
33            text,
34            start: span.start(),
35            end: span.end(),
36        }
37    }
38}
39
40#[derive(Debug)]
41enum ChangeType {
42    Struct,
43    Enum,
44    IncludeStatement,
45    UseStatement,
46    DelineatedPath,
47}
48
49/// traverses the Sway ParseTree and returns list of formatted changes
50pub fn traverse_for_changes(parse_tree: &ParseTree) -> Vec<Change> {
51    let mut changes = vec![];
52
53    for node in &parse_tree.root_nodes {
54        traverse_ast_node(node, &mut changes);
55    }
56
57    changes.sort_by(|a, b| a.start.cmp(&b.start));
58
59    changes
60}
61
62fn traverse_ast_node(ast_node: &AstNode, changes: &mut Vec<Change>) {
63    match &ast_node.content {
64        AstNodeContent::Declaration(dec) => handle_declaration(dec, ast_node, changes),
65
66        AstNodeContent::ReturnStatement(ret) => handle_return_statement(ret, changes),
67
68        AstNodeContent::Expression(expr) => handle_expression(expr, changes),
69
70        AstNodeContent::ImplicitReturnExpression(expr) => {
71            handle_implicit_return_expression(expr, changes)
72        }
73
74        AstNodeContent::UseStatement(_) => {
75            // The AST generates one root node per use statement, we must avoid duplicating them
76            // while formatting
77            let next_span = &ast_node.span;
78            match changes.last() {
79                Some(last_change) => {
80                    if last_change.start != next_span.start() {
81                        changes.push(Change::new(next_span, ChangeType::UseStatement));
82                    }
83                }
84                _ => changes.push(Change::new(next_span, ChangeType::UseStatement)),
85            }
86        }
87
88        AstNodeContent::IncludeStatement(_) => {
89            changes.push(Change::new(&ast_node.span, ChangeType::IncludeStatement))
90        }
91    }
92}
93
94fn handle_return_statement(ret: &ReturnStatement, changes: &mut Vec<Change>) {
95    handle_expression(&ret.expr, changes)
96}
97
98fn handle_declaration(dec: &Declaration, ast_node: &AstNode, changes: &mut Vec<Change>) {
99    match &dec {
100        Declaration::VariableDeclaration(var_dec) => handle_expression(&var_dec.body, changes),
101
102        Declaration::StructDeclaration(_) | Declaration::StorageDeclaration(_) => {
103            changes.push(Change::new(&ast_node.span, ChangeType::Struct))
104        }
105
106        Declaration::EnumDeclaration(_) => {
107            changes.push(Change::new(&ast_node.span, ChangeType::Enum))
108        }
109
110        Declaration::FunctionDeclaration(func) => {
111            for content in &func.body.contents {
112                traverse_ast_node(content, changes);
113            }
114        }
115
116        Declaration::ImplSelf(impl_self) => {
117            for func in &impl_self.functions {
118                for content in &func.body.contents {
119                    traverse_ast_node(content, changes);
120                }
121            }
122        }
123        Declaration::ImplTrait(impl_trait) => {
124            for func in &impl_trait.functions {
125                for content in &func.body.contents {
126                    traverse_ast_node(content, changes);
127                }
128            }
129        }
130        _ => {}
131    };
132}
133
134fn handle_expression(expr: &Expression, changes: &mut Vec<Change>) {
135    let span = &expr.span;
136    match &expr.kind {
137        ExpressionKind::Struct(_) => changes.push(Change::new(span, ChangeType::Struct)),
138        ExpressionKind::If(IfExpression {
139            condition: _,
140            then,
141            r#else,
142        }) => {
143            handle_expression(then, changes);
144
145            if let Some(else_expr) = r#else {
146                handle_expression(else_expr, changes);
147            }
148        }
149        ExpressionKind::CodeBlock(contents) => {
150            for content in &contents.contents {
151                traverse_ast_node(content, changes);
152            }
153        }
154        ExpressionKind::DelineatedPath(_) => {
155            changes.push(Change::new(span, ChangeType::DelineatedPath));
156        }
157        _ => {}
158    }
159}
160
161fn handle_implicit_return_expression(expr: &Expression, changes: &mut Vec<Change>) {
162    handle_expression(expr, changes)
163}
sway_fmt/traversal.rs

sway_fmt/
traversal.rs
