Accept block-pass arguments for optional blocks

lib/steep/type_construction.rb

@@ -4358,7 +4358,9 @@ def try_method_type(node, receiver_type:, method_name:, method_overload:, argume
             when arg.compatible?
               if arg.node
                 # Block pass (&block) is given
-                node_type, constr = constr.synthesize(arg.node, hint: arg.node_type)
+                # Passing `proc_type` instead of `node_type` because a block-pass argument is a proc, not `nil`.
+                # `nil` is still allowed by the `node_type` check below.
+                node_type, constr = constr.synthesize(arg.node, hint: arg.proc_type)
 
                 nil_given =
                   constr.check_relation(sub_type: node_type, super_type: AST::Builtin.nil_type).success? &&

lib/steep/type_inference/send_args.rb

@@ -473,16 +473,20 @@ def pair
           end
         end
 
-        def node_type
+        def proc_type
           raise unless block
 
-          type = AST::Types::Proc.new(type: block.type, block: nil, self_type: block.self_type)
+          AST::Types::Proc.new(type: block.type, block: nil, self_type: block.self_type)
+        end
+
+        def node_type
+          raise unless block
 
           if block.optional?
-            type = AST::Types::Union.build(types: [type, AST::Builtin.nil_type])
+            AST::Types::Union.build(types: [proc_type, AST::Builtin.nil_type])
+          else
+            proc_type
           end
-
-          type
         end
       end
 

sig/steep/type_inference/send_args.rbs

@@ -199,6 +199,14 @@ module Steep
 
         def pair: () -> [Parser::AST::Node, Interface::Function]?
 
+        # The type of the block, as a proc type
+        #
+        %a{pure} def proc_type: () -> AST::Types::Proc
+
+        # The type a block-pass argument node is allowed to have
+        #
+        # Same as `proc_type`, but `nil` is allowed too if the block is optional, because passing `nil` means passing no block.
+        #
         %a{pure} def node_type: () -> AST::Types::t
       end
 

sig/test/type_check_test.rbs

@@ -295,4 +295,8 @@ class TypeCheckTest < Minitest::Test
   def test_tuple_type_with_if_branch: () -> untyped
 
   def test_or_asgn_send_chain: () -> untyped
+
+  def test_block_pass_optional_block: () -> untyped
+
+  def test_block_pass_optional_block_mismatch: () -> untyped
 end

test/type_check_test.rb

@@ -4580,4 +4580,83 @@ class OrAssignBox[T]
       YAML
     )
   end
+
+  def test_block_pass_optional_block
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+          class Object
+            def stringify: (Integer value) -> String
+          end
+
+          class OptionalBlockTest
+            def sum: (?untyped init) ?{ (Integer e) -> untyped } -> untyped
+
+            def map: () ?{ (Integer element) -> String } -> Array[String]
+          end
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          def stringify(value)
+            value.to_s
+          end
+
+          test = OptionalBlockTest.new
+
+          test.sum(&:to_r)
+          test.sum(0, &:to_r)
+          test.map(&:to_s)
+          test.map(&method(:stringify))
+
+          test.map
+          test.map {|element| element.to_s }
+          test.map(&(-> (element) { element.to_s }))
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics: []
+      YAML
+    )
+  end
+
+  def test_block_pass_optional_block_mismatch
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+          class OptionalBlockTest
+            def map: () ?{ (Integer element) -> Integer } -> Array[Integer]
+          end
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          OptionalBlockTest.new.map(&:to_s)
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics:
+          - range:
+              start:
+                line: 1
+                character: 26
+              end:
+                line: 1
+                character: 32
+            severity: ERROR
+            message: |-
+              Cannot pass a value of type `^(::Integer) -> ::String` as a block-pass-argument of type `(^(::Integer) -> ::Integer | nil)`
+                ^(::Integer) -> ::String <: (^(::Integer) -> ::Integer | nil)
+                  ^(::Integer) -> ::String <: ^(::Integer) -> ::Integer
+                    ::String <: ::Integer
+                      ::Object <: ::Integer
+                        ::BasicObject <: ::Integer
+            code: Ruby::BlockTypeMismatch
+      YAML
+    )
+  end
 end