fix #25944, improve Future.then inference

pkg/analyzer/lib/src/generated/resolver.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 analyzer.src.generated.resolver;
 
 import 'dart:collection';
 
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/token.dart';
@@ skipping 4671 matching lines @@
    * of all return statements in a lambda.
    *
    * This will always be kept in sync with [_returnStack].
    */
   final List<DartType> _inferredReturn = <DartType>[];
 
   /**
    * 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)
       : _typeProvider = typeProvider;
 
   /**
    * Get the return type of the current enclosing function, if any.
    *
    * The type returned for a function is the type that is expected
    * to be used in a return or yield context.  For ordinary functions
    * this is the same as the return type of the function.  For async
    * 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;
 
   /**
    * Records the type of the expression of a return statement.
    *
    * This will be used for inferring a block bodied lambda, if no context
    * type was available.
    */
   void addReturnOrYieldType(DartType type) {
     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;
     }
   }
 
   /**
    * Match type [t1] against type [t2] as follows.
    * If `t1 = I<dynamic, ..., dynamic>`, then look for a supertype
    * of t1 of the form `K<S0, ..., Sm>` where `t2 = K<S0', ..., Sm'>`
    * If the supertype exists, use the constraints `S0 <: S0', ... Sm <: Sm'`
    * to derive a concrete instantation for I of the form `<T0, ..., Tn>`,
@@ skipping 19 matching lines @@
       if (!inferred.isBottom) {
         setType(node, inferred);
       }
     }
   }
 
   /**
    * 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);
   }
 
   /**
    * Place an info node into the error stream indicating that a
    * [type] has been inferred as the type of [node].
    */
   void recordInference(Expression node, DartType type) {
     if (!_inferenceHints) {
       return;
@@ skipping 121 matching lines @@
   }
 
   /**
    * Clear the type information assocated with [node].
    */
   static void clearType(AstNode node) {
     node?.setProperty(_typeProperty, null);
   }
 
   /**
+   * 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);
   }
 
   /**
    * Attach contextual type information [type] to [node] for use during
    * inference.
    */
   static void setTypeFromNode(AstNode innerNode, AstNode outerNode) {
     setType(innerNode, getType(outerNode));
   }
@@ skipping 996 matching lines @@
       assert(parent is PartOfDirective);
     } else {
       elementAnnotationImpl.annotationAst =
           new ConstantAstCloner().cloneNode(node);
     }
     return null;
   }
 
   @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;
       List<DartType> normalParameterTypes = callerType.normalParameterTypes;
       List<DartType> optionalParameterTypes = callerType.optionalParameterTypes;
       int normalCount = normalParameterTypes.length;
       int optionalCount = optionalParameterTypes.length;
 
       NodeList<Expression> arguments = node.arguments;
       Iterable<Expression> positional =
@@ skipping 55 matching lines @@
           node.rightHandSide, node.leftHandSide.staticType);
     }
     node.rightHandSide?.accept(this);
     node.accept(elementResolver);
     node.accept(typeAnalyzer);
     return null;
   }
 
   @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);
   }
 
   @override
   Object visitBinaryExpression(BinaryExpression node) {
     TokenType operatorType = node.operator.type;
     Expression leftOperand = node.leftOperand;
     Expression rightOperand = node.rightOperand;
     if (operatorType == TokenType.AMPERSAND_AMPERSAND) {
@@ skipping 496 matching lines @@
 
   @override
   Object visitFunctionExpression(FunctionExpression node) {
     ExecutableElement outerFunction = _enclosingFunction;
     FunctionBody outerFunctionBody = _currentFunctionBody;
     try {
       _currentFunctionBody = node.body;
       _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);
           }
         }
         super.visitFunctionExpression(node);
       } finally {
         _overrideManager.exitScope();
       }
     } finally {
       _currentFunctionBody = outerFunctionBody;
       _enclosingFunction = outerFunction;
@@ skipping 95 matching lines @@
       perBranchOverrides.add(elseOverrides);
       _overrideManager.mergeOverrides(perBranchOverrides);
     }
     return null;
   }
 
   @override
   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;
+          }
         }
       }
     }
     node.constructorName?.accept(this);
     FunctionType constructorType = node.constructorName.staticElement?.type;
     if (constructorType != null) {
       InferenceContext.setType(node.argumentList, constructorType);
     }
     node.argumentList?.accept(this);
     node.accept(elementResolver);
     node.accept(typeAnalyzer);
     return null;
   }
 
   @override
   Object visitLabel(Label node) => null;
 
   @override
   Object visitLibraryIdentifier(LibraryIdentifier node) => null;
 
   @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();
     } else if (contextType is InterfaceType) {
       InterfaceType listD =
           typeProvider.listType.instantiate([typeProvider.dynamicType]);
       targs = inferenceContext.matchTypes(listD, contextType);
     }
     if (targs != null && targs.length == 1 && !targs[0].isDynamic) {
       DartType eType = targs[0];
       InterfaceType listT = typeProvider.listType.instantiate([eType]);
       for (Expression child in node.elements) {
         InferenceContext.setType(child, eType);
       }
       InferenceContext.setType(node, listT);
     } else {
       InferenceContext.clearType(node);
     }
     super.visitListLiteral(node);
     return null;
   }
 
   @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();
     } else if (contextType is InterfaceType) {
       InterfaceType mapD = typeProvider.mapType
           .instantiate([typeProvider.dynamicType, typeProvider.dynamicType]);
       targs = inferenceContext.matchTypes(mapD, contextType);
     }
     if (targs != null && targs.length == 2 && targs.any((t) => !t.isDynamic)) {
       DartType kType = targs[0];
@@ skipping 11 matching lines @@
     return null;
   }
 
   @override
   Object visitMethodDeclaration(MethodDeclaration node) {
     ExecutableElement outerFunction = _enclosingFunction;
     FunctionBody outerFunctionBody = _currentFunctionBody;
     try {
       _currentFunctionBody = node.body;
       _enclosingFunction = node.element;
-      DartType returnType =
+      TypeContext returnType =
           _computeReturnOrYieldType(_enclosingFunction.type?.returnType);
       InferenceContext.setType(node.body, returnType);
       super.visitMethodDeclaration(node);
     } finally {
       _currentFunctionBody = outerFunctionBody;
       _enclosingFunction = outerFunction;
     }
     return null;
   }
 
@@ skipping 13 matching lines @@
     node.typeArguments?.accept(this);
     node.accept(elementResolver);
     _inferArgumentTypesFromContext(node);
     node.argumentList?.accept(this);
     node.accept(typeAnalyzer);
     return null;
   }
 
   @override
   Object visitNamedExpression(NamedExpression node) {
-    InferenceContext.setType(node.expression, InferenceContext.getType(node));
+    InferenceContext.setTypeFromNode(node.expression, node);
     return super.visitNamedExpression(node);
   }
 
   @override
   Object visitNode(AstNode node) {
     node.visitChildren(this);
     node.accept(elementResolver);
     node.accept(typeAnalyzer);
     return null;
   }
 
   @override
   Object visitParenthesizedExpression(ParenthesizedExpression node) {
-    InferenceContext.setType(node.expression, InferenceContext.getType(node));
+    InferenceContext.setTypeFromNode(node.expression, node);
     return super.visitParenthesizedExpression(node);
   }
 
   @override
   Object visitPrefixedIdentifier(PrefixedIdentifier node) {
     //
     // We visit the prefix, but do not visit the identifier because it needs to
     // be visited in the context of the prefix.
     //
     node.prefix?.accept(this);
@@ skipping 97 matching lines @@
       _overrideManager.applyOverrides(overrides);
     }
     return null;
   }
 
   @override
   Object visitTypeName(TypeName node) => null;
 
   @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) {
       (element.initializer as FunctionElementImpl).returnType =
           node.initializer.staticType;
     }
     // Note: in addition to cloning the initializers for const variables, we
     // have to clone the initializers for non-static final fields (because if
     // they occur in a class with a const constructor, they will be needed to
     // evaluate the const constructor).
@@ skipping 115 matching lines @@
         _promoteManager.setType(element, null);
       }
     }
   }
 
   /**
    * Given the declared return type of a function, compute the type of the
    * 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;
 
     // Ordinary functions just return their declared types.
     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;
   }
 
   /**
    * The given expression is the expression used to compute the iterator for a
    * for-each statement. Attempt to compute the type of objects that will be
    * assigned to the loop variable and return that type. Return `null` if the
    * type could not be determined. The [iteratorExpression] is the expression
    * that will return the Iterable being iterated over.
    */
   DartType _getIteratorElementType(Expression iteratorExpression) {
@@ skipping 52 matching lines @@
   void _inferArgumentTypesFromContext(InvocationExpression node) {
     if (!strongMode) {
       // Use propagated type inference for lambdas if not in strong mode.
       _inferFunctionExpressionsParametersTypes(node.argumentList);
       return;
     }
 
     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 &&
           originalType is FunctionType &&
           originalType.typeFormals.isNotEmpty &&
           ts is StrongTypeSystemImpl) {
         contextType = ts.inferGenericFunctionCall(typeProvider, originalType,
             DartType.EMPTY_LIST, DartType.EMPTY_LIST, returnContextType);
       }
 
@@ skipping 137 matching lines @@
       //
       // Unsound to assume that [x = "hello";] never executed after the
       // if-statement. Of course, a dead-code analysis could point out that
       // [return] here is dead.
       return _isAbruptTerminationStatement(statements[size - 1]);
     }
     return false;
   }
 
   /**
+   * 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).
    *
    * @param variable the variable to check
    * @param target the [AstNode] to check within
    * @return `true` if this variable is potentially mutated somewhere in the given ASTNode
    */
   bool _isVariableAccessedInClosure(Element variable, AstNode target) {
     _ResolverVisitor_isVariableAccessedInClosure visitor =
@@ skipping 3670 matching lines @@
       return null;
     }
     if (identical(node.staticElement, variable)) {
       if (node.inSetterContext()) {
         result = true;
       }
     }
     return null;
   }
 }