[DAS] Fixes part file handling for some create fixes

Bug: https://github.com/dart-lang/sdk/issues/61192
Change-Id: I8aa765fca2d2948ab6fe1b47a6bff6bf58de3cd3
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/447341
Auto-Submit: Felipe Morschel <git@fmorschel.dev>
Reviewed-by: Samuel Rawlins <srawlins@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Samuel Rawlins <srawlins@google.com>

pkg/analysis_server/lib/src/services/correction/dart/create_extension_member.dart

@@ -33,6 +33,7 @@ class CreateExtensionGetter extends _CreateExtensionMember {
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
+    var addStaticKeyword = inStaticContext;
     var nameNode = node;
     if (nameNode is! SimpleIdentifier) {
       return;
@@ -44,29 +45,34 @@ class CreateExtensionGetter extends _CreateExtensionMember {
     _getterName = nameNode.name;
 
     // prepare target
-    DartType? targetType;
-    ExtensionElement? extensionElement;
+    Expression? target;
     switch (nameNode.parent) {
       case PrefixedIdentifier prefixedIdentifier:
         if (prefixedIdentifier.identifier == nameNode) {
-          targetType = prefixedIdentifier.prefix.staticType;
+          target = prefixedIdentifier.prefix;
         }
       case PropertyAccess propertyAccess:
         if (propertyAccess.propertyName == nameNode) {
-          targetType = propertyAccess.realTarget.staticType;
-          if (propertyAccess.realTarget case ExtensionOverride(
-            :var element,
-          ) when targetType == null) {
-            extensionElement = element;
-            targetType = extensionElement.thisType;
-          }
+          target = propertyAccess.realTarget;
         }
-      case ExpressionFunctionBody expressionFunctionBody:
-        if (expressionFunctionBody.expression == nameNode) {
-          targetType = node.enclosingInstanceElement?.thisType;
-        }
-      default:
-        targetType = nameNode.enclosingInstanceElement?.thisType;
+    }
+
+    DartType? targetType;
+    ExtensionElement? extensionElement;
+    if (target is ExtensionOverride) {
+      targetType = target.extendedType;
+      extensionElement = target.element;
+    } else if (target == null) {
+      extensionElement = node.enclosingInstanceElement?.ifTypeOrNull();
+      targetType = extensionElement?.thisType;
+    } else {
+      // We need the type for the extension.
+      targetType = target.staticType;
+    }
+    if (targetType == null && target is SimpleIdentifier) {
+      extensionElement = target.element?.ifTypeOrNull();
+      targetType = extensionElement?.thisType;
+      addStaticKeyword = true;
     }
 
     // TODO(FMorschel): We should take into account if the target type contains
@@ -86,7 +92,7 @@ class CreateExtensionGetter extends _CreateExtensionMember {
     }
 
     void writeGetter(DartEditBuilder builder) {
-      if (inStaticContext) {
+      if (addStaticKeyword) {
         builder.write('static ');
       }
       if (fieldType != null) {
@@ -151,7 +157,7 @@ class CreateExtensionMethod extends _CreateExtensionMember {
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
-    var static = false;
+    var addStaticKeyword = inStaticContext;
     var nameNode = node;
     if (nameNode is! SimpleIdentifier) {
       return;
@@ -191,7 +197,7 @@ class CreateExtensionMethod extends _CreateExtensionMember {
     if (targetType == null && target is SimpleIdentifier) {
       extensionElement = target.element?.ifTypeOrNull();
       targetType = extensionElement?.thisType;
-      static = true;
+      addStaticKeyword = true;
     }
     if (targetType == null ||
         targetType is DynamicType ||
@@ -224,7 +230,7 @@ class CreateExtensionMethod extends _CreateExtensionMember {
     var functionType = !isInvocation ? returnType as FunctionType : null;
 
     void writeMethod(DartEditBuilder builder) {
-      if (static) {
+      if (addStaticKeyword) {
         builder.write('static ');
       }
 
@@ -356,7 +362,18 @@ class CreateExtensionOperator extends _CreateExtensionMember {
     }
 
     // We need the type for the extension.
-    var targetType = target.staticType;
+    DartType? targetType;
+    ExtensionElement? extensionElement;
+    if (target is ExtensionOverride) {
+      targetType = target.extendedType;
+      extensionElement = target.element;
+    } else {
+      targetType = target.staticType;
+    }
+    if (targetType == null && target is SimpleIdentifier) {
+      extensionElement = target.element?.ifTypeOrNull();
+      targetType = extensionElement?.thisType;
+    }
     if (targetType == null ||
         targetType is DynamicType ||
         targetType is InvalidType) {
@@ -406,14 +423,27 @@ class CreateExtensionOperator extends _CreateExtensionMember {
       builder.write(') {}');
     }
 
-    var updatedExisting = await _updateExistingExtension(builder, targetType, (
-      extension,
-      builder,
-    ) {
-      builder.insertMethod(extension, (builder) {
-        writeMethod(builder);
+    bool updatedExisting;
+    if (extensionElement != null) {
+      updatedExisting = await _updateExistingExtension2(
+        builder,
+        extensionElement,
+        (extension, builder) {
+          builder.insertMethod(extension, (builder) {
+            writeMethod(builder);
+          });
+        },
+      );
+    } else {
+      updatedExisting = await _updateExistingExtension(builder, targetType, (
+        extension,
+        builder,
+      ) {
+        builder.insertMethod(extension, (builder) {
+          writeMethod(builder);
+        });
       });
-    });
+    }
     if (updatedExisting) {
       return;
     }
@@ -441,6 +471,7 @@ class CreateExtensionSetter extends _CreateExtensionMember {
 
   @override
   Future<void> compute(ChangeBuilder builder) async {
+    var addStaticKeyword = inStaticContext;
     var nameNode = node;
     if (nameNode is! SimpleIdentifier) {
       return;
@@ -463,14 +494,28 @@ class CreateExtensionSetter extends _CreateExtensionMember {
           target = propertyAccess.realTarget;
         }
     }
-    if (target == null) {
-      return;
+
+    DartType? targetType;
+    ExtensionElement? extensionElement;
+    if (target is ExtensionOverride) {
+      targetType = target.extendedType;
+      extensionElement = target.element;
+    } else if (target == null) {
+      extensionElement = node.enclosingInstanceElement?.ifTypeOrNull();
+      targetType = extensionElement?.thisType;
+    } else {
+      // We need the type for the extension.
+      targetType = target.staticType;
+    }
+    if (targetType == null && target is SimpleIdentifier) {
+      extensionElement = target.element?.ifTypeOrNull();
+      targetType = extensionElement?.thisType;
+      addStaticKeyword = true;
     }
 
     // TODO(FMorschel): We should take into account if the target type contains
     // a setter for the same name and stop the fix from being applied.
     // We need the type for the extension.
-    var targetType = target.staticType;
     if (targetType == null ||
         targetType is DynamicType ||
         targetType is InvalidType) {
@@ -489,19 +534,33 @@ class CreateExtensionSetter extends _CreateExtensionMember {
         _setterName,
         nameGroupName: 'NAME',
         parameterType: fieldType,
+        isStatic: addStaticKeyword,
         parameterTypeGroupName: 'TYPE',
         typeParametersInScope: methodBeingCopied?.typeParameters,
       );
     }
 
-    var updatedExisting = await _updateExistingExtension(builder, targetType, (
-      extension,
-      builder,
-    ) {
-      builder.insertGetter(extension, (builder) {
-        writeSetter(builder);
+    bool updatedExisting;
+    if (extensionElement != null) {
+      updatedExisting = await _updateExistingExtension2(
+        builder,
+        extensionElement,
+        (extension, builder) {
+          builder.insertGetter(extension, (builder) {
+            writeSetter(builder);
+          });
+        },
+      );
+    } else {
+      updatedExisting = await _updateExistingExtension(builder, targetType, (
+        extension,
+        builder,
+      ) {
+        builder.insertGetter(extension, (builder) {
+          writeSetter(builder);
+        });
       });
-    });
+    }
     if (updatedExisting) {
       return;
     }
@@ -600,23 +659,34 @@ abstract class _CreateExtensionMember extends ResolvedCorrectionProducer {
     if (library.isInSdk) {
       return null;
     }
-    var path = library.library.firstFragment.source.fullName;
+    var existingExtension = await getDeclarationNodeFromElement(
+      extension,
+      includeExtensions: true,
+    );
+    if (existingExtension is! ExtensionDeclaration) {
+      return null;
+    }
+    var path =
+        existingExtension.declaredFragment?.libraryFragment.source.fullName;
+    if (path == null) {
+      // Should never happen.
+      assert(
+        false,
+        'How is path to an existing extension null? $existingExtension',
+      );
+      return null;
+    }
     var unit = await unitResult.session.getResolvedUnit(path);
     if (unit is! ResolvedUnitResult) {
       return null;
     }
-    var existingExtension = unit.unit.declarations
-        .whereType<ExtensionDeclaration>()
-        .firstWhere(
-          (declaration) => declaration.declaredFragment!.element == extension,
-        );
     var instantiated = [extension].applicableTo(
       targetLibrary: libraryElement2,
       targetType: extension.thisType as TypeImpl,
       strictCasts: true,
     );
     if (instantiated.isNotEmpty) {
-      return (path, existingExtension);
+      return (unit.path, existingExtension);
     }
     return null;
   }

pkg/analysis_server/lib/src/services/correction/dart/create_field.dart

@@ -75,13 +75,9 @@ class CreateField extends CreateFieldOrGetter {
     }
     // Prepare target `ClassDeclaration`.
     var targetFragment = targetElement.firstFragment;
-    var targetDeclarationResult = await sessionHelper.getFragmentDeclaration(
-      targetFragment,
+    var targetNode = await getDeclarationNodeFromElement(
+      targetFragment.element,
     );
-    if (targetDeclarationResult == null) {
-      return;
-    }
-    var targetNode = targetDeclarationResult.node;
     if (targetNode is! CompilationUnitMember) {
       return;
     }

pkg/analysis_server/lib/src/services/correction/dart/create_getter.dart

@@ -211,21 +211,10 @@ class CreateGetter extends CreateFieldOrGetter {
       return;
     }
     // prepare target declaration
-    var targetDeclarationResult = await sessionHelper.getFragmentDeclaration(
-      targetFragment,
+    var targetNode = await getDeclarationNodeFromElement(
+      targetFragment.element,
     );
-    if (targetDeclarationResult == null) {
-      return;
-    }
-    var targetNode = targetDeclarationResult.node;
-    if (targetNode is CompilationUnitMember) {
-      if (targetDeclarationResult.node is! ClassDeclaration &&
-          targetDeclarationResult.node is! ExtensionDeclaration &&
-          targetDeclarationResult.node is! ExtensionTypeDeclaration &&
-          targetDeclarationResult.node is! MixinDeclaration) {
-        return;
-      }
-    } else {
+    if (targetNode is! CompilationUnitMember) {
       return;
     }
     // Build method source.

pkg/analysis_server/lib/src/services/correction/dart/create_method.dart

@@ -111,7 +111,7 @@ class CreateMethod extends ResolvedCorrectionProducer {
     );
     if (targetClassElement == null) return;
 
-    var targetNode = await _declarationNodeFromElement(targetClassElement);
+    var targetNode = await getDeclarationNodeFromElement(targetClassElement);
     if (targetNode is! CompilationUnitMember) return;
 
     await _writeMethod(
@@ -167,7 +167,7 @@ class CreateMethod extends ResolvedCorrectionProducer {
       if (targetClassElement == null) return;
       targetFragment = targetClassElement.firstFragment;
 
-      targetNode = await _declarationNodeFromElement(targetClassElement);
+      targetNode = await getDeclarationNodeFromElement(targetClassElement);
       if (targetNode == null) return;
 
       // Maybe static.
@@ -189,26 +189,6 @@ class CreateMethod extends ResolvedCorrectionProducer {
     );
   }
 
-  Future<CompilationUnitMember?> _declarationNodeFromElement(
-    InterfaceElement element,
-  ) async {
-    if (element.library.isInSdk) return null;
-    if (element is MixinElement) {
-      var fragment = element.firstFragment;
-      return await getMixinDeclaration(fragment);
-    } else if (element is ClassElement) {
-      var fragment = element.firstFragment;
-      return await getClassDeclaration(fragment);
-    } else if (element is ExtensionTypeElement) {
-      var fragment = element.firstFragment;
-      return await getExtensionTypeDeclaration(fragment);
-    } else if (element is EnumElement) {
-      var fragment = element.firstFragment;
-      return await getEnumDeclaration(fragment);
-    }
-    return null;
-  }
-
   /// Inserts the new method into the source code.
   Future<void> _writeMethod(
     ChangeBuilder builder,

pkg/analysis_server/lib/src/services/correction/dart/create_method_or_function.dart

@@ -220,31 +220,14 @@ class CreateMethodOrFunction extends ResolvedCorrectionProducer {
     required bool isStatic,
   }) async {
     var name = (node as SimpleIdentifier).name;
-    // prepare environment
-    var targetSource = targetClassElement.firstFragment.libraryFragment.source;
     // prepare insert offset
-    CompilationUnitMember? targetNode;
-    List<ClassMember>? classMembers;
-    if (targetClassElement is MixinElement) {
-      var fragment = targetClassElement.firstFragment;
-      var node = targetNode = await getMixinDeclaration(fragment);
-      classMembers = node?.members;
-    } else if (targetClassElement is ClassElement) {
-      var fragment = targetClassElement.firstFragment;
-      var node = targetNode = await getClassDeclaration(fragment);
-      classMembers = node?.members;
-    } else if (targetClassElement is ExtensionTypeElement) {
-      var fragment = targetClassElement.firstFragment;
-      var node = targetNode = await getExtensionTypeDeclaration(fragment);
-      classMembers = node?.members;
-    } else if (targetClassElement is EnumElement) {
-      var fragment = targetClassElement.firstFragment;
-      var node = targetNode = await getEnumDeclaration(fragment);
-      classMembers = node?.members;
-    }
-    if (targetNode == null || classMembers == null) {
+    var targetNode = await getDeclarationNodeFromElement(targetClassElement);
+    if (targetNode == null) {
       return;
     }
+    // prepare environment
+    var targetSource = targetClassElement.firstFragment.libraryFragment.source;
+    var classMembers = targetNode.classMembers;
     var insertOffset = targetNode.end - 1;
     // prepare prefix
     var prefix = '  ';

pkg/analysis_server/lib/src/services/correction/dart/create_operator.dart

@@ -107,19 +107,7 @@ class CreateOperator extends ResolvedCorrectionProducer {
       return;
     }
     // Prepare target ClassDeclaration.
-    if (targetClassElement is MixinElement) {
-      var fragment = targetClassElement.firstFragment;
-      targetNode = await getMixinDeclaration(fragment);
-    } else if (targetClassElement is ClassElement) {
-      var fragment = targetClassElement.firstFragment;
-      targetNode = await getClassDeclaration(fragment);
-    } else if (targetClassElement is ExtensionTypeElement) {
-      var fragment = targetClassElement.firstFragment;
-      targetNode = await getExtensionTypeDeclaration(fragment);
-    } else if (targetClassElement is EnumElement) {
-      var fragment = targetClassElement.firstFragment;
-      targetNode = await getEnumDeclaration(fragment);
-    }
+    targetNode = await getDeclarationNodeFromElement(targetClassElement);
     if (targetNode == null) {
       return;
     }

pkg/analysis_server/lib/src/services/correction/dart/create_setter.dart

@@ -52,8 +52,10 @@ class CreateSetter extends ResolvedCorrectionProducer {
     InstanceElement? targetElement;
     if (target is ExtensionOverride) {
       targetElement = target.element;
-    } else if (target is Identifier && target.element is ExtensionElement) {
-      targetElement = target.element as ExtensionElement;
+    } else if (target case Identifier(
+      :InstanceElement element,
+    ) when element is! ExtensionElement) {
+      targetElement = element;
       staticModifier = true;
     } else if (target != null) {
       // prepare target interface type
@@ -81,21 +83,10 @@ class CreateSetter extends ResolvedCorrectionProducer {
       return;
     }
     // prepare target declaration
-    var targetDeclarationResult = await sessionHelper.getFragmentDeclaration(
-      targetFragment,
+    var targetNode = await getDeclarationNodeFromElement(
+      targetFragment.element,
     );
-    if (targetDeclarationResult == null) {
-      return;
-    }
-    var targetNode = targetDeclarationResult.node;
-    if (targetNode is CompilationUnitMember) {
-      if (targetDeclarationResult.node is! ClassDeclaration &&
-          targetDeclarationResult.node is! MixinDeclaration &&
-          targetDeclarationResult.node is! ExtensionDeclaration &&
-          targetDeclarationResult.node is! ExtensionTypeDeclaration) {
-        return;
-      }
-    } else {
+    if (targetNode is! CompilationUnitMember) {
       return;
     }
     // Build setter source.

pkg/analysis_server/lib/src/services/correction/error_fix_status.yaml

@@ -1595,9 +1595,7 @@ CompileTimeErrorCode.UNDEFINED_EXTENSION_GETTER:
 CompileTimeErrorCode.UNDEFINED_EXTENSION_METHOD:
   status: hasFix
 CompileTimeErrorCode.UNDEFINED_EXTENSION_OPERATOR:
-  status: needsFix
-  notes: |-
-    The fix is to create the operator in an extension.
+  status: hasFix
 CompileTimeErrorCode.UNDEFINED_EXTENSION_SETTER:
   status: hasFix
 CompileTimeErrorCode.UNDEFINED_FUNCTION:

pkg/analysis_server/lib/src/services/correction/fix_internal.dart

@@ -942,9 +942,13 @@ final _builtInNonLintGenerators = <DiagnosticCode, List<ProducerGenerator>>{
     CreateExtensionMethod.new,
     CreateMethod.method,
   ],
+  CompileTimeErrorCode.undefinedExtensionOperator: [
+    CreateExtensionOperator.new,
+  ],
   CompileTimeErrorCode.undefinedExtensionSetter: [
     ChangeTo.getterOrSetter,
     CreateSetter.new,
+    CreateExtensionSetter.new,
   ],
   CompileTimeErrorCode.undefinedFunction: [
     ChangeTo.function,

pkg/analysis_server/test/src/services/correction/fix/create_extension_member_test.dart

@@ -73,6 +73,25 @@ extension on String {
     assertLinkedGroup(change.linkedEditGroups[0], ['null']);
   }
 
+  Future<void> test_enclosing() async {
+    await resolveTestCode('''
+extension E on String {
+  void f(String s) {
+    int _ = test;
+  }
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  int get test => null;
+
+  void f(String s) {
+    int _ = test;
+  }
+}
+''');
+  }
+
   Future<void> test_existingExtension_contextType() async {
     await resolveTestCode('''
 void f() {
@@ -235,6 +254,30 @@ extension E on String {
 ''');
   }
 
+  Future<void> test_main_part() async {
+    var partPath = join(testPackageLibPath, 'part.dart');
+    newFile(partPath, '''
+part of 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part 'part.dart';
+
+void foo(Object a) {
+  int _ = E(a).myUndefinedGetter;
+}
+''');
+    await assertHasFix('''
+part of 'test.dart';
+
+extension E on Object {
+  int get myUndefinedGetter => null;
+}
+''', target: partPath);
+  }
+
   Future<void> test_namedRecord_returnType() async {
     await resolveTestCode('''
 extension E on int {
@@ -389,6 +432,58 @@ extension on String {
 ''');
   }
 
+  Future<void> test_part_main() async {
+    var mainPath = join(testPackageLibPath, 'main.dart');
+    newFile(mainPath, '''
+part 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  int _ = E(a).myUndefinedGetter;
+}
+''');
+    await assertHasFix('''
+part 'test.dart';
+
+extension E on Object {
+  int get myUndefinedGetter => null;
+}
+''', target: mainPath);
+  }
+
+  Future<void> test_part_sibling() async {
+    var part1Path = join(testPackageLibPath, 'part1.dart');
+    newFile(part1Path, '''
+part of 'main.dart';
+
+extension E on Object {
+}
+''');
+    newFile(join(testPackageLibPath, 'main.dart'), '''
+part 'part1.dart';
+part 'test.dart';
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  int _ = E(a).myUndefinedGetter;
+}
+''');
+    await assertHasFix('''
+part of 'main.dart';
+
+extension E on Object {
+  int get myUndefinedGetter => null;
+}
+''', target: part1Path);
+  }
+
   Future<void> test_record_returnType() async {
     await resolveTestCode('''
 extension E on int {
@@ -404,6 +499,26 @@ extension E on int {
 ''');
   }
 
+  Future<void> test_static() async {
+    await resolveTestCode('''
+extension E on String {
+}
+
+void f(String s) {
+  int _ = E.test;
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  static int get test => null;
+}
+
+void f(String s) {
+  int _ = E.test;
+}
+''');
+  }
+
   Future<void> test_targetType_hasTypeArguments() async {
     await resolveTestCode('''
 void f(List<int> a) {
@@ -700,6 +815,30 @@ extension on int {}
 ''');
   }
 
+  Future<void> test_main_part() async {
+    var partPath = join(testPackageLibPath, 'part.dart');
+    newFile(partPath, '''
+part of 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part 'part.dart';
+
+void foo(Object a) {
+  void Function() _ = E(a).myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part of 'test.dart';
+
+extension E on Object {
+  void myUndefinedMethod() {}
+}
+''', target: partPath);
+  }
+
   Future<void> test_multipleSameTypeExistingExtensions() async {
     await resolveTestCode('''
 void f(int i) {
@@ -1129,6 +1268,58 @@ extension on String {
 ''');
   }
 
+  Future<void> test_part_main() async {
+    var mainPath = join(testPackageLibPath, 'main.dart');
+    newFile(mainPath, '''
+part 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  void Function() _ = E(a).myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part 'test.dart';
+
+extension E on Object {
+  void myUndefinedMethod() {}
+}
+''', target: mainPath);
+  }
+
+  Future<void> test_part_sibling() async {
+    var part1Path = join(testPackageLibPath, 'part1.dart');
+    newFile(part1Path, '''
+part of 'main.dart';
+
+extension E on Object {
+}
+''');
+    newFile(join(testPackageLibPath, 'main.dart'), '''
+part 'part1.dart';
+part 'test.dart';
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  void Function() _ = E(a).myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part of 'main.dart';
+
+extension E on Object {
+  void myUndefinedMethod() {}
+}
+''', target: part1Path);
+  }
+
   Future<void> test_record_returnType() async {
     await resolveTestCode('''
 extension E on int {
@@ -1559,6 +1750,30 @@ extension on int {
 ''');
   }
 
+  Future<void> test_main_part() async {
+    var partPath = join(testPackageLibPath, 'part.dart');
+    newFile(partPath, '''
+part of 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part 'part.dart';
+
+void foo(Object a) {
+  E(a) + 0;
+}
+''');
+    await assertHasFix('''
+part of 'test.dart';
+
+extension E on Object {
+  void operator +(int other) {}
+}
+''', target: partPath);
+  }
+
   Future<void> test_nullableObject_target() async {
     await resolveTestCode('''
 void f(Object? o) {
@@ -1593,6 +1808,58 @@ extension on int? {
 ''');
   }
 
+  Future<void> test_part_main() async {
+    var mainPath = join(testPackageLibPath, 'main.dart');
+    newFile(mainPath, '''
+part 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  E(a) + 0;
+}
+''');
+    await assertHasFix('''
+part 'test.dart';
+
+extension E on Object {
+  void operator +(int other) {}
+}
+''', target: mainPath);
+  }
+
+  Future<void> test_part_sibling() async {
+    var part1Path = join(testPackageLibPath, 'part1.dart');
+    newFile(part1Path, '''
+part of 'main.dart';
+
+extension E on Object {
+}
+''');
+    newFile(join(testPackageLibPath, 'main.dart'), '''
+part 'part1.dart';
+part 'test.dart';
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  E(a) + 0;
+}
+''');
+    await assertHasFix('''
+part of 'main.dart';
+
+extension E on Object {
+  void operator +(int other) {}
+}
+''', target: part1Path);
+  }
+
   Future<void> test_returnType() async {
     await resolveTestCode('''
 void f() {
@@ -1663,6 +1930,25 @@ void f() {
     await assertNoFix();
   }
 
+  Future<void> test_enclosing() async {
+    await resolveTestCode('''
+extension E on String {
+  void f(String s) {
+    test = 0;
+  }
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  set test(int test) {}
+
+  void f(String s) {
+    test = 0;
+  }
+}
+''');
+  }
+
   Future<void> test_existingExtension() async {
     await resolveTestCode('''
 void f() {
@@ -1789,6 +2075,60 @@ extension on int {}
 ''');
   }
 
+  Future<void> test_internal_instance() async {
+    await resolveTestCode('''
+extension E on String {
+  int m(int x) => s = x;
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  set s(int s) {}
+
+  int m(int x) => s = x;
+}
+''');
+  }
+
+  Future<void> test_internal_static() async {
+    await resolveTestCode('''
+extension E on String {
+  static int m(int x) => s = x;
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  static set s(int s) {}
+
+  static int m(int x) => s = x;
+}
+''');
+  }
+
+  Future<void> test_main_part() async {
+    var partPath = join(testPackageLibPath, 'part.dart');
+    newFile(partPath, '''
+part of 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part 'part.dart';
+
+void foo(Object a) {
+  E(a).myUndefinedSetter = 0;
+}
+''');
+    await assertHasFix('''
+part of 'test.dart';
+
+extension E on Object {
+  set myUndefinedSetter(int myUndefinedSetter) {}
+}
+''', target: partPath);
+  }
+
   Future<void> test_nullableObject_target() async {
     await resolveTestCode('''
 void f(Object? o) {
@@ -1823,6 +2163,26 @@ extension on int? {
 ''');
   }
 
+  Future<void> test_override() async {
+    await resolveTestCode('''
+extension E on String {
+}
+
+void f(String s) {
+  E(s).test = '0';
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  set test(String test) {}
+}
+
+void f(String s) {
+  E(s).test = '0';
+}
+''');
+  }
+
   Future<void> test_parent_nothing() async {
     await resolveTestCode('''
 void f() {
@@ -1866,6 +2226,78 @@ extension on String {
 ''');
   }
 
+  Future<void> test_part_main() async {
+    var mainPath = join(testPackageLibPath, 'main.dart');
+    newFile(mainPath, '''
+part 'test.dart';
+
+extension E on Object {
+}
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  E(a).myUndefinedSetter = 0;
+}
+''');
+    await assertHasFix('''
+part 'test.dart';
+
+extension E on Object {
+  set myUndefinedSetter(int myUndefinedSetter) {}
+}
+''', target: mainPath);
+  }
+
+  Future<void> test_part_sibling() async {
+    var part1Path = join(testPackageLibPath, 'part1.dart');
+    newFile(part1Path, '''
+part of 'main.dart';
+
+extension E on Object {
+}
+''');
+    newFile(join(testPackageLibPath, 'main.dart'), '''
+part 'part1.dart';
+part 'test.dart';
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(Object a) {
+  E(a).myUndefinedSetter = 0;
+}
+''');
+    await assertHasFix('''
+part of 'main.dart';
+
+extension E on Object {
+  set myUndefinedSetter(int myUndefinedSetter) {}
+}
+''', target: part1Path);
+  }
+
+  Future<void> test_static() async {
+    await resolveTestCode('''
+extension E on String {
+}
+
+void f(String s) {
+  E.test = 0;
+}
+''');
+    await assertHasFix('''
+extension E on String {
+  static set test(int test) {}
+}
+
+void f(String s) {
+  E.test = 0;
+}
+''');
+  }
+
   Future<void> test_targetType_hasTypeArguments() async {
     await resolveTestCode('''
 void f(List<int> a) {

pkg/analysis_server/test/src/services/correction/fix/create_method_test.dart

@@ -235,6 +235,7 @@ mixin M {
 class CreateMethodTearoffMixinTest extends FixProcessorTest {
   @override
   FixKind get kind => DartFixKind.createMethodTearoff;
+
   Future<void> test_functionType_method_enclosingMixin_static() async {
     await resolveTestCode('''
 mixin M {
@@ -283,6 +284,85 @@ mixin M {
 useFunction(int g(double a, String b)) {}
 ''');
   }
+
+  Future<void> test_main_part() async {
+    var partPath = join(testPackageLibPath, 'part.dart');
+    newFile(partPath, '''
+part of 'test.dart';
+
+mixin M {
+}
+''');
+    await resolveTestCode('''
+part 'part.dart';
+
+void foo(M a) {
+  void Function() _ = a.myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part of 'test.dart';
+
+mixin M {
+  void myUndefinedMethod() {
+  }
+}
+''', target: partPath);
+  }
+
+  Future<void> test_part_main() async {
+    var mainPath = join(testPackageLibPath, 'main.dart');
+    newFile(mainPath, '''
+part 'test.dart';
+
+mixin M {
+}
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(M a) {
+  void Function() _ = a.myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part 'test.dart';
+
+mixin M {
+  void myUndefinedMethod() {
+  }
+}
+''', target: mainPath);
+  }
+
+  Future<void> test_part_sibling() async {
+    var part1Path = join(testPackageLibPath, 'part1.dart');
+    newFile(part1Path, '''
+part of 'main.dart';
+
+mixin M {
+}
+''');
+    newFile(join(testPackageLibPath, 'main.dart'), '''
+part 'part1.dart';
+part 'test.dart';
+''');
+    await resolveTestCode('''
+part of 'main.dart';
+
+void foo(M a) {
+  void Function() _ = a.myUndefinedMethod;
+}
+''');
+    await assertHasFix('''
+part of 'main.dart';
+
+mixin M {
+  void myUndefinedMethod() {
+  }
+}
+''', target: part1Path);
+  }
 }
 
 @reflectiveTest

pkg/analysis_server/test/src/services/correction/fix/create_setter_test.dart

@@ -212,12 +212,12 @@ void f(List p) {
 
   Future<void> test_internal_instance() async {
     await resolveTestCode('''
-extension E on String {
+class C {
   int m(int x) => s = x;
 }
 ''');
     await assertHasFix('''
-extension E on String {
+class C {
   set s(int s) {}
 
   int m(int x) => s = x;
@@ -227,12 +227,12 @@ extension E on String {
 
   Future<void> test_internal_static() async {
     await resolveTestCode('''
-extension E on String {
+class C {
   static int m(int x) => s = x;
 }
 ''');
     await assertHasFix('''
-extension E on String {
+class C {
   static set s(int s) {}
 
   static int m(int x) => s = x;
@@ -332,15 +332,8 @@ void f(String s) {
   E(s).test = '0';
 }
 ''');
-    await assertHasFix('''
-extension E on String {
-  set test(String test) {}
-}
-
-void f(String s) {
-  E(s).test = '0';
-}
-''');
+    // This should be handled by create extension member fixes
+    await assertNoFix();
   }
 
   Future<void> test_part_main() async {
@@ -509,20 +502,20 @@ void f() {
 
   Future<void> test_static() async {
     await resolveTestCode('''
-extension E on String {
+class C {
 }
 
 void f(String s) {
-  E.test = 0;
+  C.test = 0;
 }
 ''');
     await assertHasFix('''
-extension E on String {
+class C {
   static set test(int test) {}
 }
 
 void f(String s) {
-  E.test = 0;
+  C.test = 0;
 }
 ''');
   }

pkg/analysis_server_plugin/api.txt

@@ -94,6 +94,7 @@ package:analysis_server_plugin/edit/dart/correction_producer.dart:
     typeSystem (getter: TypeSystem)
     unitResult (getter: ResolvedUnitResult)
     getClassDeclaration (method: Future<ClassDeclaration?> Function(ClassFragment))
+    getDeclarationNodeFromElement (method: Future<CompilationUnitMember?> Function(InstanceElement, {bool includeExtensions}))
     getEnumDeclaration (method: Future<EnumDeclaration?> Function(EnumFragment))
     getExtensionDeclaration (method: Future<ExtensionDeclaration?> Function(ExtensionFragment))
     getExtensionTypeDeclaration (method: Future<ExtensionTypeDeclaration?> Function(ExtensionTypeFragment))
@@ -182,6 +183,7 @@ package:analyzer/dart/element/element.dart:
   EnumFragment (referenced)
   ExtensionFragment (referenced)
   ExtensionTypeFragment (referenced)
+  InstanceElement (referenced)
   InterfaceElement (referenced)
   LibraryElement (referenced)
   MixinFragment (referenced)
@@ -198,6 +200,7 @@ package:analyzer/source/source_range.dart:
 package:analyzer/src/dart/ast/ast.dart:
   AstNode (referenced)
   ClassDeclaration (referenced)
+  CompilationUnitMember (referenced)
   EnumDeclaration (referenced)
   Expression (referenced)
   ExtensionDeclaration (referenced)

pkg/analysis_server_plugin/lib/edit/dart/correction_producer.dart

@@ -435,6 +435,24 @@ abstract class ResolvedCorrectionProducer
     return null;
   }
 
+  Future<CompilationUnitMember?> getDeclarationNodeFromElement(
+    InstanceElement element, {
+    bool includeExtensions = false,
+  }) async {
+    if (element.library.isInSdk) return null;
+    return switch (element) {
+      ClassElement(:var firstFragment) => getClassDeclaration(firstFragment),
+      EnumElement(:var firstFragment) => getEnumDeclaration(firstFragment),
+      ExtensionElement(:var firstFragment) when includeExtensions =>
+        getExtensionDeclaration(firstFragment),
+      ExtensionTypeElement(:var firstFragment) => getExtensionTypeDeclaration(
+        firstFragment,
+      ),
+      MixinElement(:var firstFragment) => getMixinDeclaration(firstFragment),
+      _ => null,
+    };
+  }
+
   /// Returns the class declaration for the given [fragment], or `null` if there
   /// is no such class.
   Future<EnumDeclaration?> getEnumDeclaration(EnumFragment fragment) async {