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 4679 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>[];
 
   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 =>
       _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 inferred = _inferredReturn.last;
-    inferred = _typeSystem.getLeastUpperBound(_typeProvider, type, inferred);
-    _inferredReturn[_inferredReturn.length - 1] = inferred;
+    DartType context = _returnStack.last;
+    if (context is! FutureUnionType) {
+      DartType inferred = _inferredReturn.last;
+      inferred = _typeSystem.getLeastUpperBound(_typeProvider, type, inferred);
+      _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>`,
    * such that `I<T0, .., Tn> <: t2`
    */
@@ skipping 16 matching lines @@
       }
     } else {
       assert(false);
     }
   }
 
   /**
    * 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(getContext(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,
+   * [getContext] should be used instead.
+   */
+  static DartType getType(AstNode node) {
+    DartType t = getContext(node);
+    if (t is FutureUnionType) {
+      return t.type;
+    }
+    return t;
+  }
+
+  /**
    * Look for contextual type information attached to [node].  Returns
    * the type if found, otherwise null.
+   *
+   * If [node] has a contextual union type like `T | Future<T>` this will be
+   * returned. You can use [getType] if you prefer to only get the `T`.
    */
-  static DartType getType(AstNode node) => node?.getProperty(_typeProperty);
+  static DartType getContext(AstNode node) => node?.getProperty(_typeProperty);
+
+  /**
+   * Like [getContext] but expands a union type into a list of types.
+   */
+  static Iterable<DartType> getTypes(AstNode node) {
+    DartType t = getContext(node);
+    if (t == null) {
+      return DartType.EMPTY_LIST;
+    }
+    if (t is FutureUnionType) {
+      return t.types;
+    }
+    return <DartType>[t];
+  }
 
   /**
    * Attach contextual type information [type] to [node] for use during
    * inference.
    */
   static void setType(AstNode node, DartType type) {
-    // TODO(jmesserly): this sets the type even when it's dynamic.
-    // Can we skip that?
     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));
+    setType(innerNode, getContext(outerNode));
   }
 }
 
 /**
  * This enum holds one of four states of a field initialization state through a constructor
  * signature, not initialized, initialized in the field declaration, initialized in the field
  * formal, and finally, initialized in the initializers list.
  */
 class INIT_STATE extends Enum<INIT_STATE> {
   static const INIT_STATE NOT_INIT = const INIT_STATE('NOT_INIT', 0);
@@ skipping 1063 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);
+    DartType contextType = InferenceContext.getContext(node);
     if (contextType != null) {
-      InterfaceType futureT = typeProvider.futureType
-          .instantiate([contextType.flattenFutures(typeSystem)]);
-      InferenceContext.setType(node.expression, futureT);
+      var futureUnion =
+          FutureUnionType.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 33 matching lines @@
     } else {
       // TODO(leafp): Do downwards inference using the declared type
       // of the binary operator for other cases.
       if (operatorType == TokenType.QUESTION_QUESTION) {
         InferenceContext.setTypeFromNode(leftOperand, node);
       }
       leftOperand?.accept(this);
       if (operatorType == TokenType.QUESTION_QUESTION) {
         // Set the right side, either from the context, or using the information
         // from the left side if it is more precise.
-        DartType contextType = InferenceContext.getType(node);
+        DartType contextType = InferenceContext.getContext(node);
         DartType leftType = leftOperand?.staticType;
         if (contextType == null || contextType.isDynamic) {
           contextType = leftType;
         }
         InferenceContext.setType(rightOperand, contextType);
       }
       rightOperand?.accept(this);
     }
     node.accept(elementResolver);
     node.accept(typeAnalyzer);
@@ skipping 457 matching lines @@
       _currentFunctionBody = node.body;
       _enclosingFunction = node.element;
       _overrideManager.enterScope();
       try {
         DartType functionType = InferenceContext.getType(node);
         if (functionType is FunctionType) {
           functionType =
               matchFunctionTypeParameters(node.typeParameters, functionType);
           if (functionType is FunctionType) {
             _inferFormalParameterList(node.parameters, functionType);
-            DartType returnType =
-                _computeReturnOrYieldType(functionType.returnType);
+
+            DartType returnType;
+            if (_isFutureThenLambda(node)) {
+              var futureThenType =
+                  InferenceContext.getContext(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 = FutureUnionType.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);
@@ skipping 93 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 123 matching lines @@
    * cannot be computed from the declared return type, return null.
    */
   DartType _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 FutureUnionType(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 224 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 3685 matching lines @@
       return null;
     }
     if (identical(node.staticElement, variable)) {
       if (node.inSetterContext()) {
         result = true;
       }
     }
     return null;
   }
 }