Fix expected return type in async functions

lib/type_checker.dart

 // Copyright (c) 2016, 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 kernel.type_checker;
 
 import 'ast.dart';
 import 'class_hierarchy.dart';
 import 'core_types.dart';
 import 'type_algebra.dart';
 import 'type_environment.dart';
@@ skipping 140 matching lines @@
   }
 
   visitConstructor(Constructor node) {
     environment.returnType = null;
     environment.yieldType = null;
     node.initializers.forEach(visitInitializer);
     handleFunctionNode(node.function);
   }
 
   visitProcedure(Procedure node) {
-    environment.returnType = node.function.returnType;
+    environment.returnType = _getInternalReturnType(node.function);
     environment.yieldType = _getYieldType(node.function);
     handleFunctionNode(node.function);
   }
 
   void handleFunctionNode(FunctionNode node) {
+    var oldAsyncMarker = environment.currentAsyncMarker;
+    environment.currentAsyncMarker = node.asyncMarker;
     node.positionalParameters
         .skip(node.requiredParameterCount)
         .forEach(handleOptionalParameter);
     node.namedParameters.forEach(handleOptionalParameter);
     if (node.body != null) {
       visitStatement(node.body);
     }
+    environment.currentAsyncMarker = oldAsyncMarker;
   }
 
   void handleNestedFunctionNode(FunctionNode node) {
     var oldReturn = environment.returnType;
     var oldYield = environment.yieldType;
-    environment.returnType = node.returnType;
+    environment.returnType = _getInternalReturnType(node);
     environment.yieldType = _getYieldType(node);
     handleFunctionNode(node);
     environment.returnType = oldReturn;
     environment.yieldType = oldYield;
   }
 
   void handleOptionalParameter(VariableDeclaration parameter) {
     if (parameter.initializer != null) {
       checkAssignableExpression(parameter.initializer, parameter.type);
     }
@@ skipping 78 matching lines @@
         }
       }
       if (!found) {
         fail(argument.value, 'Unexpected named parameter: ${argument.name}');
         return const BottomType();
       }
     }
     return substitution.substituteType(function.returnType);
   }
 
+  DartType _getInternalReturnType(FunctionNode function) {
+    switch (function.asyncMarker) {
+      case AsyncMarker.Sync:
+        return function.returnType;
+
+      case AsyncMarker.Async:
+        Class container = coreTypes.futureClass;
+        DartType returnType = function.returnType;
+        if (returnType is InterfaceType && returnType.classNode == container) {
+          return returnType.typeArguments.single;
+        }
+        return const DynamicType();
+
+      case AsyncMarker.SyncStar:
+      case AsyncMarker.AsyncStar:
+      case AsyncMarker.SyncYielding:
+        return null;
+
+      default:
+        throw 'Unexpected async marker: ${function.asyncMarker}';
+    }
+  }
+
   DartType _getYieldType(FunctionNode function) {
     switch (function.asyncMarker) {
       case AsyncMarker.Sync:
       case AsyncMarker.Async:
         return null;
 
       case AsyncMarker.SyncStar:
       case AsyncMarker.AsyncStar:
         Class container = function.asyncMarker == AsyncMarker.SyncStar
             ? coreTypes.iterableClass
@@ skipping 441 matching lines @@
   visitLabeledStatement(LabeledStatement node) {
     visitStatement(node.body);
   }
 
   @override
   visitReturnStatement(ReturnStatement node) {
     if (node.expression != null) {
       if (environment.returnType == null) {
         fail(node, 'Return of a value from void method');
       } else {
-        checkAssignableExpression(node.expression, environment.returnType);
+        var type = visitExpression(node.expression);
+        if (environment.currentAsyncMarker == AsyncMarker.Async) {
+          type = environment.unfutureType(type);
+        }
+        checkAssignable(node.expression, type, environment.returnType);
       }
     }
   }
 
   @override
   visitSwitchStatement(SwitchStatement node) {
     visitExpression(node.expression);
     for (var switchCase in node.cases) {
       switchCase.expressions.forEach(visitExpression);
       visitStatement(switchCase.body);
@@ skipping 22 matching lines @@
   }
 
   @override
   visitWhileStatement(WhileStatement node) {
     checkAssignableExpression(node.condition, environment.boolType);
     visitStatement(node.body);
   }
 
   @override
   visitYieldStatement(YieldStatement node) {
-    checkAssignableExpression(node.expression, environment.yieldType);
+    if (node.isYieldStar) {
+      Class container = environment.currentAsyncMarker == AsyncMarker.AsyncStar
+          ? coreTypes.streamClass
+          : coreTypes.iterableClass;
+      var type = visitExpression(node.expression);
+      var asContainer = type is InterfaceType
+          ? hierarchy.getTypeAsInstanceOf(type, container)
+          : null;
+      if (asContainer != null) {
+        checkAssignable(node.expression, asContainer.typeArguments[0],
+            environment.yieldType);
+      } else {
+        fail(node.expression, '$type is not an instance of $container');
+      }
+    } else {
+      checkAssignableExpression(node.expression, environment.yieldType);
+    }
   }
 
   @override
   visitFieldInitializer(FieldInitializer node) {
     checkAssignableExpression(node.value, node.field.type);
   }
 
   @override
   visitRedirectingInitializer(RedirectingInitializer node) {
     handleCall(node.arguments, node.target.function,
         typeParameters: const <TypeParameter>[]);
   }
 
   @override
   visitSuperInitializer(SuperInitializer node) {
     handleCall(node.arguments, node.target.function,
         typeParameters: const <TypeParameter>[]);
   }
 
   @override
   visitLocalInitializer(LocalInitializer node) {
     visitVariableDeclaration(node.variable);
   }
 
   @override
   visitInvalidInitializer(InvalidInitializer node) {}
 }