From d53f3b9ba824c405abb708a6038368988ca6ae01 Mon Sep 17 00:00:00 2001
From: Arsen Nagdalian <arsen.nagdalian@jetbrains.com>
Date: Wed, 22 Jul 2020 04:22:36 +0300
Subject: [PATCH] [PSI] Add some useful getters to psi nodes

---
 .../src/org/jetbrains/kotlin/psi/KtContractEffect.kt | 12 ++++++++----
 .../org/jetbrains/kotlin/psi/KtContractEffectList.kt | 12 +++++++++---
 .../jetbrains/kotlin/psi/KtDeclarationWithBody.java  |  5 +++++
 .../org/jetbrains/kotlin/psi/KtNamedFunction.java    |  6 ++++++
 .../org/jetbrains/kotlin/psi/KtPropertyAccessor.java |  6 ++++++
 5 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffect.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffect.kt
index 44435bc515b..2197f6fc698 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffect.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffect.kt
@@ -6,10 +6,14 @@
 package org.jetbrains.kotlin.psi
 
 import com.intellij.lang.ASTNode
+import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.psi.stubs.KotlinContractEffectStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+import org.jetbrains.kotlin.types.checker.findCorrespondingSupertype
 
-class KtContractEffect: KtElementImplStub<KotlinContractEffectStub> { // TODO: check constructors
-    public constructor(node: ASTNode): super(node)
-    public constructor(stub: KotlinContractEffectStub): super(stub, KtStubElementTypes.CONTRACT_EFFECT)
-}
\ No newline at end of file
+class KtContractEffect: KtElementImplStub<KotlinContractEffectStub> {
+    constructor(node: ASTNode): super(node)
+    constructor(stub: KotlinContractEffectStub): super(stub, KtStubElementTypes.CONTRACT_EFFECT)
+}
+
+fun KtContractEffect.getExpression(): KtExpression = getChildOfType()!!
\ No newline at end of file
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffectList.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffectList.kt
index 2a1b18d75c4..6b982010f80 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffectList.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtContractEffectList.kt
@@ -10,6 +10,12 @@ import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 
 class KtContractEffectList : KtElementImplStub<KotlinPlaceHolderStub<KtContractEffectList>> {
-    public constructor(node: ASTNode): super(node)
-    public constructor(stub: KotlinPlaceHolderStub<KtContractEffectList>): super(stub, KtStubElementTypes.CONTRACT_EFFECT_LIST)
-}
\ No newline at end of file
+    constructor(node: ASTNode) : super(node)
+    constructor(stub: KotlinPlaceHolderStub<KtContractEffectList>) : super(stub, KtStubElementTypes.CONTRACT_EFFECT_LIST)
+}
+
+fun KtContractEffectList.getExpressions(): List<KtExpression> =
+    getStubOrPsiChildrenAsList(KtStubElementTypes.CONTRACT_EFFECT)
+        .map {
+            it.getExpression()
+        }
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java b/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java
index e424f6b2d2b..724d15e128d 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java
@@ -33,6 +33,11 @@ public interface KtDeclarationWithBody extends KtDeclaration {
     @Nullable
     String getName();
 
+    @Nullable
+    default KtContractEffectList getContractDescription() {
+        return null;
+    }
+
     boolean hasBlockBody();
 
     boolean hasBody();
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtNamedFunction.java b/compiler/psi/src/org/jetbrains/kotlin/psi/KtNamedFunction.java
index 17b960da41b..15741d2c165 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtNamedFunction.java
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtNamedFunction.java
@@ -26,6 +26,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.AstLoadingFilter;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.KtNodeTypes;
 import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.psi.psiUtil.KtPsiUtilKt;
 import org.jetbrains.kotlin.psi.stubs.KotlinFunctionStub;
@@ -245,6 +246,11 @@ public class KtNamedFunction extends KtTypeParameterListOwnerStub<KotlinFunction
         return false;
     }
 
+    @Override
+    public KtContractEffectList getContractDescription() {
+        return findChildByType(KtNodeTypes.CONTRACT_EFFECT_LIST);
+    }
+
     public boolean mayHaveContract() {
         return mayHaveContract(true);
     }
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtPropertyAccessor.java b/compiler/psi/src/org/jetbrains/kotlin/psi/KtPropertyAccessor.java
index a6483b1576d..44b96660010 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtPropertyAccessor.java
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtPropertyAccessor.java
@@ -22,6 +22,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.AstLoadingFilter;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.KtNodeTypes;
 import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.psi.stubs.KotlinPropertyAccessorStub;
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes;
@@ -137,6 +138,11 @@ public class KtPropertyAccessor extends KtDeclarationStub<KotlinPropertyAccessor
         return findChildByType(KtTokens.EQ);
     }
 
+    @Override
+    public KtContractEffectList getContractDescription() {
+        return findChildByType(KtNodeTypes.CONTRACT_EFFECT_LIST);
+    }
+
     @Override
     public boolean hasDeclaredReturnType() {
         return true;