steep 0.4.0 → 0.5.0

data/lib/steep/type_construction.rb

@@ -34,6 +34,23 @@ module Steep
       def boolean?(type)
         type.is_a?(AST::Types::Boolean)
       end
+
+      def optional(type)
+        AST::Types::Union.build(types: [type, nil_instance])
+      end
+
+      def hash_instance?(type)
+        case type
+        when AST::Types::Name
+          type.name.is_a?(TypeName::Instance) && type.name.name == ModuleName.new(name: "Hash", absolute: true)
+        else
+          false
+        end
+      end
+
+      def array_instance?(type)
+        type.is_a?(AST::Types::Name) && type.name.is_a?(TypeName::Instance) && type.name.name == ModuleName.new(name: "Array", absolute: true)
+      end
     end
 
     class MethodContext
@@ -122,8 +139,23 @@ module Steep
       @type_env = type_env
     end
 
+    def with_new_typing(typing)
+      self.class.new(
+        checker: checker,
+        source: source,
+        annotations: annotations,
+        type_env: type_env,
+        typing: typing,
+        self_type: self_type,
+        method_context: method_context,
+        block_context: block_context,
+        module_context: module_context,
+        break_context: break_context
+      )
+    end
+
     def for_new_method(method_name, node, args:, self_type:)
-      annots = source.annotations(block: node)
+      annots = source.annotations(block: node, builder: checker.builder, current_module: current_namespace)
       type_env = TypeInference::TypeEnv.new(subtyping: checker,
                                             const_env: module_context&.const_env || self.type_env.const_env)
 
@@ -131,26 +163,61 @@ module Steep
         type_env.set(const: name, type: type)
       end
 
-      self_type = annots.self_type || self_type
+      self_type = expand_alias(annots.self_type || self_type)
 
       self_interface = self_type && (self_type != Types.any || nil) && checker.resolve(self_type, with_initialize: true)
-      interface_method = self_interface&.yield_self {|interface| interface.methods[method_name] }
-      annotation_method = annotations.lookup_method_type(method_name)&.yield_self do |method_type|
+      interface_method = self_interface&.yield_self do |interface|
+        interface.methods[method_name]&.yield_self do |method|
+          if self_type.is_a?(AST::Types::Name) && method.type_name == self_type.name
+            method
+          else
+            Interface::Method.new(type_name: self_type,
+                                  name: method_name,
+                                  types: method.types,
+                                  super_method: method,
+                                  attributes: [])
+          end
+        end
+      end
+
+      annotation_method = annotations.method_type(method_name)&.yield_self do |method_type|
+        subst = Interface::Substitution.build([],
+                                              instance_type: module_context&.instance_type || AST::Types::Instance.new,
+                                              module_type: module_context&.module_type || AST::Types::Class.new,
+                                              self_type: self_type)
         Interface::Method.new(type_name: nil,
                               name: method_name,
-                              types: [checker.builder.method_type_to_method_type(method_type,
-                                                                                 current: current_namespace)],
+                              types: [method_type.subst(subst)],
                               super_method: interface_method&.super_method,
                               attributes: [])
       end
 
       if interface_method && annotation_method
+        interface_types = interface_method.types.map do |method_type|
+          subst = Interface::Substitution.build(method_type.type_params)
+          method_type.instantiate(subst)
+        end
+
+        unknowns = []
+        annotation_types = annotation_method.types.each.with_object([]) do |method_type, array|
+          fresh = method_type.type_params.map {|var| AST::Types::Var.fresh(var) }
+          unknowns.push(*fresh)
+          
+          subst = Interface::Substitution.build(method_type.type_params, fresh)
+          array << method_type.instantiate(subst)
+        end
+
+        constraints = Subtyping::Constraints.new(unknowns: unknowns)
+        interface_types.each do |type|
+          constraints.add_var *type.free_variables.to_a
+        end
+
         result = checker.check_method(method_name,
-                                      annotation_method,
-                                      interface_method,
+                                      annotation_method.with_types(annotation_types),
+                                      interface_method.with_types(interface_types),
                                       assumption: Set.new,
                                       trace: Subtyping::Trace.new,
-                                      constraints: Subtyping::Constraints.empty)
+                                      constraints: constraints)
 
         if result.failure?
           typing.add_error Errors::IncompatibleMethodTypeAnnotation.new(
@@ -168,10 +235,22 @@ module Steep
       when method && method.types.size == 1
         method_type = method.types.first
         return_type = method_type.return_type
-        var_types = TypeConstruction.parameter_types(args, method_type).transform_values {|type| absolute_type(type) }
+        var_types = TypeConstruction.parameter_types(args, method_type)
         unless TypeConstruction.valid_parameter_env?(var_types, args.reject {|arg| arg.type == :blockarg }, method_type.params)
           typing.add_error Errors::MethodArityMismatch.new(node: node)
         end
+
+        if (block_arg = args.find {|arg| arg.type == :blockarg })
+          if method_type.block
+            block_type = if method_type.block.optional?
+                           AST::Types::Union.build(types: [method_type.block.type, Types.nil_instance])
+                         else
+                           method_type.block.type
+                         end
+            var_types[block_arg.children[0].name] = block_type
+          end
+        end
+
       when method
         typing.add_error Errors::MethodDefinitionWithOverloading.new(node: node, method: method)
         return_type = union_type(*method.types.map(&:return_type))
@@ -214,9 +293,9 @@ module Steep
       end
 
       type_env = type_env.with_annotations(
-        lvar_types: annots.var_types.transform_values {|annot| absolute_type(annot.type) },
+        lvar_types: annots.lvar_types,
         ivar_types: annots.ivar_types,
-        const_types: annots.const_types.transform_values {|type| absolute_type(type) }
+        const_types: annots.const_types,
       )
 
       self.class.new(
@@ -234,23 +313,33 @@ module Steep
     end
 
     def for_module(node)
-      annots = source.annotations(block: node)
       new_module_name = ModuleName.from_node(node.children.first) or raise "Unexpected module name: #{node.children.first}"
+      new_namespace = nested_namespace(new_module_name)
 
+      annots = source.annotations(block: node, builder: checker.builder, current_module: new_namespace)
       module_type = AST::Types::Name.new_instance(name: "::Module")
 
-      implement_module_name =
-        if annots.implement_module
-          annots.implement_module.name
+      implement_module_name = yield_self do
+        if (annotation = annots.implement_module_annotation)
+          absolute_name(annotation.name.name).yield_self do |absolute_name|
+            if checker.builder.signatures.module_name?(absolute_name)
+              AST::Annotation::Implements::Module.new(name: absolute_name,
+                                                      args: annotation.name.args)
+            else
+              Steep.logger.error "Unknown module name given to @implements: #{annotation.name.name}"
+              nil
+            end
+          end
         else
-          absolute_name(new_module_name).yield_self do |module_name|
-            if checker.builder.signatures.module_name?(module_name)
-              signature = checker.builder.signatures.find_module(module_name)
-              AST::Annotation::Implements::Module.new(name: module_name,
+          absolute_name(new_module_name).yield_self do |absolute_name|
+            if checker.builder.signatures.module_name?(absolute_name)
+              signature = checker.builder.signatures.find_module(absolute_name)
+              AST::Annotation::Implements::Module.new(name: absolute_name,
                                                       args: signature.params&.variables || [])
             end
           end
         end
+      end
 
       if implement_module_name
         module_name = implement_module_name.name
@@ -258,9 +347,7 @@ module Steep
 
         abstract = checker.builder.build(TypeName::Instance.new(name: module_name))
 
-        instance_type = absolute_type(
-          AST::Types::Name.new_instance(name: module_name, args: module_args)
-        )
+        instance_type = AST::Types::Name.new_instance(name: module_name, args: module_args)
 
         unless abstract.supers.empty?
           instance_type = AST::Types::Intersection.build(
@@ -275,19 +362,18 @@ module Steep
       end
 
       if annots.instance_type
-        instance_type = absolute_type(annots.instance_type)
+        instance_type = annots.instance_type
       end
 
       if annots.module_type
-        module_type = absolute_type(annots.module_type)
+        module_type = annots.module_type
       end
 
-      new_namespace = nested_namespace(new_module_name)
       module_const_env = TypeInference::ConstantEnv.new(builder: checker.builder, current_namespace: new_namespace)
 
       module_context_ = ModuleContext.new(
         instance_type: instance_type,
-        module_type: module_type,
+        module_type: annots.self_type || module_type,
         implement_name: implement_module_name,
         current_namespace: new_namespace,
         const_env: module_const_env
@@ -313,21 +399,32 @@ module Steep
     end
 
     def for_class(node)
-      annots = source.annotations(block: node)
       new_class_name = ModuleName.from_node(node.children.first) or raise "Unexpected class name: #{node.children.first}"
+      new_namespace = nested_namespace(new_class_name)
 
-      implement_module_name =
-        if annots.implement_module
-          annots.implement_module.name
+      annots = source.annotations(block: node, builder: checker.builder, current_module: new_namespace)
+
+      implement_module_name = yield_self do
+        if (annotation = annots.implement_module_annotation)
+          absolute_name(annotation.name.name).yield_self do |absolute_name|
+            if checker.builder.signatures.class_name?(absolute_name)
+              AST::Annotation::Implements::Module.new(name: absolute_name,
+                                                      args: annotation.name.args)
+            else
+              Steep.logger.error "Unknown class name given to @implements: #{annotation.name.name}"
+              nil
+            end
+          end
         else
-          absolute_name(new_class_name).yield_self do |name|
-            if checker.builder.signatures.class_name?(name)
-              signature = checker.builder.signatures.find_class(name)
-              AST::Annotation::Implements::Module.new(name: name,
+          absolute_name(new_class_name).yield_self do |absolute_name|
+            if checker.builder.signatures.class_name?(absolute_name)
+              signature = checker.builder.signatures.find_class(absolute_name)
+              AST::Annotation::Implements::Module.new(name: absolute_name,
                                                       args: signature.params&.variables || [])
             end
           end
         end
+      end
 
       if implement_module_name
         class_name = implement_module_name.name
@@ -336,15 +433,14 @@ module Steep
         _ = checker.builder.build(TypeName::Instance.new(name: class_name))
 
         instance_type = AST::Types::Name.new_instance(name: class_name, args: class_args)
-        module_type = AST::Types::Name.new_class(name: class_name, args: [], constructor: nil)
+        module_type = AST::Types::Name.new_class(name: class_name, args: [], constructor: true)
       end
 
-      new_namespace = nested_namespace(new_class_name)
       class_const_env = TypeInference::ConstantEnv.new(builder: checker.builder, current_namespace: new_namespace)
 
       module_context = ModuleContext.new(
         instance_type: annots.instance_type || instance_type,
-        module_type: annots.module_type || module_type,
+        module_type: annots.self_type || annots.module_type || module_type,
         implement_name: implement_module_name,
         current_namespace: new_namespace,
         const_env: class_const_env
@@ -372,13 +468,13 @@ module Steep
     end
 
     def for_branch(node, truthy_vars: Set.new, type_case_override: nil)
-      annots = source.annotations(block: node)
+      annots = source.annotations(block: node, builder: checker.builder, current_module: current_namespace)
 
       type_env = self.type_env
 
       lvar_types = self.type_env.lvar_types.each.with_object({}) do |(var, type), env|
         if truthy_vars.member?(var)
-          env[var] = TypeConstruction.unwrap(type)
+          env[var] = unwrap(type)
         else
           env[var] = type
         end
@@ -399,9 +495,9 @@ module Steep
       end
 
       type_env = type_env.with_annotations(
-        lvar_types: annots.var_types.transform_values {|a| absolute_type(a.type) },
-        ivar_types: annots.ivar_types.transform_values {|ty| absolute_type(ty) },
-        const_types: annots.const_types.transform_values {|ty| absolute_type(ty) },
+        lvar_types: annots.lvar_types,
+        ivar_types: annots.ivar_types,
+        const_types: annots.const_types,
         gvar_types: {}
       ) do |var, relation, result|
         typing.add_error(
@@ -438,9 +534,15 @@ module Steep
         case node.type
         when :begin, :kwbegin
           yield_self do
-            type = each_child_node(node).map do |child|
-              synthesize(child, hint: hint)
-            end.last
+            *mid_nodes, last_node = each_child_node(node).to_a
+            mid_nodes.each do |child|
+              synthesize(child)
+            end
+            if last_node
+              type = synthesize(last_node, hint: hint)
+            else
+              type = Types.nil_instance
+            end
 
             typing.add_typing(node, type)
           end
@@ -451,9 +553,10 @@ module Steep
             rhs = node.children[1]
 
             if var.name == :_
+              synthesize(rhs, hint: Types.any)
               typing.add_typing(node, Types.any)
             else
-              type_assignment(var, rhs, node)
+              type_assignment(var, rhs, node, hint: hint)
             end
           end
 
@@ -485,7 +588,7 @@ module Steep
         when :send
           yield_self do
             if self_class?(node)
-              module_type = module_context.module_type
+              module_type = expand_alias(module_context.module_type)
               type = if module_type.is_a?(AST::Types::Name) && module_type.name.is_a?(TypeName::Class)
                        AST::Types::Name.new(name: module_type.name.updated(constructor: method_context.constructor),
                                             args: module_type.args)
@@ -501,7 +604,7 @@ module Steep
         when :csend
           yield_self do
             type = if self_class?(node)
-                     module_type = module_context.module_type
+                     module_type = expand_alias(module_context.module_type)
                      type = if module_type.is_a?(AST::Types::Name)
                               AST::Types::Name.new(name: module_type.name.updated(constructor: method_context.constructor),
                                                    args: module_type.args)
@@ -622,7 +725,11 @@ module Steep
         when :block
           yield_self do
             send_node, params, body = node.children
-            type_send(node, send_node: send_node, block_params: params, block_body: body)
+            if send_node.type == :lambda
+              type_lambda(node, block_params: params, block_body: body, type_hint: hint)
+            else
+              type_send(node, send_node: send_node, block_params: params, block_body: body)
+            end
           end
 
         when :def
@@ -636,11 +743,11 @@ module Steep
           end
 
           if node.children[2]
-            return_type = new.method_context&.return_type
+            return_type = expand_alias(new.method_context&.return_type)
             if return_type && !return_type.is_a?(AST::Types::Void)
               new.check(node.children[2], return_type) do |_, actual_type, result|
                 typing.add_error(Errors::MethodBodyTypeMismatch.new(node: node,
-                                                                    expected: return_type,
+                                                                    expected: new.method_context&.return_type,
                                                                     actual: actual_type,
                                                                     result: result))
               end
@@ -648,7 +755,7 @@ module Steep
               new.synthesize(node.children[2])
             end
           else
-            return_type = new.method_context&.return_type
+            return_type = expand_alias(new.method_context&.return_type)
             if return_type && !return_type.is_a?(AST::Types::Void)
               result = checker.check(
                 Subtyping::Relation.new(sub_type: Types.nil_instance, super_type: return_type),
@@ -656,7 +763,7 @@ module Steep
               )
               if result.failure?
                 typing.add_error(Errors::MethodBodyTypeMismatch.new(node: node,
-                                                                    expected: return_type,
+                                                                    expected: new.method_context&.return_type,
                                                                     actual: Types.nil_instance,
                                                                     result: result))
               end
@@ -681,7 +788,7 @@ module Steep
             end
 
             if node.children[3]
-              return_type = new.method_context&.return_type
+              return_type = expand_alias(new.method_context&.return_type)
               if return_type && !return_type.is_a?(AST::Types::Void)
                 new.check(node.children[3], return_type) do |return_type, actual_type, result|
                   typing.add_error(Errors::MethodBodyTypeMismatch.new(node: node,
@@ -716,18 +823,20 @@ module Steep
                              Types.array_instance(union_type(*return_types))
                            end
 
-              if method_context&.return_type && !method_context.return_type.is_a?(AST::Types::Void)
-                result = checker.check(
-                  Subtyping::Relation.new(sub_type: value_type,
-                                          super_type: method_context.return_type),
-                  constraints: Subtyping::Constraints.empty
-                )
-
-                if result.failure?
-                  typing.add_error(Errors::ReturnTypeMismatch.new(node: node,
-                                                                  expected: method_context.return_type,
-                                                                  actual: value_type,
-                                                                  result: result))
+              if (ret_type = expand_alias(method_context&.return_type))
+                unless ret_type.is_a?(AST::Types::Void)
+                  result = checker.check(
+                    Subtyping::Relation.new(sub_type: value_type,
+                                            super_type: ret_type),
+                    constraints: Subtyping::Constraints.empty
+                  )
+
+                  if result.failure?
+                    typing.add_error(Errors::ReturnTypeMismatch.new(node: node,
+                                                                    expected: method_context&.return_type,
+                                                                    actual: value_type,
+                                                                    result: result))
+                  end
                 end
               end
             end
@@ -804,7 +913,7 @@ module Steep
           yield_self do
             var = node.children[0]
             rhs = node.children[1]
-            type_assignment(var, rhs, node)
+            type_assignment(var, rhs, node, hint: hint)
           end
 
         when :restarg
@@ -837,7 +946,7 @@ module Steep
 
         when :int
           yield_self do
-            literal_type = test_literal_type(node.children[0], hint)
+            literal_type = expand_alias(hint) {|hint_| test_literal_type(node.children[0], hint_) }
 
             if literal_type
               typing.add_typing(node, literal_type)
@@ -848,7 +957,7 @@ module Steep
 
         when :sym
           yield_self do
-            literal_type = test_literal_type(node.children[0], hint)
+            literal_type = expand_alias(hint) {|hint_| test_literal_type(node.children[0], hint_) }
 
             if literal_type
               typing.add_typing(node, literal_type)
@@ -859,7 +968,7 @@ module Steep
 
         when :str
           yield_self do
-            literal_type = test_literal_type(node.children[0], hint)
+            literal_type = expand_alias(hint) {|hint_| test_literal_type(node.children[0], hint_) }
 
             if literal_type
               typing.add_typing(node, literal_type)
@@ -873,6 +982,11 @@ module Steep
 
         when :hash
           yield_self do
+            if Types.hash_instance?(hint)
+              key_hint = hint.args[0]
+              value_hint = hint.args[1]
+            end
+
             key_types = []
             value_types = []
 
@@ -880,18 +994,22 @@ module Steep
               case child.type
               when :pair
                 key, value = child.children
-                key_types << synthesize(key)
-                value_types << synthesize(value)
+                key_types << synthesize(key).yield_self do |type|
+                  select_super_type(type, key_hint)
+                end
+                value_types << synthesize(value).yield_self do |type|
+                  select_super_type(type, value_hint)
+                end
               when :kwsplat
-                splat_type = synthesize(child.children[0])
-
-                if splat_type.is_a?(AST::Types::Name) && splat_type.name == TypeName::Instance.new(name: ModuleName.parse("::Hash"))
-                  key_types << splat_type.args[0]
-                  value_types << splat_type.args[1]
-                else
-                  typing.add_error Errors::UnexpectedSplat.new(node: child, type: splat_type)
-                  key_types << Types.any
-                  value_types << Types.any
+                expand_alias(synthesize(child.children[0])) do |splat_type, original_type|
+                  if splat_type.is_a?(AST::Types::Name) && splat_type.name == TypeName::Instance.new(name: ModuleName.parse("::Hash"))
+                    key_types << splat_type.args[0]
+                    value_types << splat_type.args[1]
+                  else
+                    typing.add_error Errors::UnexpectedSplat.new(node: child, type: original_type)
+                    key_types << Types.any
+                    value_types << Types.any
+                  end
                 end
               else
                 raise "Unexpected non pair: #{child.inspect}" unless child.type == :pair
@@ -901,10 +1019,14 @@ module Steep
             key_type = key_types.empty? ? Types.any : AST::Types::Union.build(types: key_types)
             value_type = value_types.empty? ? Types.any : AST::Types::Union.build(types: value_types)
 
+            if key_types.empty? && value_types.empty? && !hint
+              typing.add_error Errors::FallbackAny.new(node: node)
+            end
+
             typing.add_typing(node, Types.hash_instance(key_type, value_type))
           end
 
-        when :dstr
+        when :dstr, :xstr
           each_child_node(node) do |child|
             synthesize(child)
           end
@@ -991,7 +1113,7 @@ module Steep
           if method_context&.method_type
             if method_context.block_type
               block_type = method_context.block_type
-              block_type.params.flat_unnamed_params.map(&:last).zip(node.children).each do |(type, node)|
+              block_type.type.params.flat_unnamed_params.map(&:last).zip(node.children).each do |(type, node)|
                 if node && type
                   check(node, type) do |_, rhs_type, result|
                     typing.add_error(Errors::IncompatibleAssignment.new(node: node,
@@ -1002,7 +1124,7 @@ module Steep
                 end
               end
 
-              typing.add_typing(node, block_type.return_type)
+              typing.add_typing(node, block_type.type.return_type)
             else
               typing.add_error(Errors::UnexpectedYield.new(node: node))
               fallback_to_any node
@@ -1015,8 +1137,13 @@ module Steep
           yield_self do
             if method_context&.method
               if method_context.super_method
-                types = method_context.super_method.types.map(&:return_type)
-                typing.add_typing(node, union_type(*types))
+                if method_context.method.incompatible?
+                  typing.add_error Errors::IncompatibleZuper.new(node: node, method: method_context.name)
+                  typing.add_typing node, Types.any
+                else
+                  types = method_context.super_method.types.map(&:return_type)
+                  typing.add_typing(node, union_type(*types))
+                end
               else
                 typing.add_error(Errors::UnexpectedSuper.new(node: node, method: method_context.name))
                 fallback_to_any node
@@ -1029,35 +1156,48 @@ module Steep
         when :array
           yield_self do
             if node.children.empty?
+              unless hint
+                typing.add_error Errors::FallbackAny.new(node: node)
+              end
               typing.add_typing(node, Types.array_instance(Types.any))
             else
-              is_tuple = hint.is_a?(AST::Types::Tuple)
-              is_tuple &&= node.children.all? {|child| child.type != :splat }
-              is_tuple &&= node.children.map.with_index do |child, index|
-                synthesize(child, hint: absolute_type(hint.types[index]))
-              end.yield_self do |types|
-                tuple = AST::Types::Tuple.new(types: types)
-                relation = Subtyping::Relation.new(sub_type: tuple, super_type: absolute_type(hint))
-                result = checker.check(relation, constraints: Subtyping::Constraints.empty)
-                result.success?
+              is_tuple = nil
+
+              expand_alias(hint) do |hint|
+                is_tuple = hint.is_a?(AST::Types::Tuple)
+                is_tuple &&= node.children.all? {|child| child.type != :splat }
+                is_tuple &&= node.children.map.with_index do |child, index|
+                  synthesize(child, hint: hint.types[index])
+                end.yield_self do |types|
+                  tuple = AST::Types::Tuple.new(types: types)
+                  relation = Subtyping::Relation.new(sub_type: tuple, super_type: hint)
+                  result = checker.check(relation, constraints: Subtyping::Constraints.empty)
+                  result.success?
+                end
               end
 
               if is_tuple
                 array_type = hint
               else
+                element_hint = expand_alias(hint) do |hint|
+                  Types.array_instance?(hint) && hint.args[0]
+                end
+
                 element_types = node.children.flat_map do |e|
                   if e.type == :splat
                     Steep.logger.info "Typing of splat in array is incompatible with Ruby; it does not use #to_a method"
                     synthesize(e.children.first).yield_self do |type|
-                      case type
-                      when AST::Types::Union
-                        type.types
-                      else
-                        [type]
+                      expand_alias(type) do |ty|
+                        case ty
+                        when AST::Types::Union
+                          ty.types
+                        else
+                          [ty]
+                        end
                       end
                     end.map do |type|
                       case
-                      when type.is_a?(AST::Types::Name) && type.name.is_a?(TypeName::Instance) && type.name.name == ModuleName.new(name: "Array", absolute: true)
+                      when Types.array_instance?(type)
                         type.args.first
                       when type.is_a?(AST::Types::Name) && type.name.is_a?(TypeName::Instance) && type.name.name == ModuleName.new(name: "Range", absolute: true)
                         type.args.first
@@ -1066,7 +1206,7 @@ module Steep
                       end
                     end
                   else
-                    [synthesize(e)]
+                    [select_super_type(synthesize(e), element_hint)]
                   end
                 end
                 array_type = Types.array_instance(AST::Types::Union.build(types: element_types))
@@ -1098,9 +1238,11 @@ module Steep
           end
 
         when :or
-          types = each_child_node(node).map {|child| synthesize(child) }
-          type = union_type(*types)
-          typing.add_typing(node, type)
+          yield_self do
+            t1, t2 = each_child_node(node).map {|child| synthesize(child) }
+            type = union_type(unwrap(t1), t2)
+            typing.add_typing(node, type)
+          end
 
         when :if
           cond, true_clause, false_clause = node.children
@@ -1129,7 +1271,7 @@ module Steep
             cond, *whens = node.children
 
             if cond
-              cond_type = synthesize(cond)
+              cond_type = expand_alias(synthesize(cond))
               if cond_type.is_a?(AST::Types::Union)
                 var_names = TypeConstruction.value_variables(cond)
                 var_types = cond_type.types.dup
@@ -1139,7 +1281,7 @@ module Steep
             pairs = whens.each.with_object([]) do |clause, pairs|
               if clause&.type == :when
                 test_types = clause.children.take(clause.children.size - 1).map do |child|
-                  synthesize(child, hint: hint)
+                  expand_alias(synthesize(child, hint: hint))
                 end
 
                 if (body = clause.children.last)
@@ -1240,6 +1382,7 @@ module Steep
               case
               when exn_classes && var_name
                 instance_types = exn_types.map do |type|
+                  type = expand_alias(type)
                   case
                   when type.is_a?(AST::Types::Name) && type.name.is_a?(TypeName::Class)
                     type.instance_type
@@ -1374,7 +1517,47 @@ module Steep
             end
           end
 
-        when :splat, :block_pass, :blockarg, :sclass
+        when :block_pass
+          yield_self do
+            value = node.children[0]
+
+            if hint.is_a?(AST::Types::Proc) && value.type == :sym
+              if hint.one_arg?
+                # Assumes Symbol#to_proc implementation
+                param_type = hint.params.required[0]
+                interface = checker.resolve(param_type, with_initialize: false)
+                method = interface.methods[value.children[0]]
+                if method
+                  return_types = method.types.flat_map do |method_type|
+                    if method_type.params.each_type.count == 0
+                      [method_type.return_type]
+                    end
+                  end
+
+                  if return_types.any?
+                    type = AST::Types::Proc.new(params: Interface::Params.empty.update(required: [param_type]),
+                                                return_type: AST::Types::Union.build(types: return_types))
+                  end
+                end
+              else
+                Steep.logger.error "Passing multiple args through Symbol#to_proc is not supported yet"
+              end
+            end
+
+            type ||= synthesize(node.children[0], hint: hint)
+            typing.add_typing node, type
+          end
+
+        when :blockarg
+          yield_self do
+            each_child_node node do |child|
+              synthesize(child)
+            end
+
+            typing.add_typing node, Types.any
+          end
+
+        when :splat, :sclass
           yield_self do
             Steep.logger.error "Unsupported node #{node.type}"
 
@@ -1404,11 +1587,12 @@ module Steep
       end
     end
 
-    def type_assignment(var, rhs, node)
+    def type_assignment(var, rhs, node, hint: nil)
       if rhs
-        rhs_type = synthesize(rhs, hint: type_env.lvar_types[var.name])
-        node_type = assign_type_to_variable(var, rhs_type, node)
-        typing.add_typing(node, node_type)
+        expand_alias(synthesize(rhs, hint: type_env.lvar_types[var.name] || hint)) do |rhs_type|
+          node_type = assign_type_to_variable(var, rhs_type, node)
+          typing.add_typing(node, node_type)
+        end
       else
         raise
         lhs_type = variable_type(var)
@@ -1433,7 +1617,7 @@ module Steep
     end
 
     def type_ivasgn(name, rhs, node)
-      rhs_type = synthesize(rhs)
+      rhs_type = synthesize(rhs, hint: type_env.get(ivar: name) { fallback_to_any(node) })
       ivar_type = type_env.assign(ivar: name, type: rhs_type) do |error|
         case error
         when Subtyping::Result::Failure
@@ -1451,7 +1635,8 @@ module Steep
 
     def type_masgn(node)
       lhs, rhs = node.children
-      rhs_type = synthesize(rhs)
+      rhs_original = synthesize(rhs)
+      rhs_type = expand_alias(rhs_original)
 
       case
       when rhs.type == :array && lhs.children.all? {|a| a.type == :lvasgn || a.type == :ivasgn } && lhs.children.size == rhs.children.size
@@ -1564,12 +1749,32 @@ module Steep
       end
     end
 
+    def type_lambda(node, block_params:, block_body:, type_hint:)
+      block_annotations = source.annotations(block: node, builder: checker.builder, current_module: current_namespace)
+      params = TypeInference::BlockParams.from_node(block_params, annotations: block_annotations)
+
+      case type_hint
+      when AST::Types::Proc
+        params_hint = type_hint.params
+        return_hint = type_hint.return_type
+      end
+
+      block_type = type_block(block_param_hint: params_hint,
+                              block_type_hint: return_hint,
+                              node_type_hint: nil,
+                              block_params: params,
+                              block_body: block_body,
+                              block_annotations: block_annotations)
+
+      typing.add_typing node, block_type
+    end
+
     def type_send(node, send_node:, block_params:, block_body:, unwrap: false)
       receiver, method_name, *arguments = send_node.children
       receiver_type = receiver ? synthesize(receiver) : self_type
 
       if unwrap
-        receiver_type = TypeConstruction.unwrap(receiver_type)
+        receiver_type = unwrap(receiver_type)
       end
 
       case receiver_type
@@ -1624,7 +1829,7 @@ module Steep
 
         if block_body && block_params
           unless typing.has_type?(block_body)
-            block_annotations = source.annotations(block: node)
+            block_annotations = source.annotations(block: node, builder: checker.builder, current_module: current_namespace)
             params = TypeInference::BlockParams.from_node(block_params, annotations: block_annotations)
             pairs = params.each.map {|param| [param, Types.any] }
 
@@ -1643,45 +1848,45 @@ module Steep
       block_type_env = type_env.dup.yield_self do |env|
         param_pairs.each do |param, type|
           if param.type
-            env.set(lvar: param.var.name, type: absolute_type(param.type))
+            env.set(lvar: param.var.name, type: param.type)
           else
-            env.set(lvar: param.var.name, type: absolute_type(type))
+            env.set(lvar: param.var.name, type: type)
           end
         end
 
         env.with_annotations(
-          lvar_types: block_annotations.var_types.transform_values {|annot| absolute_type(annot.type) },
-          ivar_types: block_annotations.ivar_types.transform_values {|type| absolute_type(type) },
-          const_types: block_annotations.const_types.transform_values {|type| absolute_type(type) }
+          lvar_types: block_annotations.lvar_types,
+          ivar_types: block_annotations.ivar_types,
+          const_types: block_annotations.const_types,
         )
       end
 
       return_type = if block_annotations.break_type
-                      union_type(method_return_type, absolute_type(block_annotations.break_type))
+                      union_type(method_return_type, block_annotations.break_type)
                     else
                       method_return_type
                     end
       Steep.logger.debug("return_type = #{return_type}")
 
-      block_context = BlockContext.new(body_type: absolute_type(block_annotations.block_type))
+      block_context = BlockContext.new(body_type: block_annotations.block_type)
       Steep.logger.debug("block_context { body_type: #{block_context.body_type} }")
 
       break_context = BreakContext.new(
-        break_type: absolute_type(block_annotations.break_type) || method_return_type,
-        next_type: absolute_type(block_annotations.block_type)
+        break_type: block_annotations.break_type || method_return_type,
+        next_type: block_annotations.block_type
       )
-      Steep.logger.debug("break_context { type: #{absolute_type(break_context.break_type)} }")
+      Steep.logger.debug("break_context { type: #{break_context.break_type} }")
 
       [self.class.new(
         checker: checker,
         source: source,
-        annotations: annotations + block_annotations,
+        annotations: annotations.merge_block_annotations(block_annotations),
         type_env: block_type_env,
         block_context: block_context,
         typing: typing,
         method_context: method_context,
         module_context: module_context,
-        self_type: absolute_type(block_annotations.self_type) || self_type,
+        self_type: block_annotations.self_type || self_type,
         break_context: break_context
       ), return_type]
     end
@@ -1689,203 +1894,291 @@ module Steep
     def type_method_call(node, receiver_type:, method:, args:, block_params:, block_body:)
       results = method.types.map do |method_type|
         Steep.logger.tagged method_type.location&.source do
-          arg_pairs = args.zip(method_type.params)
-
-          if arg_pairs
-            try_method_type(node,
-                            receiver_type: receiver_type,
-                            method_type: method_type,
-                            arg_pairs: arg_pairs,
-                            block_params: block_params,
-                            block_body: block_body)
-          else
-            Steep.logger.debug(node.inspect)
-            Errors::IncompatibleArguments.new(node: node, receiver_type: receiver_type, method_type: method_type)
-          end
+          child_typing = typing.new_child
+          arg_pairs = args.zip(method_type.params, method_type.block&.type)
+
+          type_or_error = if arg_pairs
+                            self.with_new_typing(child_typing).try_method_type(
+                              node,
+                              receiver_type: receiver_type,
+                              method_type: method_type,
+                              args: args,
+                              arg_pairs: arg_pairs,
+                              block_params: block_params,
+                              block_body: block_body,
+                              child_typing: child_typing
+                            )
+                          else
+                            Steep.logger.debug(node.inspect)
+                            Errors::IncompatibleArguments.new(node: node, receiver_type: receiver_type, method_type: method_type)
+                          end
+
+          [child_typing, type_or_error, method_type]
         end
       end
 
-      if results.all? {|result| result.is_a?(Errors::Base) }
-        results.first
+      if (success = results.find {|_, type_or_error| !type_or_error.is_a?(Errors::Base) })
+        child_typing, type, _ = success
+        child_typing.save!
+        type
       else
-        results.find do |result|
-          !result.is_a?(Errors::Base)
+        child_typing, error, method_type = results.first
+        typing.add_error error
+        child_typing.save!
+
+        case method_type.return_type
+        when AST::Types::Var
+          Types.any
+        else
+          method_type.return_type
         end
       end
     end
 
-    def try_method_type(node, receiver_type:, method_type:, arg_pairs:, block_params:, block_body:)
+    def try_method_type(node, receiver_type:, method_type:, args:, arg_pairs:, block_params:, block_body:, child_typing:)
       fresh_types = method_type.type_params.map {|x| AST::Types::Var.fresh(x) }
       fresh_vars = Set.new(fresh_types.map(&:name))
       instantiation = Interface::Substitution.build(method_type.type_params, fresh_types)
 
-      typing.new_child do |child_typing|
-        construction = self.class.new(
-          checker: checker,
-          source: source,
-          annotations: annotations,
-          type_env: type_env,
-          typing: child_typing,
-          self_type: self_type,
-          method_context: method_context,
-          block_context: block_context,
-          module_context: module_context,
-          break_context: break_context
-        )
+      construction = self.class.new(
+        checker: checker,
+        source: source,
+        annotations: annotations,
+        type_env: type_env,
+        typing: child_typing,
+        self_type: self_type,
+        method_context: method_context,
+        block_context: block_context,
+        module_context: module_context,
+        break_context: break_context
+      )
 
-        method_type.instantiate(instantiation).yield_self do |method_type|
-          constraints = Subtyping::Constraints.new(unknowns: fresh_types.map(&:name))
-          variance = Subtyping::VariableVariance.from_method_type(method_type)
-          occurence = Subtyping::VariableOccurence.from_method_type(method_type)
+      method_type.instantiate(instantiation).yield_self do |method_type|
+        constraints = Subtyping::Constraints.new(unknowns: fresh_types.map(&:name))
+        variance = Subtyping::VariableVariance.from_method_type(method_type)
+        occurence = Subtyping::VariableOccurence.from_method_type(method_type)
 
-          arg_pairs.each do |(arg_node, param_type)|
-            arg_type = if arg_node.type == :splat
-                         type = construction.synthesize(arg_node.children[0])
-                         child_typing.add_typing(arg_node, Types.array_instance(type))
-                       else
-                         construction.synthesize(arg_node, hint: param_type)
-                       end
+        arg_pairs.each do |(arg_node, param_type)|
+          arg_type = if arg_node.type == :splat
+                       type = construction.synthesize(arg_node.children[0])
+                       child_typing.add_typing(arg_node, type)
+                     else
+                       construction.synthesize(arg_node, hint: param_type)
+                     end
 
-            relation = Subtyping::Relation.new(
-              sub_type: arg_type,
-              super_type: param_type.subst(instantiation)
-            )
+          relation = Subtyping::Relation.new(
+            sub_type: arg_type,
+            super_type: param_type.subst(instantiation)
+          )
 
-            checker.check(relation, constraints: constraints).else do |result|
-              return Errors::ArgumentTypeMismatch.new(
-                node: arg_node,
-                receiver_type: receiver_type,
-                expected: relation.super_type,
-                actual: relation.sub_type
-              )
-            end
+          checker.check(relation, constraints: constraints).else do |result|
+            return Errors::ArgumentTypeMismatch.new(
+              node: arg_node,
+              receiver_type: receiver_type,
+              expected: relation.super_type,
+              actual: relation.sub_type
+            )
           end
+        end
 
-          if method_type.block && block_params
-            block_annotations = source.annotations(block: node)
+        if block_params && method_type.block
+          block_annotations = source.annotations(block: node, builder: checker.builder, current_module: current_namespace)
+          block_params_ = TypeInference::BlockParams.from_node(block_params, annotations: block_annotations)
+          block_param_hint = block_params_.params_type(hint: method_type.block.type.params)
 
-            params = TypeInference::BlockParams.from_node(block_params, annotations: block_annotations)
-            block_param_pairs = params.zip(method_type.block.params)
+          relation = Subtyping::Relation.new(
+            sub_type: AST::Types::Proc.new(params: block_param_hint, return_type: AST::Types::Any.new),
+            super_type: method_type.block.type
+          )
 
-            unless block_param_pairs
-              return Errors::IncompatibleBlockParameters.new(
-                node: node,
-                method_type: method_type
-              )
-            end
+          checker.check(relation, constraints: constraints).else do |result|
+            return Errors::IncompatibleBlockParameters.new(
+              node: node,
+              method_type: method_type
+            )
+          end
+        end
 
-            block_param_pairs.each do |param, type|
-              if param.type
-                relation = Subtyping::Relation.new(
-                  sub_type: type,
-                  super_type: absolute_type(param.type)
-                )
-
-                checker.check(relation, constraints: constraints).else do |result|
-                  return Errors::BlockParameterTypeMismatch.new(
-                    node: param.node,
-                    expected: type,
-                    actual: param.type
-                  )
+        case
+        when method_type.block && block_params
+          Steep.logger.debug "block is okay: method_type=#{method_type}"
+          Steep.logger.debug "Constraints = #{constraints}"
+
+          begin
+            method_type.subst(constraints.solution(checker, variance: variance, variables: occurence.params)).yield_self do |method_type|
+              block_type = construction.type_block(block_param_hint: method_type.block.type.params,
+                                                   block_type_hint: method_type.block.type.return_type,
+                                                   node_type_hint: method_type.return_type,
+                                                   block_params: block_params_,
+                                                   block_body: block_body,
+                                                   block_annotations: block_annotations)
+
+              relation = Subtyping::Relation.new(sub_type: block_type.return_type,
+                                                 super_type: method_type.block.type.return_type)
+
+              result = checker.check(relation, constraints: constraints)
+
+              case result
+              when Subtyping::Result::Success
+                method_type.return_type.subst(constraints.solution(checker, variance: variance, variables: fresh_vars)).yield_self do |ret_type|
+                  if block_annotations.break_type
+                    AST::Types::Union.new(types: [block_annotations.break_type, ret_type])
+                  else
+                    ret_type
+                  end
                 end
+
+              when Subtyping::Result::Failure
+                Errors::BlockTypeMismatch.new(node: node,
+                                              expected: method_type.block.type,
+                                              actual: block_type,
+                                              result: result)
               end
             end
 
-            if block_annotations.block_type
-              relation = Subtyping::Relation.new(
-                sub_type: absolute_type(block_annotations.block_type),
-                super_type: method_type.block.return_type
-              )
-
-              checker.check(relation, constraints: constraints).else do |result|
-                typing.add_error Errors::BlockTypeMismatch.new(node: node,
-                                                               expected: method_type.block.return_type,
-                                                               actual: absolute_type(block_annotations.block_type),
-                                                               result: result)
+          rescue Subtyping::Constraints::UnsatisfiableConstraint => exn
+            Errors::UnsatisfiableConstraint.new(node: node,
+                                                method_type: method_type,
+                                                var: exn.var,
+                                                sub_type: exn.sub_type,
+                                                super_type: exn.super_type,
+                                                result: exn.result)
+          end
+
+        when method_type.block && args.block_pass_arg
+          begin
+            method_type.subst(constraints.solution(checker, variance: variance, variables: occurence.params)).yield_self do |method_type|
+              block_type = synthesize(args.block_pass_arg, hint: method_type.block.type)
+              relation = Subtyping::Relation.new(sub_type: block_type,
+                                                 super_type: method_type.block.yield_self do |expected_block|
+                                                   if expected_block.optional?
+                                                     Types.optional(expected_block.type)
+                                                   else
+                                                     expected_block.type
+                                                   end
+                                                 end)
+              result = checker.check(relation, constraints: constraints)
+
+              case result
+              when Subtyping::Result::Success
+                method_type.return_type.subst(constraints.solution(checker, variance: variance, variables: fresh_vars))
+
+              when Subtyping::Result::Failure
+                Errors::BlockTypeMismatch.new(node: node,
+                                              expected: method_type.block.type,
+                                              actual: block_type,
+                                              result: result)
               end
             end
+
+          rescue Subtyping::Constraints::UnsatisfiableConstraint => exn
+            Errors::UnsatisfiableConstraint.new(node: node,
+                                                method_type: method_type,
+                                                var: exn.var,
+                                                sub_type: exn.sub_type,
+                                                super_type: exn.super_type,
+                                                result: exn.result)
           end
 
-          case
-          when method_type.block && block_params && block_body
-            Steep.logger.debug "block is okay: method_type=#{method_type}"
-            Steep.logger.debug "Constraints = #{constraints}"
+        when (!method_type.block || method_type.block.optional?) && !block_params && !block_body && !args.block_pass_arg
+          # OK, without block
+          method_type.subst(constraints.solution(checker, variance: variance, variables: fresh_vars)).return_type
 
-            begin
-              method_type.subst(constraints.solution(checker, variance: variance, variables: occurence.params)).yield_self do |method_type|
-                block_annotations = source.annotations(block: node)
+        when !method_type.block && (block_params || args.block_pass_arg)
+          Errors::UnexpectedBlockGiven.new(
+            node: node,
+            method_type: method_type
+          )
 
-                block_param_pairs = TypeInference::BlockParams.from_node(block_params, annotations: block_annotations).yield_self do |params|
-                  params.zip(method_type.block.params)
-                end
+        when method_type.block && !method_type.block.optional? && !block_params && !block_body && !args.block_pass_arg
+          Errors::RequiredBlockMissing.new(
+            node: node,
+            method_type: method_type
+          )
 
-                for_block, return_type = for_block(block_annotations: block_annotations,
-                                                   param_pairs: block_param_pairs,
-                                                   method_return_type: method_type.return_type,
-                                                   typing: child_typing)
+        else
+          raise "Unexpected case condition"
+        end
+      end
+    end
 
-                each_child_node(block_params) do |p|
-                  for_block.synthesize(p)
-                end
+    def type_block(block_param_hint:, block_type_hint:, node_type_hint:, block_params:, block_body:, block_annotations:)
+      block_param_pairs = block_param_hint && block_params.zip(block_param_hint)
 
-                block_body_type = for_block.synthesize(block_body)
+      param_types_hash = {}
+      if block_param_pairs
+        block_param_pairs.each do |param, type|
+          var_name = param.var.name
+          param_types_hash[var_name] = type
+        end
+      else
+        block_params.each do |param|
+          var_name = param.var.name
+          param_types_hash[var_name] = param.type || Types.any
+        end
+      end
 
-                unless method_type.block.return_type.is_a?(AST::Types::Void)
-                  result = checker.check(Subtyping::Relation.new(
-                    sub_type: block_annotations.block_type || block_body_type,
-                    super_type: method_type.block.return_type
-                  ), constraints: constraints)
+      block_type_env = type_env.dup.tap do |env|
+        param_types_hash.each do |name, type|
+          env.set(lvar: name, type: type)
+        end
+
+        block_annotations.lvar_types.each do |name, type|
+          env.set(lvar: name, type: type)
+        end
+      end
 
-                  if result.success?
-                    return_type.subst(constraints.solution(checker, variance: variance, variables: fresh_vars)).tap do
-                      child_typing.save!
+      break_type = if block_annotations.break_type
+                      union_type(node_type_hint, block_annotations.break_type)
+                    else
+                      node_type_hint
                     end
-                  else
-                    typing.add_error Errors::BlockTypeMismatch.new(node: node,
-                                                                   expected: method_type.block.return_type,
-                                                                   actual: block_annotations.block_type || block_body_type,
-                                                                   result: result)
-                    return_type
-                  end
-                else
-                  child_typing.save!
-                  return_type
-                end
-              end
+      Steep.logger.debug("return_type = #{break_type}")
 
-            rescue Subtyping::Constraints::UnsatisfiableConstraint => exn
-              typing.add_error Errors::UnsatisfiableConstraint.new(node: node,
-                                                                   method_type: method_type,
-                                                                   var: exn.var,
-                                                                   sub_type: exn.sub_type,
-                                                                   super_type: exn.super_type,
-                                                                   result: exn.result
-              )
-              fallback_any_rec node
-            end
+      block_context = BlockContext.new(body_type: block_annotations.block_type)
+      Steep.logger.debug("block_context { body_type: #{block_context.body_type} }")
 
-          when !method_type.block && !block_params && !block_body
-            # OK, without block
-            method_type.subst(constraints.solution(checker, variance: variance, variables: fresh_vars)).return_type.tap do
-              child_typing.save!
-            end
+      break_context = BreakContext.new(
+        break_type: break_type,
+        next_type: block_context.body_type
+      )
+      Steep.logger.debug("break_context { type: #{break_context.break_type} }")
 
-          when !method_type.block && block_params && block_body
-            Errors::UnexpectedBlockGiven.new(
-              node: node,
-              method_type: method_type
-            )
+      for_block_body = self.class.new(
+        checker: checker,
+        source: source,
+        annotations: annotations.merge_block_annotations(block_annotations),
+        type_env: block_type_env,
+        block_context: block_context,
+        typing: typing,
+        method_context: method_context,
+        module_context: module_context,
+        self_type: block_annotations.self_type || self_type,
+        break_context: break_context
+      )
 
-          when method_type.block && !block_params && !block_body
-            Errors::RequiredBlockMissing.new(
-              node: node,
-              method_type: method_type
-            )
+      if block_body
+        return_type = if (body_type = block_context.body_type)
+                        for_block_body.check(block_body, body_type) do |expected, actual, result|
+                          typing.add_error Errors::BlockTypeMismatch.new(node: block_body,
+                                                                         expected: expected,
+                                                                         actual: actual,
+                                                                         result: result)
 
-          else
-            raise "Unexpected case condition"
-          end
-        end
+                        end
+                        body_type
+                      else
+                        for_block_body.synthesize(block_body, hint: block_type_hint)
+                      end
+      else
+        return_type = Types.any
+      end
+
+      AST::Types::Proc.new(
+        params: block_param_hint || block_params.params_type,
+        return_type: return_type
+      ).tap do |type|
+        Steep.logger.debug "block_type == #{type}"
       end
     end
 
@@ -2137,7 +2430,7 @@ module Steep
         case
         when module_context.defined_instance_methods.include?(method_name)
           # ok
-        when annotations.dynamics[method_name]&.instance_method?
+        when annotations.instance_dynamics.include?(method_name)
           # ok
         else
           typing.add_error Errors::MethodDefinitionMissing.new(node: node,
@@ -2150,7 +2443,7 @@ module Steep
         case
         when module_context.defined_module_methods.include?(method_name)
           # ok
-        when annotations.dynamics[method_name]&.module_method?
+        when annotations.module_dynamics.include?(method_name)
           # ok
         else
           typing.add_error Errors::MethodDefinitionMissing.new(node: node,
@@ -2160,13 +2453,15 @@ module Steep
         end
       end
 
-      annotations.dynamics.each do |method_name, annotation|
-        case
-        when annotation.module_method? && expected_module_method_names.member?(method_name)
-          # ok
-        when annotation.instance_method? && expected_instance_method_names.member?(method_name)
-          # ok
-        else
+      annotations.instance_dynamics.each do |method_name|
+        unless expected_instance_method_names.member?(method_name)
+          typing.add_error Errors::UnexpectedDynamicMethod.new(node: node,
+                                                               module_name: module_name.name,
+                                                               method_name: method_name)
+        end
+      end
+      annotations.module_dynamics.each do |method_name|
+        unless expected_module_method_names.member?(method_name)
           typing.add_error Errors::UnexpectedDynamicMethod.new(node: node,
                                                                module_name: module_name.name,
                                                                method_name: method_name)
@@ -2234,13 +2529,15 @@ module Steep
       typing.type_of(node: node)
     end
 
-    def self.unwrap(type)
-      case
-      when type.is_a?(AST::Types::Union)
-        types = type.types.reject {|type| type.is_a?(AST::Types::Nil) }
-        AST::Types::Union.build(types: types)
-      else
-        type
+    def unwrap(type)
+      expand_alias(type) do |expanded|
+        case
+        when expanded.is_a?(AST::Types::Union)
+          types = expanded.types.reject {|type| type.is_a?(AST::Types::Nil) }
+          AST::Types::Union.build(types: types)
+        else
+          type
+        end
       end
     end
 
@@ -2272,6 +2569,18 @@ module Steep
       end
     end
 
+    def expand_alias(type)
+      if type
+        expanded = checker.expand_alias(type)
+      end
+
+      if block_given?
+        yield expanded, type
+      else
+        expanded
+      end
+    end
+
     def test_literal_type(literal, hint)
       case hint
       when AST::Types::Literal
@@ -2284,5 +2593,26 @@ module Steep
         end
       end
     end
+
+    def select_super_type(sub_type, super_type)
+      if super_type
+        result = checker.check(
+          Subtyping::Relation.new(sub_type: sub_type, super_type: super_type),
+          constraints: Subtyping::Constraints.empty
+        )
+
+        if result.success?
+          super_type
+        else
+          if block_given?
+            yield result
+          else
+            sub_type
+          end
+        end
+      else
+        sub_type
+      end
+    end
   end
 end