Don't resolve target expression to latest defs syntactically below it in the PSI (PY-7541)

Resolving to such elements may become a source of resolve loops.

python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -373,6 +373,11 @@ public class PyPsiUtils {
     return expression instanceof PyStringLiteralExpression ? ((PyStringLiteralExpression)expression).getStringValue() : null;
   }
 
+  public static boolean isBefore(@NotNull final PsiElement element, @NotNull final PsiElement element2) {
+    // TODO: From RubyPsiUtil, should be moved to PsiTreeUtil
+    return element.getTextOffset() <= element2.getTextOffset();
+  }
+
   private static abstract class TopLevelVisitor extends PyRecursiveElementVisitor {
     public void visitPyElement(final PyElement node) {
       super.visitPyElement(node);

python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java

@@ -227,6 +227,19 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
           else if (owner == originalOwner && !scope.isGlobal(referencedName)) {
             final ResolveResultList latest = resolveToLatestDefs(owner, myElement, referencedName);
             if (!latest.isEmpty()) {
+              if (myElement instanceof PyTargetExpression) {
+                final RatedResolveResult result = latest.get(0);
+                final PsiElement element = result.getElement();
+                if (element instanceof PyTargetExpression) {
+                  if (PyPsiUtils.isBefore(element, myElement)) {
+                    return latest;
+                  }
+                  else {
+                    ret.poke(myElement, getRate(myElement));
+                    return ret;
+                  }
+                }
+              }
               return latest;
             }
             if (owner instanceof PyClass) {

python/testData/resolve/LoopToLowerReassignment.py

@@ -0,0 +1,5 @@
+def f():
+    while True:
+        foo = 1
+#        <ref>
+        foo = 2

python/testData/resolve/LoopToUpperReassignment.py

@@ -0,0 +1,5 @@
+def f():
+    while True:
+        foo = 1
+        foo = 2
+#        <ref>

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -7,6 +7,7 @@ import com.intellij.psi.ResolveResult;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.fixtures.PyResolveTestCase;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 import com.jetbrains.python.psi.resolve.ImportedResolveResult;
 
@@ -483,4 +484,23 @@ public class PyResolveTest extends PyResolveTestCase {
   public void testReferenceInDocstring() {
     assertResolvesTo(PyClass.class, "datetime");
   }
+
+  // PY-7541
+  public void testLoopToUpperReassignment() {
+    final PsiReference ref = findReferenceByMarker();
+    final PsiElement source = ref.getElement();
+    final PsiElement target = ref.resolve();
+    assertNotNull(target);
+    assertTrue(source != target);
+    assertTrue(PyPsiUtils.isBefore(target, source));
+  }
+
+  // PY-7541
+  public void testLoopToLowerReassignment() {
+    final PsiReference ref = findReferenceByMarker();
+    final PsiElement source = ref.getElement();
+    final PsiElement target = ref.resolve();
+    assertNotNull(target);
+    assertTrue(source == target);
+  }
 }