Add await/async when a method with 'await' in its body is extracted.

pkg/analysis_server/lib/src/services/refactoring/extract_method.dart

 // Copyright (c) 2014, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 library services.src.refactoring.extract_method;
 
 import 'dart:async';
 
 import 'package:analysis_server/src/protocol_server.dart' hide Element;
 import 'package:analysis_server/src/services/correction/name_suggestion.dart';
 import 'package:analysis_server/src/services/correction/selection_analyzer.dart';
 import 'package:analysis_server/src/services/correction/source_range.dart';
 import 'package:analysis_server/src/services/correction/statement_analyzer.dart';
 import 'package:analysis_server/src/services/correction/status.dart';
 import 'package:analysis_server/src/services/correction/util.dart';
 import 'package:analysis_server/src/services/refactoring/naming_conventions.dart';
 import 'package:analysis_server/src/services/refactoring/refactoring.dart';
 import 'package:analysis_server/src/services/refactoring/refactoring_internal.dart';
 import 'package:analysis_server/src/services/refactoring/rename_class_member.dart';
 import 'package:analysis_server/src/services/refactoring/rename_unit_member.dart';
 import 'package:analysis_server/src/services/search/element_visitors.dart';
 import 'package:analysis_server/src/services/search/search_engine.dart';
 import 'package:analyzer/src/generated/ast.dart';
 import 'package:analyzer/src/generated/element.dart';
+import 'package:analyzer/src/generated/engine.dart';
 import 'package:analyzer/src/generated/java_core.dart';
 import 'package:analyzer/src/generated/resolver.dart' show ExitDetector;
 import 'package:analyzer/src/generated/scanner.dart';
 import 'package:analyzer/src/generated/source.dart';
 
 const String _TOKEN_SEPARATOR = '\uFFFF';
 
 /**
  * Returns the "normalized" version of the given source, which is reconstructed
  * from tokens, so ignores all the comments and spaces.
@@ skipping 20 matching lines @@
 class ExtractMethodRefactoringImpl extends RefactoringImpl
     implements ExtractMethodRefactoring {
   static const ERROR_EXITS =
       'Selected statements contain a return statement, but not all possible '
       'execuion flows exit. Semantics may not be preserved.';
 
   final SearchEngine searchEngine;
   final CompilationUnit unit;
   final int selectionOffset;
   final int selectionLength;
+  AnalysisContext context;
   CompilationUnitElement unitElement;
   LibraryElement libraryElement;
   SourceRange selectionRange;
   CorrectionUtils utils;
   Set<LibraryElement> librariesToImport = new Set<LibraryElement>();
 
   String returnType;
+  String variableType;
   String name;
   bool extractAll = true;
   bool canCreateGetter = false;
   bool createGetter = false;
   final List<String> names = <String>[];
   final List<int> offsets = <int>[];
   final List<int> lengths = <int>[];
 
   /**
    * The map of local names to their visibility ranges.
    */
   Map<String, List<SourceRange>> _localNames = <String, List<SourceRange>>{};
 
   /**
    * The set of names that are referenced without any qualifier.
    */
   Set<String> _unqualifiedNames = new Set<String>();
 
   Set<String> _excludedNames = new Set<String>();
   List<RefactoringMethodParameter> _parameters = <RefactoringMethodParameter>[];
   Map<String, RefactoringMethodParameter> _parametersMap =
       <String, RefactoringMethodParameter>{};
   Map<String, List<SourceRange>> _parameterReferencesMap =
       <String, List<SourceRange>>{};
+  bool _hasAwait = false;
   DartType _returnType;
   String _returnVariableName;
   AstNode _parentMember;
   Expression _selectionExpression;
   FunctionExpression _selectionFunctionExpression;
   List<Statement> _selectionStatements;
   List<_Occurrence> _occurrences = [];
   bool _staticContext = false;
 
   ExtractMethodRefactoringImpl(this.searchEngine, this.unit,
       this.selectionOffset, this.selectionLength) {
     unitElement = unit.element;
     libraryElement = unitElement.library;
+    context = libraryElement.context;
     selectionRange = new SourceRange(selectionOffset, selectionLength);
     utils = new CorrectionUtils(unit);
   }
 
   @override
   List<RefactoringMethodParameter> get parameters => _parameters;
 
   @override
   void set parameters(List<RefactoringMethodParameter> parameters) {
     _parameters = parameters.toList();
@@ skipping 56 matching lines @@
   @override
   Future<RefactoringStatus> checkInitialConditions() {
     RefactoringStatus result = new RefactoringStatus();
     // selection
     result.addStatus(_checkSelection());
     if (result.hasFatalError) {
       return new Future.value(result);
     }
     // prepare parts
     result.addStatus(_initializeParameters());
+    _initializeHasAwait();
     _initializeReturnType();
     // occurrences
     _initializeOccurrences();
     _prepareOffsetsLengths();
     // getter
     canCreateGetter = _computeCanCreateGetter();
     createGetter =
         canCreateGetter && _isExpressionForGetter(_selectionExpression);
     // names
     _prepareExcludedNames();
@@ skipping 24 matching lines @@
       if (!extractAll && !occurence.isSelection) {
         continue;
       }
       // prepare invocation source
       String invocationSource;
       if (_selectionFunctionExpression != null) {
         invocationSource = name;
       } else {
         StringBuffer sb = new StringBuffer();
         // may be returns value
-        if (_selectionStatements != null && returnType != 'void') {
+        if (_selectionStatements != null && variableType != null) {
           // single variable assignment / return statement
           if (_returnVariableName != null) {
             String occurrenceName =
                 occurence._parameterOldToOccurrenceName[_returnVariableName];
             // may be declare variable
             if (!_parametersMap.containsKey(_returnVariableName)) {
-              if (returnType.isEmpty) {
+              if (variableType.isEmpty) {
                 sb.write('var ');
               } else {
-                sb.write(returnType);
+                sb.write(variableType);
                 sb.write(' ');
               }
             }
             // assign the return value
             sb.write(occurrenceName);
             sb.write(' = ');
           } else {
             sb.write('return ');
           }
         }
+        // await
+        if (_hasAwait) {
+          sb.write('await ');
+        }
         // invocation itself
         sb.write(name);
         if (!createGetter) {
           sb.write('(');
           bool firstParameter = true;
           for (RefactoringMethodParameter parameter in _parameters) {
             // may be comma
             if (firstParameter) {
               firstParameter = false;
             } else {
@@ skipping 35 matching lines @@
       String declarationSource = null;
       {
         String returnExpressionSource = _getMethodBodySource();
         // closure
         if (_selectionFunctionExpression != null) {
           declarationSource = '${name}${returnExpressionSource}';
           if (_selectionFunctionExpression.body is ExpressionFunctionBody) {
             declarationSource += ';';
           }
         }
+        // optional 'async' body modifier
+        String asyncKeyword = _hasAwait ? ' async' : '';
         // expression
         if (_selectionExpression != null) {
           // add return type
           if (returnType.isNotEmpty) {
             annotations += '$returnType ';
           }
           // just return expression
-          declarationSource =
-              '${annotations}${signature} => ${returnExpressionSource};';
+          declarationSource = '$annotations$signature$asyncKeyword => ';
+          declarationSource += '$returnExpressionSource;';
         }
         // statements
         if (_selectionStatements != null) {
           if (returnType.isNotEmpty) {
             annotations += returnType + ' ';
           }
-          declarationSource = '${annotations}${signature} {${eol}';
+          declarationSource = '$annotations$signature$asyncKeyword {$eol';
           declarationSource += returnExpressionSource;
           if (_returnVariableName != null) {
             declarationSource +=
                 '${prefix}  return ${_returnVariableName};$eol';
           }
           declarationSource += '${prefix}}';
         }
       }
       // insert declaration
       if (declarationSource != null) {
@@ skipping 194 matching lines @@
     replaceEdits = replaceEdits.reversed.toList();
     String source = SourceEdit.applySequence(originalSource, replaceEdits);
     pattern.normalizedSource = _getNormalizedSource(source);
     return pattern;
   }
 
   String _getTypeCode(DartType type) {
     return utils.getTypeSource(type, librariesToImport);
   }
 
+  void _initializeHasAwait() {
+    _HasAwaitVisitor visitor = new _HasAwaitVisitor();
+    if (_selectionExpression != null) {
+      _selectionExpression.accept(visitor);
+    } else if (_selectionStatements != null) {
+      _selectionStatements.forEach((statement) => statement.accept(visitor));
+    }
+    _hasAwait = visitor.result;
+  }
+
   /**
    * Fills [_occurrences] field.
    */
   void _initializeOccurrences() {
     _occurrences.clear();
     // prepare selection
     _SourcePattern selectionPattern = _getSourcePattern(selectionRange);
     Map<String, String> patternToSelectionName =
         _inverseMap(selectionPattern.originalToPatternNames);
     // prepare an enclosing parent - class or unit
@@ skipping 57 matching lines @@
       result.addFatalError(format(
           'Ambiguous return value: Selected block contains more than one '
           'assignment to local variables. Affected variables are:\n\n{0}',
           sb.toString().trim()));
     }
     // done
     return result;
   }
 
   void _initializeReturnType() {
+    InterfaceType futureType = context.typeProvider.futureType;
     if (_selectionFunctionExpression != null) {
+      variableType = '';
       returnType = '';
     } else if (_returnType == null) {
-      returnType = 'void';
+      variableType = null;
+      if (_hasAwait) {
+        returnType = _getTypeCode(futureType);
+      } else {
+        returnType = 'void';
+      }
+    } else if (_returnType.isDynamic) {
+      variableType = '';
+      if (_hasAwait) {
+        returnType = _getTypeCode(futureType);
+      } else {
+        returnType = '';
+      }
     } else {
-      returnType = _getTypeCode(_returnType);
-    }
-    if (returnType == 'dynamic') {
-      returnType = '';
+      variableType = _getTypeCode(_returnType);
+      if (_hasAwait) {
+        if (_returnType.element != futureType.element) {
+          returnType = _getTypeCode(futureType.substitute4([_returnType]));
+        }
+      } else {
+        returnType = variableType;
+      }
     }
   }
 
   /**
    * Checks if the given [VariableElement] is declared in [selectionRange].
    */
   bool _isDeclaredInSelection(VariableElement element) {
     return selectionRange.contains(element.nameOffset);
   }
 
@@ skipping 246 matching lines @@
           patternName = '__refVar${pattern.originalToPatternNames.length}';
           pattern.originalToPatternNames[originalName] = patternName;
         }
         replaceEdits.add(new SourceEdit(nodeRange.offset - partRange.offset,
             nodeRange.length, patternName));
       }
     }
   }
 }
 
+class _HasAwaitVisitor extends GeneralizingAstVisitor {
+  bool result = false;
+
+  @override
+  visitAwaitExpression(AwaitExpression node) {
+    result = true;
+  }
+
+  @override
+  visitForEachStatement(ForEachStatement node) {
+    if (node.awaitKeyword != null) {
+      result = true;
+    }
+    super.visitForEachStatement(node);
+  }
+}
+
 class _HasReturnStatementVisitor extends RecursiveAstVisitor {
   bool hasReturn = false;
 
   @override
   visitBlockFunctionBody(BlockFunctionBody node) {}
 
   @override
   visitReturnStatement(ReturnStatement node) {
     hasReturn = true;
   }
@@ skipping 247 matching lines @@
       return false;
     }
     for (int i = 0; i < parameterTypes.length; i++) {
       if (other.parameterTypes[i] != parameterTypes[i]) {
         return false;
       }
     }
     return true;
   }
 }