added visitor with type parameters (TypeInferrerVisitor now implements it);

added expected type (pass it down the psi tree);
added information about is-checks to type enrichment

idea/src/org/jetbrains/jet/lang/psi/JetAnnotatedExpression.java

@@ -17,10 +17,15 @@ public class JetAnnotatedExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitAnnotatedExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitAnnotatedExpression(this, data);
+    }
+
     public JetExpression getBaseExpression() {
         return findChildByClass(JetExpression.class);
     }

idea/src/org/jetbrains/jet/lang/psi/JetAnnotation.java

@@ -15,10 +15,15 @@ public class JetAnnotation extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitAnnotation(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitAnnotation(this, data);
+    }
+
     public List<JetAnnotationEntry> getEntries() {
         return findChildrenByType(JetNodeTypes.ANNOTATION_ENTRY);
     }

idea/src/org/jetbrains/jet/lang/psi/JetAnnotationEntry.java

@@ -17,10 +17,16 @@ public class JetAnnotationEntry extends JetElement implements JetCall {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitAnnotationEntry(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitAnnotationEntry(this, data);
+    }
+
+
     @Nullable @IfNotParsed
     public JetTypeReference getTypeReference() {
         return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetArrayAccessExpression.java

@@ -29,10 +29,15 @@ public class JetArrayAccessExpression extends JetReferenceExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitArrayAccessExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitArrayAccessExpression(this, data);
+    }
+
     @NotNull
     public JetExpression getArrayExpression() {
         JetExpression baseExpression = findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetBinaryExpression.java

@@ -16,10 +16,15 @@ public class JetBinaryExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBinaryExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBinaryExpression(this, data);
+    }
+
     @NotNull
     public JetExpression getLeft() {
         JetExpression left = findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetBinaryExpressionWithTypeRHS.java

@@ -15,10 +15,15 @@ public class JetBinaryExpressionWithTypeRHS extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBinaryWithTypeRHSExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBinaryWithTypeRHSExpression(this, data);
+    }
+
     @NotNull
     public JetExpression getLeft() {
         JetExpression left = findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetBindingPattern.java

@@ -27,4 +27,9 @@ public class JetBindingPattern extends JetPattern {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBindingPattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBindingPattern(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetBlockExpression.java

@@ -15,10 +15,15 @@ public class JetBlockExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBlockExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBlockExpression(this, data);
+    }
+
     @NotNull
     public List<JetElement> getStatements() {
         return Arrays.asList(findChildrenByClass(JetElement.class));

idea/src/org/jetbrains/jet/lang/psi/JetBlockStringTemplateEntry.java

@@ -15,4 +15,9 @@ public class JetBlockStringTemplateEntry extends JetStringTemplateEntryWithExpre
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBlockStringTemplateEntry(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBlockStringTemplateEntry(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetBreakExpression.java

@@ -12,7 +12,12 @@ public class JetBreakExpression extends JetLabelQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitBreakExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitBreakExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetCallExpression.java

@@ -17,10 +17,15 @@ public class JetCallExpression extends JetExpression implements JetCall {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitCallExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitCallExpression(this, data);
+    }
+
     @Nullable
     public JetExpression getCalleeExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetCatchClause.java

@@ -16,10 +16,15 @@ public class JetCatchClause extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitCatchSection(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitCatchSection(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetParameterList getParameterList() {
         return (JetParameterList) findChildByType(JetNodeTypes.VALUE_PARAMETER_LIST);

idea/src/org/jetbrains/jet/lang/psi/JetClass.java

@@ -38,6 +38,11 @@ public class JetClass extends JetTypeParameterListOwner implements JetClassOrObj
         visitor.visitClass(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitClass(this, data);
+    }
+
     @Nullable
     public JetParameterList getPrimaryConstructorParameterList() {
         return (JetParameterList) findChildByType(JetNodeTypes.VALUE_PARAMETER_LIST);

idea/src/org/jetbrains/jet/lang/psi/JetClassBody.java

@@ -29,6 +29,11 @@ public class JetClassBody extends JetElement {
         visitor.visitClassBody(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitClassBody(this, data);
+    }
+
     @NotNull
     public List<JetClassInitializer> getAnonymousInitializers() {
         return findChildrenByType(JetNodeTypes.ANONYMOUS_INITIALIZER);

idea/src/org/jetbrains/jet/lang/psi/JetClassInitializer.java

@@ -16,6 +16,11 @@ public class JetClassInitializer extends JetDeclaration {
         visitor.visitAnonymousInitializer(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitAnonymousInitializer(this, data);
+    }
+
     @NotNull
     public JetExpression getBody() {
         JetExpression body = findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetClassObject.java

@@ -18,6 +18,11 @@ public class JetClassObject extends JetDeclaration {
         visitor.visitClassObject(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitClassObject(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetObjectDeclaration getObjectDeclaration() {
         return (JetObjectDeclaration) findChildByType(JetNodeTypes.OBJECT_DECLARATION);

idea/src/org/jetbrains/jet/lang/psi/JetConstantExpression.java

@@ -12,7 +12,12 @@ public class JetConstantExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitConstantExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitConstantExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetConstructor.java

@@ -22,6 +22,11 @@ public class JetConstructor extends JetDeclaration implements JetDeclarationWith
         visitor.visitConstructor(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitConstructor(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetParameterList getParameterList() {
         return (JetParameterList) findChildByType(JetNodeTypes.VALUE_PARAMETER_LIST);
@@ -54,6 +59,11 @@ public class JetConstructor extends JetDeclaration implements JetDeclarationWith
         return findChildByType(JetTokens.EQ) == null;
     }
 
+    @Override
+    public boolean hasDeclaredReturnType() {
+        return true;
+    }
+
     @NotNull
     @Override
     public JetElement asElement() {

idea/src/org/jetbrains/jet/lang/psi/JetContinueExpression.java

@@ -12,8 +12,12 @@ public class JetContinueExpression extends JetLabelQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitContinueExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitContinueExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetDeclarationWithBody.java

@@ -17,6 +17,8 @@ public interface JetDeclarationWithBody {
 
     boolean hasBlockBody();
 
+    boolean hasDeclaredReturnType();
+
     @NotNull
     JetElement asElement();
 }

idea/src/org/jetbrains/jet/lang/psi/JetDecomposerPattern.java

@@ -26,4 +26,9 @@ public class JetDecomposerPattern extends JetPattern {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDecomposerPattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDecomposerPattern(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetDelegationSpecifier.java

@@ -14,10 +14,15 @@ public class JetDelegationSpecifier extends JetElement{
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationSpecifier(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationSpecifier(this, data);
+    }
+
     @Nullable
     public JetTypeReference getTypeReference() {
         return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetDelegationSpecifierList.java

@@ -15,10 +15,15 @@ public class JetDelegationSpecifierList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationSpecifierList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationSpecifierList(this, data);
+    }
+
     public List<JetDelegationSpecifier> getDelegationSpecifiers() {
         return PsiTreeUtil.getChildrenOfTypeAsList(this, JetDelegationSpecifier.class);
     }

idea/src/org/jetbrains/jet/lang/psi/JetDelegatorByExpressionSpecifier.java

@@ -13,10 +13,15 @@ public class JetDelegatorByExpressionSpecifier extends JetDelegationSpecifier {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationByExpressionSpecifier(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationByExpressionSpecifier(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getDelegateExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetDelegatorToSuperCall.java

@@ -17,10 +17,15 @@ public class JetDelegatorToSuperCall extends JetDelegationSpecifier implements J
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationToSuperCallSpecifier(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationToSuperCallSpecifier(this, data);
+    }
+
     @Nullable
     public JetValueArgumentList getValueArgumentList() {
         return (JetValueArgumentList) findChildByType(JetNodeTypes.VALUE_ARGUMENT_LIST);

idea/src/org/jetbrains/jet/lang/psi/JetDelegatorToSuperClass.java

@@ -12,7 +12,12 @@ public class JetDelegatorToSuperClass extends JetDelegationSpecifier {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationToSuperClassSpecifier(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationToSuperClassSpecifier(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetDelegatorToThisCall.java

@@ -18,10 +18,15 @@ public class JetDelegatorToThisCall extends JetDelegationSpecifier implements Je
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDelegationToThisCall(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDelegationToThisCall(this, data);
+    }
+
     @Nullable
     public JetValueArgumentList getValueArgumentList() {
         return (JetValueArgumentList) findChildByType(JetNodeTypes.VALUE_ARGUMENT_LIST);

idea/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.java

@@ -14,10 +14,15 @@ public class JetDoWhileExpression extends JetLoopExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDoWhileExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDoWhileExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getCondition() {
         return findExpressionUnder(JetNodeTypes.CONDITION);

idea/src/org/jetbrains/jet/lang/psi/JetDotQualifiedExpression.java

@@ -12,7 +12,12 @@ public class JetDotQualifiedExpression extends JetQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitDotQualifiedExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitDotQualifiedExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetElement.java

@@ -51,4 +51,8 @@ public class JetElement extends ASTWrapperPsiElement {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitJetElement(this);
     }
+    
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitJetElement(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetEnumEntry.java

@@ -50,4 +50,9 @@ public class JetEnumEntry extends JetClass {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitEnumEntry(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitEnumEntry(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetEscapeStringTemplateEntry.java

@@ -17,6 +17,11 @@ public class JetEscapeStringTemplateEntry extends JetStringTemplateEntry {
         visitor.visitEscapeStringTemplateEntry(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitEscapeStringTemplateEntry(this, data);
+    }
+
     public String getUnescapedValue() {
         return StringUtil.unescapeStringCharacters(getText());
     }

idea/src/org/jetbrains/jet/lang/psi/JetExpression.java

@@ -13,10 +13,15 @@ public class JetExpression extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitExpression(this, data);
+    }
+
     protected JetExpression findExpressionUnder(JetNodeType type) {
         JetContainerNode containerNode = (JetContainerNode) findChildByType(type);
         if (containerNode == null) return null;

idea/src/org/jetbrains/jet/lang/psi/JetExpressionPattern.java

@@ -21,4 +21,9 @@ public class JetExpressionPattern extends JetPattern {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitExpressionPattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitExpressionPattern(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetExtendedVisitor.java

@@ -0,0 +1,422 @@
+package org.jetbrains.jet.lang.psi;
+
+import com.intellij.psi.PsiElementVisitor;
+
+/**
+ * @author svtk
+ */
+public class JetExtendedVisitor<R, D> extends PsiElementVisitor {
+    public R visitJetElement(JetElement element, D data) {
+        visitElement(element);
+        return null;
+    }
+
+    public R visitDeclaration(JetDeclaration dcl, D data) {
+        return visitExpression(dcl, data);
+    }
+
+    public R visitNamespace(JetNamespace namespace, D data) {
+        return visitDeclaration(namespace, data);
+    }
+
+    public R visitClass(JetClass klass, D data) {
+        return visitNamedDeclaration(klass, data);
+    }
+
+    public R visitClassObject(JetClassObject classObject, D data) {
+        return visitDeclaration(classObject, data);
+    }
+
+    public R visitConstructor(JetConstructor constructor, D data) {
+        return visitDeclaration(constructor, data);
+    }
+
+    public R visitNamedFunction(JetNamedFunction function, D data) {
+        return visitNamedDeclaration(function, data);
+    }
+
+    public R visitProperty(JetProperty property, D data) {
+        return visitNamedDeclaration(property, data);
+    }
+
+    public R visitTypedef(JetTypedef typedef, D data) {
+        return visitNamedDeclaration(typedef, data);
+    }
+
+    public R visitJetFile(JetFile file, D data) {
+        visitFile(file);
+        return null;
+    }
+
+    public R visitImportDirective(JetImportDirective importDirective, D data) {
+        return visitJetElement(importDirective, data);
+    }
+
+    public R visitClassBody(JetClassBody classBody, D data) {
+        return visitJetElement(classBody, data);
+    }
+
+    public R visitNamespaceBody(JetNamespaceBody body, D data) {
+        return visitJetElement(body, data);
+    }
+
+    public R visitModifierList(JetModifierList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitAnnotation(JetAnnotation annotation, D data) {
+        return visitJetElement(annotation, data);
+    }
+
+    public R visitAnnotationEntry(JetAnnotationEntry annotationEntry, D data) {
+        return visitJetElement(annotationEntry, data);
+    }
+
+    public R visitTypeParameterList(JetTypeParameterList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitTypeParameter(JetTypeParameter parameter, D data) {
+        return visitNamedDeclaration(parameter, data);
+    }
+
+    public R visitEnumEntry(JetEnumEntry enumEntry, D data) {
+        return visitClass(enumEntry, data);
+    }
+
+    public R visitParameterList(JetParameterList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitParameter(JetParameter parameter, D data) {
+        return visitNamedDeclaration(parameter, data);
+    }
+
+    public R visitDelegationSpecifierList(JetDelegationSpecifierList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitDelegationSpecifier(JetDelegationSpecifier specifier, D data) {
+        return visitJetElement(specifier, data);
+    }
+
+    public R visitDelegationByExpressionSpecifier(JetDelegatorByExpressionSpecifier specifier, D data) {
+        return visitDelegationSpecifier(specifier, data);
+    }
+
+    public R visitDelegationToSuperCallSpecifier(JetDelegatorToSuperCall call, D data) {
+        return visitDelegationSpecifier(call, data);
+    }
+
+    public R visitDelegationToSuperClassSpecifier(JetDelegatorToSuperClass specifier, D data) {
+        return visitDelegationSpecifier(specifier, data);
+    }
+
+    public R visitDelegationToThisCall(JetDelegatorToThisCall thisCall, D data) {
+        return visitDelegationSpecifier(thisCall, data);
+    }
+
+    public R visitTypeReference(JetTypeReference typeReference, D data) {
+        return visitJetElement(typeReference, data);
+    }
+
+    public R visitValueArgumentList(JetValueArgumentList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitArgument(JetValueArgument argument, D data) {
+        return visitJetElement(argument, data);
+    }
+
+    public R visitExpression(JetExpression expression, D data) {
+        return visitJetElement(expression, data);
+    }
+
+    public R visitLoopExpression(JetLoopExpression loopExpression, D data) {
+        return visitExpression(loopExpression, data);
+    }
+
+    public R visitConstantExpression(JetConstantExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitSimpleNameExpression(JetSimpleNameExpression expression, D data) {
+        return visitReferenceExpression(expression, data);
+    }
+
+    public R visitReferenceExpression(JetReferenceExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitTupleExpression(JetTupleExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitPrefixExpression(JetPrefixExpression expression, D data) {
+        return visitUnaryExpression(expression, data);
+    }
+
+    public R visitPostfixExpression(JetPostfixExpression expression, D data) {
+        return visitUnaryExpression(expression, data);
+    }
+
+    public R visitUnaryExpression(JetUnaryExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitBinaryExpression(JetBinaryExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+//    public R visitNewExpression(JetNewExpression expression, D data) {
+//        return visitExpression(expression, data);
+//    }
+//
+    public R visitReturnExpression(JetReturnExpression expression, D data) {
+        return visitLabelQualifiedExpression(expression, data);
+    }
+
+    public R visitLabelQualifiedExpression(JetLabelQualifiedExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitThrowExpression(JetThrowExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitBreakExpression(JetBreakExpression expression, D data) {
+        return visitLabelQualifiedExpression(expression, data);
+    }
+
+    public R visitContinueExpression(JetContinueExpression expression, D data) {
+        return visitLabelQualifiedExpression(expression, data);
+    }
+
+    public R visitIfExpression(JetIfExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitWhenExpression(JetWhenExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitTryExpression(JetTryExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitForExpression(JetForExpression expression, D data) {
+        return visitLoopExpression(expression, data);
+    }
+
+    public R visitWhileExpression(JetWhileExpression expression, D data) {
+        return visitLoopExpression(expression, data);
+    }
+
+    public R visitDoWhileExpression(JetDoWhileExpression expression, D data) {
+        return visitLoopExpression(expression, data);
+    }
+
+    public R visitFunctionLiteralExpression(JetFunctionLiteralExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitAnnotatedExpression(JetAnnotatedExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitCallExpression(JetCallExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitArrayAccessExpression(JetArrayAccessExpression expression, D data) {
+        return visitReferenceExpression(expression, data);
+    }
+
+    public R visitQualifiedExpression(JetQualifiedExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitHashQualifiedExpression(JetHashQualifiedExpression expression, D data) {
+        return visitQualifiedExpression(expression, data);
+    }
+
+    public R visitDotQualifiedExpression(JetDotQualifiedExpression expression, D data) {
+        return visitQualifiedExpression(expression, data);
+    }
+
+    public R visitPredicateExpression(JetPredicateExpression expression, D data) {
+        return visitQualifiedExpression(expression, data);
+    }
+
+    public R visitSafeQualifiedExpression(JetSafeQualifiedExpression expression, D data) {
+        return visitQualifiedExpression(expression, data);
+    }
+
+    public R visitObjectLiteralExpression(JetObjectLiteralExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitRootNamespaceExpression(JetRootNamespaceExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitBlockExpression(JetBlockExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitCatchSection(JetCatchClause catchClause, D data) {
+        return visitJetElement(catchClause, data);
+    }
+
+    public R visitFinallySection(JetFinallySection finallySection, D data) {
+        return visitJetElement(finallySection, data);
+    }
+
+    public R visitTypeArgumentList(JetTypeArgumentList typeArgumentList, D data) {
+        return visitJetElement(typeArgumentList, data);
+    }
+
+    public R visitThisExpression(JetThisExpression expression, D data) {
+        return visitLabelQualifiedExpression(expression, data);
+    }
+
+    public R visitParenthesizedExpression(JetParenthesizedExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitInitializerList(JetInitializerList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitAnonymousInitializer(JetClassInitializer initializer, D data) {
+        return visitDeclaration(initializer, data);
+    }
+
+    public R visitPropertyAccessor(JetPropertyAccessor accessor, D data) {
+        return visitDeclaration(accessor, data);
+    }
+
+    public R visitTypeConstraintList(JetTypeConstraintList list, D data) {
+        return visitJetElement(list, data);
+    }
+
+    public R visitTypeConstraint(JetTypeConstraint constraint, D data) {
+        return visitJetElement(constraint, data);
+    }
+
+    private R visitTypeElement(JetTypeElement type, D data) {
+        return visitJetElement(type, data);
+    }
+
+    public R visitUserType(JetUserType type, D data) {
+        return visitTypeElement(type, data);
+    }
+
+    public R visitTupleType(JetTupleType type, D data) {
+        return visitTypeElement(type, data);
+    }
+
+    public R visitFunctionType(JetFunctionType type, D data) {
+        return visitTypeElement(type, data);
+    }
+
+    public R visitSelfType(JetSelfType type, D data) {
+        return visitTypeElement(type, data);
+    }
+
+    public R visitBinaryWithTypeRHSExpression(JetBinaryExpressionWithTypeRHS expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitStringTemplateExpression(JetStringTemplateExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitNamedDeclaration(JetNamedDeclaration declaration, D data) {
+        return visitDeclaration(declaration, data);
+    }
+
+    public R visitNullableType(JetNullableType nullableType, D data) {
+        return visitTypeElement(nullableType, data);
+    }
+
+    public R visitTypeProjection(JetTypeProjection typeProjection, D data) {
+        return visitJetElement(typeProjection, data);
+    }
+
+    public R visitWhenEntry(JetWhenEntry jetWhenEntry, D data) {
+        return visitJetElement(jetWhenEntry, data);
+    }
+
+    public R visitIsExpression(JetIsExpression expression, D data) {
+        return visitExpression(expression, data);
+    }
+
+    public R visitWhenConditionCall(JetWhenConditionCall condition, D data) {
+        return visitJetElement(condition, data);
+    }
+
+    public R visitWhenConditionIsPattern(JetWhenConditionIsPattern condition, D data) {
+        return visitJetElement(condition, data);
+    }
+
+    public R visitWhenConditionInRange(JetWhenConditionInRange condition, D data) {
+        return visitJetElement(condition, data);
+    }
+
+    public R visitTypePattern(JetTypePattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    public R visitWildcardPattern(JetWildcardPattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    public R visitExpressionPattern(JetExpressionPattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    public R visitTuplePattern(JetTuplePattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    private R visitPattern(JetPattern pattern, D data) {
+        return visitJetElement(pattern, data);
+    }
+
+    public R visitDecomposerPattern(JetDecomposerPattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    public R visitObjectDeclaration(JetObjectDeclaration declaration, D data) {
+        return visitNamedDeclaration(declaration, data);
+    }
+
+    public R visitBindingPattern(JetBindingPattern pattern, D data) {
+        return visitPattern(pattern, data);
+    }
+
+    public R visitStringTemplateEntry(JetStringTemplateEntry entry, D data) {
+        return visitJetElement(entry, data);
+    }
+
+    public R visitStringTemplateEntryWithExpression(JetStringTemplateEntryWithExpression entry, D data) {
+        return visitStringTemplateEntry(entry, data);
+    }
+
+    public R visitBlockStringTemplateEntry(JetBlockStringTemplateEntry entry, D data) {
+        return visitStringTemplateEntryWithExpression(entry, data);
+    }
+
+    public R visitSimpleNameStringTemplateEntry(JetSimpleNameStringTemplateEntry entry, D data) {
+        return visitStringTemplateEntryWithExpression(entry, data);
+    }
+
+    public R visitLiteralStringTemplateEntry(JetLiteralStringTemplateEntry entry, D data) {
+        return visitStringTemplateEntry(entry, data);
+    }
+
+    public R visitEscapeStringTemplateEntry(JetEscapeStringTemplateEntry entry, D data) {
+        return visitStringTemplateEntry(entry, data);
+    }
+}

idea/src/org/jetbrains/jet/lang/psi/JetFinallySection.java

@@ -13,10 +13,15 @@ public class JetFinallySection extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitFinallySection(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitFinallySection(this, data);
+    }
+
     public JetBlockExpression getFinalExpression() {
         return (JetBlockExpression) findChildByType(JetNodeTypes.BLOCK);
     }

idea/src/org/jetbrains/jet/lang/psi/JetForExpression.java

@@ -14,10 +14,15 @@ public class JetForExpression extends JetLoopExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitForExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitForExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetParameter getLoopParameter() {
         return (JetParameter) findChildByType(JetNodeTypes.LOOP_PARAMETER);

idea/src/org/jetbrains/jet/lang/psi/JetFunction.java

@@ -36,6 +36,11 @@ abstract public class JetFunction extends JetTypeParameterListOwner implements J
         return findChildByClass(JetExpression.class);
     }
 
+    @Override
+    public boolean hasDeclaredReturnType() {
+        return getReturnTypeRef() != null;
+    }
+
     @Nullable
     public JetTypeReference getReceiverTypeRef() {
         PsiElement child = getFirstChild();

idea/src/org/jetbrains/jet/lang/psi/JetFunctionLiteralExpression.java

@@ -13,10 +13,15 @@ public class JetFunctionLiteralExpression extends JetExpression implements JetDe
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitFunctionLiteralExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitFunctionLiteralExpression(this, data);
+    }
+
     @NotNull
     public JetFunctionLiteral getFunctionLiteral() {
         return (JetFunctionLiteral) findChildByType(JetNodeTypes.FUNCTION_LITERAL);
@@ -89,6 +94,11 @@ public class JetFunctionLiteralExpression extends JetExpression implements JetDe
         return getFunctionLiteral().hasBlockBody();
     }
 
+    @Override
+    public boolean hasDeclaredReturnType() {
+        return getFunctionLiteral().getReturnTypeRef() != null;
+    }
+
     @NotNull
     @Override
     public JetElement asElement() {

idea/src/org/jetbrains/jet/lang/psi/JetFunctionType.java

@@ -40,10 +40,15 @@ public class JetFunctionType extends JetTypeElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitFunctionType(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitFunctionType(this, data);
+    }
+
     @Nullable
     public JetParameterList getParameterList() {
         return (JetParameterList) findChildByType(JetNodeTypes.VALUE_PARAMETER_LIST);

idea/src/org/jetbrains/jet/lang/psi/JetHashQualifiedExpression.java

@@ -12,7 +12,12 @@ public class JetHashQualifiedExpression extends JetQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitHashQualifiedExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitHashQualifiedExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetIfExpression.java

@@ -14,10 +14,15 @@ public class JetIfExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitIfExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitIfExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getCondition() {
         return findExpressionUnder(JetNodeTypes.CONDITION);

idea/src/org/jetbrains/jet/lang/psi/JetImportDirective.java

@@ -19,6 +19,11 @@ public class JetImportDirective extends JetElement {
         visitor.visitImportDirective(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitImportDirective(this, data);
+    }
+
     public boolean isAbsoluteInRootNamespace() {
         return findChildByType(JetTokens.NAMESPACE_KEYWORD) != null;
     }

idea/src/org/jetbrains/jet/lang/psi/JetInitializerList.java

@@ -15,10 +15,15 @@ public class JetInitializerList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitInitializerList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitInitializerList(this, data);
+    }
+
     @NotNull
     public List<JetDelegationSpecifier> getInitializers() {
         return PsiTreeUtil.getChildrenOfTypeAsList(this, JetDelegationSpecifier.class);

idea/src/org/jetbrains/jet/lang/psi/JetIsExpression.java

@@ -15,10 +15,15 @@ public class JetIsExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitIsExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitIsExpression(this, data);
+    }
+
     @NotNull
     public JetExpression getLeftHandSide() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetLiteralStringTemplateEntry.java

@@ -15,4 +15,9 @@ public class JetLiteralStringTemplateEntry extends JetStringTemplateEntry {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitLiteralStringTemplateEntry(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitLiteralStringTemplateEntry(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetModifierList.java

@@ -23,6 +23,11 @@ public class JetModifierList extends JetElement {
         visitor.visitModifierList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitModifierList(this, data);
+    }
+
     @NotNull
     public List<JetAnnotation> getAnnotations() {
         return findChildrenByType(JetNodeTypes.ANNOTATION);

idea/src/org/jetbrains/jet/lang/psi/JetNamedFunction.java

@@ -18,6 +18,11 @@ public class JetNamedFunction extends JetFunction {
         visitor.visitNamedFunction(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitNamedFunction(this, data);
+    }
+
     public boolean hasTypeParameterListBeforeFunctionName() {
         JetTypeParameterList typeParameterList = getTypeParameterList();
         if (typeParameterList == null) {

idea/src/org/jetbrains/jet/lang/psi/JetNamespace.java

@@ -23,6 +23,11 @@ public class JetNamespace extends JetNamedDeclaration {
         visitor.visitNamespace(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitNamespace(this, data);
+    }
+
     public List<JetDeclaration> getDeclarations() {
         PsiElement body = findChildByType(JetNodeTypes.NAMESPACE_BODY);
         return PsiTreeUtil.getChildrenOfTypeAsList(body != null ? body : this, JetDeclaration.class);

idea/src/org/jetbrains/jet/lang/psi/JetNamespaceBody.java

@@ -22,4 +22,9 @@ public class JetNamespaceBody extends JetElement {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitNamespaceBody(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitNamespaceBody(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetNullableType.java

@@ -20,10 +20,15 @@ public class JetNullableType extends JetTypeElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitNullableType(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitNullableType(this, data);
+    }
+
     @NotNull
     public JetTypeElement getInnerType() {
         return findChildByClass(JetTypeElement.class);

idea/src/org/jetbrains/jet/lang/psi/JetObjectDeclaration.java

@@ -79,10 +79,15 @@ public class JetObjectDeclaration extends JetNamedDeclaration implements JetClas
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitObjectDeclaration(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitObjectDeclaration(this, data);
+    }
+
     public boolean isObjectLiteral() {
         return getNameAsDeclaration() == null;
     }

idea/src/org/jetbrains/jet/lang/psi/JetObjectLiteralExpression.java

@@ -13,10 +13,15 @@ public class JetObjectLiteralExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitObjectLiteralExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitObjectLiteralExpression(this, data);
+    }
+
     @NotNull
     public JetObjectDeclaration getObjectDeclaration() {
         return (JetObjectDeclaration) findChildByType(JetNodeTypes.OBJECT_DECLARATION);

idea/src/org/jetbrains/jet/lang/psi/JetParameter.java

@@ -15,10 +15,15 @@ public class JetParameter extends JetNamedDeclaration {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitParameter(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitParameter(this, data);
+    }
+
     @Nullable
     public JetTypeReference getTypeReference() {
         return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetParameterList.java

@@ -15,10 +15,15 @@ public class JetParameterList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitParameterList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitParameterList(this, data);
+    }
+
     public List<JetParameter> getParameters() {
         return findChildrenByType(JetNodeTypes.VALUE_PARAMETER);
     }

idea/src/org/jetbrains/jet/lang/psi/JetParenthesizedExpression.java

@@ -13,10 +13,15 @@ public class JetParenthesizedExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitParenthesizedExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitParenthesizedExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetPostfixExpression.java

@@ -19,8 +19,12 @@ public class JetPostfixExpression extends JetUnaryExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitPostfixExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitPostfixExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetPredicateExpression.java

@@ -12,7 +12,12 @@ public class JetPredicateExpression extends JetQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitPredicateExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitPredicateExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetPrefixExpression.java

@@ -14,10 +14,15 @@ public class JetPrefixExpression extends JetUnaryExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitPrefixExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitPrefixExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getBaseExpression() {
         PsiElement expression = getOperationSign().getNextSibling();

idea/src/org/jetbrains/jet/lang/psi/JetProperty.java

@@ -24,6 +24,11 @@ public class JetProperty extends JetTypeParameterListOwner {
         visitor.visitProperty(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitProperty(this, data);
+    }
+
     public boolean isVar() {
         return getNode().findChildByType(JetTokens.VAR_KEYWORD) != null;
     }

idea/src/org/jetbrains/jet/lang/psi/JetPropertyAccessor.java

@@ -17,10 +17,15 @@ public class JetPropertyAccessor extends JetDeclaration implements JetDeclaratio
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitPropertyAccessor(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitPropertyAccessor(this, data);
+    }
+
     public boolean isSetter() {
         return findChildByType(JetTokens.SET_KEYWORD) != null;
     }
@@ -49,6 +54,11 @@ public class JetPropertyAccessor extends JetDeclaration implements JetDeclaratio
         return findChildByType(JetTokens.EQ) == null;
     }
 
+    @Override
+    public boolean hasDeclaredReturnType() {
+        return false;
+    }
+
     @NotNull
     @Override
     public JetElement asElement() {

idea/src/org/jetbrains/jet/lang/psi/JetReturnExpression.java

@@ -13,10 +13,15 @@ public class JetReturnExpression extends JetLabelQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitReturnExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitReturnExpression(this, data);
+    }
+
     @Nullable
     public JetExpression getReturnedExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetRootNamespaceExpression.java

@@ -12,7 +12,12 @@ public class JetRootNamespaceExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitRootNamespaceExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitRootNamespaceExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetSafeQualifiedExpression.java

@@ -12,7 +12,12 @@ public class JetSafeQualifiedExpression extends JetQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitSafeQualifiedExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitSafeQualifiedExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetSelfType.java

@@ -21,7 +21,12 @@ public class JetSelfType extends JetTypeElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitSelfType(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitSelfType(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetSimpleNameExpression.java

@@ -157,8 +157,12 @@ public class JetSimpleNameExpression extends JetReferenceExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitSimpleNameExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitSimpleNameExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetSimpleNameStringTemplateEntry.java

@@ -17,4 +17,9 @@ public class JetSimpleNameStringTemplateEntry extends JetStringTemplateEntryWith
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitSimpleNameStringTemplateEntry(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitSimpleNameStringTemplateEntry(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetStringTemplateEntryWithExpression.java

@@ -15,4 +15,9 @@ public abstract class JetStringTemplateEntryWithExpression extends JetStringTemp
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitStringTemplateEntryWithExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitStringTemplateEntryWithExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetStringTemplateExpression.java

@@ -12,10 +12,15 @@ public class JetStringTemplateExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitStringTemplateExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitStringTemplateExpression(this, data);
+    }
+
     @NotNull
     public JetStringTemplateEntry[] getEntries() {
         return findChildrenByClass(JetStringTemplateEntry.class);

idea/src/org/jetbrains/jet/lang/psi/JetThisExpression.java

@@ -15,10 +15,15 @@ public class JetThisExpression extends JetLabelQualifiedExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitThisExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitThisExpression(this, data);
+    }
+
     /**
      * class A : B, C {
      *     override fun foo() {

idea/src/org/jetbrains/jet/lang/psi/JetThrowExpression.java

@@ -13,10 +13,15 @@ public class JetThrowExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitThrowExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitThrowExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getThrownExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetTryExpression.java

@@ -16,10 +16,15 @@ public class JetTryExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTryExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTryExpression(this, data);
+    }
+
     @NotNull
     public JetBlockExpression getTryBlock() {
         return (JetBlockExpression) findChildByType(JetNodeTypes.BLOCK);

idea/src/org/jetbrains/jet/lang/psi/JetTupleExpression.java

@@ -15,10 +15,15 @@ public class JetTupleExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTupleExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTupleExpression(this, data);
+    }
+
     public List<JetExpression> getEntries() {
         return PsiTreeUtil.getChildrenOfTypeAsList(this, JetExpression.class);
     }

idea/src/org/jetbrains/jet/lang/psi/JetTuplePattern.java

@@ -19,6 +19,11 @@ public class JetTuplePattern extends JetPattern {
         visitor.visitTuplePattern(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTuplePattern(this, data);
+    }
+
     @NotNull
     public List<JetTuplePatternEntry> getEntries() {
         return findChildrenByType(JetNodeTypes.TUPLE_PATTERN_ENTRY);

idea/src/org/jetbrains/jet/lang/psi/JetTupleType.java

@@ -21,10 +21,15 @@ public class JetTupleType extends JetTypeElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTupleType(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTupleType(this, data);
+    }
+
     @NotNull
     public List<JetTypeReference> getComponentTypeRefs() {
         return findChildrenByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetTypeArgumentList.java

@@ -15,10 +15,15 @@ public class JetTypeArgumentList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeArgumentList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeArgumentList(this, data);
+    }
+
     public List<JetTypeProjection> getArguments() {
         return findChildrenByType(JetNodeTypes.TYPE_PROJECTION);
     }

idea/src/org/jetbrains/jet/lang/psi/JetTypeConstraint.java

@@ -16,10 +16,15 @@ public class JetTypeConstraint extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeConstraint(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeConstraint(this, data);
+    }
+
     public boolean isClassObjectContraint() {
         return findChildByType(JetTokens.CLASS_KEYWORD) != null &&
                 findChildByType(JetTokens.OBJECT_KEYWORD) != null;

idea/src/org/jetbrains/jet/lang/psi/JetTypeConstraintList.java

@@ -15,10 +15,15 @@ public class JetTypeConstraintList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeConstraintList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeConstraintList(this, data);
+    }
+
     @NotNull
     public List<JetTypeConstraint> getConstraints() {
         return findChildrenByType(JetNodeTypes.TYPE_CONSTRAINT);

idea/src/org/jetbrains/jet/lang/psi/JetTypeParameter.java

@@ -16,10 +16,15 @@ public class JetTypeParameter extends JetNamedDeclaration {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeParameter(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeParameter(this, data);
+    }
+
     @NotNull
     public Variance getVariance() {
         JetModifierList modifierList = getModifierList();

idea/src/org/jetbrains/jet/lang/psi/JetTypeParameterList.java

@@ -19,7 +19,12 @@ public class JetTypeParameterList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeParameterList(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeParameterList(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetTypePattern.java

@@ -21,4 +21,9 @@ public class JetTypePattern extends JetPattern {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypePattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypePattern(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetTypeProjection.java

@@ -32,6 +32,11 @@ public class JetTypeProjection extends JetDeclaration {
         visitor.visitTypeProjection(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeProjection(this, data);
+    }
+
     @Nullable
     public JetTypeReference getTypeReference() {
         return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetTypeReference.java

@@ -18,10 +18,15 @@ public class JetTypeReference extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitTypeReference(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypeReference(this, data);
+    }
+
     public List<JetAnnotation> getAttributeAnnotations() {
         return findChildrenByType(JetNodeTypes.ANNOTATION);
     }

idea/src/org/jetbrains/jet/lang/psi/JetTypedef.java

@@ -18,6 +18,11 @@ public class JetTypedef extends JetTypeParameterListOwner {
         visitor.visitTypedef(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitTypedef(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetTypeReference getTypeReference() {
         return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);

idea/src/org/jetbrains/jet/lang/psi/JetUserType.java

@@ -27,6 +27,11 @@ public class JetUserType extends JetTypeElement {
         visitor.visitUserType(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitUserType(this, data);
+    }
+
     public JetTypeArgumentList getTypeArgumentList() {
         return (JetTypeArgumentList) findChildByType(JetNodeTypes.TYPE_ARGUMENT_LIST);
     }

idea/src/org/jetbrains/jet/lang/psi/JetValueArgument.java

@@ -18,6 +18,11 @@ public class JetValueArgument extends JetElement implements ValueArgumentPsi {
         visitor.visitArgument(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitArgument(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getArgumentExpression() {
         return findChildByClass(JetExpression.class);

idea/src/org/jetbrains/jet/lang/psi/JetValueArgumentList.java

@@ -15,10 +15,15 @@ public class JetValueArgumentList extends JetElement {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitValueArgumentList(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitValueArgumentList(this, data);
+    }
+
     public List<JetValueArgument> getArguments() {
         return findChildrenByType(JetNodeTypes.VALUE_ARGUMENT);
     }

idea/src/org/jetbrains/jet/lang/psi/JetWhenConditionCall.java

@@ -32,4 +32,9 @@ public class JetWhenConditionCall extends JetWhenCondition {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitWhenConditionCall(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhenConditionCall(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetWhenConditionInRange.java

@@ -39,6 +39,11 @@ public class JetWhenConditionInRange extends JetWhenCondition {
         visitor.visitWhenConditionInRange(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhenConditionInRange(this, data);
+    }
+
     public JetSimpleNameExpression getOperationReference() {
         return (JetSimpleNameExpression) findChildByType(JetNodeTypes.OPERATION_REFERENCE);
     }

idea/src/org/jetbrains/jet/lang/psi/JetWhenConditionIsPattern.java

@@ -26,4 +26,9 @@ public class JetWhenConditionIsPattern extends JetWhenCondition {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitWhenConditionIsPattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhenConditionIsPattern(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetWhenEntry.java

@@ -27,6 +27,11 @@ public class JetWhenEntry extends JetElement {
         visitor.visitWhenEntry(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhenEntry(this, data);
+    }
+
     @NotNull
     public JetWhenCondition[] getConditions() {
         return findChildrenByClass(JetWhenCondition.class);

idea/src/org/jetbrains/jet/lang/psi/JetWhenExpression.java

@@ -26,7 +26,12 @@ public class JetWhenExpression extends JetExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitWhenExpression(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhenExpression(this, data);
+    }
 }

idea/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java

@@ -14,10 +14,15 @@ public class JetWhileExpression extends JetLoopExpression {
     }
 
     @Override
-    public void accept(JetVisitor visitor) {
+    public void accept(@NotNull JetVisitor visitor) {
         visitor.visitWhileExpression(this);
     }
 
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWhileExpression(this, data);
+    }
+
     @Nullable @IfNotParsed
     public JetExpression getCondition() {
         return findExpressionUnder(JetNodeTypes.CONDITION);

idea/src/org/jetbrains/jet/lang/psi/JetWildcardPattern.java

@@ -15,4 +15,9 @@ public class JetWildcardPattern extends JetPattern {
     public void accept(@NotNull JetVisitor visitor) {
         visitor.visitWildcardPattern(this);
     }
+
+    @Override
+    public <R, D> R visit(@NotNull JetExtendedVisitor<R, D> visitor, D data) {
+        return visitor.visitWildcardPattern(this, data);
+    }
 }

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

@@ -119,7 +119,7 @@ public class BindingTraceContext implements BindingContext, BindingTrace {
 
     @Override
     public void recordCompileTimeValue(@NotNull JetExpression expression, @NotNull CompileTimeConstant<?> value) {
-        safePut(compileTimeValues, expression, value);
+        safePut(compileTimeValues, expression, value, true);
     }
 
     @Override
@@ -152,8 +152,13 @@ public class BindingTraceContext implements BindingContext, BindingTrace {
     }
 
     private <K, V> void safePut(Map<K, V> map, K key, V value) {
+        safePut(map, key, value, false);
+    }
+
+    private <K, V> void safePut(Map<K, V> map, K key, V value, boolean canBeEquals) {
         V oldValue = map.put(key, value);
-        assert oldValue == null || oldValue == value : (key instanceof PsiElement ? key.toString() + " \"" + ((PsiElement) key).getText() + "\"" : key.toString()) + " -> " + oldValue + " and " + value;
+        assert oldValue == null || oldValue == value || (!canBeEquals || oldValue.equals(value)) :
+                (key instanceof PsiElement ? key.toString() + " \"" + ((PsiElement) key).getText() + "\"" : key.toString()) + " -> " + oldValue + " and " + value;
     }
 
     @Override

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

@@ -733,7 +733,7 @@ public class TopDownAnalyzer {
             JetFlowInformationProvider flowInformationProvider = classDescriptorResolver.computeFlowData(declaration, bodyExpression);
             JetTypeInferrer.Services typeInferrer = semanticServices.getTypeInferrerServices(traceForConstructors, flowInformationProvider);
 
-            typeInferrer.checkFunctionReturnType(functionInnerScope, declaration, descriptor, JetStandardClasses.getUnitType());
+            typeInferrer.checkFunctionReturnType(functionInnerScope, declaration, JetStandardClasses.getUnitType());
         }
     }
 

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

@@ -93,7 +93,7 @@ public class TypeResolver {
                                     String errorMessage = (expectedArgumentCount == 0 ? "No" : expectedArgumentCount) + " type arguments expected";
                                     if (actualArgumentCount == 0) {
                                         trace.getErrorHandler().genericError(type.getNode(), errorMessage);
-                                    } else if (expectedArgumentCount == 0) {
+                                    } else {
                                         trace.getErrorHandler().genericError(type.getTypeArgumentList().getNode(), errorMessage);
                                     }
                                 } else {

idea/src/org/jetbrains/jet/lang/resolve/constants/IntValue.java

@@ -45,4 +45,21 @@ public class IntValue implements CompileTimeConstant<Integer> {
     public String toString() {
         return value + ".int";
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        IntValue intValue = (IntValue) o;
+
+        if (value != intValue.value) return false;
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return value;
+    }
 }

idea/src/org/jetbrains/jet/lang/resolve/constants/StringValue.java

@@ -36,4 +36,21 @@ public class StringValue implements CompileTimeConstant<String> {
     public String toString() {
         return "\"" + value + "\"";
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        StringValue that = (StringValue) o;
+
+        if (value != null ? !value.equals(that.value) : that.value != null) return false;
+
+        return true;
+    }
+
+    @Override
+    public int hashCode() {
+        return value != null ? value.hashCode() : 0;
+    }
 }

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

@@ -67,7 +67,21 @@ public class DataFlowInfo {
 
     @NotNull
     public List<JetType> getPossibleTypes(VariableDescriptor variableDescriptor) {
-        return typeInfo.get(variableDescriptor);
+        List<JetType> types = typeInfo.get(variableDescriptor);
+        NullabilityFlags nullabilityFlags = nullabilityInfo.get(variableDescriptor);
+        if (nullabilityFlags == null || nullabilityFlags.canBeNull()) {
+            return types;
+        }
+        List<JetType> enrichedTypes = Lists.newArrayListWithCapacity(types.size());
+        for (JetType type: types) {
+            if (type.isNullable()) {
+                enrichedTypes.add(TypeUtils.makeNotNullable(type));
+            }
+            else {
+                enrichedTypes.add(type);
+            }
+        }
+        return enrichedTypes;
     }
 
     public DataFlowInfo equalsToNull(@NotNull VariableDescriptor variableDescriptor, boolean notNull) {
@@ -76,7 +90,9 @@ public class DataFlowInfo {
 
     private ImmutableMap<VariableDescriptor, NullabilityFlags> getEqualsToNullMap(VariableDescriptor variableDescriptor, boolean notNull) {
         Map<VariableDescriptor, NullabilityFlags> builder = Maps.newHashMap(nullabilityInfo);
-        builder.put(variableDescriptor, new NullabilityFlags(!notNull, notNull));
+        NullabilityFlags nullabilityFlags = nullabilityInfo.get(variableDescriptor);
+        boolean varNotNull = notNull || (nullabilityFlags != null && !nullabilityFlags.canBeNull);
+        builder.put(variableDescriptor, new NullabilityFlags(!varNotNull, varNotNull));
         return ImmutableMap.copyOf(builder);
     }
 
@@ -85,7 +101,9 @@ public class DataFlowInfo {
         Map<VariableDescriptor, NullabilityFlags> builder = Maps.newHashMap(nullabilityInfo);
         for (VariableDescriptor variableDescriptor : variableDescriptors) {
             if (variableDescriptor != null) {
-                builder.put(variableDescriptor, new NullabilityFlags(!notNull, notNull));
+                NullabilityFlags nullabilityFlags = nullabilityInfo.get(variableDescriptor);
+                boolean varNotNull = notNull || (nullabilityFlags != null && !nullabilityFlags.canBeNull);
+                builder.put(variableDescriptor, new NullabilityFlags(!varNotNull, varNotNull));
             }
         }
         return ImmutableMap.copyOf(builder);