Fix lambda typing

lib/steep/diagnostic/ruby.rb

@@ -669,6 +669,19 @@ def header_line
         end
       end
 
+      class ProcTypeExpected < Base
+        attr_reader :type
+
+        def initialize(node:, type:)
+          super(node: node)
+          @type = type
+        end
+
+        def header_line
+          "Proc type is expected but `#{type.to_s}` is specified"
+        end
+      end
+
       class UnsupportedSyntax < Base
         attr_reader :message
 

lib/steep/type_construction.rb

@@ -2727,10 +2727,24 @@ def type_masgn(node)
       end
     end
 
+    def optional_proc?(type)
+      if type.is_a?(AST::Types::Union)
+        if type.types.size == 2
+          if type.types.find {|t| t.is_a?(AST::Types::Nil) }
+            if proc_type = type.types.find {|t| t.is_a?(AST::Types::Proc) }
+              proc_type
+            end
+          end
+        end
+      end
+    end
+
     def type_lambda(node, params_node:, body_node:, type_hint:)
       block_annotations = source.annotations(block: node, factory: checker.factory, current_module: current_namespace)
       params = TypeInference::BlockParams.from_node(params_node, annotations: block_annotations)
 
+      type_hint = deep_expand_alias(type_hint) if type_hint
+
       case type_hint
       when AST::Types::Proc
         params_hint = type_hint.type.params
@@ -2749,6 +2763,13 @@ def type_lambda(node, params_node:, body_node:, type_hint:)
 
       block_constr.typing.add_context_for_body(node, context: block_constr.context)
 
+      default_proc_function =
+        Interface::Function.new(
+          params: Interface::Function::Params.empty,
+          return_type: AST::Builtin.any_type,
+          location: nil
+        )
+
       params.params.each do |param|
         _, block_constr = block_constr.synthesize(param.node, hint: param.type)
       end
@@ -2759,17 +2780,29 @@ def type_lambda(node, params_node:, body_node:, type_hint:)
             case block_param_type
             when AST::Types::Proc
               Interface::Block.new(type: block_param_type.type, optional: false)
-            when AST::Types::Union
-              if block_param_type.types.size == 2
-                if block_param_type.types.find {|t| t.is_a?(AST::Types::Nil) }
-                  if proc_type = block_param_type.types.find {|t| t.is_a?(AST::Types::Proc) }
-                    Interface::Block.new(type: proc_type.type, optional: true)
-                  end
-                end
+            else
+              if proc_type = optional_proc?(block_param_type)
+                Interface::Block.new(type: proc_type.type, optional: true)
+              else
+                block_constr.typing.add_error(
+                  Diagnostic::Ruby::ProcTypeExpected.new(
+                    node: block_param.node,
+                    type: block_param_type
+                  )
+                )
+
+                Interface::Block.new(
+                  type: Interface::Function.new(
+                    params: Interface::Function::Params.empty,
+                    return_type: AST::Builtin.any_type,
+                    location: nil
+                  ),
+                  optional: false
+                )
               end
             end
           else
-            type_hint.block
+            block_hint
           end
         end
 

smoke/diagnostics/proc_type_expected.rb

@@ -0,0 +1,3 @@
+-> (&block) do
+  # @type var block: Integer
+end

smoke/diagnostics/test_expectations.yml

@@ -332,6 +332,18 @@
     severity: ERROR
     message: Type `::Integer` does not have method `foo`
     code: Ruby::NoMethod
+- file: proc_type_expected.rb
+  diagnostics:
+  - range:
+      start:
+        line: 1
+        character: 4
+      end:
+        line: 1
+        character: 10
+    severity: ERROR
+    message: Proc type is expected but `::Integer` is specified
+    code: Ruby::ProcTypeExpected
 - file: required_block_missing.rb
   diagnostics:
   - range:

test/type_construction_test.rb

@@ -8420,6 +8420,73 @@ def foo(&block)
     end
   end
 
+  def test_lambda_with_block_alias
+    with_checker(<<-RBS) do |checker|
+type callback[T] = ^() { (T) -> void } -> T
+    RBS
+
+      source = parse_ruby(<<-'RUBY')
+# @type var foo: callback[Integer]
+foo = -> (&block) { block[80]; 123 }
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _, _ = construction.synthesize(source.node)
+
+        assert_no_error typing
+      end
+    end
+  end
+
+  def test_lambda_with_block_non_proc
+    with_checker(<<-RBS) do |checker|
+    RBS
+
+      source = parse_ruby(<<-'RUBY')
+# @type var foo: Integer
+foo = -> (&block) do
+  block
+end
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _, _ = construction.synthesize(source.node)
+
+        assert_typing_error(typing, size: 1) do |errors|
+          assert_any!(errors) do |error|
+            assert_instance_of Diagnostic::Ruby::IncompatibleAssignment, error
+          end
+        end
+      end
+    end
+  end
+
+  def test_lambda_with_block_non_proc_arg
+    with_checker(<<-RBS) do |checker|
+    RBS
+
+      source = parse_ruby(<<-'RUBY')
+foo = -> (&block) do
+  # @type var block: Integer
+  block+1
+end
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _, _ = construction.synthesize(source.node)
+
+        assert_equal parse_type("^() { () -> untyped } -> ::Integer"), type
+
+        assert_typing_error(typing, size: 1) do |errors|
+          assert_any!(errors) do |error|
+            assert_instance_of Diagnostic::Ruby::ProcTypeExpected, error
+            assert_equal parse_type("::Integer"), error.type
+          end
+        end
+      end
+    end
+  end
+
   def test_flat_map
     with_checker(<<-RBS) do |checker|
 class FlatMap