From 45abd4580619f219d68b93170e6ac588704db7bf Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Mon, 21 Mar 2011 19:43:18 +0100
Subject: [PATCH] correct isReferenceTo() for variables used in nested
 functions (part of PY-3118)

---
 .../codeInsight/controlflow/ScopeOwner.java   |  1 -
 .../python/psi/impl/PyReferenceImpl.java      | 23 +++++++++++++++++--
 python/testData/findUsages/NestedFunctions.py | 10 ++++++++
 .../jetbrains/python/PyFindUsagesTest.java    |  9 ++++++++
 4 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 python/testData/findUsages/NestedFunctions.py

diff --git a/python/src/com/jetbrains/python/codeInsight/controlflow/ScopeOwner.java b/python/src/com/jetbrains/python/codeInsight/controlflow/ScopeOwner.java
index 64fb3e0bd765..f7d1a8631e63 100644
--- a/python/src/com/jetbrains/python/codeInsight/controlflow/ScopeOwner.java
+++ b/python/src/com/jetbrains/python/codeInsight/controlflow/ScopeOwner.java
@@ -1,6 +1,5 @@
 package com.jetbrains.python.codeInsight.controlflow;
 
-import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
 import com.jetbrains.python.psi.PyElement;
 
 /**
diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceImpl.java
index 9a774f3c432e..15221bbdd2df 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyReferenceImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceImpl.java
@@ -314,8 +314,27 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
         if (element instanceof PyParameter || element instanceof PyTargetExpression) {
           PsiElement ourContainer = PsiTreeUtil.getParentOfType(getElement(), PsiNamedElement.class, PyLambdaExpression.class, PyComprehensionElement.class);
           PsiElement theirContainer = PsiTreeUtil.getParentOfType(element, PsiNamedElement.class, PyLambdaExpression.class, PyComprehensionElement.class);
-          if (ourContainer != null && ourContainer == theirContainer) {
-            return true;
+          if (ourContainer != null) {
+            if (ourContainer == theirContainer) {
+              return true;
+            }
+            if (PsiTreeUtil.isAncestor(theirContainer, ourContainer, true)) {
+              ScopeOwner ourScopeOwner = PsiTreeUtil.getParentOfType(getElement(), ScopeOwner.class);
+              ScopeOwner theirScopeOwner = PsiTreeUtil.getParentOfType(element, ScopeOwner.class);
+              if (ourScopeOwner != theirScopeOwner) {
+                boolean shadowsName = false;
+                while(ourScopeOwner != theirScopeOwner && ourScopeOwner != null) {
+                  if (ControlFlowCache.getScope(ourScopeOwner).containsDeclaration(elementName)) {
+                    shadowsName = true;
+                    break;
+                  }
+                  ourScopeOwner = PsiTreeUtil.getParentOfType(ourScopeOwner, ScopeOwner.class);
+                }
+                if (!shadowsName) {
+                  return true;
+                }
+              }
+            }
           }
         }
         final PsiElement resolveResult = resolve();
diff --git a/python/testData/findUsages/NestedFunctions.py b/python/testData/findUsages/NestedFunctions.py
new file mode 100644
index 000000000000..ff39441de007
--- /dev/null
+++ b/python/testData/findUsages/NestedFunctions.py
@@ -0,0 +1,10 @@
+xyzzy = False
+
+def foo():
+    if xyzzy:
+        b<caret>ar = {}
+    else:
+        bar = []
+    def x():
+        for y in enumerate(bar):
+            pass
diff --git a/python/testSrc/com/jetbrains/python/PyFindUsagesTest.java b/python/testSrc/com/jetbrains/python/PyFindUsagesTest.java
index 0f39e31369ed..e455303ca21d 100644
--- a/python/testSrc/com/jetbrains/python/PyFindUsagesTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFindUsagesTest.java
@@ -64,6 +64,15 @@ public class PyFindUsagesTest extends PyLightFixtureTestCase {
     assertUsages(usages, "import <caret>re", "<caret>re.compile");
   }
 
+  public void testNestedFunctions() {  // PY-3118
+    final Collection<UsageInfo> usages = doTest();
+    assertUsages(usages, "<caret>bar = {}", "<caret>bar = []", "enumerate(<caret>bar)");
+  }
+
+  private Collection<UsageInfo> doTest() {
+    return myFixture.testFindUsages("findUsages/" + getTestName(false) + ".py");
+  }
+
   private void assertUsages(Collection<UsageInfo> usages, String... usageTexts) {
     assertEquals(usageTexts.length, usages.size());
     List<UsageInfo> sortedUsages = new ArrayList<UsageInfo>(usages);