Add migration for renaming existing panic identifiers to r#panic (#7231)

## Description

This PR:
- implements the [Rename existing "panic"
identifiers](https://github.com/FuelLabs/sway/issues/6765) migration
step.
- further extends the migration infrastructure by covering visiting
additional expressions and declarations.

As with all previous migrations, a pragmatic approach is taken between
the invested effort and the gained coverage. The goal is to support
migrating the use cases most probable in real-life code. Non-supported
cases are documented in code comments.

The PR is the last step in fully implementing #6765.

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [ ] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [ ] I have added tests that prove my fix is effective or that my
feature works.
- [ ] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

forc-plugins/forc-migrate/src/cli/shared.rs

@@ -96,14 +96,11 @@ pub(crate) fn compile_package<'a>(
     let manifest_dir = build_instructions.manifest_dir()?;
     let manifest = ManifestFile::from_dir(manifest_dir.clone())?;
     let pkg_manifest = get_pkg_manifest_file(&manifest)?;
+    let pkg_name = pkg_manifest.project_name();
 
     println_action_green(
         "Compiling",
-        &format!(
-            "{} ({})",
-            pkg_manifest.project_name(),
-            manifest.dir().to_string_lossy()
-        ),
+        &format!("{} ({})", pkg_name, manifest.dir().to_string_lossy()),
     );
 
     let member_manifests = manifest.member_manifests()?;
@@ -149,6 +146,7 @@ pub(crate) fn compile_package<'a>(
     };
 
     return Ok(ProgramInfo {
+        pkg_name: pkg_name.to_string(),
         lexed_program: programs.lexed,
         ty_program,
         engines,

forc-plugins/forc-migrate/src/migrations/error_type.rs

@@ -0,0 +1,389 @@
+#![allow(deprecated)]
+
+use std::{sync::Arc, vec};
+
+use crate::{internal_error, migrations::MutProgramInfo, modifying::*, visiting::*};
+use anyhow::{bail, Ok, Result};
+use sway_ast::{
+    assignable::ElementAccess, expr, Assignable, Expr, ItemFn, ItemStruct, StatementLet,
+};
+use sway_core::language::{
+    ty::{
+        TyExpression, TyExpressionVariant, TyFunctionDecl, TyReassignmentTarget, TyStructDecl,
+        TyVariableDecl,
+    },
+    CallPathType,
+};
+use sway_types::{Ident, Span, Spanned};
+
+use super::{ContinueMigrationProcess, DryRun, MigrationStep, MigrationStepKind};
+
+// NOTE: We assume idiomatic usage of the identifier `panic`. This means we support
+//       its migration only if it is used as a function name, struct field, or variable name.
+//       E.g., renaming `panic` in `struct panic { ... }` is not supported,
+//       as it is not an idiomatic usage.
+
+// NOTE: We don't have infrastructure in place for searching for usages of a symbol.
+//       Ideally, if we had it, we would use such infrastructure to rename symbol usages
+//       when its definition get renamed.
+//       Luckily, for this particular migration, it is sufficient to visit specific expression,
+//       like, e.g., function calls, and rename them.
+
+// NOTE: We don't support renaming modules named `panic`. The reason is that we have the `str`
+//       module in the standard library, signaling that using keywords as module names is acceptable.
+
+#[allow(dead_code)]
+pub(super) const RENAME_EXISTING_PANIC_IDENTIFIERS_TO_R_PANIC_STEP: MigrationStep = MigrationStep {
+    title: "Rename existing `panic` identifiers to `r#panic`",
+    duration: 0,
+    kind: MigrationStepKind::CodeModification(
+        rename_existing_panic_identifiers_to_r_panic_step,
+        &[],
+        ContinueMigrationProcess::IfNoManualMigrationActionsNeeded,
+    ),
+    help: &[
+        "Migration will rename existing `panic` identifiers in struct fields,",
+        "function names and arguments, and variable names to `r#panic`.",
+        " ",
+        "E.g., `let panic = 42;` will become `let r#panic = 42;`.",
+    ],
+};
+
+fn rename_existing_panic_identifiers_to_r_panic_step(
+    program_info: &mut MutProgramInfo,
+    dry_run: DryRun,
+) -> Result<Vec<Span>> {
+    struct Visitor;
+    impl TreesVisitorMut<Span> for Visitor {
+        fn visit_fn_decl(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_fn: &mut ItemFn,
+            _ty_fn: Option<Arc<TyFunctionDecl>>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            // First, let's check the arguments.
+            for lexed_arg in lexed_fn.fn_signature.arguments.inner.args_mut() {
+                let arg_name = match &mut lexed_arg.pattern {
+                    sway_ast::Pattern::Var { name, .. } => name,
+                    // A valid identifier in a function argument pattern can only be a variable,
+                    // never an enum variant. So we know that this `ident` is a variable.
+                    sway_ast::Pattern::AmbiguousSingleIdent(ident) => ident,
+                    _ => continue,
+                };
+
+                if arg_name.as_raw_ident_str() != "panic" {
+                    continue;
+                }
+
+                output.push(arg_name.span());
+
+                if ctx.dry_run == DryRun::Yes {
+                    continue;
+                }
+
+                *arg_name = Ident::new_with_raw(arg_name.span(), true);
+            }
+
+            // Then, the function name.
+            if lexed_fn.fn_signature.name.as_raw_ident_str() != "panic" {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            output.push(lexed_fn.fn_signature.name.span());
+
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            modify(lexed_fn).set_name("r#panic");
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_struct_decl(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_struct: &mut ItemStruct,
+            _ty_struct: Option<Arc<TyStructDecl>>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            for lexed_field in lexed_struct.fields.inner.iter_mut() {
+                let field_name = &mut lexed_field.value.name;
+
+                if field_name.as_raw_ident_str() != "panic" {
+                    continue;
+                }
+
+                output.push(field_name.span());
+
+                if ctx.dry_run == DryRun::Yes {
+                    continue;
+                }
+
+                *field_name = Ident::new_with_raw(field_name.span(), true);
+            }
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_fn_call(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_fn_call: &mut Expr,
+            ty_fn_call: Option<&TyExpression>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            // We report the occurrences only if it is not a dry-run.
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let Some(ty_fn_call) = ty_fn_call else {
+                // Without the typed function call, we cannot proceed
+                // because we cannot check if the function is actually defined in the current package.
+                return Ok(InvalidateTypedElement::No);
+            };
+
+            let Expr::FuncApp { func, args: _ } = lexed_fn_call else {
+                bail!(internal_error("`lexed_fn_call` is not an `Expr::FuncApp`."));
+            };
+
+            let Expr::Path(path) = &mut **func else {
+                // We are interested only in function calls that are paths.
+                // Only such calls can be renamed.
+                return Ok(InvalidateTypedElement::No);
+            };
+
+            let last_segment = path.last_segment_mut();
+
+            if last_segment.name.as_raw_ident_str() != "panic" {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            // Check if the function is actually defined in the current package.
+            let TyExpressionVariant::FunctionApplication { fn_ref, .. } = &ty_fn_call.expression
+            else {
+                bail!(internal_error(
+                    "`ty_fn_call` is not a `TyExpressionVariant::FunctionApplication`."
+                ));
+            };
+
+            let ty_fn = ctx.engines.de().get_function(fn_ref.id());
+            // We need the full path to the function to ensure it is defined in the current package.
+            if ty_fn.call_path.callpath_type != CallPathType::Full {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let Some(fn_pkg_name) = ty_fn.call_path.prefixes.first() else {
+                return Ok(InvalidateTypedElement::No);
+            };
+
+            if fn_pkg_name.as_str() != ctx.pkg_name {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            output.push(last_segment.span());
+
+            modify(last_segment).set_name("r#panic");
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_method_call(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_method_call: &mut Expr,
+            ty_method_call: Option<&TyExpression>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            // We report the occurrences only if it is not a dry-run.
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let Some(ty_method_call) = ty_method_call else {
+                // Without the typed method call, we cannot proceed
+                // because we cannot check if the method is actually defined in the current package.
+                return Ok(InvalidateTypedElement::No);
+            };
+
+            let Expr::MethodCall {
+                path_seg, args: _, ..
+            } = lexed_method_call
+            else {
+                bail!(internal_error(
+                    "`lexed_method_call` is not an `Expr::MethodCall`."
+                ));
+            };
+
+            if path_seg.name.as_raw_ident_str() != "panic" {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            // Check if the method is actually defined in the current package.
+            let TyExpressionVariant::FunctionApplication { fn_ref, .. } =
+                &ty_method_call.expression
+            else {
+                bail!(internal_error(
+                    "`ty_method_call` is not a `TyExpressionVariant::FunctionApplication`."
+                ));
+            };
+
+            let ty_method = ctx.engines.de().get_function(fn_ref.id());
+            // We need the full path to the function to ensure it is defined in the current package.
+            if ty_method.call_path.callpath_type != CallPathType::Full {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let Some(fn_pkg_name) = ty_method.call_path.prefixes.first() else {
+                return Ok(InvalidateTypedElement::No);
+            };
+
+            if fn_pkg_name.as_str() != ctx.pkg_name {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            output.push(path_seg.span());
+
+            modify(path_seg).set_name("r#panic");
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_statement_let(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_let: &mut StatementLet,
+            _ty_var_decl: Option<&TyVariableDecl>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            let var_name = match &mut lexed_let.pattern {
+                sway_ast::Pattern::Var { name, .. } => name,
+                // A valid identifier in a variable name pattern can only be a variable,
+                // never an enum variant. So we know that this `ident` is a variable.
+                sway_ast::Pattern::AmbiguousSingleIdent(ident) => ident,
+                _ => {
+                    // NOTE: We don't support renaming `panic` in patterns other than variables,
+                    //       e.g., in deconstruction patterns.
+                    return Ok(InvalidateTypedElement::No);
+                }
+            };
+
+            if var_name.as_raw_ident_str() != "panic" {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            output.push(var_name.span());
+
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            *var_name = Ident::new_with_raw(var_name.span(), true);
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_expr(
+            &mut self,
+            ctx: &VisitingContext,
+            lexed_expr: &mut Expr,
+            _ty_expr: Option<&TyExpression>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            // We report the occurrences only if it is not a dry-run.
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let var_names = match lexed_expr {
+                Expr::Path(path) if path.suffix.is_empty() => vec![&mut path.prefix.name],
+                Expr::Struct { fields, .. } => fields
+                    .inner
+                    .iter_mut()
+                    .map(|field| &mut field.field_name)
+                    .collect(),
+                Expr::FieldProjection { name, .. } => vec![name],
+                _ => vec![],
+            };
+
+            for var_name in var_names
+                .into_iter()
+                .filter(|n| n.as_raw_ident_str() == "panic")
+            {
+                output.push(var_name.span());
+
+                *var_name = Ident::new_with_raw(var_name.span(), true);
+            }
+
+            Ok(InvalidateTypedElement::No)
+        }
+
+        fn visit_reassignment(
+            &mut self,
+            ctx: &VisitingContext,
+            _lexed_op: &mut expr::ReassignmentOp,
+            lexed_lhs: &mut Assignable,
+            _ty_lhs: Option<&TyReassignmentTarget>,
+            _lexed_rhs: &mut Expr,
+            _ty_rhs: Option<&TyExpression>,
+            output: &mut Vec<Span>,
+        ) -> Result<InvalidateTypedElement> {
+            // On the LHS, we support renaming `panic` only in these cases:
+            // - Variable names, e.g., `let panic = 42;`
+            // - Single field access, e.g., `let x.panic = 42;`
+            // But occurrences in, e.g., `foo[panic].x = 42;` will not be renamed.
+            // Full traversal of reassignments' LHS will be done as a part of migration
+            // infrastructure in the future.
+
+            // We report the occurrences only if it is not a dry-run.
+            if ctx.dry_run == DryRun::Yes {
+                return Ok(InvalidateTypedElement::No);
+            }
+
+            let var_names = match lexed_lhs {
+                Assignable::ElementAccess(element_access) => match element_access {
+                    ElementAccess::Var(name) => vec![name],
+                    ElementAccess::FieldProjection {
+                        target: element_access,
+                        name,
+                        ..
+                    } => {
+                        let mut names = vec![name];
+                        if let ElementAccess::Var(name) = &mut **element_access {
+                            names.push(name)
+                        };
+                        names
+                    }
+                    ElementAccess::TupleFieldProjection {
+                        target: element_access,
+                        ..
+                    }
+                    | ElementAccess::Index {
+                        target: element_access,
+                        ..
+                    } => match &mut **element_access {
+                        ElementAccess::Var(name) => vec![name],
+                        _ => vec![],
+                    },
+                    _ => vec![],
+                },
+                Assignable::Deref { .. } => vec![],
+            };
+
+            for var_name in var_names
+                .into_iter()
+                .filter(|n| n.as_raw_ident_str() == "panic")
+            {
+                output.push(var_name.span());
+
+                *var_name = Ident::new_with_raw(var_name.span(), true);
+            }
+
+            Ok(InvalidateTypedElement::No)
+        }
+    }
+
+    ProgramVisitorMut::visit_program(program_info, dry_run, &mut Visitor {})
+}

forc-plugins/forc-migrate/src/migrations/mod.rs

@@ -9,6 +9,7 @@
 //! the migration tool.
 
 mod demo;
+mod error_type;
 mod merge_core_std;
 mod partial_eq;
 mod references;
@@ -34,6 +35,8 @@ use sway_types::Span;
 use crate::internal_error;
 
 pub(crate) struct ProgramInfo<'a> {
+    /// The name of the current package being migrated.
+    pub pkg_name: String,
     pub lexed_program: Arc<LexedProgram>,
     pub ty_program: Arc<TyProgram>,
     pub engines: &'a Engines,
@@ -44,6 +47,8 @@ pub(crate) struct ProgramInfo<'a> {
 /// [TyProgram] and the [Engines]. It is used in migrations
 /// that modify the source code by altering the lexed program.
 pub(crate) struct MutProgramInfo<'a> {
+    /// The name of the current package being migrated.
+    pub pkg_name: &'a str,
     pub lexed_program: &'a mut LexedProgram,
     pub ty_program: &'a TyProgram,
     pub engines: &'a Engines,
@@ -52,7 +57,11 @@ pub(crate) struct MutProgramInfo<'a> {
 impl ProgramInfo<'_> {
     pub(crate) fn as_mut(&mut self) -> MutProgramInfo {
         MutProgramInfo {
-            lexed_program: Arc::get_mut(&mut self.lexed_program).unwrap(),
+            pkg_name: &self.pkg_name,
+            // Because the `ProgramsCacheEntry` clones the `programs`, the compilation will always
+            // result in two strong `Arc` references to the `lexed_program`.
+            // Therefore, we must use `Arc::make_mut` to get the copy-on-write behavior.
+            lexed_program: Arc::make_mut(&mut self.lexed_program),
             ty_program: &self.ty_program,
             engines: self.engines,
         }
@@ -485,4 +494,7 @@ fn assert_migration_steps_consistency(migration_steps: MigrationSteps) {
 
 /// The list of the migration steps, grouped by the Sway feature that causes
 /// the breaking changes behind the migration steps.
-const MIGRATION_STEPS: MigrationSteps = &[];
+const MIGRATION_STEPS: MigrationSteps = &[(
+    Feature::ErrorType,
+    &[error_type::RENAME_EXISTING_PANIC_IDENTIFIERS_TO_R_PANIC_STEP],
+)];

forc-plugins/forc-migrate/src/visiting/mod.rs

@@ -7,9 +7,10 @@ use std::sync::Arc;
 
 use duplicate::duplicate_item;
 use sway_ast::{
-    expr::{LoopControlFlow, ReassignmentOpVariant},
-    CodeBlockContents, Expr, IfCondition, IfExpr, ItemFn, ItemImpl, ItemImplItem, ItemKind,
-    ItemUse, PathExprSegment, Statement, StatementLet,
+    expr::{LoopControlFlow, ReassignmentOp, ReassignmentOpVariant},
+    keywords::*,
+    Assignable, CodeBlockContents, Expr, IfCondition, IfExpr, ItemFn, ItemImpl, ItemImplItem,
+    ItemKind, ItemStruct, ItemUse, PathExprSegment, Statement, StatementLet,
 };
 use sway_core::{
     decl_engine::DeclEngine,
@@ -17,8 +18,9 @@ use sway_core::{
         lexed::LexedModule,
         ty::{
             self, TyAstNodeContent, TyCodeBlock, TyDecl, TyExpression, TyExpressionVariant,
-            TyFunctionDecl, TyImplSelfOrTrait, TyIntrinsicFunctionKind, TyModule, TySideEffect,
-            TySideEffectVariant, TyTraitItem, TyUseStatement, TyVariableDecl,
+            TyFunctionDecl, TyImplSelfOrTrait, TyIntrinsicFunctionKind, TyModule,
+            TyReassignmentTarget, TySideEffect, TySideEffectVariant, TyStructDecl, TyTraitItem,
+            TyUseStatement, TyVariableDecl,
         },
         CallPath,
     },
@@ -32,6 +34,8 @@ use crate::{
 };
 
 pub(crate) struct VisitingContext<'a> {
+    /// The name of the current package being migrated.
+    pub pkg_name: &'a str,
     pub engines: &'a Engines,
     pub dry_run: DryRun,
 }
@@ -97,7 +101,16 @@ pub(crate) trait __TreesVisitor<O> {
     ) -> Result<InvalidateTypedElement> {
         Ok(InvalidateTypedElement::No)
     }
-    fn visit_fn(
+    fn visit_struct_decl(
+        &mut self,
+        ctx: &VisitingContext,
+        lexed_struct: __ref_type([ItemStruct]),
+        ty_struct: Option<Arc<TyStructDecl>>,
+        output: &mut Vec<O>,
+    ) -> Result<InvalidateTypedElement> {
+        Ok(InvalidateTypedElement::No)
+    }
+    fn visit_fn_decl(
         &mut self,
         ctx: &VisitingContext,
         lexed_fn: __ref_type([ItemFn]),
@@ -169,6 +182,32 @@ pub(crate) trait __TreesVisitor<O> {
     ) -> Result<InvalidateTypedElement> {
         Ok(InvalidateTypedElement::No)
     }
+    #[allow(clippy::too_many_arguments)]
+    fn visit_reassignment(
+        &mut self,
+        ctx: &VisitingContext,
+        lexed_op: __ref_type([ReassignmentOp]),
+        lexed_lhs: __ref_type([Assignable]),
+        ty_lhs: Option<&TyReassignmentTarget>,
+        lexed_rhs: __ref_type([Expr]),
+        ty_rhs: Option<&TyExpression>,
+        output: &mut Vec<O>,
+    ) -> Result<InvalidateTypedElement> {
+        Ok(InvalidateTypedElement::No)
+    }
+    #[allow(clippy::too_many_arguments)]
+    fn visit_binary_op(
+        &mut self,
+        ctx: &VisitingContext,
+        op: &'static str,
+        lexed_lhs: __ref_type([Expr]),
+        ty_lhs: Option<&TyExpression>,
+        lexed_rhs: __ref_type([Expr]),
+        ty_rhs: Option<&TyExpression>,
+        output: &mut Vec<O>,
+    ) -> Result<InvalidateTypedElement> {
+        Ok(InvalidateTypedElement::No)
+    }
 }
 
 #[allow(dead_code)]
@@ -190,6 +229,8 @@ impl __ProgramVisitor {
         V: __TreesVisitor<O>,
     {
         let ctx = VisitingContext {
+            #[allow(clippy::needless_borrow)] // Clippy lint false positive. Actually, a Clippy bug.
+            pkg_name: &program_info.pkg_name,
             engines: program_info.engines,
             dry_run,
         };
@@ -257,15 +298,35 @@ impl __ProgramVisitor {
                             .find_map(|node| match &node.content {
                                 TyAstNodeContent::SideEffect(TySideEffect {
                                     side_effect: TySideEffectVariant::UseStatement(ty_use),
-                                }) => Some(ty_use),
+                                }) if ty_use.span == item_use.span() => Some(ty_use),
                                 _ => None,
                             })
                     });
 
                     visitor.visit_use(ctx, item_use, ty_use, output)?;
                 }
-                ItemKind::Struct(_item_struct) => {
-                    // TODO: Implement visiting `struct`.
+                ItemKind::Struct(item_struct) => {
+                    let ty_struct_decl = ty_module.and_then(|ty_module| {
+                        ty_module
+                            .all_nodes
+                            .iter()
+                            .find_map(|node| match &node.content {
+                                TyAstNodeContent::Declaration(TyDecl::StructDecl(
+                                    ty_struct_decl,
+                                )) => {
+                                    let ty_struct_decl =
+                                        ctx.engines.de().get_struct(&ty_struct_decl.decl_id);
+                                    if ty_struct_decl.span == item_struct.span() {
+                                        Some(ty_struct_decl)
+                                    } else {
+                                        None
+                                    }
+                                }
+                                _ => None,
+                            })
+                    });
+
+                    visitor.visit_struct_decl(ctx, item_struct, ty_struct_decl, output)?;
                 }
                 ItemKind::Enum(_item_enum) => {
                     // TODO: Implement visiting `enum`.
@@ -288,7 +349,7 @@ impl __ProgramVisitor {
                             })
                     });
 
-                    Self::visit_fn(ctx, item_fn, ty_fn, visitor, output)?;
+                    Self::visit_fn_decl(ctx, item_fn, ty_fn, visitor, output)?;
                 }
                 ItemKind::Trait(_item_trait) => {
                     // TODO: Implement visiting `trait`.
@@ -335,7 +396,7 @@ impl __ProgramVisitor {
         Ok(())
     }
 
-    fn visit_fn<V, O>(
+    fn visit_fn_decl<V, O>(
         ctx: &VisitingContext,
         lexed_fn: __ref_type([ItemFn]),
         ty_fn: Option<Arc<TyFunctionDecl>>,
@@ -345,7 +406,7 @@ impl __ProgramVisitor {
     where
         V: __TreesVisitor<O>,
     {
-        let ty_fn = match visitor.visit_fn(ctx, lexed_fn, ty_fn.clone(), output)? {
+        let ty_fn = match visitor.visit_fn_decl(ctx, lexed_fn, ty_fn.clone(), output)? {
             InvalidateTypedElement::Yes => None,
             InvalidateTypedElement::No => ty_fn,
         };
@@ -393,7 +454,7 @@ impl __ProgramVisitor {
                         })
                     });
 
-                    Self::visit_fn(ctx, item_fn, ty_item_fn, visitor, output)?;
+                    Self::visit_fn_decl(ctx, item_fn, ty_item_fn, visitor, output)?;
                 }
                 ItemImplItem::Const(_item_const) => {
                     // TODO: Implement visiting `associated consts`.
@@ -517,6 +578,38 @@ impl __ProgramVisitor {
         Ok(())
     }
 
+    fn visit_binary_op<V, O>(
+        ctx: &VisitingContext,
+        op: &'static str,
+        lexed_lhs: __ref_type([Expr]),
+        lexed_rhs: __ref_type([Expr]),
+        visitor: &mut V,
+        output: &mut Vec<O>,
+    ) -> Result<()>
+    where
+        V: __TreesVisitor<O>,
+    {
+        // TODO: Implement getting typed LHS and RHS when visiting operands' expressions.
+        //       We need to properly handle the desugaring.
+        //       E.g., `x + func(1, 2);`
+        //       will be desugared into `add(x, func(1, 2));`
+        //       When visiting the operands in the lexed tree, in the typed tree
+        //       we need to skip the operator method call, like `add` in the above example,
+        //       and provide the typed arguments instead.
+        let ty_lhs = None;
+        let ty_rhs = None;
+
+        match visitor.visit_binary_op(ctx, op, lexed_lhs, ty_lhs, lexed_rhs, ty_rhs, output)? {
+            InvalidateTypedElement::No => (ty_lhs, ty_rhs),
+            InvalidateTypedElement::Yes => (None, None),
+        };
+
+        Self::visit_expr(ctx, lexed_lhs, ty_lhs, visitor, output)?;
+        Self::visit_expr(ctx, lexed_rhs, ty_rhs, visitor, output)?;
+
+        Ok(())
+    }
+
     fn visit_expr<V, O>(
         ctx: &VisitingContext,
         lexed_expr: __ref_type([Expr]),
@@ -550,9 +643,7 @@ impl __ProgramVisitor {
             Expr::Error(..) => {
                 bail!(internal_error("`Expr::Error` cannot happen, because `forc migrate` analyzes only successfully compiled programs."));
             }
-            Expr::Path(_path_expr) => {
-                // TODO: Implement visiting `path_expr`.
-            }
+            Expr::Path(_path_expr) => {}
             Expr::Literal(_literal) => {}
             Expr::AbiCast { args, .. } => {
                 let ty_abi_cast_expr = ty_expr
@@ -690,8 +781,13 @@ impl __ProgramVisitor {
 
                 Self::visit_args(ctx, lexed_expr, ty_expr, false, visitor, output)?;
             }
-            Expr::Index { target: _, arg: _ } => {
-                // TODO: Implement visiting `array[index]`.
+            Expr::Index { target, arg } => {
+                // TODO: Implement extracting typed elements for `array[index]`.
+                let ty_target = None;
+                let ty_arg = None;
+
+                Self::visit_expr(ctx, target.__as_ref(), ty_target, visitor, output)?;
+                Self::visit_expr(ctx, arg.inner.__as_ref(), ty_arg, visitor, output)?;
             }
             Expr::MethodCall {
                 target: _,
@@ -711,19 +807,25 @@ impl __ProgramVisitor {
                 Self::visit_args(ctx, lexed_expr, ty_expr, true, visitor, output)?;
             }
             Expr::FieldProjection {
-                target: _,
+                target,
                 dot_token: _,
                 name: _,
             } => {
-                // TODO: Implement visiting `struct.field`.
+                // TODO: Implement extracting typed target for `struct.field`.
+                let ty_target = None;
+
+                Self::visit_expr(ctx, target.__as_ref(), ty_target, visitor, output)?;
             }
             Expr::TupleFieldProjection {
-                target: _,
+                target,
                 dot_token: _,
                 field: _,
                 field_span: _,
             } => {
-                // TODO: Implement visiting `tuple.0`.
+                // TODO: Implement extracting typed target for `tuple.N`.
+                let ty_target = None;
+
+                Self::visit_expr(ctx, target.__as_ref(), ty_target, visitor, output)?;
             }
             Expr::Ref {
                 ampersand_token: _,
@@ -740,174 +842,327 @@ impl __ProgramVisitor {
             }
             Expr::Not {
                 bang_token: _,
-                expr: _,
+                expr,
             } => {
-                // TODO: Implement visiting `not`.
+                // TODO: Implement extracting typed expressions when visiting `not`.
+                let ty_expr = None;
+                Self::visit_expr(ctx, expr.__as_ref(), ty_expr, visitor, output)?;
             }
             Expr::Mul {
-                lhs: _,
+                lhs,
                 star_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `mul`.
+                Self::visit_binary_op(
+                    ctx,
+                    <StarToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Div {
-                lhs: _,
+                lhs,
                 forward_slash_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `div`.
+                Self::visit_binary_op(
+                    ctx,
+                    <ForwardSlashToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Pow {
-                lhs: _,
+                lhs,
                 double_star_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `pow`.
+                Self::visit_binary_op(
+                    ctx,
+                    <DoubleStarToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Modulo {
-                lhs: _,
+                lhs,
                 percent_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `modulo`.
+                Self::visit_binary_op(
+                    ctx,
+                    <PercentToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Add {
-                lhs: _,
+                lhs,
                 add_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `add`.
+                Self::visit_binary_op(
+                    ctx,
+                    <AddToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Sub {
-                lhs: _,
+                lhs,
                 sub_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `sub`.
+                Self::visit_binary_op(
+                    ctx,
+                    <SubToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Shl {
-                lhs: _,
+                lhs,
                 shl_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `<<`.
+                Self::visit_binary_op(
+                    ctx,
+                    <ShlToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Shr {
-                lhs: _,
+                lhs,
                 shr_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `>>`.
+                Self::visit_binary_op(
+                    ctx,
+                    <ShrToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::BitAnd {
-                lhs: _,
+                lhs,
                 ampersand_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `bit_and`.
+                Self::visit_binary_op(
+                    ctx,
+                    <AmpersandToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::BitXor {
-                lhs: _,
+                lhs,
                 caret_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `bit_xor`.
+                Self::visit_binary_op(
+                    ctx,
+                    <CaretToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::BitOr {
-                lhs: _,
+                lhs,
                 pipe_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `bit_or`.
+                Self::visit_binary_op(
+                    ctx,
+                    <PipeToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Equal {
-                lhs: _,
+                lhs,
                 double_eq_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `==`.
+                Self::visit_binary_op(
+                    ctx,
+                    <DoubleEqToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::NotEqual {
-                lhs: _,
+                lhs,
                 bang_eq_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `!=`.
+                Self::visit_binary_op(
+                    ctx,
+                    <BangEqToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::LessThan {
-                lhs: _,
+                lhs,
                 less_than_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `<`.
+                Self::visit_binary_op(
+                    ctx,
+                    <LessThanToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::GreaterThan {
-                lhs: _,
+                lhs,
                 greater_than_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `>`.
+                Self::visit_binary_op(
+                    ctx,
+                    <GreaterThanToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::LessThanEq {
-                lhs: _,
+                lhs,
                 less_than_eq_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `<=`.
+                Self::visit_binary_op(
+                    ctx,
+                    <LessThanEqToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::GreaterThanEq {
-                lhs: _,
+                lhs,
                 greater_than_eq_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `>=`.
+                Self::visit_binary_op(
+                    ctx,
+                    <GreaterThanEqToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::LogicalAnd {
-                lhs: _,
+                lhs,
                 double_ampersand_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `logical_and`.
+                Self::visit_binary_op(
+                    ctx,
+                    <DoubleAmpersandToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::LogicalOr {
-                lhs: _,
+                lhs,
                 double_pipe_token: _,
-                rhs: _,
+                rhs,
             } => {
-                // TODO: Implement visiting `logical_or`.
+                Self::visit_binary_op(
+                    ctx,
+                    <DoublePipeToken as Token>::AS_STR,
+                    lhs.__as_ref(),
+                    rhs.__as_ref(),
+                    visitor,
+                    output,
+                )?;
             }
             Expr::Reassignment {
-                assignable: _,
+                assignable,
                 reassignment_op,
                 expr,
             } => {
+                let ty_reassignment = ty_expr
+                    .map(|ty_expr| match &ty_expr.expression {
+                        ty::TyExpressionVariant::Reassignment(ty_reassignment) => {
+                            Ok(ty_reassignment.as_ref())
+                        }
+                        _ => bail!(invalid_ty_expression_variant(
+                            "Reassignment",
+                            "Reassignment"
+                        )),
+                    })
+                    .transpose()?;
+                let ty_lhs = ty_reassignment.map(|ty_reassignment| &ty_reassignment.lhs);
+                let ty_rhs = ty_reassignment.map(|ty_reassignment| &ty_reassignment.rhs);
+
+                let (_ty_lhs, ty_rhs) = match visitor.visit_reassignment(
+                    ctx,
+                    reassignment_op,
+                    assignable,
+                    ty_lhs,
+                    expr.__as_ref(),
+                    ty_rhs,
+                    output,
+                )? {
+                    InvalidateTypedElement::Yes => (None, None),
+                    InvalidateTypedElement::No => (ty_lhs, ty_rhs),
+                };
+
                 match reassignment_op.variant {
                     ReassignmentOpVariant::Equals => {
-                        let ty_reassignment_rhs = ty_expr
-                            .map(|ty_expr| match &ty_expr.expression {
-                                ty::TyExpressionVariant::Reassignment(ty_reassignment) => {
-                                    Ok(&ty_reassignment.rhs)
-                                }
-                                _ => bail!(invalid_ty_expression_variant(
-                                    "Reassignment",
-                                    "Reassignment"
-                                )),
-                            })
-                            .transpose()?;
-
-                        Self::visit_expr(ctx, expr, ty_reassignment_rhs, visitor, output)?;
+                        // TODO: Implement visiting expressions in the reassignment LHS.
+                        Self::visit_expr(ctx, expr, ty_rhs, visitor, output)?;
                     }
                     _ => {
                         // TODO: Implement getting `ty_expr` when visiting `compound reassignments`.
                         //       We need to properly handle the desugaring.
                         //       E.g., `x += func(1, 2);`
-                        //       will be desugared into `x = x.add(func(1, 2));`
+                        //       will be desugared into `x = add(x, func(1, 2));`
                         //       When visiting the RHS in the lexed tree, we need to skip the
-                        //       operator method call in the typed tree.
+                        //       operator method call in the typed tree, and provide the
+                        //       typed arguments instead.
                         //       To provide visiting without loosing the information about compound
                         //       reassignment, we will need to have a dedicated `visit_reassignment`
                         //       method.
+                        // TODO: Implement visiting expressions in the reassignment LHS.
                         Self::visit_expr(ctx, expr, None, visitor, output)?;
                     }
                 }
@@ -929,8 +1184,8 @@ impl __ProgramVisitor {
     where
         V: __TreesVisitor<O>,
     {
-        match &lexed_if.condition {
-            IfCondition::Expr(_) => {
+        match __ref([lexed_if.condition]) {
+            IfCondition::Expr(lexed_condition) => {
                 let ty_if = ty_if_expr
                     .map(|ty_expr| match &ty_expr.expression {
                         ty::TyExpressionVariant::IfExp {
@@ -945,7 +1200,7 @@ impl __ProgramVisitor {
                         _ => bail!(invalid_ty_expression_variant("IfExpr", "If")),
                     })
                     .transpose()?;
-                let _ty_if_condition = ty_if.map(|ty_if| ty_if.0);
+                let ty_if_condition = ty_if.map(|ty_if| ty_if.0);
                 let ty_if_then = ty_if
                     .map(|ty_if| match &ty_if.1.expression {
                         ty::TyExpressionVariant::CodeBlock(ty_code_block) => Ok(ty_code_block),
@@ -957,7 +1212,7 @@ impl __ProgramVisitor {
                     .transpose()?;
                 let ty_if_else = ty_if.and_then(|ty_if| ty_if.2);
 
-                // TODO: Implement visiting `if_condition`.
+                visitor.visit_expr(ctx, lexed_condition.__as_ref(), ty_if_condition, output)?;
 
                 Self::visit_block(
                     ctx,

sway-ast/src/item/mod.rs

@@ -131,6 +131,38 @@ pub enum FnArgs {
     },
 }
 
+impl FnArgs {
+    /// Returns all the [FnArg]s, from the function signature defined by `self`.
+    ///
+    /// If the `self` is [FnArgs::NonStatic], the first `self` argument is not
+    /// returned, because it is not an [FnArg].
+    pub fn args(&self) -> Vec<&FnArg> {
+        match self {
+            Self::Static(punctuated) => punctuated.iter().collect(),
+            Self::NonStatic { args_opt, .. } => args_opt
+                .as_ref()
+                .map_or(vec![], |(_comma_token, punctuated)| {
+                    punctuated.iter().collect()
+                }),
+        }
+    }
+
+    /// Returns all the [FnArg]s, from the function signature defined by `self`.
+    ///
+    /// If the `self` is [FnArgs::NonStatic], the first `self` argument is not
+    /// returned, because it is not an [FnArg].
+    pub fn args_mut(&mut self) -> Vec<&mut FnArg> {
+        match self {
+            Self::Static(punctuated) => punctuated.iter_mut().collect(),
+            Self::NonStatic { args_opt, .. } => args_opt
+                .as_mut()
+                .map_or(vec![], |(_comma_token, punctuated)| {
+                    punctuated.iter_mut().collect()
+                }),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Serialize)]
 pub struct FnArg {
     pub pattern: Pattern,

swayfmt/src/utils/language/expr/struct_field.rs

@@ -22,7 +22,7 @@ impl Format for ExprStructField {
         formatted_code: &mut FormattedCode,
         formatter: &mut Formatter,
     ) -> Result<(), FormatterError> {
-        write!(formatted_code, "{}", self.field_name.as_str())?;
+        write!(formatted_code, "{}", self.field_name.as_raw_ident_str())?;
         if let Some((_colon_token, expr)) = &self.expr_opt {
             formatter.with_shape(
                 formatter

swayfmt/src/utils/language/punctuated.rs

@@ -213,7 +213,7 @@ impl Format for TypeField {
         write!(
             formatted_code,
             "{}{} ",
-            self.name.as_str(),
+            self.name.as_raw_ident_str(),
             ColonToken::AS_STR,
         )?;
         self.ty.format(formatted_code, formatter)?;