[DAS] Fixes part file handling for some create fixes

Bug: #61192
Change-Id: I8aa765fca2d2948ab6fe1b47a6bff6bf58de3cd3
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/447341
Auto-Submit: Felipe Morschel <[email protected]>
Reviewed-by: Samuel Rawlins <[email protected]>
Reviewed-by: Brian Wilkerson <[email protected]>
Commit-Queue: Samuel Rawlins <[email protected]>

Author: FMorschel

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;
-          }
-        }
-      case ExpressionFunctionBody expressionFunctionBody:
-        if (expressionFunctionBody.expression == nameNode) {
-          targetType = node.enclosingInstanceElement?.thisType;
+          target = propertyAccess.realTarget;
         }
-      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 = '  ';