Typing for When-expressions

idea/src/org/jetbrains/jet/JetNodeTypes.java

@@ -115,8 +115,8 @@ public interface JetNodeTypes {
     JetNodeType TYPE_PATTERN              = new JetNodeType("TYPE_PATTERN");
     JetNodeType EXPRESSION_PATTERN        = new JetNodeType("EXPRESSION_PATTERN");
     JetNodeType BINDING_PATTERN           = new JetNodeType("BINDING_PATTERN");
-    JetNodeType WHEN                      = new JetNodeType("WHEN");
+    JetNodeType WHEN                      = new JetNodeType("WHEN", JetWhenExpression.class);
-    JetNodeType WHEN_ENTRY                = new JetNodeType("WHEN_ENTRY");
+    JetNodeType WHEN_ENTRY                = new JetNodeType("WHEN_ENTRY", JetWhenEntry.class);
     JetNodeType WHEN_CONDITION            = new JetNodeType("WHEN_CONDITION");
     JetNodeType TUPLE_PATTERN_ENTRY       = new JetNodeType("TUPLE_PATTERN_ENTRY");
     JetNodeType NULLABLE_TYPE             = new JetNodeType("NULLABLE_TYPE", JetNullableType.class);

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

@@ -186,6 +186,10 @@ public class JetVisitor extends PsiElementVisitor {
         visitExpression(expression);
     }
 
+    public void visitWhenExpression(JetWhenExpression expression) {
+        visitExpression(expression);
+    }
+
     public void visitTryExpression(JetTryExpression expression) {
         visitExpression(expression);
     }
@@ -321,4 +325,8 @@ public class JetVisitor extends PsiElementVisitor {
     public void visitTypeProjection(JetTypeProjection typeProjection) {
         visitJetElement(typeProjection);
     }
+
+    public void visitWhenEntry(JetWhenEntry jetWhenEntry) {
+        visitJetElement(jetWhenEntry);
+    }
 }

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

@@ -0,0 +1,40 @@
+package org.jetbrains.jet.lang.psi;
+
+import com.intellij.lang.ASTNode;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.JetNodeTypes;
+import org.jetbrains.jet.lexer.JetTokens;
+
+/**
+ * @author abreslav
+ */
+public class JetWhenEntry extends JetElement {
+    public JetWhenEntry(@NotNull ASTNode node) {
+        super(node);
+    }
+
+    public boolean isElse() {
+        return findChildByType(JetTokens.ELSE_KEYWORD) != null;
+    }
+
+    public boolean isElseContinue() {
+        return isElse() && (findChildByType(JetTokens.DOUBLE_ARROW) == null) && (findChildByType(JetTokens.CONTINUE_KEYWORD) != null);
+    }
+
+    @Nullable
+    public JetExpression getExpression() {
+        return findChildByClass(JetExpression.class);
+    }
+
+    @Nullable
+    public JetWhenExpression getSubWhen() {
+        // TODO: this may be a WHEN that goes after "=>"
+        return (JetWhenExpression) findChildByType(JetNodeTypes.WHEN);
+    }
+
+    @Override
+    public void accept(@NotNull JetVisitor visitor) {
+        visitor.visitWhenEntry(this);
+    }
+}

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

@@ -0,0 +1,26 @@
+package org.jetbrains.jet.lang.psi;
+
+import com.intellij.lang.ASTNode;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.JetNodeTypes;
+
+import java.util.List;
+
+/**
+ * @author abreslav
+ */
+public class JetWhenExpression extends JetExpression {
+    public JetWhenExpression(@NotNull ASTNode node) {
+        super(node);
+    }
+
+    @NotNull
+    public List<JetWhenEntry> getEntries() {
+        return findChildrenByType(JetNodeTypes.WHEN_ENTRY);
+    }
+
+    @Override
+    public void accept(JetVisitor visitor) {
+        visitor.visitWhenExpression(this);
+    }
+}

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

@@ -119,7 +119,7 @@ public class JetTypeChecker {
             @Override
             public void visitIfExpression(JetIfExpression expression) {
                 // TODO : check condition type
-                // TODO : change types according to is and nullability according to null comparisons
+                // TODO : change types according to is and nullability checks
                 JetExpression elseBranch = expression.getElse();
                 if (elseBranch == null) {
                     // TODO : type-check the branch
@@ -131,6 +131,15 @@ public class JetTypeChecker {
                 }
             }
 
+            @Override
+            public void visitWhenExpression(JetWhenExpression expression) {
+                // TODO : what if there're sub-whens?
+                // TODO : what if there's "else continue"?
+                List<Type> expressions = new ArrayList<Type>();
+                collectAllReturnTypes(expression, scope, expressions);
+                result[0] = commonSupertype(expressions);
+            }
+
             @Override
             public void visitTupleExpression(JetTupleExpression expression) {
                 List<JetExpression> entries = expression.getEntries();
@@ -150,6 +159,19 @@ public class JetTypeChecker {
         return result[0];
     }
 
+    private void collectAllReturnTypes(JetWhenExpression whenExpression, JetScope scope, List<Type> result) {
+        for (JetWhenEntry entry : whenExpression.getEntries()) {
+            if (entry.getSubWhen() != null) {
+                collectAllReturnTypes(entry.getSubWhen(), scope, result);
+            } else {
+                JetExpression resultExpression = entry.getExpression();
+                if (resultExpression != null) {
+                    result.add(getType(scope, resultExpression));
+                }
+            }
+        }
+    }
+
     public Type commonSupertype(Collection<Type> types) {
         Collection<Type> typeSet = new HashSet<Type>(types);
         assert !typeSet.isEmpty();

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

@@ -89,6 +89,16 @@ public class JetTypeCheckerTest extends LightDaemonAnalyzerTestCase {
         assertType("if (true) 1 else '1'", "Any");
     }
 
+    public void testWhen() throws Exception {
+        assertType("when (1) { is 1 => 2; } ", "Int");
+        assertType("when (1) { is 1 => 2; is 1 => '2'} ", "Any");
+        assertType("when (1) { is 1 => 2; is 1 => '2'; is 1 => null} ", "Any?");
+        assertType("when (1) { is 1 => 2; is 1 => '2'; else => null} ", "Any?");
+        assertType("when (1) { is 1 => 2; is 1 => '2'; else continue} ", "Any");
+        assertType("when (1) { is 1 => 2; is 1 => '2'; is 1 when(e) {is 1 => null}} ", "Any?");
+        assertType("when (1) { is 1 => 2; is 1 => '2'; is 1 => when(e) {is 1 => null}} ", "Any?");
+    }
+
     public void testCommonSupertypes() throws Exception {
         assertCommonSupertype("Int", "Int", "Int");