Fixed #2988: Union(...) restriction doesn't work

crystal-lang · Jul 13, 2016 · 351fe7f · 351fe7f · bcardiff · Jul 13, 2016
1 parent 4dfc24a
commit 351fe7f
Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/spec/compiler/type_inference/union_spec.cr b/spec/compiler/type_inference/union_spec.cr
@@ -155,4 +155,14 @@ describe "Type inference: union" do
       nil.as(Void?)
       )) { nil_type }
   end
+
+  it "can use Union in type restriction (#2988)" do
+    assert_type(%(
+      def foo(x : Union(Int32, String))
+        x
+      end
+
+      {foo(1), foo("hi")}
+      )) { tuple_of([int32, string]) }
+  end
 end
diff --git a/src/compiler/crystal/semantic/restrictions.cr b/src/compiler/crystal/semantic/restrictions.cr
@@ -342,6 +342,12 @@ module Crystal
     end
 
     def restrict(other : Generic, context)
+      # Special case: consider `Union(X, Y, ...)` the same as `X | Y | ...`
+      generic_class = context.type_lookup.lookup_type other.name
+      if generic_class.is_a?(GenericUnionType)
+        return restrict(Union.new(other.type_vars), context)
+      end
+
       parents.try &.each do |parent|
         next if parent.is_a?(NonGenericModuleType)
 
@@ -686,7 +692,7 @@ module Crystal
 
     def restrict(other : Generic, context)
       generic_module = context.type_lookup.lookup_type other.name
-      return nil unless generic_module == @module
+      return super unless generic_module == @module
 
       generic_module = generic_module.as(GenericModuleType)
       return nil unless generic_module.type_vars.size == @module.type_vars.size
@@ -728,7 +734,7 @@ module Crystal
 
     def restrict(other : Generic, context)
       generic_class = context.type_lookup.lookup_type other.name
-      return nil unless generic_class == @extended_class
+      return super unless generic_class == @extended_class
 
       generic_class = generic_class.as(GenericClassType)
       return nil unless generic_class.type_vars.size == type_vars.size