Initial PyReachingDefs implementation

python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -313,11 +313,11 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
       myBuilder.addPendingEdge(node, myBuilder.prevInstruction);
     }
 
-    final ArrayList<Instruction> rescueInstructions = new ArrayList<Instruction>();
+    final ArrayList<Instruction> exceptInstructions = new ArrayList<Instruction>();
     for (PyExceptPart exceptPart : node.getExceptParts()) {
       myBuilder.prevInstruction = lastBlockInstruction;
-      final Instruction rescueInstruction = myBuilder.startNode(exceptPart);
-      rescueInstructions.add(rescueInstruction);
+      final Instruction exceptInstruction = myBuilder.startNode(exceptPart);
+      exceptInstructions.add(exceptInstruction);
       exceptPart.accept(this);
       myBuilder.addPendingEdge(node, myBuilder.prevInstruction);
     }
@@ -340,7 +340,7 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
         // handle raise instructions inside compound statement
         if (pendingElement instanceof PyRaiseStatement && PsiTreeUtil.isAncestor(tryPart, pendingElement, false)) {
-          for (Instruction rescueInstruction : rescueInstructions) {
+          for (Instruction rescueInstruction : exceptInstructions) {
             myBuilder.addEdge(instruction, rescueInstruction);
           }
           return;
@@ -367,6 +367,23 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
   @Override
   public void visitPyListCompExpression(final PyListCompExpression node) {
-    super.visitPyListCompExpression(node);
+    myBuilder.startNode(node);
+    for (ComprhIfComponent component : node.getIfComponents()) {
+      final PyExpression condition = component.getTest();
+      condition.accept(this);
+      final Instruction head = myBuilder.prevInstruction;
+      final Instruction prevInstruction = myBuilder.startConditionalNode(condition, condition, true);
+      // restore head
+      myBuilder.prevInstruction = head;
+      myBuilder.addPendingEdge(node, myBuilder.startConditionalNode(condition, condition, false)); // false condition
+      myBuilder.prevInstruction = prevInstruction;
+    }
+
+    for (ComprhForComponent forComponent : node.getForComponents()) {
+      forComponent.getIteratedList().accept(this);
+      forComponent.getIteratorVariable().accept(this);
+    }
+
+    node.getResultExpression().accept(this);
   }
 }

python/src/com/jetbrains/python/codeInsight/dataflow/PyReachingDefsDfaInstance.java

@@ -0,0 +1,74 @@
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.codeInsight.controlflow.Instruction;
+import com.intellij.codeInsight.dataflow.DFAMap;
+import com.intellij.codeInsight.dataflow.DfaInstance;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.codeInsight.controlflow.WriteInstruction;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Map;
+
+/**
+ * @author oleg
+ */
+public class PyReachingDefsDfaInstance implements DfaInstance<ScopeVariable> {
+  // Use this its own map, because check in PyReachingDefsDfaSemilattice is important
+  public static final DFAMap<ScopeVariable> INITIAL_MAP = new DFAMap<ScopeVariable>();
+
+  public DFAMap<ScopeVariable> fun(final DFAMap<ScopeVariable> map, Instruction instruction) {
+    final PsiElement element = instruction.getElement();
+    //if (element == null || element.getUserData(ReferenceCompletionUtil.REFERENCE_BEING_COMPLETED)!=null){
+    //  return map;
+    //}
+
+    // Scope reduction
+    final DFAMap<ScopeVariable> reducedMap = new DFAMap<ScopeVariable>();
+    for (Map.Entry<String, ScopeVariable> entry : map.entrySet()) {
+      final ScopeVariable value = entry.getValue();
+      if (element != null && PsiTreeUtil.isAncestor(value.getScope(), element, false)){
+        reducedMap.put(entry.getKey(), value);
+      }
+    }
+
+    return processReducedMap(reducedMap, instruction, element);
+  }
+
+  private static DFAMap<ScopeVariable> processReducedMap(final DFAMap<ScopeVariable> map,
+                                                         final Instruction instruction,
+                                                         final PsiElement element) {
+    if (instruction instanceof WriteInstruction) {
+      final WriteInstruction wInstruction = (WriteInstruction)instruction;
+      final String name = wInstruction.getName();
+      final ScopeVariable variable = map.get(name);
+      // If parameter
+      if (UsageAnalyzer.isParameter(element)) {
+        final PsiElement scope = ScopeUtil.getScopeElement(element);
+        final ScopeVariable scopeVariable = new ScopeVariableImpl(name, true, scope, element);
+        map.put(name, scopeVariable);
+      } else {
+        final ScopeVariableImpl scopeVariable;
+        final boolean isParameter = variable != null && variable.isParameter();
+        if (variable == null) {
+          final PsiElement scope = ScopeUtil.getScopeElement(element);
+          scopeVariable = new ScopeVariableImpl(name, isParameter, scope, element);
+        } else {
+          scopeVariable = new ScopeVariableImpl(name, isParameter, variable.getScope(),
+                                                variable.getDeclarations());
+        }
+        map.put(name, scopeVariable);
+      }
+    }
+    return map;
+  }
+
+  @NotNull
+  public DFAMap<ScopeVariable> initial() {
+    return INITIAL_MAP;
+  }
+
+  public boolean isForward() {
+    return true;
+  }
+}

python/src/com/jetbrains/python/codeInsight/dataflow/PyReachingDefsSemilattice.java

@@ -0,0 +1,79 @@
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.codeInsight.dataflow.DFAMap;
+import com.intellij.codeInsight.dataflow.Semilattice;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.containers.HashSet;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Set;
+
+/**
+ * @author oleg
+ */
+public class PyReachingDefsSemilattice implements Semilattice<ScopeVariable> {
+  public boolean eq(DFAMap<ScopeVariable> e1, DFAMap<ScopeVariable> e2) {
+    if (e1 == PyReachingDefsDfaInstance.INITIAL_MAP && e2 != PyReachingDefsDfaInstance.INITIAL_MAP ||
+        e2 == PyReachingDefsDfaInstance.INITIAL_MAP && e1 != PyReachingDefsDfaInstance.INITIAL_MAP) {
+      return false;
+    }
+    return e1.equals(e2);
+  }
+
+  public DFAMap<ScopeVariable> join(ArrayList<DFAMap<ScopeVariable>> ins) {
+    if (ins.isEmpty()) {
+      return DFAMap.empty();
+    }
+    if (ins.size() == 1) {
+      return ins.get(0);
+    }
+
+    final Set<String> resultNames = getResultNames(ins);
+    if (resultNames == null || resultNames.isEmpty()) {
+      return new DFAMap<ScopeVariable>();
+    }
+
+    final DFAMap<ScopeVariable> result = new DFAMap<ScopeVariable>();
+    for (String name : resultNames) {
+
+      boolean isParameter = true;
+      Set<PsiElement> declarations = new HashSet<PsiElement>();
+      PsiElement scope = null;
+
+      // iterating over all maps
+      for (DFAMap<ScopeVariable> map : ins) {
+        final ScopeVariable variable = map.get(name);
+        if (variable == null) {
+          continue;
+        }
+        final PsiElement valueScope = variable.getScope();
+        if (scope == null || scope == valueScope) {
+          scope = valueScope;
+          declarations.addAll(variable.getDeclarations());
+        } else
+        if (PsiTreeUtil.isAncestor(valueScope, scope, false)){
+          scope = valueScope;
+          declarations = new HashSet<PsiElement>(variable.getDeclarations());
+        }
+      }
+      final ScopeVariable scopeVariable = new ScopeVariableImpl(name, isParameter, scope, declarations);
+      result.put(name, scopeVariable);
+    }
+    return result;
+  }
+
+  @Nullable
+  private static Set<String> getResultNames(final ArrayList<DFAMap<ScopeVariable>> ins) {
+    // Compute intersection of all the names
+    Set<String> names2Include = null;
+    for (DFAMap<ScopeVariable> map : ins) {
+      if (map == PyReachingDefsDfaInstance.INITIAL_MAP) {
+        continue;
+      }
+      names2Include = map.intersectKeys(names2Include);
+    }
+    return names2Include;
+  }
+}

python/src/com/jetbrains/python/codeInsight/dataflow/ScopeUtil.java

@@ -0,0 +1,42 @@
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.codeInsight.controlflow.ControlFlowOwner;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyExceptPartNavigator;
+import com.jetbrains.python.psi.impl.PyForStatementNavigator;
+import com.jetbrains.python.psi.impl.PyListCompExpressionNavigator;
+
+/**
+ * @author oleg
+ */
+public class ScopeUtil {
+  public static PsiElement getScopeElement(final PsiElement element) {
+    if (element instanceof PyNamedParameter){
+      final PyFunction function = PsiTreeUtil.getParentOfType(element, PyFunction.class, false);
+      if (function != null){
+        return function;
+      }
+    }
+
+    final PyExceptPart exceptPart = PyExceptPartNavigator.getPyExceptPartByTarget(element);
+    if (exceptPart != null){
+      return exceptPart;
+    }
+
+    final PyForStatement forStatement = PyForStatementNavigator.getPyForStatementByIterable(element);
+    if (forStatement != null){
+      return forStatement;
+    }
+
+    final PyListCompExpression listCompExpression = PyListCompExpressionNavigator.getPyListCompExpressionByVariable(element);
+    if (listCompExpression != null){
+      return listCompExpression;
+    }
+
+    final ControlFlowOwner owner = PsiTreeUtil.getParentOfType(element, ControlFlowOwner.class, false);
+    assert owner != null : "element should have not null controlflow owner";
+    return owner;
+  }
+}

python/src/com/jetbrains/python/codeInsight/dataflow/ScopeVariable.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2000-2008 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.psi.PsiElement;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+
+/**
+ * @author oleg
+ */
+public interface ScopeVariable {
+  @NotNull
+  String getName();
+
+  @NotNull
+  Collection<PsiElement> getDeclarations();
+
+  PsiElement getScope();
+
+  boolean isParameter();
+}

python/src/com/jetbrains/python/codeInsight/dataflow/ScopeVariableImpl.java

@@ -0,0 +1,47 @@
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.PyElement;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+import java.util.Collections;
+
+/**
+ * @author oleg
+ */
+public class ScopeVariableImpl implements ScopeVariable {
+  private String myName;
+  private final Collection<PsiElement> myDeclarations;
+  private PsiElement myScope;
+  private boolean isParameter;
+
+  public ScopeVariableImpl(final String name, final boolean parameter, final PsiElement scope, final Collection<PsiElement> declarations) {
+    myName = name;
+    myDeclarations = declarations;
+    myScope = scope;
+    isParameter = parameter;
+  }
+
+  public ScopeVariableImpl(final String name, final boolean parameter, final PsiElement scope, PsiElement declaration) {
+    this(name, parameter, scope, Collections.singletonList(declaration));
+  }
+
+  @NotNull
+  public String getName() {
+    return myName;
+  }
+
+  @NotNull
+  public Collection<PsiElement> getDeclarations() {
+    return myDeclarations;
+  }
+
+  public PsiElement getScope() {
+    return myScope;
+  }
+
+  public boolean isParameter() {
+    return isParameter;
+  }
+}

python/src/com/jetbrains/python/codeInsight/dataflow/UsageAnalyzer.java

@@ -0,0 +1,34 @@
+package com.jetbrains.python.codeInsight.dataflow;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.PyParameter;
+import com.jetbrains.python.psi.impl.PyExceptPartNavigator;
+import com.jetbrains.python.psi.impl.PyForStatementNavigator;
+import com.jetbrains.python.psi.impl.PyListCompExpressionNavigator;
+
+/**
+ * @author oleg
+ */
+public class UsageAnalyzer {
+  public static boolean isParameter(final PsiElement element) {
+    // Except Block
+    if (PyExceptPartNavigator.getPyExceptPartByTarget(element) != null){
+      return true;
+    }
+
+    // For iterable statement
+    if (PyForStatementNavigator.getPyForStatementByIterable(element) != null){
+      return true;
+    }
+
+    // List comprehension expression
+   if (PyListCompExpressionNavigator.getPyListCompExpressionByVariable(element) != null){
+     return true;
+   }
+
+    if (element instanceof PyParameter) {
+      return true;
+    }
+    return false;
+  }
+}

python/src/com/jetbrains/python/psi/PyListCompExpression.java

@@ -27,5 +27,4 @@ public interface PyListCompExpression extends PyExpression, NameDefiner {
   PyExpression getResultExpression();
   List<ComprhForComponent> getForComponents();
   List<ComprhIfComponent> getIfComponents();
-
 }

python/src/com/jetbrains/python/psi/impl/PyExceptPartNavigator.java

@@ -0,0 +1,25 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.PyExceptPart;
+import com.jetbrains.python.psi.PyExpression;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author oleg
+ */
+public class PyExceptPartNavigator {
+  @Nullable
+  public static PyExceptPart getPyExceptPartByTarget(final PsiElement element){
+    final PyExceptPart pyExceptPart = PsiTreeUtil.getParentOfType(element, PyExceptPart.class, false);
+    if (pyExceptPart == null){
+      return null;
+    }
+    final PyExpression expr = pyExceptPart.getTarget();
+    if (expr != null && PsiTreeUtil.isAncestor(expr, element, false)){
+      return pyExceptPart;
+    }
+    return null;
+  }
+}

python/src/com/jetbrains/python/psi/impl/PyForStatementNavigator.java

@@ -0,0 +1,25 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyForStatement;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author oleg
+ */
+public class PyForStatementNavigator {
+  @Nullable
+  public static PyForStatement getPyForStatementByIterable(final PsiElement element){
+    final PyForStatement forStatement = PsiTreeUtil.getParentOfType(element, PyForStatement.class, false);
+    if (forStatement == null){
+      return null;
+    }
+    final PyExpression target = forStatement.getForPart().getTarget();
+    if (target != null && PsiTreeUtil.isAncestor(target, element, false)){
+      return forStatement;
+    }
+    return null;
+  }
+}

python/src/com/jetbrains/python/psi/impl/PyListCompExpressionNavigator.java

@@ -0,0 +1,28 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.ComprhForComponent;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyListCompExpression;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author oleg
+ */
+public class PyListCompExpressionNavigator {
+  @Nullable
+  public static PyListCompExpression getPyListCompExpressionByVariable(final PsiElement element){
+    final PyListCompExpression listCompExpression = PsiTreeUtil.getParentOfType(element, PyListCompExpression.class, false);
+    if (listCompExpression == null){
+      return null;
+    }
+    for (ComprhForComponent component : listCompExpression.getForComponents()) {
+      final PyExpression variable = component.getIteratorVariable();
+      if (variable != null && PsiTreeUtil.isAncestor(variable, element, false)){
+        return listCompExpression;
+      }
+    }
+    return null;
+  }
+}

python/testData/codeInsight/controlflow/function.txt

@@ -4,12 +4,16 @@
 3(4) WRITE ACCESS: collapse
 4(5) element: PyExpressionStatement
 5(6) element: PyAssignmentStatement
-6(7) WRITE ACCESS: method
-7(8) WRITE ACCESS: methodList
-8(9) element: PyAssignmentStatement
-9(10) WRITE ACCESS: s
-10(11) WRITE ACCESS: s
-11(12) WRITE ACCESS: processFunc
-12(13) element: PyPrintStatement
-13(14) WRITE ACCESS: method
-14() element: null
+6(7,8) element: PyListCompExpression
+7(9) element: PyCallExpression: callable. Condition: callable(getattr(object, method)):true
+8(10) element: PyCallExpression: callable. Condition: callable(getattr(object, method)):false
+9(10) WRITE ACCESS: method
+10(11) WRITE ACCESS: methodList
+11(12) element: PyAssignmentStatement
+12(13) WRITE ACCESS: s
+13(14) WRITE ACCESS: s
+14(15) WRITE ACCESS: processFunc
+15(16) element: PyPrintStatement
+16(17) element: PyListCompExpression
+17(18) WRITE ACCESS: method
+18() element: null

python/testData/codeInsight/controlflow/listcomp.py

@@ -1 +1 @@
-[k for k, v in params.items()]
+[k for k, v in params.items() if foo]

python/testData/codeInsight/controlflow/listcomp.txt

@@ -0,0 +1,8 @@
+0(1) element: null
+1(2) element: PyExpressionStatement
+2(3,4) element: PyListCompExpression
+3(5) element: PyReferenceExpression: foo. Condition: foo:true
+4(7) element: PyReferenceExpression: foo. Condition: foo:false
+5(6) WRITE ACCESS: k
+6(7) WRITE ACCESS: v
+7() element: null

python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java

@@ -66,6 +66,10 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testListComp() throws Exception {
+    doTest();
+  }
+
   public void testFunction() throws Exception {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");