From 2121f707abf1ea69aa4531c06a9e93ef25a14bd5 Mon Sep 17 00:00:00 2001
From: Natalia Ukhorskaya <natalia.ukhorskaya@jetbrains.com>
Date: Mon, 20 Jul 2015 16:37:36 +0300
Subject: [PATCH] Fix completion for codeFragments on functions with expression
 body

---
 .../resolve/util/codeFragmentResolveUtils.kt  | 27 ++++++-------------
 .../codeFragments/withoutBodyFunction.kt      |  1 +
 .../withoutBodyFunction.kt.fragment           |  1 +
 .../codeFragments/withoutBodyProperty.kt      |  6 +++++
 .../withoutBodyProperty.kt.fragment           |  1 +
 ...CodeFragmentHighlightingTestGenerated.java | 12 +++++++++
 6 files changed, 29 insertions(+), 19 deletions(-)
 create mode 100644 idea/testData/checker/codeFragments/withoutBodyFunction.kt
 create mode 100644 idea/testData/checker/codeFragments/withoutBodyFunction.kt.fragment
 create mode 100644 idea/testData/checker/codeFragments/withoutBodyProperty.kt
 create mode 100644 idea/testData/checker/codeFragments/withoutBodyProperty.kt.fragment

diff --git a/idea/ide-common/src/org/jetbrains/kotlin/resolve/util/codeFragmentResolveUtils.kt b/idea/ide-common/src/org/jetbrains/kotlin/resolve/util/codeFragmentResolveUtils.kt
index b9d79298474..9dc91eb9da8 100644
--- a/idea/ide-common/src/org/jetbrains/kotlin/resolve/util/codeFragmentResolveUtils.kt
+++ b/idea/ide-common/src/org/jetbrains/kotlin/resolve/util/codeFragmentResolveUtils.kt
@@ -30,27 +30,16 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 //TODO: this code should be moved into debugger which should set correct context for its code fragment
 private fun JetExpression.correctContextForExpression(): JetExpression {
-    when (this) {
-        is JetBlockExpression -> {
-            return this.getStatements().lastOrNull() ?: this
-        }
-        is JetFunctionLiteral -> {
-            return this.getBodyExpression()?.getStatements()?.lastOrNull() ?: this
-        }
+    return when (this) {
+        is JetProperty -> this.getDelegateExpressionOrInitializer()
+        is JetFunctionLiteral -> this.getBodyExpression()?.getStatements()?.lastOrNull()
+        is JetDeclarationWithBody -> this.getBodyExpression()
+        is JetBlockExpression -> this.getStatements().lastOrNull()
         else -> {
-            val previousExpression = this.siblings(forward = false, withItself = false).firstIsInstanceOrNull<JetExpression>()
-            if (previousExpression is JetExpression) {
-                return previousExpression
-            }
-            else {
-                val parentExpression = this.parents.firstIsInstanceOrNull<JetExpression>()
-                if (parentExpression is JetExpression) {
-                    return parentExpression
-                }
-            }
+            this.siblings(forward = false, withItself = false).firstIsInstanceOrNull<JetExpression>()
+                    ?: this.parents.firstIsInstanceOrNull<JetExpression>()
         }
-    }
-    return this
+    } ?: this
 }
 
 public fun JetCodeFragment.getScopeAndDataFlowForAnalyzeFragment(
diff --git a/idea/testData/checker/codeFragments/withoutBodyFunction.kt b/idea/testData/checker/codeFragments/withoutBodyFunction.kt
new file mode 100644
index 00000000000..2d5499e583f
--- /dev/null
+++ b/idea/testData/checker/codeFragments/withoutBodyFunction.kt
@@ -0,0 +1 @@
+<caret>fun foo(aaaB: Int) = aaaB
\ No newline at end of file
diff --git a/idea/testData/checker/codeFragments/withoutBodyFunction.kt.fragment b/idea/testData/checker/codeFragments/withoutBodyFunction.kt.fragment
new file mode 100644
index 00000000000..6924940e079
--- /dev/null
+++ b/idea/testData/checker/codeFragments/withoutBodyFunction.kt.fragment
@@ -0,0 +1 @@
+aaaB
\ No newline at end of file
diff --git a/idea/testData/checker/codeFragments/withoutBodyProperty.kt b/idea/testData/checker/codeFragments/withoutBodyProperty.kt
new file mode 100644
index 00000000000..efb72646336
--- /dev/null
+++ b/idea/testData/checker/codeFragments/withoutBodyProperty.kt
@@ -0,0 +1,6 @@
+package withoutBodyProperties
+
+class A {
+    val aaaB = 1
+    <caret>val aaaC = 1
+}
diff --git a/idea/testData/checker/codeFragments/withoutBodyProperty.kt.fragment b/idea/testData/checker/codeFragments/withoutBodyProperty.kt.fragment
new file mode 100644
index 00000000000..8d9b093ecf4
--- /dev/null
+++ b/idea/testData/checker/codeFragments/withoutBodyProperty.kt.fragment
@@ -0,0 +1 @@
+aaaB + aaaC
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/CodeFragmentHighlightingTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/CodeFragmentHighlightingTestGenerated.java
index 346da628466..15da139a3e2 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/CodeFragmentHighlightingTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/debugger/evaluate/CodeFragmentHighlightingTestGenerated.java
@@ -132,6 +132,18 @@ public class CodeFragmentHighlightingTestGenerated extends AbstractCodeFragmentH
             String fileName = JetTestUtils.navigationMetadata("idea/testData/checker/codeFragments/smartCasts.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("withoutBodyFunction.kt")
+        public void testWithoutBodyFunction() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/checker/codeFragments/withoutBodyFunction.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("withoutBodyProperty.kt")
+        public void testWithoutBodyProperty() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/checker/codeFragments/withoutBodyProperty.kt");
+            doTest(fileName);
+        }
     }
 
     @TestMetadata("idea/testData/checker/codeFragments/imports")