crystal-lang · Mar 23, 2017 · Mar 23, 2017 · Mar 23, 2017
diff --git a/spec/compiler/semantic/macro_spec.cr b/spec/compiler/semantic/macro_spec.cr
@@ -775,6 +775,22 @@ describe "Semantic: macro" do
       )) { int32.metaclass }
   end
 
+  it "finds generic type argument of included module with self" do
+    assert_type(%(
+      module Bar(T)
+        def t
+          {{ T }}
+        end
+      end
+
+      class Foo(U)
+        include Bar(self)
+      end
+
+      Foo(Int32).new.t
+      )) { generic_class("Foo", int32).metaclass }
+  end
+
   it "finds free type vars" do
     assert_type(%(
       module Foo(T)

diff --git a/spec/compiler/semantic/module_spec.cr b/spec/compiler/semantic/module_spec.cr
@@ -226,15 +226,183 @@ describe "Semantic: module" do
         end
       end
 
-      class Baz(X)
+      class Bar(U)
+        include Foo(self)
+      end
+
+      Bar(Int32).new.foo
+      ") { generic_class("Bar", int32).metaclass }
+  end
+
+  it "includes generic module with self, and inherits it" do
+    assert_type("
+      module Foo(T)
+        def foo
+          T
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz < Bar(Int32)
+      end
+
+      Baz.new.foo
+      ") { types["Baz"].metaclass }
+  end
+
+  it "includes generic module with self (check argument type, success)" do
+    assert_type("
+      module Foo(T)
+        def foo(x : T)
+          x
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      Bar(Int32).new.foo Bar(Int32).new
+      ") { generic_class("Bar", int32) }
+  end
+
+  it "includes generic module with self (check argument superclass type, success)" do
+    assert_type("
+      module Foo(T)
+        def foo(x : T)
+          x
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz < Bar(Int32)
+      end
+
+      Bar(Int32).new.foo Baz.new
+      ") { types["Baz"] }
+  end
+
+  it "includes generic module with self (check argument type, error)" do
+    assert_error "
+      module Foo(T)
+        def foo(x : T)
+          x
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz1 < Bar(Int32)
+      end
+
+      class Baz2 < Bar(Int32)
+      end
+
+      Baz1.new.foo Baz2.new
+      ", "no overload matches"
+  end
+
+  it "includes generic module with self (check argument superclass type, error)" do
+    assert_error "
+      module Foo(T)
+        def foo(x : T)
+          x
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz < Bar(Int32)
+      end
+
+      Baz.new.foo Bar(Int32).new
+      ", "no overload matches"
+  end
+
+  it "includes generic module with self (check return type, success)" do
+    assert_type("
+      module Foo(T)
+        def foo : T
+          Bar(Int32).new
+        end
       end
 
       class Bar(U)
         include Foo(self)
       end
 
       Bar(Int32).new.foo
-      ") { generic_class("Bar", int32).metaclass }
+      ") { generic_class("Bar", int32) }
+  end
+
+  it "includes generic module with self (check return subclass type, success)" do
+    assert_type("
+      module Foo(T)
+        def foo : T
+          Baz.new
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz < Bar(Int32)
+      end
+
+      Bar(Int32).new.foo
+      ") { types["Baz"] }
+  end
+
+  it "includes generic module with self (check return type, error)" do
+    assert_error "
+      module Foo(T)
+        def foo : T
+          Bar(Int32).new
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz < Bar(Int32)
+      end
+
+      Baz.new.foo
+      ", "type must be Baz, not Bar(Int32)"
+  end
+
+  it "includes generic module with self (check return subclass type, error)" do
+    assert_error "
+      module Foo(T)
+        def foo : T
+          Baz2.new
+        end
+      end
+
+      class Bar(U)
+        include Foo(self)
+      end
+
+      class Baz1 < Bar(Int32)
+      end
+
+      class Baz2 < Bar(Int32)
+      end
+
+      Baz1.new.foo
+      ", "type must be Baz1, not Baz2"
   end
 
   it "includes module but can't access metaclass methods" do

diff --git a/src/compiler/crystal/macros/interpreter.cr b/src/compiler/crystal/macros/interpreter.cr
@@ -418,6 +418,9 @@ module Crystal
         end
 
         TypeNode.new(matched_type)
+      when Self
+        target = @scope == @program.class_type ? @scope : @scope.instance_type
+        TypeNode.new(target)
       when ASTNode
         matched_type
       else

diff --git a/src/compiler/crystal/semantic/main_visitor.cr b/src/compiler/crystal/semantic/main_visitor.cr
@@ -177,6 +177,8 @@ module Crystal
         # It's different if from a virtual type we do `v.class.new`
         # because the class could be any in the hierarchy.
         node.type = check_type_in_type_args(type.remove_alias_if_simple).devirtualize
+      when Self
+        node.type = check_type_in_type_args(the_self(node).remove_alias_if_simple)
       when ASTNode
         type.accept self unless type.type?
         node.syntax_replacement = type
@@ -310,12 +312,7 @@ module Crystal
     end
 
     def visit(node : Self)
-      the_self = (@scope || current_type)
-      if the_self.is_a?(Program)
-        node.raise "there's no self in this scope"
-      end
-
-      node.type = the_self.instance_type
+      node.type = the_self(node).instance_type
     end
 
     def visit(node : Var)
@@ -3124,6 +3121,14 @@ module Crystal
       end
     end
 
+    def the_self(node)
+      the_self = (@scope || current_type)
+      if the_self.is_a?(Program)
+        node.raise "there's no self in this scope"
+      end
+      the_self
+    end
+
     def visit(node : When | Unless | Until | MacroLiteral | OpAssign)
       raise "BUG: #{node.class_desc} node '#{node}' (#{node.location}) should have been eliminated in normalize"
     end

diff --git a/src/compiler/crystal/semantic/semantic_visitor.cr b/src/compiler/crystal/semantic/semantic_visitor.cr
@@ -220,8 +220,8 @@ abstract class Crystal::SemanticVisitor < Crystal::Visitor
     end
   end
 
-  def lookup_type(node : ASTNode, free_vars = nil)
-    current_type.lookup_type(node, free_vars: free_vars, allow_typeof: false)
+  def lookup_type(node : ASTNode, free_vars = nil, lazy_self = false)
+    current_type.lookup_type(node, free_vars: free_vars, allow_typeof: false, lazy_self: lazy_self)
   end
 
   def check_outside_exp(node, op)

diff --git a/src/compiler/crystal/semantic/top_level_visitor.cr b/src/compiler/crystal/semantic/top_level_visitor.cr
@@ -811,7 +811,7 @@ class Crystal::TopLevelVisitor < Crystal::SemanticVisitor
   def include_in(current_type, node, kind)
     node_name = node.name
 
-    type = lookup_type(node_name)
+    type = lookup_type(node_name, lazy_self: true)
     case type
     when GenericModuleType
       node.raise "wrong number of type vars for #{type} (given 0, expected #{type.type_vars.size})"