Infer block-pass symbols (#793)

* BlockSymbol link

* First iteration of block-pass symbol resolution

* Chain::Call#block

* Superfluous block tracking

* Refactor NodeChainer

* Fix chain block check

* Sync with #780

* BlockSymbol link

* First iteration of block-pass symbol resolution

* Change spec to test Call#block

* Resolve generics from BlockSymbol

Author: castwide

lib/solargraph/parser/parser_gem/node_chainer.rb

@@ -11,11 +11,11 @@ class NodeChainer
 
         # @param node [Parser::AST::Node]
         # @param filename [String, nil]
-        # @param in_block [Boolean]
-        def initialize node, filename = nil, in_block = false
+        # @param parent [Parser::AST::Node, nil]
+        def initialize node, filename = nil, parent = nil
           @node = node
           @filename = filename
-          @in_block = in_block ? 1 : 0
+          @parent = parent
         end
 
         # @return [Source::Chain]
@@ -27,10 +27,10 @@ def chain
         class << self
           # @param node [Parser::AST::Node]
           # @param filename [String, nil]
-          # @param in_block [Boolean]
+          # @param parent [Parser::AST::Node, nil]
           # @return [Source::Chain]
-          def chain node, filename = nil, in_block = false
-            NodeChainer.new(node, filename, in_block).chain
+          def chain node, filename = nil, parent = nil
+            NodeChainer.new(node, filename, parent).chain
           end
 
           # @param code [String]
@@ -52,53 +52,39 @@ def generate_links n
           return generate_links(n.children[0]) if n.type == :splat
           result = []
           if n.type == :block
-            @in_block += 1
-            result.concat generate_links(n.children[0])
-            @in_block -= 1
+            result.concat NodeChainer.chain(n.children[0], @filename, n).links
           elsif n.type == :send
             if n.children[0].is_a?(::Parser::AST::Node)
               result.concat generate_links(n.children[0])
-              args = []
-              n.children[2..-1].each do |c|
-                args.push NodeChainer.chain(c)
-              end
-              result.push Chain::Call.new(n.children[1].to_s, args, @in_block > 0 || block_passed?(n))
+              result.push Chain::Call.new(n.children[1].to_s, node_args(n), passed_block(n))
             elsif n.children[0].nil?
               args = []
               n.children[2..-1].each do |c|
-                args.push NodeChainer.chain(c)
+                args.push NodeChainer.chain(c, @filename, n)
               end
-              result.push Chain::Call.new(n.children[1].to_s, args, @in_block > 0 || block_passed?(n))
+              result.push Chain::Call.new(n.children[1].to_s, node_args(n), passed_block(n))
             else
               raise "No idea what to do with #{n}"
             end
           elsif n.type == :csend
             if n.children[0].is_a?(::Parser::AST::Node)
               result.concat generate_links(n.children[0])
-              args = []
-              n.children[2..-1].each do |c|
-                args.push NodeChainer.chain(c)
-              end
-              result.push Chain::QCall.new(n.children[1].to_s, args, @in_block > 0 || block_passed?(n))
+              result.push Chain::QCall.new(n.children[1].to_s, node_args(n))
             elsif n.children[0].nil?
-              args = []
-              n.children[2..-1].each do |c|
-                args.push NodeChainer.chain(c)
-              end
-              result.push Chain::QCall.new(n.children[1].to_s, args, @in_block > 0 || block_passed?(n))
+              result.push Chain::QCall.new(n.children[1].to_s, node_args(n))
             else
               raise "No idea what to do with #{n}"
             end
           elsif n.type == :self
             result.push Chain::Head.new('self')
           elsif n.type == :zsuper
-            result.push Chain::ZSuper.new('super', @in_block > 0 || block_passed?(n))
+            result.push Chain::ZSuper.new('super')
           elsif n.type == :super
-            args = n.children.map { |c| NodeChainer.chain(c) }
-            result.push Chain::Call.new('super', args, @in_block > 0 || block_passed?(n))
+            args = n.children.map { |c| NodeChainer.chain(c, @filename, n) }
+            result.push Chain::Call.new('super', args)
           elsif n.type == :yield
-            args = n.children.map { |c| NodeChainer.chain(c) }
-            result.push Chain::Call.new('yield', args, @in_block > 0 || block_passed?(n))
+            args = n.children.map { |c| NodeChainer.chain(c, @filename, n) }
+            result.push Chain::Call.new('yield', args)
           elsif n.type == :const
             const = unpack_name(n)
             result.push Chain::Constant.new(const)
@@ -118,7 +104,7 @@ def generate_links n
           elsif n.type == :and
             result.concat generate_links(n.children.last)
           elsif n.type == :or
-            result.push Chain::Or.new([NodeChainer.chain(n.children[0], @filename), NodeChainer.chain(n.children[1], @filename)])
+            result.push Chain::Or.new([NodeChainer.chain(n.children[0], @filename, n), NodeChainer.chain(n.children[1], @filename, n)])
           elsif [:begin, :kwbegin].include?(n.type)
             result.concat generate_links(n.children.last)
           elsif n.type == :block_pass
@@ -128,7 +114,11 @@ def generate_links n
               # added in Ruby 3.1 - https://bugs.ruby-lang.org/issues/11256
               result.push Chain::BlockVariable.new(nil)
             else
-              result.push Chain::BlockVariable.new("&#{block_variable_name_node.children[0].to_s}")
+              if block_variable_name_node.type == :sym
+                result.push Chain::BlockSymbol.new("#{block_variable_name_node.children[0].to_s}")
+              else
+                result.push Chain::BlockVariable.new("&#{block_variable_name_node.children[0].to_s}")
+              end
             end
           elsif n.type == :hash
             result.push Chain::Hash.new('::Hash', hash_is_splatted?(n))
@@ -150,9 +140,16 @@ def hash_is_splatted? node
           true
         end
 
-        # @param node [Parser::AST::Node]
-        def block_passed? node
-          node.children.last.is_a?(::Parser::AST::Node) && node.children.last.type == :block_pass
+        def passed_block node
+          return unless node == @node && @parent&.type == :block
+
+          NodeChainer.chain(@parent.children[2], @filename)
+        end
+
+        def node_args node
+          node.children[2..-1].map do |child|
+            NodeChainer.chain(child, @filename, node)
+          end
         end
       end
     end

lib/solargraph/pin/base_variable.rb

@@ -50,7 +50,7 @@ def probe api_map
             clip = api_map.clip_at(location.filename, pos)
             # Use the return node for inference. The clip might infer from the
             # first node in a method call instead of the entire call.
-            chain = Parser.chain(node, nil, clip.in_block?)
+            chain = Parser.chain(node, nil, nil)
             result = chain.infer(api_map, closure, clip.locals)
             types.push result unless result.undefined?
           end

lib/solargraph/pin/method.rb

@@ -346,7 +346,7 @@ def param_type_from_name(tag, name)
 
       # @return [ComplexType]
       def generate_complex_type
-        tags = docstring.tags(:return).map(&:types).flatten.reject(&:nil?)
+        tags = docstring.tags(:return).map(&:types).flatten.compact
         return ComplexType::UNDEFINED if tags.empty?
         ComplexType.try_parse *tags
       end

lib/solargraph/source.rb

@@ -435,23 +435,16 @@ def string_nodes_in n
     def inner_tree_at node, position, stack
       return if node.nil?
       here = Range.from_node(node)
-      if here.contain?(position) || colonized(here, position, node)
+      if here.contain?(position)
         stack.unshift node
         node.children.each do |c|
           next unless Parser.is_ast_node?(c)
-          next if !Parser.rubyvm? && c.loc.expression.nil?
+          next if c.loc.expression.nil?
           inner_tree_at(c, position, stack)
         end
       end
     end
 
-    def colonized range, position, node
-      node.type == :COLON2 &&
-        range.ending.line == position.line &&
-        range.ending.character == position.character - 2 &&
-        code[Position.to_offset(code, Position.new(position.line, position.character - 2)), 2] == '::'
-    end
-
     protected
 
     # @return [String]

lib/solargraph/source/chain.rb

@@ -21,6 +21,7 @@ class Chain
       autoload :Head,             'solargraph/source/chain/head'
       autoload :Or,               'solargraph/source/chain/or'
       autoload :BlockVariable,    'solargraph/source/chain/block_variable'
+      autoload :BlockSymbol,      'solargraph/source/chain/block_symbol'
       autoload :ZSuper,           'solargraph/source/chain/z_super'
       autoload :Hash,             'solargraph/source/chain/hash'
       autoload :Array,            'solargraph/source/chain/array'
@@ -168,13 +169,15 @@ def infer_first_defined pins, context, api_map, locals
           @@inference_depth -= 1
         end
         return ComplexType::UNDEFINED if possibles.empty?
+
         type = if possibles.length > 1
           sorted = possibles.map { |t| t.rooted? ? "::#{t}" : t.to_s }.sort { |a, _| a == 'nil' ? 1 : 0 }
           ComplexType.parse(*sorted)
         else
           ComplexType.parse(possibles.map(&:to_s).join(', '))
         end
         return type if context.nil? || context.return_type.undefined?
+
         type.self_to(context.return_type.tag)
       end
 

lib/solargraph/source/chain/block_symbol.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Solargraph
+  class Source
+    class Chain
+      class BlockSymbol < Link
+        def resolve api_map, name_pin, locals
+          [Pin::ProxyType.anonymous(ComplexType.try_parse('Proc'))]
+        end
+      end
+    end
+  end
+end

lib/solargraph/source/chain/call.rb

@@ -10,18 +10,21 @@ class Call < Link
         # @return [::Array<Chain>]
         attr_reader :arguments
 
+        # @return [Chain, nil]
+        attr_reader :block
+
         # @param word [String]
         # @param arguments [::Array<Chain>]
-        # @param with_block [Boolean] True if the chain is inside a block
-        # @param head [Boolean] True if the call is the start of its chain
-        def initialize word, arguments = [], with_block = false
+        # @param block [Chain, nil]
+        def initialize word, arguments = [], block = nil
           @word = word
           @arguments = arguments
-          @with_block = with_block
+          @block = block
+          fix_block_pass
         end
 
         def with_block?
-          @with_block
+          !!@block
         end
 
         # @param api_map [ApiMap]
@@ -74,16 +77,7 @@ def inferred_pins pins, api_map, context, locals
                 param = ol.parameters[idx]
                 atype = nil
                 if param.nil?
-                  last_arg = idx == arguments.length - 1
-                  match = if ol.parameters.any?(&:restarg?)
-                            true
-                          elsif last_arg && ol.block?
-                            # block argument that isn't declared as an arg as well - let's add that here
-                            atypes[idx] ||= arg.infer(api_map, Pin::ProxyType.anonymous(context), locals)
-                            atypes[idx].namespace == 'Proc'
-                          else
-                            false
-                          end
+                  match = ol.parameters.any?(&:restarg?)
                   break
                 end
                 atype = atypes[idx] ||= arg.infer(api_map, Pin::ProxyType.anonymous(context), locals)
@@ -95,10 +89,16 @@ def inferred_pins pins, api_map, context, locals
                 end
               end
               if match
-                new_signature_pin = ol.resolve_generics_from_context_until_complete(ol.generics, atypes)
-                new_return_type = new_signature_pin.return_type
-                type = with_params(new_return_type.self_to(context.to_s), context).qualify(api_map, context.namespace) if new_return_type.defined?
-                type ||= ComplexType::UNDEFINED
+                # @todo Functional but dodgy generic resolution from block inference
+                if block && ol.block && ol.block.return_type.name == 'generic' && ol.return_type.to_s.include?(ol.block.return_type.to_s)
+                  blocktype = block_call_type(api_map, context)
+                  type = ComplexType.parse(ol.return_type.to_s.gsub("generic<#{ol.generics.first}>", blocktype.to_s))
+                else
+                  new_signature_pin = ol.resolve_generics_from_context_until_complete(ol.generics, atypes)
+                  new_return_type = new_signature_pin.return_type
+                  type = with_params(new_return_type.self_to(context.to_s), context).qualify(api_map, context.namespace) if new_return_type.defined?
+                  type ||= ComplexType::UNDEFINED
+                end
               end
               break if type.defined?
             end
@@ -234,6 +234,21 @@ def with_params type, context
           return type unless type.to_s.include?('$')
           ComplexType.try_parse(type.to_s.gsub('$', context.value_types.map(&:tag).join(', ')).gsub('<>', ''))
         end
+
+        def fix_block_pass
+          argument = @arguments.last&.links&.first
+          @block = @arguments.pop if argument.is_a?(BlockSymbol) || argument.is_a?(BlockVariable)
+        end
+
+        def block_call_type(api_map, context)
+          return nil unless with_block?
+
+          # @todo Handle BlockVariable and literal blocks
+          if block.links.first.is_a?(BlockSymbol)
+            callee = api_map.get_path_pins("#{context.subtypes.first}##{block.links.first.word}").first
+            callee&.return_type || ComplexType::UNDEFINED
+          end
+        end
       end
     end
   end

lib/solargraph/source/source_chainer.rb

@@ -14,10 +14,10 @@ class SourceChainer
 
       class << self
         # @param source [Source]
-        # @param position [Position]
+        # @param position [Position, Array(Integer, Integer)]
         # @return [Source::Chain]
         def chain source, position
-          new(source, position).chain
+          new(source, Solargraph::Position.normalize(position)).chain
         end
       end
 
@@ -41,8 +41,6 @@ def chain
           parent = nil
           if !source.repaired? && source.parsed? && source.synchronized?
             tree = source.tree_at(position.line, position.column)
-            # node, parent = source.tree_at(position.line, position.column)[0..2]
-            tree.shift while tree.length > 1 && tree.first.type == :SCOPE
             node, parent = tree[0..2]
           elsif source.parsed? && source.repaired? && end_of_phrase == '.'
             node, parent = source.tree_at(fixed_position.line, fixed_position.column)[0..2]
@@ -60,7 +58,7 @@ def chain
         end
         return Chain.new([Chain::UNDEFINED_CALL]) if node.nil? || (node.type == :sym && !phrase.start_with?(':'))
         # chain = NodeChainer.chain(node, source.filename, parent && parent.type == :block)
-        chain = Parser.chain(node, source.filename, parent && [:ITER, :block].include?(parent.type))
+        chain = Parser.chain(node, source.filename, parent)
         if source.repaired? || !source.parsed? || !source.synchronized?
           if end_of_phrase.strip == '.'
             chain.links.push Chain::UNDEFINED_CALL

spec/parser/node_chainer_spec.rb

@@ -109,7 +109,26 @@ class Foo
       foo(bar, &baz)
     ))
     chain = Solargraph::Parser.chain(source.node)
-    expect(chain.links.first.arguments.length).to eq(2)
+    expect(chain.links.first.arguments.length).to eq(1)
+    expect(chain.links.first).to be_with_block
+  end
+
+  it 'tracks block-pass symbols' do
+    source = Solargraph::Source.load_string(%(
+      foo(&:bar)
+    ))
+    chain = Solargraph::Parser.chain(source.node)
+    expect(chain.links.first.block).to be_a(Solargraph::Source::Chain)
+    expect(chain.links.first.block.links.first).to be_a(Solargraph::Source::Chain::BlockSymbol)
+  end
+
+  it 'tracks block-pass symbols' do
+    source = Solargraph::Source.load_string(%(
+      foo(&:bar)
+    ))
+    chain = Solargraph::Parser.chain(source.node)
+    arg = chain.links.first.block.links.first
+    expect(arg).to be_a(Solargraph::Source::Chain::BlockSymbol)
   end
 
   # feature added in Ruby 3.1

spec/source/source_chainer_spec.rb

@@ -326,4 +326,14 @@ def strings; end
     expect(chain.node.type).to be(:send)
     expect(chain.node.children[1]).to be(:s)
   end
+
+  it 'adds blocks to calls' do
+    source = Solargraph::Source.load_string(%(
+      x.y do
+        z
+      end
+    ))
+    chain = Solargraph::Source::SourceChainer.chain(source, [1, 9])
+    expect(chain.links.last.block).to be
+  end
 end