Return values collected from CF data

idea/src/org/jetbrains/jet/lang/JetSemanticServices.java

@@ -2,6 +2,7 @@ package org.jetbrains.jet.lang;
 
 import com.intellij.openapi.project.Project;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.cfg.JetFlowInformationProvider;
 import org.jetbrains.jet.lang.resolve.ClassDescriptorResolver;
 import org.jetbrains.jet.lang.resolve.JetScope;
 import org.jetbrains.jet.lang.resolve.OverloadResolver;
@@ -44,8 +45,8 @@ public class JetSemanticServices {
     }
 
     @NotNull
-    public JetTypeInferrer getTypeInferrer(BindingTrace trace) {
-        return new JetTypeInferrer(trace, this);
+    public JetTypeInferrer getTypeInferrer(BindingTrace trace, JetFlowInformationProvider flowInformationProvider) {
+        return new JetTypeInferrer(trace, flowInformationProvider, this);
     }
 
     @NotNull

idea/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilder.java

@@ -47,7 +47,7 @@ public interface JetControlFlowBuilder {
     JetElement getCurrentSubroutine();
 
     void returnValue(@NotNull JetElement subroutine);
-    void returnNoValue(@NotNull JetElement subroutine);
+    void returnNoValue(@NotNull JetElement expression, @NotNull JetElement subroutine);
 
     void writeNode(@NotNull JetElement assignment, @NotNull JetElement lValue);
 

idea/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilderAdapter.java

@@ -110,8 +110,8 @@ public class JetControlFlowBuilderAdapter implements JetControlFlowBuilder {
     }
 
     @Override
-    public void returnNoValue(@NotNull JetElement subroutine) {
-        builder.returnNoValue(subroutine);
+    public void returnNoValue(@NotNull JetElement expression, @NotNull JetElement subroutine) {
+        builder.returnNoValue(expression, subroutine);
     }
 
     @Override

idea/src/org/jetbrains/jet/lang/cfg/JetControlFlowProcessor.java

@@ -366,7 +366,7 @@ public class JetControlFlowProcessor {
             }
             if (subroutine != null) {
                 if (returnedExpression == null) {
-                    builder.returnNoValue(subroutine);
+                    builder.returnNoValue(expression, subroutine);
                 }
                 else {
                     builder.returnValue(subroutine);

idea/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java

+package org.jetbrains.jet.lang.cfg;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.psi.JetElement;
+import org.jetbrains.jet.lang.psi.JetExpression;
+import org.jetbrains.jet.lang.psi.JetFunction;
+
+import java.util.Collection;
+
+/**
+ * @author abreslav
+ */
+public interface JetFlowInformationProvider {
+    JetFlowInformationProvider ERROR = new JetFlowInformationProvider() {
+        @Override
+        public void collectReturnedInformation(@NotNull JetFunction function, Collection<JetExpression> returnedExpressions, Collection<JetElement> elementsReturningUnit) {
+            throw new UnsupportedOperationException();
+        }
+    };
+
+    void collectReturnedInformation(@NotNull JetFunction function, Collection<JetExpression> returnedExpressions, Collection<JetElement> elementsReturningUnit);
+}

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/FunctionLiteralValueInstruction.java

 package org.jetbrains.jet.lang.cfg.pseudocode;
 
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.psi.JetElement;
+import org.jetbrains.jet.lang.psi.JetFunctionLiteralExpression;
 
 /**
 * @author abreslav
 */
-public class FunctionLiteralValueInstruction extends InstructionWithNext {
+public class FunctionLiteralValueInstruction extends ReadValueInstruction {
 
     private Pseudocode body;
 
-    public FunctionLiteralValueInstruction(@NotNull JetElement element) {
-        super(element);
+    public FunctionLiteralValueInstruction(@NotNull JetFunctionLiteralExpression expression) {
+        super(expression);
     }
 
     public Pseudocode getBody() {

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/InstructionVisitor.java

@@ -9,7 +9,7 @@ public class InstructionVisitor {
     }
 
     public void visitFunctionLiteralValue(FunctionLiteralValueInstruction instruction) {
-        visitInstructionWithNext(instruction);
+        visitReadValue(instruction);
     }
 
     public void visitUnconditionalJump(UnconditionalJumpInstruction instruction) {

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowDataTraceFactory.java

@@ -10,11 +10,11 @@ public interface JetControlFlowDataTraceFactory {
     JetControlFlowDataTraceFactory EMPTY = new JetControlFlowDataTraceFactory() {
         @NotNull
         @Override
-        public JetControlFlowDataTrace createTrace(JetElement element) {
-            return JetControlFlowDataTrace.EMPTY;
+        public JetPseudocodeTrace createTrace(JetElement element) {
+            return JetPseudocodeTrace.EMPTY;
         }
     };
 
     @NotNull
-    JetControlFlowDataTrace createTrace(JetElement element);
+    JetPseudocodeTrace createTrace(JetElement element);
 }

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowInstructionsGenerator.java

@@ -7,6 +7,7 @@ import org.jetbrains.jet.lang.cfg.Label;
 import org.jetbrains.jet.lang.psi.JetBlockExpression;
 import org.jetbrains.jet.lang.psi.JetElement;
 import org.jetbrains.jet.lang.psi.JetExpression;
+import org.jetbrains.jet.lang.psi.JetFunctionLiteralExpression;
 
 import java.util.*;
 
@@ -21,9 +22,9 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
     private int labelCount = 0;
 
     private final Stack<JetControlFlowInstructionsGeneratorWorker> builders = new Stack<JetControlFlowInstructionsGeneratorWorker>();
-    private final JetControlFlowDataTrace trace;
+    private final JetPseudocodeTrace trace;
 
-    public JetControlFlowInstructionsGenerator(JetControlFlowDataTrace trace) {
+    public JetControlFlowInstructionsGenerator(JetPseudocodeTrace trace) {
         super(null);
         this.trace = trace;
     }
@@ -60,7 +61,7 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
         JetControlFlowInstructionsGeneratorWorker worker = popBuilder(subroutine);
         if (functionLiteral) {
             JetControlFlowInstructionsGeneratorWorker builder = builders.peek();
-            FunctionLiteralValueInstruction instruction = new FunctionLiteralValueInstruction(subroutine);
+            FunctionLiteralValueInstruction instruction = new FunctionLiteralValueInstruction((JetFunctionLiteralExpression) subroutine);
             instruction.setBody(worker.getPseudocode());
             builder.add(instruction);
         }
@@ -155,8 +156,8 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
         }
 
         @Override
-        public void returnNoValue(@NotNull JetElement subroutine) {
-            add(new ReturnNoValueInstruction(getExitPoint(subroutine)));
+        public void returnNoValue(@NotNull JetElement expression, @NotNull JetElement subroutine) {
+            add(new ReturnNoValueInstruction(expression, getExitPoint(subroutine)));
         }
 
         @Override

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowDataTrace.java → idea/src/org/jetbrains/jet/lang/cfg/pseudocode/JetPseudocodeTrace.java

@@ -6,9 +6,9 @@ import org.jetbrains.jet.lang.psi.JetElement;
 /**
  * @author abreslav
  */
-public interface JetControlFlowDataTrace {
+public interface JetPseudocodeTrace {
 
-    JetControlFlowDataTrace EMPTY = new JetControlFlowDataTrace() {
+    JetPseudocodeTrace EMPTY = new JetPseudocodeTrace() {
         @Override
         public void recordControlFlowData(@NotNull JetElement element, @NotNull Pseudocode pseudocode) {
         }

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/Pseudocode.java

@@ -52,6 +52,8 @@ public class Pseudocode {
     private final List<PseudocodeLabel> labels = new ArrayList<PseudocodeLabel>();
 
     private final JetElement correspondingElement;
+    private SubroutineExitInstruction exitInstruction;
+    private boolean postPrecessed = false;
 
     public Pseudocode(JetElement correspondingElement) {
         this.correspondingElement = correspondingElement;
@@ -69,6 +71,16 @@ public class Pseudocode {
 
     public void addInstruction(Instruction instruction) {
         instructions.add(instruction);
+        if (instruction instanceof SubroutineExitInstruction) {
+            SubroutineExitInstruction exitInstruction = (SubroutineExitInstruction) instruction;
+            assert this.exitInstruction == null;
+            this.exitInstruction = exitInstruction;
+        }
+    }
+
+    @NotNull
+    public SubroutineExitInstruction getExitInstruction() {
+        return exitInstruction;
     }
 
     public void bindLabel(Label label) {
@@ -76,6 +88,8 @@ public class Pseudocode {
     }
 
     public void postProcess() {
+        if (postPrecessed) return;
+        postPrecessed = true;
         for (int i = 0, instructionsSize = instructions.size(); i < instructionsSize; i++) {
             Instruction instruction = instructions.get(i);
             final int currentPosition = i;

idea/src/org/jetbrains/jet/lang/cfg/pseudocode/ReturnNoValueInstruction.java

 package org.jetbrains.jet.lang.cfg.pseudocode;
 
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.cfg.Label;
+import org.jetbrains.jet.lang.psi.JetElement;
 
 /**
  * @author abreslav
  */
 public class ReturnNoValueInstruction extends AbstractJumpInstruction {
 
-    public ReturnNoValueInstruction(Label targetLabel) {
+    private final JetElement element;
+
+    public ReturnNoValueInstruction(@NotNull JetElement element, Label targetLabel) {
         super(targetLabel);
+        this.element = element;
+    }
+
+    public JetElement getElement() {
+        return element;
     }
 
     @Override

idea/src/org/jetbrains/jet/lang/resolve/ClassDescriptorResolver.java

@@ -4,11 +4,15 @@ import com.intellij.lang.ASTNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.JetSemanticServices;
+import org.jetbrains.jet.lang.cfg.JetFlowInformationProvider;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.*;
 import org.jetbrains.jet.lexer.JetTokens;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 
 /**
  * @author abreslav
@@ -305,7 +309,7 @@ public class ClassDescriptorResolver {
                 type = ErrorUtils.createErrorType("No type, no body");
             } else {
                 // TODO : ??? Fix-point here: what if we have something like "val a = foo {a.bar()}"
-                type = semanticServices.getTypeInferrer(trace).getType(scope, initializer, false);
+                type = semanticServices.getTypeInferrer(trace, JetFlowInformationProvider.ERROR).getType(scope, initializer, false);
             }
         } else {
             type = typeResolver.resolveType(scope, propertyTypeRef);

idea/src/org/jetbrains/jet/lang/resolve/TopDownAnalyzer.java

 package org.jetbrains.jet.lang.resolve;
 
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.JetSemanticServices;
 import org.jetbrains.jet.lang.cfg.JetControlFlowProcessor;
-import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTrace;
-import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTraceFactory;
-import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowInstructionsGenerator;
+import org.jetbrains.jet.lang.cfg.JetFlowInformationProvider;
+import org.jetbrains.jet.lang.cfg.pseudocode.*;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.*;
 
@@ -25,7 +23,6 @@ public class TopDownAnalyzer {
     private final JetSemanticServices semanticServices;
     private final ClassDescriptorResolver classDescriptorResolver;
     private final BindingTrace trace;
-    private final JetTypeInferrer typeInferrer;
     private final JetControlFlowDataTraceFactory flowDataTraceFactory;
     private boolean readyToProcessExpressions = false;
 
@@ -33,7 +30,6 @@ public class TopDownAnalyzer {
         this.semanticServices = semanticServices;
         this.classDescriptorResolver = new ClassDescriptorResolver(semanticServices, bindingTrace);
         this.trace = bindingTrace;
-        this.typeInferrer = semanticServices.getTypeInferrer(trace);
         this.flowDataTraceFactory = flowDataTraceFactory;
     }
 
@@ -41,7 +37,6 @@ public class TopDownAnalyzer {
         this.semanticServices = semanticServices;
         this.classDescriptorResolver = new ClassDescriptorResolver(semanticServices, bindingTrace);
         this.trace = bindingTrace;
-        this.typeInferrer = semanticServices.getTypeInferrer(trace);
         this.flowDataTraceFactory = JetControlFlowDataTraceFactory.EMPTY;
     }
 
@@ -101,7 +96,7 @@ public class TopDownAnalyzer {
                         if (importDirective.isAllUnder()) {
                             JetExpression importedReference = importDirective.getImportedReference();
                             if (importedReference != null) {
-                                JetType type = typeInferrer.getType(namespaceScope, importedReference, false);
+                                JetType type = semanticServices.getTypeInferrer(trace, JetFlowInformationProvider.ERROR).getType(namespaceScope, importedReference, false);
                                 if (type != null) {
                                     namespaceScope.importScope(type.getMemberScope());
                                 }
@@ -263,71 +258,137 @@ public class TopDownAnalyzer {
             WritableScope declaringScope = declaringScopes.get(declaration);
             assert declaringScope != null;
 
-            WritableScope parameterScope = semanticServices.createWritableScope(declaringScope, descriptor);
-            for (TypeParameterDescriptor typeParameter : descriptor.getTypeParameters()) {
-                parameterScope.addTypeParameterDescriptor(typeParameter);
-            }
-            for (ValueParameterDescriptor valueParameterDescriptor : descriptor.getUnsubstitutedValueParameters()) {
-                parameterScope.addPropertyDescriptor(valueParameterDescriptor);
-            }
-            parameterScope.addLabeledDeclaration(descriptor);
 
             assert declaration instanceof JetFunction || declaration instanceof JetConstructor;
             JetDeclarationWithBody declarationWithBody = (JetDeclarationWithBody) declaration;
 
             JetExpression bodyExpression = declarationWithBody.getBodyExpression();
 
-            if (bodyExpression != null) {
-                JetControlFlowDataTrace controlFlowDataTrace = flowDataTraceFactory.createTrace(declaration);
-                JetControlFlowInstructionsGenerator instructionsGenerator = new JetControlFlowInstructionsGenerator(controlFlowDataTrace);
-                new JetControlFlowProcessor(semanticServices, trace, instructionsGenerator).generate(declaration, bodyExpression);
-                controlFlowDataTrace.close();
-
-                boolean preferBlock = true;
-                FunctionDescriptorImpl functionDescriptorImpl = null;
-                if (declaration instanceof JetFunction) {
-                    JetFunction jetFunction = (JetFunction) declaration;
-                    preferBlock = jetFunction.hasBlockBody();
-                    functionDescriptorImpl = (FunctionDescriptorImpl) descriptor;
-                }
+            if (declaration instanceof JetFunction) {
+                JetFunction function = (JetFunction) declaration;
+                FunctionDescriptorImpl functionDescriptorImpl = (FunctionDescriptorImpl) descriptor;
+                if (bodyExpression != null) {
+                    JetFlowInformationProvider flowInformationProvider = computeFlowData(declaration, bodyExpression);
+                    JetTypeInferrer typeInferrer = semanticServices.getTypeInferrer(trace, flowInformationProvider);
 
-                JetType returnType = resolveExpression(parameterScope, bodyExpression, preferBlock, controlFlowDataTrace);
+                    assert readyToProcessExpressions : "Must be ready collecting types";
 
-                if (declaration instanceof JetFunction) {
-                    JetFunction function = (JetFunction) declaration;
                     if (function.getReturnTypeRef() != null) {
-                        if (returnType != null) {
-                            if (!semanticServices.getTypeChecker().isConvertibleTo(returnType, descriptor.getUnsubstitutedReturnType())) {
-                                semanticServices.getErrorHandler().typeMismatch(bodyExpression, descriptor.getUnsubstitutedReturnType(), returnType);
-                            }
-                        }
+                        typeInferrer.checkFunctionReturnType(declaringScope, function, descriptor);
                     }
                     else {
-                        JetType safeReturnType = returnType;
-                        if (safeReturnType == null) {
-                            safeReturnType = ErrorUtils.createErrorType("Unable to infer body type");
+                        JetType returnType = typeInferrer.getFunctionReturnType(declaringScope, function, descriptor);
+                        if (returnType == null) {
+                            returnType = ErrorUtils.createErrorType("Unable to infer body type");
                         }
-                        functionDescriptorImpl.setUnsubstitutedReturnType(safeReturnType);
+                        functionDescriptorImpl.setUnsubstitutedReturnType(returnType);
                     }
                 }
-            }
-            else {
-                if (declaration instanceof JetFunction) {
-                    JetFunction function = (JetFunction) declaration;
+                else {
                     if (function.getReturnTypeRef() == null) {
                         semanticServices.getErrorHandler().genericError(function.getNode(), "This function must either declare a return type or have a body element");
                         ((FunctionDescriptorImpl) descriptor).setUnsubstitutedReturnType(ErrorUtils.createErrorType("No type, no body"));
                     }
                 }
             }
+            else if (declaration instanceof JetConstructor) {
+                if (bodyExpression != null) {
+                    computeFlowData(declaration, bodyExpression);
+                    JetFlowInformationProvider flowInformationProvider = computeFlowData(declaration, bodyExpression);
+                    JetTypeInferrer typeInferrer = semanticServices.getTypeInferrer(trace, flowInformationProvider);
+                    typeInferrer.getType(FunctionDescriptorUtil.getFunctionInnerScope(declaringScope, descriptor, semanticServices), bodyExpression, true);
+                }
+
+            }
             assert descriptor.getUnsubstitutedReturnType() != null;
         }
     }
 
-    @Nullable
-    private JetType resolveExpression(@NotNull JetScope scope, JetExpression expression, boolean preferBlock, JetControlFlowDataTrace controlFlowDataTrace) {
-        assert readyToProcessExpressions : "Must be ready collecting types";
-        return typeInferrer.getType(scope, expression, preferBlock);
+    private JetFlowInformationProvider computeFlowData(@NotNull JetDeclaration declaration, @NotNull JetExpression bodyExpression) {
+        final JetPseudocodeTrace pseudocodeTrace = flowDataTraceFactory.createTrace(declaration);
+        final Map<JetElement, Pseudocode> pseudocodeMap = new HashMap<JetElement, Pseudocode>();
+        JetPseudocodeTrace wrappedTrace = new JetPseudocodeTrace() {
+            @Override
+            public void recordControlFlowData(@NotNull JetElement element, @NotNull Pseudocode pseudocode) {
+                pseudocodeTrace.recordControlFlowData(element, pseudocode);
+                pseudocodeMap.put(element, pseudocode);
+            }
+
+            @Override
+            public void close() {
+                pseudocodeTrace.close();
+                for (Pseudocode pseudocode : pseudocodeMap.values()) {
+                    pseudocode.postProcess();
+                }
+            }
+        };
+        JetControlFlowInstructionsGenerator instructionsGenerator = new JetControlFlowInstructionsGenerator(wrappedTrace);
+        new JetControlFlowProcessor(semanticServices, trace, instructionsGenerator).generate(declaration, bodyExpression);
+        wrappedTrace.close();
+        return new JetFlowInformationProvider() {
+            @Override
+            public void collectReturnedInformation(@NotNull JetFunction function, Collection<JetExpression> returnedExpressions, Collection<JetElement> elementsReturningUnit) {
+                Pseudocode pseudocode = pseudocodeMap.get(function);
+                assert pseudocode != null;
+
+                SubroutineExitInstruction exitInstruction = pseudocode.getExitInstruction();
+                processPreviousInstructions(exitInstruction, returnedExpressions, elementsReturningUnit);
+            }
+        };
+    }
+
+    private void processPreviousInstructions(Instruction previousFor, final Collection<JetExpression> returnedExpressions, final Collection<JetElement> elementsReturningUnit) {
+        Collection<Instruction> previousInstructions = previousFor.getPreviousInstructions();
+        InstructionVisitor visitor = new InstructionVisitor() {
+            @Override
+            public void visitReadValue(ReadValueInstruction instruction) {
+                returnedExpressions.add((JetExpression) instruction.getElement());
+            }
+
+            @Override
+            public void visitReturnValue(ReturnValueInstruction instruction) {
+                processPreviousInstructions(instruction, returnedExpressions, elementsReturningUnit);
+            }
+
+            @Override
+            public void visitReturnNoValue(ReturnNoValueInstruction instruction) {
+                elementsReturningUnit.add(instruction.getElement());
+            }
+
+            @Override
+            public void visitSubroutineEnter(SubroutineEnterInstruction instruction) {
+                elementsReturningUnit.add(instruction.getSubroutine());
+            }
+
+            @Override
+            public void visitUnsupportedElementInstruction(UnsupportedElementInstruction instruction) {
+                semanticServices.getErrorHandler().genericError(instruction.getElement().getNode(), "Unsupported by control-flow builder");
+            }
+
+            @Override
+            public void visitWriteValue(WriteValueInstruction writeValueInstruction) {
+                elementsReturningUnit.add(writeValueInstruction.getElement());
+            }
+
+            @Override
+            public void visitJump(AbstractJumpInstruction instruction) {
+                processPreviousInstructions(instruction, returnedExpressions, elementsReturningUnit);
+            }
+
+            @Override
+            public void visitInstruction(Instruction instruction) {
+                if (instruction instanceof JetElementInstruction) {
+                    JetElementInstruction elementInstruction = (JetElementInstruction) instruction;
+                    semanticServices.getErrorHandler().genericError(elementInstruction.getElement().getNode(), "Unsupported by control-flow builder " + elementInstruction.getElement());
+                }
+                else {
+                    throw new UnsupportedOperationException(instruction.toString());
+                }
+            }
+        };
+        for (Instruction previousInstruction : previousInstructions) {
+            previousInstruction.accept(visitor);
+        }
     }
 
 }

idea/src/org/jetbrains/jet/lang/types/FunctionDescriptorUtil.java

@@ -2,6 +2,9 @@ package org.jetbrains.jet.lang.types;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.JetSemanticServices;
+import org.jetbrains.jet.lang.resolve.JetScope;
+import org.jetbrains.jet.lang.resolve.WritableScope;
 
 import java.util.*;
 
@@ -116,4 +119,17 @@ public class FunctionDescriptorUtil {
         );
         return substitutedDescriptor;
     }
+
+    @NotNull
+    public static JetScope getFunctionInnerScope(@NotNull JetScope outerScope, @NotNull FunctionDescriptor descriptor, @NotNull JetSemanticServices semanticServices) {
+        WritableScope parameterScope = semanticServices.createWritableScope(outerScope, descriptor);
+        for (TypeParameterDescriptor typeParameter : descriptor.getTypeParameters()) {
+            parameterScope.addTypeParameterDescriptor(typeParameter);
+        }
+        for (ValueParameterDescriptor valueParameterDescriptor : descriptor.getUnsubstitutedValueParameters()) {
+            parameterScope.addPropertyDescriptor(valueParameterDescriptor);
+        }
+        parameterScope.addLabeledDeclaration(descriptor);
+        return parameterScope;
+    }
 }

idea/src/org/jetbrains/jet/lang/types/JetTypeInferrer.java

@@ -6,6 +6,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
 import org.jetbrains.jet.lang.JetSemanticServices;
+import org.jetbrains.jet.lang.cfg.JetFlowInformationProvider;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.*;
 import org.jetbrains.jet.lexer.JetTokens;
@@ -60,16 +61,20 @@ public class JetTypeInferrer {
         assignmentOperationCounterparts.put(JetTokens.MINUSEQ, JetTokens.MINUS);
     }
 
+    private final Map<JetExpression, JetType> typeCache = new HashMap<JetExpression, JetType>();
+
     private final BindingTrace trace;
     private final JetSemanticServices semanticServices;
     private final TypeResolver typeResolver;
     private final ClassDescriptorResolver classDescriptorResolver;
+    private final JetFlowInformationProvider flowInformationProvider;
 
-    public JetTypeInferrer(BindingTrace trace, JetSemanticServices semanticServices) {
-        this.trace = trace;
+    public JetTypeInferrer(@NotNull BindingTrace trace, @NotNull JetFlowInformationProvider flowInformationProvider, @NotNull JetSemanticServices semanticServices) {
+        this.trace = new CachedBindingTrace(trace);
         this.semanticServices = semanticServices;
         this.typeResolver = new TypeResolver(trace, semanticServices);
         this.classDescriptorResolver = semanticServices.getClassDescriptorResolver(trace);
+        this.flowInformationProvider = flowInformationProvider;
     }
 
     @NotNull
@@ -107,7 +112,6 @@ public class JetTypeInferrer {
             @NotNull JetReferenceExpression reference,
             @NotNull String name,
             @NotNull JetType receiverType,
-
             @NotNull List<JetType> argumentTypes,
             boolean reportUnresolved) {
         OverloadDomain overloadDomain = semanticServices.getOverloadResolver().getOverloadDomain(receiverType, scope, name);
@@ -264,21 +268,75 @@ public class JetTypeInferrer {
         };
     }
 
+    @NotNull
+    public JetType getFunctionReturnType(@NotNull JetScope outerScope, JetFunction function, FunctionDescriptor functionDescriptor) {
+        Map<JetElement, JetType> typeMap = getReturnedExpressions(outerScope, function, functionDescriptor);
+        Collection<JetType> types = typeMap.values();
+        return types.isEmpty() ? JetStandardClasses.getNothingType() : semanticServices.getTypeChecker().commonSupertype(types);
+    }
+
+    private JetType getCachedType(@NotNull JetExpression expression) {
+//        assert typeCache.containsKey(expression) : "No type cached for " + expression.getText();
+        return typeCache.get(expression);
+    }
+
+    public void checkFunctionReturnType(@NotNull JetScope outerScope, @NotNull JetFunction function, @NotNull FunctionDescriptor functionDescriptor) {
+        Map<JetElement, JetType> typeMap = getReturnedExpressions(outerScope, function, functionDescriptor);
+        if (typeMap.isEmpty()) {
+            return; // The function returns Nothing
+        }
+        JetType expectedReturnType = functionDescriptor.getUnsubstitutedReturnType();
+        for (Map.Entry<JetElement, JetType> entry : typeMap.entrySet()) {
+            JetType actualType = entry.getValue();
+            JetElement element = entry.getKey();
+            if (!semanticServices.getTypeChecker().isConvertibleTo(actualType, expectedReturnType)) {
+                if (element instanceof JetExpression) {
+                    JetExpression expression = (JetExpression) element;
+                    semanticServices.getErrorHandler().typeMismatch(expression, expectedReturnType, actualType);
+                }
+                else {
+                    semanticServices.getErrorHandler().genericError(element.getNode(), "This function must return a value of type " + expectedReturnType);
+                }
+            }
+        }
+    }
+
+    private Map<JetElement, JetType> getReturnedExpressions(JetScope outerScope, JetFunction function, FunctionDescriptor functionDescriptor) {
+        JetExpression bodyExpression = function.getBodyExpression();
+        assert bodyExpression != null;
+        JetScope functionInnerScope = FunctionDescriptorUtil.getFunctionInnerScope(outerScope, functionDescriptor, semanticServices);
+        getType(functionInnerScope, bodyExpression, function.hasBlockBody());
+        Collection<JetExpression> returnedExpressions = new ArrayList<JetExpression>();
+        Collection<JetElement> elementsReturningUnit = new ArrayList<JetElement>();
+        flowInformationProvider.collectReturnedInformation(function, returnedExpressions, elementsReturningUnit);
+        Map<JetElement,JetType> typeMap = new HashMap<JetElement, JetType>();
+        for (JetExpression returnedExpression : returnedExpressions) {
+            JetType cachedType = getCachedType(returnedExpression);
+            if (cachedType != null) {
+                typeMap.put(returnedExpression, cachedType);
+            }
+        }
+        for (JetElement jetElement : elementsReturningUnit) {
+            typeMap.put(jetElement, JetStandardClasses.getUnitType());
+        }
+        return typeMap;
+    }
+
     @Nullable
-    private JetType getBlockReturnedType(@NotNull JetScope outerScope, @NotNull List<JetElement> block, @NotNull LabeledJumpDomain jumpDomain) {
+    private JetType getBlockReturnedType(@NotNull JetScope outerScope, @NotNull List<JetElement> block) {
         if (block.isEmpty()) {
             return JetStandardClasses.getUnitType();
         }
 
         DeclarationDescriptor containingDescriptor = outerScope.getContainingDeclaration();
         WritableScope scope = semanticServices.createWritableScope(outerScope, containingDescriptor);
-        return getBlockReturnedTypeWithWritableScope(scope, block, jumpDomain);
+        return getBlockReturnedTypeWithWritableScope(scope, block);
     }
 
-    private JetType getBlockReturnedTypeWithWritableScope(@NotNull WritableScope scope, @NotNull List<? extends JetElement> block, @NotNull LabeledJumpDomain jumpDomain) {
+    private JetType getBlockReturnedTypeWithWritableScope(@NotNull WritableScope scope, @NotNull List<? extends JetElement> block) {
         assert !block.isEmpty();
 
-        TypeInferrerVisitorWithWritableScope blockLevelVisitor = new TypeInferrerVisitorWithWritableScope(scope, true, jumpDomain);
+        TypeInferrerVisitorWithWritableScope blockLevelVisitor = new TypeInferrerVisitorWithWritableScope(scope, true);
 
         JetType result = null;
         for (JetElement statement : block) {
@@ -305,18 +363,12 @@ public class JetTypeInferrer {
     private class TypeInferrerVisitor extends JetVisitor {
         private final JetScope scope;
         private final boolean preferBlock;
-        private final LabeledJumpDomain jumpDomain;
 
         protected JetType result;
 
-        private TypeInferrerVisitor(@NotNull JetScope scope, boolean preferBlock, @NotNull LabeledJumpDomain jumpDomain) {
+        private TypeInferrerVisitor(@NotNull JetScope scope, boolean preferBlock) {
             this.scope = scope;
             this.preferBlock = preferBlock;
-            this.jumpDomain = jumpDomain;
-        }
-
-        private TypeInferrerVisitor(JetScope scope, boolean preferBlock) {
-            this(scope, preferBlock, LabeledJumpDomain.EMPTY);
         }
 
         @Nullable
@@ -379,7 +431,7 @@ public class JetTypeInferrer {
         public void visitFunctionLiteralExpression(JetFunctionLiteralExpression expression) {
             if (preferBlock && !expression.hasParameterSpecification()) {
                 trace.recordBlock(expression);
-                result = getBlockReturnedType(scope, expression.getBody(), LabeledJumpDomain.ERROR);
+                result = getBlockReturnedType(scope, expression.getBody());
                 return;
             }
 
@@ -416,7 +468,7 @@ public class JetTypeInferrer {
                     writableScope.addPropertyDescriptor(propertyDescriptor);
                 }
                 writableScope.setThisType(receiverType);
-                returnType = getBlockReturnedType(writableScope, body, LabeledJumpDomain.ERROR);
+                returnType = getBlockReturnedType(writableScope, body);
             }
             JetType effectiveReceiverType = receiverTypeRef == null ? null : receiverType;
             JetType safeReturnType = returnType == null ? ErrorUtils.createErrorType("<return type>") : returnType;
@@ -477,7 +529,6 @@ public class JetTypeInferrer {
             else {
                 returnedType = JetStandardClasses.getUnitType();
             }
-            jumpDomain.registerReturn(expression, returnedType);
 
             result = JetStandardClasses.getNothingType();
         }
@@ -485,15 +536,11 @@ public class JetTypeInferrer {
         @Override
         public void visitBreakExpression(JetBreakExpression expression) {
             result = JetStandardClasses.getNothingType();
-
-            jumpDomain.registerBreakOrContinue(expression);
         }
 
         @Override
         public void visitContinueExpression(JetContinueExpression expression) {
             result = JetStandardClasses.getNothingType();
-
-            jumpDomain.registerBreakOrContinue(expression);
         }
 
         @Override
@@ -603,7 +650,7 @@ public class JetTypeInferrer {
 
         @Override
         public void visitBlockExpression(JetBlockExpression expression) {
-            result = getBlockReturnedType(scope, expression.getStatements(), jumpDomain);
+            result = getBlockReturnedType(scope, expression.getStatements());
         }
 
         @Override
@@ -692,7 +739,7 @@ public class JetTypeInferrer {
                 if (!function.hasParameterSpecification()) {
                     WritableScope writableScope = semanticServices.createWritableScope(scope, scope.getContainingDeclaration());
                     conditionScope = writableScope;
-                    getBlockReturnedTypeWithWritableScope(writableScope, function.getBody(), LabeledJumpDomain.ERROR); // TODO
+                    getBlockReturnedTypeWithWritableScope(writableScope, function.getBody());
                     trace.recordBlock(function);
                 } else {
                     getType(scope, body, true);
@@ -701,7 +748,7 @@ public class JetTypeInferrer {
             else if (body != null) {
                 WritableScope writableScope = semanticServices.createWritableScope(scope, scope.getContainingDeclaration());
                 conditionScope = writableScope;
-                getBlockReturnedTypeWithWritableScope(writableScope, Collections.singletonList(body), LabeledJumpDomain.ERROR); // TODO
+                getBlockReturnedTypeWithWritableScope(writableScope, Collections.singletonList(body));
             }
             checkCondition(conditionScope, expression.getCondition());
             result = JetStandardClasses.getUnitType();
@@ -1240,8 +1287,8 @@ public class JetTypeInferrer {
     private class TypeInferrerVisitorWithWritableScope extends TypeInferrerVisitor {
         private final WritableScope scope;
 
-        public TypeInferrerVisitorWithWritableScope(@NotNull WritableScope scope, boolean preferBlock, @NotNull LabeledJumpDomain jumpDomain) {
-            super(scope, preferBlock, jumpDomain);
+        public TypeInferrerVisitorWithWritableScope(@NotNull WritableScope scope, boolean preferBlock) {
+            super(scope, preferBlock);
             this.scope = scope;
         }
 
@@ -1358,30 +1405,44 @@ public class JetTypeInferrer {
         }
     }
 
-    private interface LabeledJumpDomain {
-        LabeledJumpDomain EMPTY = new LabeledJumpDomain() {
-            @Override
-            public void registerReturn(@NotNull JetReturnExpression expression, JetType returnedExpressionType) {
-            }
+    private class CachedBindingTrace extends BindingTrace {
+        private final BindingTrace originalTrace;
 
-            @Override
-            public void registerBreakOrContinue(@NotNull JetLabelQualifiedExpression expression) {
-            }
-        };
+        public CachedBindingTrace(BindingTrace originalTrace) {
+            this.originalTrace = originalTrace;
+        }
 
-        LabeledJumpDomain ERROR = new LabeledJumpDomain() {
-            @Override
-            public void registerReturn(@NotNull JetReturnExpression expression, JetType returnedExpressionType) {
-//                throw new UnsupportedOperationException();
-            }
+        public void recordExpressionType(@NotNull JetExpression expression, @NotNull JetType type) {
+            originalTrace.recordExpressionType(expression, type);
+            typeCache.put(expression, type);
+        }
 
-            @Override
-            public void registerBreakOrContinue(@NotNull JetLabelQualifiedExpression expression) {
-//                throw new UnsupportedOperationException();
-            }
-        };
+        public void recordReferenceResolution(@NotNull JetReferenceExpression expression, @NotNull DeclarationDescriptor descriptor) {
+            originalTrace.recordReferenceResolution(expression, descriptor);
+        }
+
+        public void recordLabelResolution(@NotNull JetReferenceExpression expression, @NotNull PsiElement element) {
+            originalTrace.recordLabelResolution(expression, element);
+        }
+
+        public void recordDeclarationResolution(@NotNull PsiElement declaration, @NotNull DeclarationDescriptor descriptor) {
+            originalTrace.recordDeclarationResolution(declaration, descriptor);
+        }
+
+        public void recordTypeResolution(@NotNull JetTypeReference typeReference, @NotNull JetType type) {
+            originalTrace.recordTypeResolution(typeReference, type);
+        }
+
+        public void setToplevelScope(JetScope toplevelScope) {
+            originalTrace.setToplevelScope(toplevelScope);
+        }
 
-        void registerReturn(@NotNull JetReturnExpression expression, JetType returnedExpressionType);
-        void registerBreakOrContinue(@NotNull JetLabelQualifiedExpression expression);
+        public void recordBlock(JetFunctionLiteralExpression expression) {
+            originalTrace.recordBlock(expression);
+        }
+
+        public void removeReferenceResolution(@NotNull JetReferenceExpression referenceExpression) {
+            originalTrace.removeReferenceResolution(referenceExpression);
+        }
     }
 }

idea/testData/cfg/EmptyFunction.instructions

+== empty ==
+fun empty() {}
+---------------------
+l0:
+    <START>
+l1:
+    <END>
+=====================

idea/testData/cfg/EmptyFunction.jet

+fun empty() {}
\ No newline at end of file

idea/testData/cfg/ShortFunction.instructions

+== short ==
+fun short() = 1
+---------------------
+l0:
+    <START>
+    r(1)
+l1:
+    <END>
+=====================

idea/testData/cfg/ShortFunction.jet

+fun short() = 1

idea/tests/org/jetbrains/jet/cfg/JetControlFlowTest.java

@@ -10,7 +10,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.JetTestCaseBase;
 import org.jetbrains.jet.lang.ErrorHandler;
 import org.jetbrains.jet.lang.cfg.pseudocode.Instruction;
-import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTrace;
+import org.jetbrains.jet.lang.cfg.pseudocode.JetPseudocodeTrace;
 import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTraceFactory;
 import org.jetbrains.jet.lang.cfg.pseudocode.Pseudocode;
 import org.jetbrains.jet.lang.psi.JetElement;
@@ -42,7 +42,7 @@ public class JetControlFlowTest extends JetTestCaseBase {
         JetFile file = (JetFile) getFile();
 
         final Map<JetElement, Pseudocode> data = new LinkedHashMap<JetElement, Pseudocode>();
-        final JetControlFlowDataTrace jetControlFlowDataTrace = new JetControlFlowDataTrace() {
+        final JetPseudocodeTrace pseudocodeTrace = new JetPseudocodeTrace() {
 
             @Override
             public void recordControlFlowData(@NotNull JetElement element, @NotNull Pseudocode pseudocode) {
@@ -51,6 +51,9 @@ public class JetControlFlowTest extends JetTestCaseBase {
 
             @Override
             public void close() {
+                for (Pseudocode pseudocode : data.values()) {
+                    pseudocode.postProcess();
+                }
             }
 
         };
@@ -58,8 +61,8 @@ public class JetControlFlowTest extends JetTestCaseBase {
         AnalyzingUtils.analyzeNamespace(file.getRootNamespace(), ErrorHandler.THROW_EXCEPTION, new JetControlFlowDataTraceFactory() {
             @NotNull
             @Override
-            public JetControlFlowDataTrace createTrace(JetElement element) {
-                return jetControlFlowDataTrace;
+            public JetPseudocodeTrace createTrace(JetElement element) {
+                return pseudocodeTrace;
             }
         });
 
@@ -75,9 +78,6 @@ public class JetControlFlowTest extends JetTestCaseBase {
 
     private void processCFData(String name, Map<JetElement, Pseudocode> data) throws IOException {
         Collection<Pseudocode> pseudocodes = data.values();
-        for (Pseudocode pseudocode : pseudocodes) {
-            pseudocode.postProcess();
-        }
 
         StringBuilder instructionDump = new StringBuilder();
         int i = 0;

idea/tests/org/jetbrains/jet/types/JetTypeCheckerTest.java

@@ -6,6 +6,7 @@ import com.intellij.openapi.project.Project;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.ErrorHandler;
 import org.jetbrains.jet.lang.JetSemanticServices;
+import org.jetbrains.jet.lang.cfg.JetFlowInformationProvider;
 import org.jetbrains.jet.lang.psi.JetChangeUtil;
 import org.jetbrains.jet.lang.psi.JetClass;
 import org.jetbrains.jet.lang.psi.JetExpression;
@@ -455,14 +456,14 @@ public class JetTypeCheckerTest extends LightDaemonAnalyzerTestCase {
     private void assertType(String expression, JetType expectedType) {
         Project project = getProject();
         JetExpression jetExpression = JetChangeUtil.createExpression(project, expression);
-        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY).getType(classDefinitions.BASIC_SCOPE, jetExpression, false);
+        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY, JetFlowInformationProvider.ERROR).getType(classDefinitions.BASIC_SCOPE, jetExpression, false);
         assertTrue(type + " != " + expectedType, JetTypeImpl.equalTypes(type, expectedType));
     }
 
     private void assertErrorType(String expression) {
         Project project = getProject();
         JetExpression jetExpression = JetChangeUtil.createExpression(project, expression);
-        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY).safeGetType(classDefinitions.BASIC_SCOPE, jetExpression, false);
+        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY, JetFlowInformationProvider.ERROR).safeGetType(classDefinitions.BASIC_SCOPE, jetExpression, false);
         assertTrue("Error type expected but " + type + " returned", ErrorUtils.isErrorType(type));
     }
 
@@ -485,7 +486,7 @@ public class JetTypeCheckerTest extends LightDaemonAnalyzerTestCase {
     private void assertType(JetScope scope, String expression, String expectedTypeStr) {
         Project project = getProject();
         JetExpression jetExpression = JetChangeUtil.createExpression(project, expression);
-        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY).getType(scope, jetExpression, false);
+        JetType type = semanticServices.getTypeInferrer(BindingTrace.DUMMY, JetFlowInformationProvider.ERROR).getType(scope, jetExpression, false);
         JetType expectedType = expectedTypeStr == null ? null : makeType(expectedTypeStr);
         assertEquals(expectedType, type);
     }