Pseudocode: Store lazily computed value usage map in the pseudocode

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/Pseudocode.java

@@ -58,4 +58,7 @@ public interface Pseudocode {
 
     @Nullable
     PseudoValue getElementValue(@Nullable JetElement element);
+
+    @NotNull
+    List<? extends Instruction> getUsages(@Nullable PseudoValue value);
 }

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/PseudocodeImpl.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.lang.cfg.pseudocode;
 
 import com.google.common.collect.*;
+import com.intellij.openapi.util.NotNullLazyValue;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.cfg.Label;
@@ -89,6 +90,15 @@ public class PseudocodeImpl implements Pseudocode {
 
     private final Map<JetElement, PseudoValue> elementsToValues = new HashMap<JetElement, PseudoValue>();
 
+    private final NotNullLazyValue<Map<PseudoValue, List<? extends Instruction>>> valueUsages =
+            new NotNullLazyValue<Map<PseudoValue, List<? extends Instruction>>>() {
+                @NotNull
+                @Override
+                protected Map<PseudoValue, List<? extends Instruction>> compute() {
+                    return PseudocodePackage.collectValueUsages(PseudocodeImpl.this);
+                }
+            };
+
     private Pseudocode parent = null;
     private Set<LocalFunctionDeclarationInstruction> localDeclarations = null;
     //todo getters
@@ -251,6 +261,13 @@ public class PseudocodeImpl implements Pseudocode {
         return elementsToValues.get(element);
     }
 
+    @NotNull
+    @Override
+    public List<? extends Instruction> getUsages(@Nullable PseudoValue value) {
+        List<? extends Instruction> result = valueUsages.getValue().get(value);
+        return result != null ? result : Collections.<Instruction>emptyList();
+    }
+
     /*package*/ void bindElementToValue(@NotNull JetElement element, @NotNull PseudoValue value) {
         elementsToValues.put(element, value);
     }

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/pseudocodeUtil.kt

@@ -64,18 +64,6 @@ val PseudoValue.implicitReturnValue: Boolean
         return false
     }
 
-fun Pseudocode.collectValueUsages(): Map<PseudoValue, List<Instruction>> {
-    val map = HashMap<PseudoValue, MutableList<Instruction>>()
-    traverseFollowingInstructions(getEnterInstruction(), HashSet(), TraversalOrder.FORWARD) {
-        for (value in it.inputValues) {
-            map.getOrPut(value){ ArrayList() }.add(it)
-        }
-        true
-    }
-
-    return map
-}
-
 fun getReceiverTypePredicate(resolvedCall: ResolvedCall<*>, receiverValue: ReceiverValue): TypePredicate? {
     val callableDescriptor = resolvedCall.getResultingDescriptor()
     if (callableDescriptor == null) return null
@@ -96,11 +84,8 @@ fun getReceiverTypePredicate(resolvedCall: ResolvedCall<*>, receiverValue: Recei
     return null
 }
 
-fun getExpectedTypePredicate(
-        value: PseudoValue,
-        valueUsageMap: Map<PseudoValue, List<Instruction>>,
-        bindingContext: BindingContext
-): TypePredicate {
+fun getExpectedTypePredicate(value: PseudoValue, bindingContext: BindingContext): TypePredicate {
+    val pseudocode = value.createdAt.owner
     val typePredicates = HashSet<TypePredicate?>()
 
     fun addSubtypesOf(jetType: JetType?) = typePredicates.add(jetType?.getSubtypesPredicate())
@@ -112,7 +97,7 @@ fun getExpectedTypePredicate(
             addSubtypesOf(functionDescriptor?.getReturnType())
         }
 
-        valueUsageMap[value]?.forEach {
+        pseudocode.getUsages(value).forEach {
             when (it) {
                 is ReturnValueInstruction -> {
                     val functionDescriptor = (it.element as JetReturnExpression).getTargetFunctionDescriptor(bindingContext)
@@ -173,4 +158,16 @@ fun getExpectedTypePredicate(
 
     addTypePredicates(value)
     return and(typePredicates.filterNotNull())
+}
+
+private fun Pseudocode.collectValueUsages(): Map<PseudoValue, List<Instruction>> {
+    val map = HashMap<PseudoValue, MutableList<Instruction>>()
+    traverseFollowingInstructions(getEnterInstruction(), HashSet(), TraversalOrder.FORWARD) {
+        for (value in it.inputValues) {
+            map.getOrPut(value){ ArrayList() }.add(it)
+        }
+        true
+    }
+
+    return map
 }

compiler/tests/org/jetbrains/jet/cfg/AbstractPseudoValueTest.kt

@@ -22,14 +22,12 @@ import org.jetbrains.jet.lang.psi.JetElement
 import org.jetbrains.jet.lang.psi.JetTreeVisitorVoid
 import org.jetbrains.jet.lang.resolve.BindingContext
 import java.util.*
-import org.jetbrains.jet.lang.cfg.pseudocode.collectValueUsages
 import org.jetbrains.jet.lang.cfg.pseudocode.TypePredicate
 import org.jetbrains.jet.lang.cfg.pseudocode.getExpectedTypePredicate
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.eval.InstructionWithValue
 
 public abstract class AbstractPseudoValueTest : AbstractPseudocodeTest() {
     override fun dumpInstructions(pseudocode: PseudocodeImpl, out: StringBuilder, bindingContext: BindingContext) {
-        val valueUsageMap = pseudocode.collectValueUsages()
         val expectedTypePredicateMap = HashMap<PseudoValue, TypePredicate>()
 
         fun getElementToValueMap(pseudocode: PseudocodeImpl): Map<JetElement, PseudoValue> {
@@ -51,7 +49,7 @@ public abstract class AbstractPseudoValueTest : AbstractPseudocodeTest() {
                 element?.getText()?.replaceAll("\\s+", " ") ?: ""
 
         fun valueDecl(value: PseudoValue): String {
-            val typePredicate = expectedTypePredicateMap.getOrPut(value) { getExpectedTypePredicate(value, valueUsageMap, bindingContext) }
+            val typePredicate = expectedTypePredicateMap.getOrPut(value) { getExpectedTypePredicate(value, bindingContext) }
             return "${value.debugName}: $typePredicate"
         }
 

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -373,8 +373,6 @@ private fun ExtractionData.inferParametersInfo(
 
     val extractedDescriptorToParameter = HashMap<DeclarationDescriptor, MutableParameter>()
 
-    val valueUsageMap = pseudocode.collectValueUsages()
-
     for (refInfo in getBrokenReferencesInfo(createTemporaryCodeBlock())) {
         val (originalRef, originalDeclaration, originalDescriptor, resolvedCall) = refInfo.resolveResult
         val ref = refInfo.refExpr
@@ -437,7 +435,7 @@ private fun ExtractionData.inferParametersInfo(
                 if (!parameterType.processTypeIfExtractable(typeParameters, nonDenotableTypes)) continue
 
                 val parameterTypePredicate =
-                        pseudocode.getElementValue(originalRef)?.let { getExpectedTypePredicate(it, valueUsageMap, bindingContext) } ?: AllTypes
+                        pseudocode.getElementValue(originalRef)?.let { getExpectedTypePredicate(it, bindingContext) } ?: AllTypes
 
                 val parameter = extractedDescriptorToParameter.getOrPut(descriptorToExtract) {
                     val parameterName =