ruby · soutaro · Sep 19, 2019 · Sep 23, 2019
diff --git a/lib/ruby/signature/method_type.rb b/lib/ruby/signature/method_type.rb
@@ -4,29 +4,34 @@ class MethodType
       class Block
         attr_reader :type
         attr_reader :required
+        attr_reader :self_type
 
-        def initialize(type:, required:)
+        def initialize(type:, required:, self_type:)
           @type = type
           @required = required
+          @self_type = self_type
         end
 
         def ==(other)
           other.is_a?(Block) &&
             other.type == type &&
-            other.required == required
+            other.required == required &&
+            other.self_type == self_type
         end
 
         def to_json(*a)
           {
             type: type,
-            required: required
+            required: required,
+            self_type: self_type
           }.to_json(*a)
         end
 
         def sub(s)
           self.class.new(
             type: type.sub(s),
-            required: required
+            required: required,
+            self_type: self_type&.sub(s)
           )
         end
       end
@@ -87,7 +92,9 @@ def map_type(&block)
           type_params: type_params,
           type: type.map_type(&block),
           block: self.block&.yield_self do |b|
-            Block.new(type: b.type.map_type(&block), required: b.required)
+            Block.new(type: b.type.map_type(&block),
+                      required: b.required,
+                      self_type: b.self_type && b.self_type.map_type(&block))
           end,
           location: location
         )
@@ -105,11 +112,15 @@ def each_type(&block)
       end
 
       def to_s
+        self_type = block&.self_type&.yield_self do |type|
+          " @ #{type.to_s}"
+        end
+
         s = case
             when block && block.required
-              "(#{type.param_to_s}) { (#{block.type.param_to_s}) -> #{block.type.return_to_s} } -> #{type.return_to_s}"
+              "(#{type.param_to_s}) { (#{block.type.param_to_s}) -> #{block.type.return_to_s} }#{self_type} -> #{type.return_to_s}"
             when block
-              "(#{type.param_to_s}) ?{ (#{block.type.param_to_s}) -> #{block.type.return_to_s} } -> #{type.return_to_s}"
+              "(#{type.param_to_s}) ?{ (#{block.type.param_to_s}) -> #{block.type.return_to_s} }#{self_type} -> #{type.return_to_s}"
             else
               "(#{type.param_to_s}) -> #{type.return_to_s}"
             end

diff --git a/lib/ruby/signature/parser.y b/lib/ruby/signature/parser.y
@@ -5,6 +5,7 @@ class Ruby::Signature::Parser
         tSTRING tSYMBOL tINTEGER tWRITE_ATTR
         kLPAREN kRPAREN kLBRACKET kRBRACKET kLBRACE kRBRACE
         kVOID kNIL kANY kTOP kBOT kSELF kSELFQ kINSTANCE kCLASS kBOOL kSINGLETON kTYPE kDEF kMODULE kSUPER
+        kATMARK
         kPRIVATE kPUBLIC kALIAS
         kCOLON kCOLON2 kCOMMA kBAR kAMP kHAT kARROW kQUESTION kEXCLAMATION kSTAR kSTAR2 kFATARROW kEQ kDOT kLT
         kINTERFACE kEND kINCLUDE kEXTEND kATTRREADER kATTRWRITER kATTRACCESSOR tOPERATOR tQUOTEDMETHOD
@@ -459,13 +460,19 @@ rule
 
   block_opt:
       { result = nil }
-    | kLBRACE function_type kRBRACE {
-        block = MethodType::Block.new(type: val[1].value, required: true)
-        result = LocatedValue.new(value: block, location: val[0].location + val[2].location)
+    | kLBRACE function_type kRBRACE self_type_opt {
+        block = MethodType::Block.new(type: val[1].value, required: true, self_type: val[3])
+        result = LocatedValue.new(value: block, location: val[0].location + (val[3] || val[2]).location)
       }
-    | kQUESTION kLBRACE function_type kRBRACE {
-        block = MethodType::Block.new(type: val[2].value, required: false)
-        result = LocatedValue.new(value: block, location: val[0].location + val[3].location)
+    | kQUESTION kLBRACE function_type kRBRACE self_type_opt {
+        block = MethodType::Block.new(type: val[2].value, required: false, self_type: val[4])
+        result = LocatedValue.new(value: block, location: val[0].location + (val[4] || val[3]).location)
+      }
+
+  self_type_opt:
+      { result = nil }
+    | kATMARK simple_type {
+        result = val[1]
       }
 
   def_name:
@@ -1173,6 +1180,7 @@ PUNCTS = {
   "/" => :tOPERATOR,
   "`" => :tOPERATOR,
   "%" => :tOPERATOR,
+  "@" => :kATMARK
 }
 PUNCTS_RE = Regexp.union(*PUNCTS.keys)
 
@@ -1229,6 +1237,10 @@ def next_token
   when input.scan(/:\w+\b/)
     s = input.matched.yield_self {|s| s[1, s.length] }.to_sym
     new_token(:tSYMBOL, s)
+  when input.scan(/@[a-zA-Z_]\w*/)
+    new_token(:tIVAR, input.matched.to_sym)
+  when input.scan(/@@[a-zA-Z_]\w*/)
+    new_token(:tCLASSVAR, input.matched.to_sym)
   when input.scan(PUNCTS_RE)
     new_token(PUNCTS[input.matched])
   when input.scan(/(::)?([A-Z]\w*::)+/)
@@ -1239,10 +1251,6 @@ def next_token
     new_token(:tUKEYWORD, input.matched.chop.to_sym)
   when input.scan(/\$[A-Za-z_]\w*/)
     new_token(:tGLOBALIDENT)
-  when input.scan(/@[a-zA-Z_]\w*/)
-    new_token(:tIVAR, input.matched.to_sym)
-  when input.scan(/@@[a-zA-Z_]\w*/)
-    new_token(:tCLASSVAR, input.matched.to_sym)
   when input.scan(/_[a-zA-Z]\w*\b/)
     new_token(:tINTERFACEIDENT)
   when input.scan(/[A-Z]\w*\b/)

diff --git a/lib/ruby/signature/scaffold/rbi.rb b/lib/ruby/signature/scaffold/rbi.rb
@@ -308,6 +308,7 @@ def method_type(args_node, type_node, variables:)
 
         def parse_params(args_node, args, method_type, variables:)
           vars = (node_to_hash(each_arg(args).to_a[0]) || {}).transform_values {|value| type_of(value, variables: variables) }
+          nodes = (node_to_hash(each_arg(args).to_a[0]) || {})
 
           required_positionals = []
           optional_positionals = []
@@ -374,15 +375,19 @@ def parse_params(args_node, args, method_type, variables:)
 
           method_block = nil
           if block
-            if (type = vars[block])
+            if (type = vars[block]) && (node = nodes[block])
               if type.is_a?(Types::Proc)
-                method_block = MethodType::Block.new(required: true, type: type.type)
+                self_type = proc_bind(node, variables: variables)
+                method_block = MethodType::Block.new(required: true,
+                                                     type: type.type,
+                                                     self_type: self_type)
               else
                 STDERR.puts "Unexpected block type: #{type}"
                 PP.pp args_node, STDERR
                 method_block = MethodType::Block.new(
                   required: true,
-                  type: Types::Function.empty(Types::Bases::Any.new(location: nil))
+                  type: Types::Function.empty(Types::Bases::Any.new(location: nil)),
+                  self_type: nil
                 )
               end
             end
@@ -479,7 +484,17 @@ def proc_type?(type_node)
           else
             type_node.type == :CALL && proc_type?(type_node.children[0])
           end
+        end
 
+        def proc_bind(node, variables:)
+          if call_node?(node, name: :bind, receiver: -> (_) { true })
+            type_node = each_arg(node.children[2]).to_a[0]
+            type_of(type_node, variables: variables)
+          else
+            if node.type == :CALL
+              proc_bind(node.children[0], variables: variables)
+            end
+          end
         end
 
         def call_node?(node, name:, receiver: -> (node) { node.type == :CONST && node.children[0] == :T }, args: -> (node) { true })

diff --git a/test/ruby/signature/method_type_parsing_test.rb b/test/ruby/signature/method_type_parsing_test.rb
@@ -30,4 +30,18 @@ def test_method_type_eof_re_error
 
     assert_equal "}", error.error_value
   end
+
+  def test_self_type
+    Parser.parse_method_type("[A] () { () -> A } @ Integer -> A").yield_self do |type|
+      assert_equal "Integer", type.block.self_type.to_s
+    end
+
+    Parser.parse_method_type("[A] () { () -> A } @ singleton(Integer) -> A").yield_self do |type|
+      assert_equal "singleton(Integer)", type.block.self_type.to_s
+    end
+
+    Parser.parse_method_type("[A] () { () -> A } @ self -> A").yield_self do |type|
+      assert_equal "self", type.block.self_type.to_s
+    end
+  end
 end
diff --git a/test/ruby/signature/rbi_scaffold_test.rb b/test/ruby/signature/rbi_scaffold_test.rb
@@ -405,6 +405,29 @@ class Dir
     assert_write parser.decls, <<-EOF
 class Dir
   include Enumerable
+end
+    EOF
+  end
+
+  def test_bind_proc
+    parser = RBI.new
+
+    parser.parse <<-EOF
+class Hello
+  sig do
+    type_parameters(:U)
+    .params(
+        blk: T.proc.bind(T.untyped).params().returns(T.type_parameter(:U)),
+    )
+    .returns(T.type_parameter(:U))
+  end
+  def instance_eval(arg0=T.unsafe(nil), filename=T.unsafe(nil), lineno=T.unsafe(nil), &blk); end
+end
+    EOF
+
+    assert_write parser.decls, <<-EOF
+class Hello
+  def instance_eval: [U] () { () -> U } @ any -> U
 end
     EOF
   end

diff --git a/test/ruby/signature/writer_test.rb b/test/ruby/signature/writer_test.rb
@@ -126,6 +126,14 @@ def initialize: () -> void
 
 class Bar
   def self.new: (String) -> Bar
+end
+    SIG
+  end
+
+  def test_block_self
+    assert_writer <<-SIG
+class Foo
+  def instance_eval: [A] () { () -> A } @ self -> A
 end
     SIG
   end