Drop psi builder cache on pushing or popping parsing scope (PY-36167)

Since builder could cache something related to the current scope (e.g. token type)

GitOrigin-RevId: ffabda058dac874bd9187fa27d30b823f782fcf6

python/python-psi-impl/src/com/jetbrains/python/parsing/ParsingContext.java

@@ -42,11 +42,15 @@ public class ParsingContext {
 
   @NotNull
   public ParsingScope popScope() {
-    return myScopes.pop();
+    final ParsingScope prevScope = myScopes.pop();
+    resetBuilderCache(prevScope);
+    return prevScope;
   }
 
   public void pushScope(@NotNull ParsingScope scope) {
+    final ParsingScope prevScope = getScope();
     myScopes.push(scope);
+    resetBuilderCache(prevScope);
   }
 
   @NotNull
@@ -77,4 +81,12 @@ public class ParsingContext {
   public ParsingScope emptyParsingScope() {
     return new ParsingScope();
   }
+
+  private void resetBuilderCache(@NotNull ParsingScope prevScope) {
+    if (!getScope().equals(prevScope)) {
+      // builder could cache some data based on the previous scope (e.g. token type)
+      // and we have to reset its cache if scope has changed
+      getBuilder().mark().rollbackTo();
+    }
+  }
 }

python/python-psi-impl/src/com/jetbrains/python/parsing/ParsingScope.java

@@ -15,6 +15,8 @@
  */
 package com.jetbrains.python.parsing;
 
+import java.util.Objects;
+
 /**
  * @author vlan
  */
@@ -82,4 +84,22 @@ public class ParsingScope {
     result.mySuite = mySuite;
     return result;
   }
+
+  @Override
+  public boolean equals(Object o) {
+    if (o == this) return true;
+    if (!(o instanceof ParsingScope)) return false;
+
+    final ParsingScope scope = (ParsingScope)o;
+    return myFunction == scope.myFunction &&
+           myClass == scope.myClass &&
+           mySuite == scope.mySuite &&
+           myAfterSemicolon == scope.myAfterSemicolon &&
+           myAsync == scope.myAsync;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(myFunction, myClass, mySuite, myAfterSemicolon, myAsync);
+  }
 }

python/python-psi-impl/test/com/jetbrains/python/parsing/PythonParsingTest.java

@@ -925,6 +925,11 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest(LanguageLevel.PYTHON38);
   }
 
+  // PY-36167
+  public void testFunctionWithPassAndAwaitAfterInPy36() {
+    doTest(LanguageLevel.PYTHON36);
+  }
+
   public void doTest() {
     doTest(LanguageLevel.PYTHON26);
   }

python/testData/psi/FunctionWithPassAndAwaitAfterInPy36.py

@@ -0,0 +1,3 @@
+async def a():
+    def e(): pass
+    await undefined

python/testData/psi/FunctionWithPassAndAwaitAfterInPy36.txt

@@ -0,0 +1,32 @@
+PyFile:FunctionWithPassAndAwaitAfterInPy36.py
+  PyFunction('a')
+    PsiElement(Py:ASYNC_KEYWORD)('async')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:DEF_KEYWORD)('def')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('a')
+    PyParameterList
+      PsiElement(Py:LPAR)('(')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n    ')
+    PyStatementList
+      PyFunction('e')
+        PsiElement(Py:DEF_KEYWORD)('def')
+        PsiWhiteSpace(' ')
+        PsiElement(Py:IDENTIFIER)('e')
+        PyParameterList
+          PsiElement(Py:LPAR)('(')
+          PsiElement(Py:RPAR)(')')
+        PsiElement(Py:COLON)(':')
+        PsiWhiteSpace(' ')
+        PyStatementList
+          PyPassStatement
+            PsiElement(Py:PASS_KEYWORD)('pass')
+      PsiWhiteSpace('\n    ')
+      PyExpressionStatement
+        PyPrefixExpression
+          PsiElement(Py:AWAIT_KEYWORD)('await')
+          PsiWhiteSpace(' ')
+          PyReferenceExpression: undefined
+            PsiElement(Py:IDENTIFIER)('undefined')