[PSI] Add some useful getters to psi nodes

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)
-}
+class KtContractEffect: KtElementImplStub<KotlinContractEffectStub> {
+    constructor(node: ASTNode): super(node)
+    constructor(stub: KotlinContractEffectStub): super(stub, KtStubElementTypes.CONTRACT_EFFECT)
+}
+
+fun KtContractEffect.getExpression(): KtExpression = getChildOfType()!!

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)
-}
+    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()
+        }

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();

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);
     }

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;