use correct location for use statements

Author: giacomocavalieri

compiler-core/src/ast.rs

@@ -3569,7 +3569,7 @@ impl UntypedStatement {
         match self {
             Statement::Expression(expression) => expression.location(),
             Statement::Assignment(assignment) => assignment.location,
-            Statement::Use(use_) => use_.location,
+            Statement::Use(use_) => use_.location.merge(&use_.call.location()),
             Statement::Assert(assert) => assert.location,
         }
     }
@@ -3598,7 +3598,9 @@ impl TypedStatement {
         match self {
             Statement::Expression(expression) => expression.location(),
             Statement::Assignment(assignment) => assignment.location,
-            Statement::Use(use_) => use_.location,
+            // A use statement covers the entire block: `use_.location` covers
+            // just the use's first line and not what comes after it.
+            Statement::Use(use_) => use_.location.merge(&use_.call.location()),
             Statement::Assert(assert) => assert.location,
         }
     }

compiler-core/src/language_server/code_action.rs

@@ -8626,12 +8626,7 @@ impl<'a> ExtractFunction<'a> {
                 let location = statements
                     .first()
                     .location()
-                    .merge(&match &statements.last() {
-                        ast::Statement::Expression(expression) => expression.location(),
-                        ast::Statement::Assignment(assignment) => assignment.location,
-                        ast::Statement::Assert(assert) => assert.location,
-                        ast::Statement::Use(use_) => use_.call.location(),
-                    });
+                    .merge(&statements.last().location());
 
                 self.extract_code_in_tail_position(
                     *full_location,
@@ -8972,12 +8967,7 @@ impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
         statements: &'ast [TypedStatement],
     ) {
         let last_statement_location = self.last_statement_location;
-        self.last_statement_location = statements.last().map(|last| match last {
-            ast::Statement::Expression(_)
-            | ast::Statement::Assignment(_)
-            | ast::Statement::Assert(_) => last.location(),
-            ast::Statement::Use(use_) => last.location().merge(&use_.call.location()),
-        });
+        self.last_statement_location = statements.last().map(|last| last.location());
 
         ast::visit::visit_typed_expr_block(self, location, statements);
 
@@ -9002,12 +8992,7 @@ impl<'ast> ast::visit::Visit<'ast> for ExtractFunction<'ast> {
     }
 
     fn visit_typed_statement(&mut self, statement: &'ast TypedStatement) {
-        let statement_location = match statement {
-            ast::Statement::Expression(expression) => expression.location(),
-            ast::Statement::Assignment(assignment) => assignment.location,
-            ast::Statement::Assert(assert) => assert.location,
-            ast::Statement::Use(use_) => use_.location.merge(&use_.call.location()),
-        };
+        let statement_location = statement.location();
 
         if self.can_extract(statement_location) {
             let is_in_tail_position =