python/src/com/jetbrains/python/refactoring/PyDefUseUtil.java

/*
 * Copyright 2000-2013 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.refactoring;

import com.google.common.collect.Sets;
import com.intellij.codeInsight.controlflow.ControlFlow;
import com.intellij.codeInsight.controlflow.ControlFlowUtil;
import com.intellij.codeInsight.controlflow.Instruction;
import com.intellij.openapi.util.Comparing;
import com.intellij.psi.PsiElement;
import com.intellij.util.Function;
import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator;
import com.intellij.psi.util.QualifiedName;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * @author Dennis.Ushakov
 */
public class PyDefUseUtil {
  private PyDefUseUtil() {
  }

  @NotNull
  public static List<ReadWriteInstruction> getLatestDefs(ScopeOwner block, String varName, PsiElement anchor, boolean acceptTypeAssertions) {
    final ControlFlow controlFlow = ControlFlowCache.getControlFlow(block);
    final Instruction[] instructions = controlFlow.getInstructions();
    final PyAugAssignmentStatement augAssignment = PyAugAssignmentStatementNavigator.getStatementByTarget(anchor);
    if (augAssignment != null) {
      anchor = augAssignment;
    }
    int instr = ControlFlowUtil.findInstructionNumberByElement(instructions, anchor);
    if (instr < 0) {
      return Collections.emptyList();
    }
    if (anchor instanceof PyTargetExpression) {
      Collection<Instruction> pred = instructions[instr].allPred();
      if (!pred.isEmpty()) {
        instr = pred.iterator().next().num();
      }
    }
    final Collection<ReadWriteInstruction> result = getLatestDefs(varName, instructions, instr, acceptTypeAssertions);
    return new ArrayList<ReadWriteInstruction>(result);
  }

  private static Collection<ReadWriteInstruction> getLatestDefs(final String varName, final Instruction[] instructions, final int instr,
                                                                final boolean acceptTypeAssertions) {
    final Collection<ReadWriteInstruction> result = new LinkedHashSet<ReadWriteInstruction>();
    ControlFlowUtil.iteratePrev(instr, instructions,
                                new Function<Instruction, ControlFlowUtil.Operation>() {
                                  @Override
                                  public ControlFlowUtil.Operation fun(Instruction instruction) {
                                    if (instruction instanceof ReadWriteInstruction) {
                                      final ReadWriteInstruction rwInstruction = (ReadWriteInstruction)instruction;
                                      final PsiElement element = instruction.getElement();
                                      final String name = elementName(element);
                                      final ReadWriteInstruction.ACCESS access = rwInstruction.getAccess();
                                      if ((access.isWriteAccess() || (acceptTypeAssertions && access.isAssertTypeAccess())) &&
                                          Comparing.strEqual(name, varName)) {
                                        result.add(rwInstruction);
                                        return ControlFlowUtil.Operation.CONTINUE;
                                      }
                                    }
                                    return ControlFlowUtil.Operation.NEXT;
                                  }
                                });
    return result;
  }

  @Nullable
  private static String elementName(PsiElement element) {
    if (element instanceof PyImportElement) {
      return ((PyImportElement) element).getVisibleName();
    }
    if (element instanceof PyReferenceExpression) {
      final QualifiedName qname = ((PyReferenceExpression)element).asQualifiedName();
      if (qname != null) {
        return qname.toString();
      }
    }
    return element instanceof PyElement ? ((PyElement)element).getName() : null;
  }

  @NotNull
  public static PsiElement[] getPostRefs(ScopeOwner block, PyTargetExpression var, PyExpression anchor) {
    final ControlFlow controlFlow = ControlFlowCache.getControlFlow(block);
    final Instruction[] instructions = controlFlow.getInstructions();
    final int instr = ControlFlowUtil.findInstructionNumberByElement(instructions, anchor);
    if (instr < 0) {
      return new PyElement[0];
    }
    final boolean[] visited = new boolean[instructions.length];
    final Collection<PyElement> result = Sets.newHashSet();
    for (Instruction instruction : instructions[instr].allSucc()) {
      getPostRefs(var, instructions, instruction.num(), visited, result);
    }
    return result.toArray(new PyElement[result.size()]);
  }

  private static void getPostRefs(PyTargetExpression var,
                                  Instruction[] instructions,
                                  int instr,
                                  boolean[] visited,
                                  Collection<PyElement> result) {
    // TODO: Use ControlFlowUtil.process() for forwards CFG traversal
    if (visited[instr]) return;
    visited[instr] = true;
    if (instructions[instr] instanceof ReadWriteInstruction) {
      final ReadWriteInstruction instruction = (ReadWriteInstruction)instructions[instr];
      final PsiElement element = instruction.getElement();
      String name = elementName(element);
      if (Comparing.strEqual(name, var.getName())) {
        final ReadWriteInstruction.ACCESS access = instruction.getAccess();
        if (access.isWriteAccess()) {
          return;
        }
        result.add((PyElement)instruction.getElement());
      }
    }
    for (Instruction instruction : instructions[instr].allSucc()) {
      getPostRefs(var, instructions, instruction.num(), visited, result);
    }
  }

  public static class InstructionNotFoundException extends RuntimeException {
  }
}