fix #25944, improve Future.then inference

pkg/analyzer/lib/src/generated/resolver.dart

Index: pkg/analyzer/lib/src/generated/resolver.dart
diff --git a/pkg/analyzer/lib/src/generated/resolver.dart b/pkg/analyzer/lib/src/generated/resolver.dart
index afa7ccec98ee9c48746663cc1c1a9893c0b0b541..b432e0dbb5de86b584e4e12355c5cf2b2c585399 100644
--- a/pkg/analyzer/lib/src/generated/resolver.dart
+++ b/pkg/analyzer/lib/src/generated/resolver.dart
@@ -4689,9 +4689,7 @@ class InferenceContext {
    * A stack of return types for all of the enclosing
    * functions and methods.
    */
-  // TODO(leafp) Handle the implicit union type for Futures
-  // https://github.com/dart-lang/sdk/issues/25322
-  final List<DartType> _returnStack = <DartType>[];
+  final List<TypeContext> _returnStack = <TypeContext>[];
 
   InferenceContext._(this._errorReporter, TypeProvider typeProvider,
       this._typeSystem, this._inferenceHints)
@@ -4706,7 +4704,7 @@ class InferenceContext {
    * functions returning Future<T> and for generator functions
    * returning Stream<T> or Iterable<T>, this is T.
    */
-  DartType get returnContext =>
+  TypeContext get returnContext =>
       _returnStack.isNotEmpty ? _returnStack.last : null;
 
   /**
@@ -4719,10 +4717,10 @@ class InferenceContext {
     if (_returnStack.isEmpty) {
       return;
     }
-    DartType context = _returnStack.last;
-    if (context == null || context.isDynamic) {
+    TypeContext context = _returnStack.last;
+    if (context == null || context is DartType && context.isDynamic) {
       DartType inferred = _inferredReturn.last;
-      inferred = _typeSystem.getLeastUpperBound(_typeProvider, type, inferred);
+      inferred = _typeSystem.getLeastUpperBound(_typeProvider, inferred, type);
       _inferredReturn[_inferredReturn.length - 1] = inferred;
     }
   }
@@ -4762,8 +4760,7 @@ class InferenceContext {
    * Push a block function body's return type onto the return stack.
    */
   void pushReturnContext(BlockFunctionBody node) {
-    DartType returnType = getType(node);
-    _returnStack.add(returnType);
+    _returnStack.add(getType(node));
     _inferredReturn.add(BottomTypeImpl.instance);
   }
 
@@ -4905,16 +4902,46 @@ class InferenceContext {
   }
 
   /**
+   * Look for a single contextual type attached to [node], and returns the type
+   * if found, otherwise null.
+   *
+   * If [node] has a contextual union type like `T | Future<T>` this will
+   * simplify it to only return `T`. If the caller can handle a union type,
+   * [getType] should be used instead.
+   */
+  static DartType getNonFutureType(AstNode node) {

[Leaf, 2016/08/05 22:32:35]

I don't love this name.  It seems to be that it would be nicer to keep the names
matching what gets returned.  So this would be `getType` (it returns a type),
and what's now called `getType below would maybe become `getContext`?  This
matches better with returnContext which yields a TypeContext.  `setType` should
probably become `setContext` if you make this change, I guess.

[Jennifer Messerly, 2016/08/08 21:59:26]

Great idea! Yeah I didn't like those names either... Done!

+    TypeContext t = getType(node);
+    if (t is FutureUnionTypeContext) {
+      return t.type;
+    }
+    return t;
+  }
+
+  /**
    * Look for contextual type information attached to [node].  Returns
    * the type if found, otherwise null.
    */
-  static DartType getType(AstNode node) => node?.getProperty(_typeProperty);
+  static TypeContext getType(AstNode node) => node?.getProperty(_typeProperty);
+
+  /**
+   * Like [getType] but expands a union type into a list of types.
+   */
+  static Iterable<DartType> getTypes(AstNode node) {
+    TypeContext t = getType(node);
+    if (t == null) {
+      return DartType.EMPTY_LIST;
+    }
+    if (t is FutureUnionTypeContext) {
+      return t.types;
+    }
+    return <DartType>[t];
+  }
 
   /**
    * Attach contextual type information [type] to [node] for use during
    * inference.
    */
-  static void setType(AstNode node, DartType type) {
+  static void setType(AstNode node, TypeContext type) {
     node?.setProperty(_typeProperty, type);
   }
 
@@ -5931,7 +5958,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitArgumentList(ArgumentList node) {
-    DartType callerType = InferenceContext.getType(node);
+    DartType callerType = InferenceContext.getNonFutureType(node);
     if (callerType is FunctionType) {
       Map<String, DartType> namedParameterTypes =
           callerType.namedParameterTypes;
@@ -6007,13 +6034,11 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitAwaitExpression(AwaitExpression node) {
-    // TODO(leafp): Handle the implicit union type here
-    // https://github.com/dart-lang/sdk/issues/25322
-    DartType contextType = InferenceContext.getType(node);
+    TypeContext contextType = InferenceContext.getType(node);
     if (contextType != null) {
-      InterfaceType futureT = typeProvider.futureType
-          .instantiate([contextType.flattenFutures(typeSystem)]);
-      InferenceContext.setType(node.expression, futureT);
+      var futureUnion =
+          FutureUnionTypeContext.from(contextType, typeProvider, typeSystem);
+      InferenceContext.setType(node.expression, futureUnion);
     }
     return super.visitAwaitExpression(node);
   }
@@ -6530,14 +6555,30 @@ class ResolverVisitor extends ScopedVisitor {
       _enclosingFunction = node.element;
       _overrideManager.enterScope();
       try {
-        DartType functionType = InferenceContext.getType(node);
+        DartType functionType = InferenceContext.getNonFutureType(node);
         if (functionType is FunctionType) {
           functionType =
               matchFunctionTypeParameters(node.typeParameters, functionType);
           if (functionType is FunctionType) {
             _inferFormalParameterList(node.parameters, functionType);
-            DartType returnType =
-                _computeReturnOrYieldType(functionType.returnType);
+
+            TypeContext returnType;
+            if (_isFutureThenLambda(node)) {
+              var futureThenType =
+                  InferenceContext.getType(node.parent) as FunctionType;
+
+              // Pretend the return type of Future<T>.then<S> first parameter is
+              //
+              //     T -> (S | Future<S>)
+              //
+              // We can't represent this in Dart so we populate it here during
+              // inference.
+              returnType = FutureUnionTypeContext.from(
+                  futureThenType.returnType, typeProvider, typeSystem);
+            } else {
+              returnType = _computeReturnOrYieldType(functionType.returnType);
+            }
+
             InferenceContext.setType(node.body, returnType);
           }
         }
@@ -6653,30 +6694,37 @@ class ResolverVisitor extends ScopedVisitor {
   Object visitInstanceCreationExpression(InstanceCreationExpression node) {
     TypeName classTypeName = node.constructorName.type;
     if (classTypeName.typeArguments == null) {
-      DartType contextType = InferenceContext.getType(node);
-      if (contextType is InterfaceType &&
-          contextType.typeArguments != null &&
-          contextType.typeArguments.length > 0) {
-        // TODO(jmesserly): for generic methods we use the
-        // StrongTypeSystemImpl.inferGenericFunctionCall, which appears to
-        // be a tad more powerful than matchTypes.
-        //
-        // For example it can infer this case:
-        //
-        //     class E<S, T> extends A<C<S>, T> { ... }
-        //     A<C<int>, String> a0 = /*infer<int, String>*/new E("hello");
-        //
-        // See _inferArgumentTypesFromContext in this file for use of it.
-        List<DartType> targs =
-            inferenceContext.matchTypes(classTypeName.type, contextType);
-        if (targs != null && targs.any((t) => !t.isDynamic)) {
-          ClassElement classElement = classTypeName.type.element;
-          InterfaceType rawType = classElement.type;
-          InterfaceType fullType =
-              rawType.substitute2(targs, rawType.typeArguments);
-          // The element resolver uses the type on the constructor name, so
-          // infer it first
-          typeAnalyzer.inferConstructorName(node.constructorName, fullType);
+      // Given a union of context types ` T0 | T1 | ... | Tn`, find the first
+      // valid instantiation `new C<Ti>`, if it exists.
+      // TODO(jmesserly): if we support union types for real, `new C<Ti | Tj>`
+      // will become a valid possibility. Right now the only allowed union is
+      // `T | Future<T>` so we can take a simple approach.
+      for (var contextType in InferenceContext.getTypes(node)) {
+        if (contextType is InterfaceType &&
+            contextType.typeArguments != null &&
+            contextType.typeArguments.isNotEmpty) {
+          // TODO(jmesserly): for generic methods we use the
+          // StrongTypeSystemImpl.inferGenericFunctionCall, which appears to
+          // be a tad more powerful than matchTypes.
+          //
+          // For example it can infer this case:
+          //
+          //     class E<S, T> extends A<C<S>, T> { ... }
+          //     A<C<int>, String> a0 = /*infer<int, String>*/new E("hello");
+          //
+          // See _inferArgumentTypesFromContext in this file for use of it.
+          List<DartType> targs =
+              inferenceContext.matchTypes(classTypeName.type, contextType);
+          if (targs != null && targs.any((t) => !t.isDynamic)) {
+            ClassElement classElement = classTypeName.type.element;
+            InterfaceType rawType = classElement.type;
+            InterfaceType fullType =
+                rawType.substitute2(targs, rawType.typeArguments);
+            // The element resolver uses the type on the constructor name, so
+            // infer it first
+            typeAnalyzer.inferConstructorName(node.constructorName, fullType);
+            break;
+          }
         }
       }
     }
@@ -6699,7 +6747,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitListLiteral(ListLiteral node) {
-    DartType contextType = InferenceContext.getType(node);
+    DartType contextType = InferenceContext.getNonFutureType(node);
     List<DartType> targs = null;
     if (node.typeArguments != null) {
       targs = node.typeArguments.arguments.map((t) => t.type).toList();
@@ -6724,7 +6772,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitMapLiteral(MapLiteral node) {
-    DartType contextType = InferenceContext.getType(node);
+    DartType contextType = InferenceContext.getNonFutureType(node);
     List<DartType> targs = null;
     if (node.typeArguments != null) {
       targs = node.typeArguments.arguments.map((t) => t.type).toList();
@@ -6756,7 +6804,7 @@ class ResolverVisitor extends ScopedVisitor {
     try {
       _currentFunctionBody = node.body;
       _enclosingFunction = node.element;
-      DartType returnType =
+      TypeContext returnType =
           _computeReturnOrYieldType(_enclosingFunction.type?.returnType);
       InferenceContext.setType(node.body, returnType);
       super.visitMethodDeclaration(node);
@@ -6790,7 +6838,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitNamedExpression(NamedExpression node) {
-    InferenceContext.setType(node.expression, InferenceContext.getType(node));
+    InferenceContext.setTypeFromNode(node.expression, node);
     return super.visitNamedExpression(node);
   }
 
@@ -6804,7 +6852,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitParenthesizedExpression(ParenthesizedExpression node) {
-    InferenceContext.setType(node.expression, InferenceContext.getType(node));
+    InferenceContext.setTypeFromNode(node.expression, node);
     return super.visitParenthesizedExpression(node);
   }
 
@@ -6922,7 +6970,7 @@ class ResolverVisitor extends ScopedVisitor {
 
   @override
   Object visitVariableDeclaration(VariableDeclaration node) {
-    InferenceContext.setType(node.initializer, InferenceContext.getType(node));
+    InferenceContext.setTypeFromNode(node.initializer, node);
     super.visitVariableDeclaration(node);
     VariableElement element = node.element;
     if (element.initializer != null && node.initializer != null) {
@@ -7058,7 +7106,7 @@ class ResolverVisitor extends ScopedVisitor {
    * values which should be returned or yielded as appropriate.  If a type
    * cannot be computed from the declared return type, return null.
    */
-  DartType _computeReturnOrYieldType(DartType declaredType) {
+  TypeContext _computeReturnOrYieldType(DartType declaredType) {
     bool isGenerator = _enclosingFunction.isGenerator;
     bool isAsynchronous = _enclosingFunction.isAsynchronous;
 
@@ -7066,21 +7114,22 @@ class ResolverVisitor extends ScopedVisitor {
     if (!isGenerator && !isAsynchronous) {
       return declaredType;
     }
-    if (isGenerator) {
-      if (declaredType is! InterfaceType) {
-        return null;
+    if (declaredType is InterfaceType) {
+      if (isGenerator) {
+        // If it's sync* we expect Iterable<T>
+        // IF it's async* we expect Stream<T>
+        InterfaceType rawType = isAsynchronous
+            ? typeProvider.streamDynamicType
+            : typeProvider.iterableDynamicType;
+        // Match the types to instantiate the type arguments if possible
+        List<DartType> typeArgs =
+            inferenceContext.matchTypes(rawType, declaredType);
+        return (typeArgs?.length == 1) ? typeArgs[0] : null;
       }
-      // If it's synchronous, we expect Iterable<T>, otherwise Stream<T>
-      InterfaceType rawType = isAsynchronous
-          ? typeProvider.streamDynamicType
-          : typeProvider.iterableDynamicType;
-      // Match the types to instantiate the type arguments if possible
-      List<DartType> typeArgs =
-          inferenceContext.matchTypes(rawType, declaredType);
-      return (typeArgs?.length == 1) ? typeArgs[0] : null;
+      // async functions expect `Future<T> | T`
+      return new FutureUnionTypeContext(declaredType, typeProvider, typeSystem);
     }
-    // Must be asynchronous to reach here, so strip off any layers of Future
-    return declaredType.flattenFutures(typeSystem);
+    return declaredType;
   }
 
   /**
@@ -7153,7 +7202,7 @@ class ResolverVisitor extends ScopedVisitor {
     DartType contextType = node.staticInvokeType;
     if (contextType is FunctionType) {
       DartType originalType = node.function.staticType;
-      DartType returnContextType = InferenceContext.getType(node);
+      DartType returnContextType = InferenceContext.getNonFutureType(node);
       TypeSystem ts = typeSystem;
       if (returnContextType != null &&
           node.typeArguments == null &&
@@ -7311,6 +7360,19 @@ class ResolverVisitor extends ScopedVisitor {
   }
 
   /**
+   * Returns true if this expression is being passed to `Future.then`.
+   *
+   * If so we will apply special typing rules in strong mode, to handle the
+   * implicit union of `S | Future<S>`
+   */
+  bool _isFutureThenLambda(FunctionExpression node) {
+    Element element = node.staticParameterElement?.enclosingElement;
+    return element is MethodElement &&
+        element.name == 'then' &&
+        element.enclosingElement.type.isDartAsyncFuture;
+  }
+
+  /**
    * Return `true` if the given variable is accessed within a closure in the given
    * [AstNode] and also mutated somewhere in variable scope. This information is only
    * available for local variables (including parameters).