Pseudocode: Add explicit type predicates to magic instructions

annotations/com/intellij/util/annotations.xml

@@ -14,4 +14,10 @@
   <item name='com.intellij.util.Range T getTo()'>
     <annotation name='org.jetbrains.annotations.NotNull'/>
   </item>
+  <item name='com.intellij.util.SmartFMap com.intellij.util.SmartFMap&lt;K,V&gt; emptyMap()'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
+  <item name='com.intellij.util.SmartFMap com.intellij.util.SmartFMap&lt;K,V&gt; plus(K, V)'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
 </root>

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilder.java

@@ -20,6 +20,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.cfg.pseudocode.PseudoValue;
 import org.jetbrains.jet.lang.cfg.pseudocode.Pseudocode;
+import org.jetbrains.jet.lang.cfg.pseudocode.TypePredicate;
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.eval.AccessTarget;
 import org.jetbrains.jet.lang.descriptors.ReceiverParameterDescriptor;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
@@ -115,6 +116,7 @@ public interface JetControlFlowBuilder {
             @NotNull JetElement instructionElement,
             @Nullable JetElement valueElement,
             @NotNull List<PseudoValue> inputValues,
+            @NotNull Map<PseudoValue, TypePredicate> expectedTypes,
             boolean synthetic
     );
 

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilderAdapter.java

@@ -20,6 +20,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.cfg.pseudocode.PseudoValue;
 import org.jetbrains.jet.lang.cfg.pseudocode.Pseudocode;
+import org.jetbrains.jet.lang.cfg.pseudocode.TypePredicate;
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.eval.AccessTarget;
 import org.jetbrains.jet.lang.descriptors.ReceiverParameterDescriptor;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
@@ -71,9 +72,10 @@ public abstract class JetControlFlowBuilderAdapter implements JetControlFlowBuil
             @NotNull JetElement instructionElement,
             @Nullable JetElement valueElement,
             @NotNull List<PseudoValue> inputValues,
+            @NotNull Map<PseudoValue, TypePredicate> expectedTypes,
             boolean synthetic
     ) {
-        return getDelegateBuilder().magic(instructionElement, valueElement, inputValues, synthetic);
+        return getDelegateBuilder().magic(instructionElement, valueElement, inputValues, expectedTypes, synthetic);
     }
 
     @NotNull

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowProcessor.java

@@ -22,16 +22,12 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.SmartFMap;
-import com.intellij.util.containers.ContainerUtil;
 import kotlin.Function0;
 import kotlin.Function1;
 import kotlin.KotlinPackage;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowInstructionsGenerator;
+import org.jetbrains.jet.lang.cfg.pseudocode.*;
-import org.jetbrains.jet.lang.cfg.pseudocode.PseudoValue;
-import org.jetbrains.jet.lang.cfg.pseudocode.Pseudocode;
-import org.jetbrains.jet.lang.cfg.pseudocode.PseudocodeImpl;
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.eval.AccessTarget;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.psi.*;
@@ -188,12 +184,24 @@ public class JetControlFlowProcessor {
 
         @NotNull
         private PseudoValue createSyntheticValue(@NotNull JetElement instructionElement) {
-            return builder.magic(instructionElement, null, Collections.<PseudoValue>emptyList(), true);
+            return createSyntheticValue(instructionElement, Collections.<JetElement>emptyList());
+        }
+
+        @NotNull
+        private PseudoValue createSyntheticValue(@NotNull JetElement instructionElement, @NotNull List<? extends JetElement> from) {
+            List<PseudoValue> values = elementsToValues(from);
+            return builder.magic(instructionElement, null, values, defaultTypeMap(values), true);
         }
 
         @NotNull
         private PseudoValue createNonSyntheticValue(@NotNull List<? extends JetElement> from, @NotNull JetElement to) {
-            return builder.magic(to, to, elementsToValues(from), false);
+            List<PseudoValue> values = elementsToValues(from);
+            return builder.magic(to, to, values, defaultTypeMap(values), false);
+        }
+
+        @NotNull
+        private Map<PseudoValue, TypePredicate> defaultTypeMap(List<PseudoValue> values) {
+            return PseudocodePackage.expectedTypeFor(AllTypes.instance$, values);
         }
 
         private void mergeValues(@NotNull List<JetExpression> from, @NotNull JetExpression to) {
@@ -218,10 +226,10 @@ public class JetControlFlowProcessor {
         }
 
         private List<PseudoValue> elementsToValues(List<? extends JetElement> from) {
+            if (from.isEmpty()) return Collections.emptyList();
             return KotlinPackage.filterNotNull(
-                    KotlinPackage.mapTo(
+                    KotlinPackage.map(
                             from,
-                            new LinkedHashSet<PseudoValue>(),
                             new Function1<JetElement, PseudoValue>() {
                                 @Override
                                 public PseudoValue invoke(JetElement element) {
@@ -862,12 +870,17 @@ public class JetControlFlowProcessor {
         private void writeLoopParameterAssignment(JetForExpression expression) {
             JetParameter loopParameter = expression.getLoopParameter();
             JetMultiDeclaration multiDeclaration = expression.getMultiParameter();
-
             JetExpression loopRange = expression.getLoopRange();
+
+            JetType loopRangeType = trace.get(BindingContext.EXPRESSION_TYPE, loopRange);
+            TypePredicate loopRangeTypeSet = loopRangeType != null ? new SingleType(loopRangeType) : AllTypes.instance$;
+            PseudoValue loopRangeValue = builder.getBoundValue(loopRange);
+
             PseudoValue value = builder.magic(
                     loopRange != null ? loopRange : expression,
                     null,
-                    Collections.singletonList(builder.getBoundValue(loopRange)),
+                    Collections.singletonList(loopRangeValue),
+                    Collections.singletonMap(loopRangeValue, loopRangeTypeSet),
                     true
             );
 
@@ -1058,7 +1071,10 @@ public class JetControlFlowProcessor {
                 inputExpressions.add(calleeExpression);
                 inputExpressions.add(generateAndGetReceiverIfAny(expression));
 
-                createNonSyntheticValue(inputExpressions, calleeExpression != null ? calleeExpression : expression);
+                createNonSyntheticValue(
+                        inputExpressions,
+                        calleeExpression != null ? calleeExpression : expression
+                );
             }
 
             copyValue(calleeExpression, expression);
@@ -1106,16 +1122,16 @@ public class JetControlFlowProcessor {
             generateInstructions(initializer, NOT_IN_CONDITION);
             for (JetMultiDeclarationEntry entry : declaration.getEntries()) {
                 builder.declareVariable(entry);
+
                 ResolvedCall<FunctionDescriptor> resolvedCall = trace.get(BindingContext.COMPONENT_RESOLVED_CALL, entry);
                 PseudoValue writtenValue = resolvedCall != null
                                            ? builder.call(
-                        entry,
+                                                entry,
-                        resolvedCall,
+                                                resolvedCall,
-                        getReceiverValues(resolvedCall, false),
+                                                getReceiverValues(resolvedCall, false),
-                        Collections.<PseudoValue, ValueParameterDescriptor>emptyMap()
+                                                Collections.<PseudoValue, ValueParameterDescriptor>emptyMap()
-                )
+                                            )
-                                           : builder.magic(entry, null,
+                                           : createSyntheticValue(entry, Collections.singletonList(initializer));
-                                                           ContainerUtil.createMaybeSingletonList(builder.getBoundValue(initializer)), true);
                 if (generateWriteForEntries) {
                     generateInitializer(entry, writtenValue != null ? writtenValue : createSyntheticValue(entry));
                 }
@@ -1209,9 +1225,7 @@ public class JetControlFlowProcessor {
                     JetWhenCondition condition = conditions[i];
                     condition.accept(conditionVisitor, context);
                     if (i + 1 < conditions.length) {
-                        PseudoValue conditionValue = builder.magic(
+                        PseudoValue conditionValue = createSyntheticValue(condition, Arrays.asList(subjectExpression, condition));
-                                condition, null, elementsToValues(Arrays.asList(subjectExpression, condition)), true
-                        );
                         builder.nondeterministicJump(bodyLabel, expression, conditionValue);
                     }
                 }
@@ -1221,9 +1235,7 @@ public class JetControlFlowProcessor {
                     PseudoValue conditionValue = null;
                     JetWhenCondition lastCondition = KotlinPackage.lastOrNull(conditions);
                     if (lastCondition != null) {
-                        conditionValue = builder.magic(
+                        conditionValue = createSyntheticValue(lastCondition, Arrays.asList(subjectExpression, lastCondition));
-                                lastCondition, null, elementsToValues(Arrays.asList(subjectExpression, lastCondition)), true
-                        );
                     }
                     builder.nondeterministicJump(nextLabel, expression, conditionValue);
                 }

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

@@ -420,7 +420,8 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
         @NotNull
         @Override
         public PseudoValue loadStringTemplate(@NotNull JetStringTemplateExpression expression, @NotNull List<PseudoValue> inputValues) {
-            return inputValues.isEmpty() ? read(expression) : magic(expression, expression, inputValues, false);
+            if (inputValues.isEmpty()) return read(expression);
+            return magic(expression, expression, inputValues, PseudocodePackage.expectedTypeFor(AllTypes.instance$, inputValues), false);
         }
 
         @NotNull
@@ -429,10 +430,12 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
                 @NotNull JetElement instructionElement,
                 @Nullable JetElement valueElement,
                 @NotNull List<PseudoValue> inputValues,
+                @NotNull Map<PseudoValue, TypePredicate> expectedTypes,
                 boolean synthetic
         ) {
-            MagicInstruction instruction =
+            MagicInstruction instruction = MagicInstruction.object$.create(
-                    MagicInstruction.object$.create(instructionElement, valueElement, getCurrentScope(), synthetic, inputValues, valueFactory);
+                    instructionElement, valueElement, getCurrentScope(), synthetic, inputValues, expectedTypes, valueFactory
+            );
             add(instruction);
             return instruction.getOutputValue();
         }
@@ -489,7 +492,21 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
                 @NotNull PredefinedOperation operation,
                 @NotNull List<PseudoValue> inputValues
         ) {
-            return magic(expression, expression, inputValues, false);
+            Map<PseudoValue, TypePredicate> expectedTypes;
+            switch(operation) {
+                case AND:
+                case OR:
+                    SingleType onlyBoolean = new SingleType(KotlinBuiltIns.getInstance().getBooleanType());
+                    expectedTypes = PseudocodePackage.expectedTypeFor(onlyBoolean, inputValues);
+                    break;
+                case NOT_NULL_ASSERTION:
+                    expectedTypes = PseudocodePackage.expectedTypeFor(AllTypes.instance$, inputValues);
+                    break;
+                default:
+                    throw new IllegalArgumentException("Invalid operation: " + operation);
+            }
+
+            return magic(expression, expression, inputValues, expectedTypes, false);
         }
 
         @Override

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2014 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 org.jetbrains.jet.lang.cfg.pseudocode
+
+import org.jetbrains.jet.lang.types.JetType
+import org.jetbrains.jet.lang.types.checker.JetTypeChecker
+import org.jetbrains.jet.renderer.DescriptorRenderer
+import com.intellij.util.SmartFMap
+import org.jetbrains.jet.lang.types.TypeUtils
+
+public trait TypePredicate: (JetType) -> Boolean {
+    [suppress("PARAMETER_NAME_CHANGED_ON_OVERRIDE")]
+    override fun invoke(typeToCheck: JetType): Boolean
+}
+
+public data class SingleType(val targetType: JetType): TypePredicate {
+    override fun invoke(typeToCheck: JetType): Boolean = JetTypeChecker.INSTANCE.equalTypes(typeToCheck, targetType)
+    override fun toString(): String = targetType.render()
+}
+
+public object AllTypes : TypePredicate {
+    override fun invoke(typeToCheck: JetType): Boolean = true
+
+    override fun toString(): String = ""
+}
+
+private fun JetType.render(): String = DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(this)
+
+public fun <T> TypePredicate.expectedTypeFor(keys: Iterable<T>): Map<T, TypePredicate> =
+        keys.fold(SmartFMap.emptyMap<T, TypePredicate>()) { (map, key) -> map.plus(key, this) }

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/instructions/eval/operationInstructions.kt

@@ -26,6 +26,7 @@ import org.jetbrains.jet.lang.cfg.pseudocode.instructions.InstructionVisitor
 import org.jetbrains.jet.lang.cfg.pseudocode.instructions.InstructionVisitorWithResult
 import org.jetbrains.jet.lang.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor
+import org.jetbrains.jet.lang.cfg.pseudocode.TypePredicate
 
 public abstract class OperationInstruction protected(
         element: JetElement,
@@ -102,8 +103,13 @@ public class MagicInstruction(
         element: JetElement,
         lexicalScope: LexicalScope,
         val synthetic: Boolean,
-        inputValues: List<PseudoValue>
+        inputValues: List<PseudoValue>,
+        val expectedTypes: Map<PseudoValue, TypePredicate>
 ) : OperationInstruction(element, lexicalScope, inputValues), StrictlyValuedOperationInstruction {
+    protected fun setResult(valueElement: JetElement?, factory: PseudoValueFactory?): OperationInstruction {
+        return setResult(factory?.newValue(valueElement, this))
+    }
+
     override fun accept(visitor: InstructionVisitor) {
         visitor.visitMagic(this)
     }
@@ -113,7 +119,7 @@ public class MagicInstruction(
     }
 
     override fun createCopy() =
-            MagicInstruction(element, lexicalScope, synthetic, inputValues).setResult(resultValue)
+            MagicInstruction(element, lexicalScope, synthetic, inputValues, expectedTypes).setResult(resultValue)
 
     override fun toString() = renderInstruction("magic", render(element))
 
@@ -124,8 +130,11 @@ public class MagicInstruction(
                 lexicalScope: LexicalScope,
                 synthetic: Boolean,
                 inputValues: List<PseudoValue>,
+                expectedTypes: Map<PseudoValue, TypePredicate>,
                 factory: PseudoValueFactory
-        ): MagicInstruction = MagicInstruction(element, lexicalScope, synthetic, inputValues).setResult(factory, valueElement) as MagicInstruction
+        ): MagicInstruction = MagicInstruction(
+                element, lexicalScope, synthetic, inputValues, expectedTypes
+        ).setResult(factory, valueElement) as MagicInstruction
     }
 }