[PSI] Introduce "context" soft keyword and parse context receivers

compiler/psi/src/org/jetbrains/kotlin/KtNodeTypes.java

@@ -39,6 +39,7 @@ public interface KtNodeTypes {
     IElementType SCRIPT_INITIALIZER     = new KtNodeType("SCRIPT_INITIALIZER", KtScriptInitializer.class);
     IElementType SECONDARY_CONSTRUCTOR  = KtStubElementTypes.SECONDARY_CONSTRUCTOR;
     IElementType PRIMARY_CONSTRUCTOR    = KtStubElementTypes.PRIMARY_CONSTRUCTOR;
+    IElementType CONTEXT_RECEIVER       = KtStubElementTypes.CONTEXT_RECEIVER;
 
     IElementType TYPE_PARAMETER_LIST                = KtStubElementTypes.TYPE_PARAMETER_LIST;
     IElementType TYPE_PARAMETER                     = KtStubElementTypes.TYPE_PARAMETER;

compiler/psi/src/org/jetbrains/kotlin/lexer/KtTokens.java

@@ -153,6 +153,7 @@ public interface KtTokens {
     KtKeywordToken SET_KEYWORD       = KtKeywordToken.softKeyword("set");
     KtKeywordToken CONSTRUCTOR_KEYWORD = KtKeywordToken.softKeyword("constructor");
     KtKeywordToken INIT_KEYWORD        = KtKeywordToken.softKeyword("init");
+    KtKeywordToken CONTEXT_KEYWORD     = KtKeywordToken.softKeyword("context");
 
     KtModifierKeywordToken ABSTRACT_KEYWORD  = KtModifierKeywordToken.softKeywordModifier("abstract");
     KtModifierKeywordToken ENUM_KEYWORD      = KtModifierKeywordToken.softKeywordModifier("enum");
@@ -221,7 +222,7 @@ public interface KtTokens {
                                              DATA_KEYWORD, INLINE_KEYWORD, NOINLINE_KEYWORD, TAILREC_KEYWORD, EXTERNAL_KEYWORD,
                                              ANNOTATION_KEYWORD, CROSSINLINE_KEYWORD, CONST_KEYWORD, OPERATOR_KEYWORD, INFIX_KEYWORD,
                                              SUSPEND_KEYWORD, HEADER_KEYWORD, IMPL_KEYWORD, EXPECT_KEYWORD, ACTUAL_KEYWORD,
-                                             VALUE_KEYWORD
+                                             VALUE_KEYWORD, CONTEXT_KEYWORD
     );
 
     /*

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParsing.java

@@ -423,6 +423,10 @@ public class KotlinParsing extends AbstractKotlinParsing {
         }
         PsiBuilder.Marker decl = mark();
 
+        if (at(CONTEXT_KEYWORD)) {
+            parseContextReceiver();
+        }
+
         ModifierDetector detector = new ModifierDetector();
         parseModifierList(detector, DEFAULT, TokenSet.EMPTY);
 
@@ -612,6 +616,41 @@ public class KotlinParsing extends AbstractKotlinParsing {
         return false;
     }
 
+    /*
+     * contextReceiver
+     *   : "context" "(" typeReference{","}+ ")"
+     */
+    private void parseContextReceiver() {
+        assert _at(CONTEXT_KEYWORD);
+        PsiBuilder.Marker contextReceiver = mark();
+        advance(); // CONTEXT_KEYWORD
+        if (at(LPAR)) {
+            advance(); // LPAR
+            while (true) {
+                if (at(COMMA)) {
+                    errorAndAdvance("Expecting a type reference");
+                }
+                parseTypeRef();
+                if (at(RPAR)) {
+                    advance();
+                    break;
+                }
+                if (at(COMMA)) {
+                    advance();
+                }
+                else {
+                    if (!at(RPAR)) {
+                        error("Expecting comma or ')'");
+                    }
+                }
+            }
+            contextReceiver.done(CONTEXT_RECEIVER);
+        } else {
+            errorWithRecovery("Expecting context receivers", TokenSet.EMPTY);
+            contextReceiver.drop();
+        }
+    }
+
     /*
      * fileAnnotationList
      *   : ("[" "file:" annotationEntry+ "]")*
@@ -1178,6 +1217,10 @@ public class KotlinParsing extends AbstractKotlinParsing {
         }
         PsiBuilder.Marker decl = mark();
 
+        if (at(CONTEXT_KEYWORD)) {
+            parseContextReceiver();
+        }
+
         ModifierDetector detector = new ModifierDetector();
         parseModifierList(detector, DEFAULT, TokenSet.EMPTY);
 

compiler/psi/src/org/jetbrains/kotlin/psi/KtContextReceiver.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.psi
+
+import com.intellij.lang.ASTNode
+import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+
+class KtContextReceiver : KtElementImplStub<KotlinPlaceHolderStub<KtContextReceiver>> {
+    constructor(node: ASTNode) : super(node)
+    constructor(stub: KotlinPlaceHolderStub<KtContextReceiver>) : super(stub, KtStubElementTypes.CONTEXT_RECEIVER)
+
+    override fun <R : Any?, D : Any?> accept(visitor: KtVisitor<R, D>, data: D): R {
+        return visitor.visitContextReceiver(this, data)
+    }
+
+    fun typeReferences(): List<KtTypeReference> = findChildrenByType(KtNodeTypes.TYPE_REFERENCE)
+}

compiler/psi/src/org/jetbrains/kotlin/psi/KtVisitor.java

@@ -158,6 +158,10 @@ public class KtVisitor<R, D> extends PsiElementVisitor {
         return visitSuperTypeListEntry(specifier, data);
     }
 
+    public R visitContextReceiver(@NotNull KtContextReceiver contextReceiver, D data) {
+        return visitKtElement(contextReceiver, data);
+    }
+
     public R visitConstructorDelegationCall(@NotNull KtConstructorDelegationCall call, D data) {
         return visitKtElement(call, data);
     }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/elements/KtStubElementTypes.java

@@ -143,6 +143,9 @@ public interface KtStubElementTypes {
     KtPlaceHolderStubElementType<KtConstructorCalleeExpression> CONSTRUCTOR_CALLEE =
             new KtPlaceHolderStubElementType<>("CONSTRUCTOR_CALLEE", KtConstructorCalleeExpression.class);
 
+    KtPlaceHolderStubElementType<KtContextReceiver> CONTEXT_RECEIVER =
+            new KtPlaceHolderStubElementType<>("CONTEXT_RECEIVER", KtContextReceiver.class);
+
     KtConstantExpressionElementType NULL                = new KtConstantExpressionElementType("NULL");
     KtConstantExpressionElementType BOOLEAN_CONSTANT    = new KtConstantExpressionElementType("BOOLEAN_CONSTANT");
     KtConstantExpressionElementType FLOAT_CONSTANT      = new KtConstantExpressionElementType("FLOAT_CONSTANT");