lambda: smart completion: substitute return type; respect one line lambda as return statement (IDEA-90941; IDEA-90944)

java/java-impl/src/com/intellij/codeInsight/ExpectedTypesProvider.java

@@ -330,29 +330,52 @@ public class ExpectedTypesProvider {
       return null;
     }
 
-    @Override public void visitReturnStatement(PsiReturnStatement statement) {
-      final PsiLambdaExpression lambdaExpression = PsiTreeUtil.getParentOfType(statement, PsiLambdaExpression.class);
-      final PsiMethod interfaceMethod = lambdaExpression != null ? LambdaUtil.getFunctionalInterfaceMethod(lambdaExpression.getFunctionalInterfaceType()) : null;
-      final PsiMethod scopeMethod = interfaceMethod != null ? interfaceMethod : PsiTreeUtil.getParentOfType(statement, PsiMethod.class);
+    @Override
+    public void visitLambdaExpression(PsiLambdaExpression lambdaExpression) {
+      super.visitLambdaExpression(lambdaExpression);
+      final PsiType functionalInterfaceType = lambdaExpression.getFunctionalInterfaceType();
+      final PsiMethod scopeMethod = LambdaUtil.getFunctionalInterfaceMethod(functionalInterfaceType);
       if (scopeMethod != null) {
-        PsiType type = scopeMethod.getReturnType();
-        if (type != null) {
-          ExpectedTypeInfoImpl info = createInfoImpl(type, ExpectedTypeInfo.TYPE_OR_SUBTYPE, type,
-                                                     TailType.SEMICOLON);
-          if (PropertyUtil.isSimplePropertyAccessor(scopeMethod)) {
-            info.expectedName = new NullableComputable<String>() {
-              @Override
-              public String compute() {
-                return PropertyUtil.getPropertyName(scopeMethod);
-              }
-            };
-          }
+        visitMethodReturnType(scopeMethod, LambdaUtil.getFunctionalInterfaceReturnType(functionalInterfaceType));
+      }
+    }
 
-          myResult = new ExpectedTypeInfo[]{info};
-        }
-        else {
-          myResult = ExpectedTypeInfo.EMPTY_ARRAY;
+    @Override public void visitReturnStatement(PsiReturnStatement statement) {
+      final PsiMethod method;
+      final PsiType type;
+      final PsiLambdaExpression lambdaExpression = PsiTreeUtil.getParentOfType(statement, PsiLambdaExpression.class);
+      if (lambdaExpression != null) {
+        final PsiType functionalInterfaceType = lambdaExpression.getFunctionalInterfaceType();
+        method = LambdaUtil.getFunctionalInterfaceMethod(functionalInterfaceType);
+        type = LambdaUtil.getFunctionalInterfaceReturnType(functionalInterfaceType);
+      }
+      else {
+        method = PsiTreeUtil.getParentOfType(statement, PsiMethod.class);
+        type = method != null ? method.getReturnType() : null;
+      }
+
+      if (method != null) {
+        visitMethodReturnType(method, type);
+      }
+    }
+
+    private void visitMethodReturnType(final PsiMethod scopeMethod, PsiType type) {
+      if (type != null) {
+        ExpectedTypeInfoImpl info = createInfoImpl(type, ExpectedTypeInfo.TYPE_OR_SUBTYPE, type,
+                                                   TailType.SEMICOLON);
+        if (PropertyUtil.isSimplePropertyAccessor(scopeMethod)) {
+          info.expectedName = new NullableComputable<String>() {
+            @Override
+            public String compute() {
+              return PropertyUtil.getPropertyName(scopeMethod);
+            }
+          };
         }
+
+        myResult = new ExpectedTypeInfo[]{info};
+      }
+      else {
+        myResult = ExpectedTypeInfo.EMPTY_ARRAY;
       }
     }
 

java/java-tests/testData/codeInsight/completion/smartType/ExpectedReturnType1-out.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2000-2012 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+interface I {
+  boolean foo(String s);
+}
+public class Test {
+  public int m() {
+    I i = s -> s.contains();
+    return 0;
+  }
+}

java/java-tests/testData/codeInsight/completion/smartType/ExpectedReturnType1.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2000-2012 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+interface I {
+  boolean foo(String s);
+}
+public class Test {
+  public int m() {
+    I i = s -> s.<caret>
+    return 0;
+  }
+}

java/java-tests/testData/codeInsight/completion/smartType/ExpectedReturnTypeWithSubstitution.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2000-2012 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+public class Test {
+  public interface Some<T> {
+    T test(T p1, T p2);
+  }
+
+  class Foo {
+    public Foo create() {return null;}
+    public Foo create1() {return null;}
+  }
+
+  {
+    Some<Foo> some = (i1, i2) -> {
+      return i1.<caret>
+    };
+  }
+}

java/java-tests/testSrc/com/intellij/codeInsight/completion/SmartType18CompletionTest.java

@@ -43,6 +43,14 @@ public class SmartType18CompletionTest extends LightFixtureCompletionTestCase {
     doTest();
   }
 
+  public void testExpectedReturnTypeWithSubstitution() {
+    doTest();
+  }
+
+  public void testExpectedReturnType1() {
+    doTest();
+  }
+
   private void doTest() {
     configureByFile("/" + getTestName(false) + ".java");
     assertNotNull(myItems);