Support type narrowing by Module#<

lib/steep/ast/types/logic.rb

@@ -68,6 +68,12 @@ def initialize(location: nil)
           end
         end
 
+        class ArgIsAncestor < Base
+          def initialize(location: nil)
+            @location = location
+          end
+        end
+
         class Env < Base
           attr_reader :truthy, :falsy, :type
 

lib/steep/interface/builder.rb

@@ -766,6 +766,15 @@ def replace_primitive_method(method_name, method_def, method_type)
                 )
               )
             end
+          when :<, :<=
+            case defined_in
+            when RBS::BuiltinNames::Module.name
+              return method_type.with(
+                type: method_type.type.with(
+                  return_type: AST::Types::Logic::ArgIsAncestor.new(location: method_type.type.return_type.location)
+                )
+              )
+            end
           end
         end
 

lib/steep/type_inference/logic_type_interpreter.rb

@@ -367,6 +367,32 @@ def evaluate_method_call(env:, type:, receiver:, arguments:)
               [truthy_result, falsy_result]
             end
           end
+
+        when AST::Types::Logic::ArgIsAncestor
+          if receiver && (arg = arguments[0])
+            receiver_type = typing.type_of(node: receiver)
+            arg_type = factory.deep_expand_alias(typing.type_of(node: arg))
+
+            if arg_type.is_a?(AST::Types::Name::Singleton)
+              truthy_type = arg_type
+              falsy_type = receiver_type
+              truthy_env, falsy_env = refine_node_type(
+                env: env,
+                node: receiver,
+                truthy_type: truthy_type,
+                falsy_type: falsy_type
+              )
+
+              truthy_result = Result.new(type: TRUE, env: truthy_env, unreachable: false)
+              truthy_result.unreachable! unless truthy_type
+
+              falsy_result = Result.new(type: FALSE, env: falsy_env, unreachable: false)
+              falsy_result.unreachable! unless falsy_type
+
+              [truthy_result, falsy_result]
+            end
+          end
+
         when AST::Types::Logic::Not
           if receiver
             truthy_result, falsy_result = evaluate_node(env: env, node: receiver)

sig/steep/ast/types.rbs

@@ -6,7 +6,7 @@ module Steep
            | Intersection  | Record | Tuple | Union
            | Name::Alias | Name::Instance | Name::Interface | Name::Singleton
            | Proc | Var
-           | Logic::Not | Logic::ReceiverIsNil | Logic::ReceiverIsNotNil | Logic::ReceiverIsArg | Logic::ArgIsReceiver | Logic::ArgEqualsReceiver | Logic::Env
+           | Logic::Not | Logic::ReceiverIsNil | Logic::ReceiverIsNotNil | Logic::ReceiverIsArg | Logic::ArgIsReceiver | Logic::ArgEqualsReceiver | Logic::ArgIsAncestor | Logic::Env
 
       # Variables and special types that is subject for substitution
       #

sig/steep/ast/types/logic.rbs

@@ -52,6 +52,11 @@ module Steep
           def initialize: (?location: untyped?) -> void
         end
 
+        # A type for `Class#<` or `Class#<=` call results.
+        class ArgIsAncestor < Base
+          def initialize: (?location: untyped?) -> void
+        end
+
         # A type with truthy/falsy type environment.
         class Env < Base
           attr_reader truthy: TypeInference::TypeEnv

test/logic_type_interpreter_test.rb

@@ -281,6 +281,46 @@ def test_call_arg_equals_receiver
     end
   end
 
+  def test_call_arg_is_ancestor
+    with_checker() do |checker|
+      source = parse_ruby("klass = Foo; klass < String")
+
+      node = source.node.children[1]
+
+      call = TypeInference::MethodCall::Typed.new(
+        node: node,
+        context: TypeInference::MethodCall::TopLevelContext.new,
+        method_name: :<,
+        receiver_type: parse_type("singleton(::Object)"),
+        actual_method_type: parse_method_type("(::Module) -> bool?").yield_self do |method_type|
+          method_type.with(
+            type: method_type.type.with(
+              return_type: AST::Types::Logic::ArgIsAncestor.new()
+            )
+          )
+        end,
+        method_decls: [],
+        return_type: AST::Types::Logic::ArgIsAncestor.new()
+      )
+
+      typing = Typing.new(source: source, root_context: nil)
+      typing.add_typing(dig(node), AST::Types::Logic::ArgIsAncestor.new(), nil)
+      typing.add_typing(dig(node, 0), parse_type("singleton(::Object)"), nil)
+      typing.add_typing(dig(node, 2), parse_type("singleton(::String)"), nil)
+
+      env = type_env
+        .assign_local_variable(:klass, parse_type("singleton(::Object)"), nil)
+
+      interpreter = LogicTypeInterpreter.new(subtyping: checker, typing: typing, config: config)
+      truthy_result, falsy_result = interpreter.eval(env: env, node: node)
+
+      assert_equal parse_type("true"), truthy_result.type
+      assert_equal parse_type("false"), falsy_result.type
+      assert_equal parse_type("singleton(::String)"), truthy_result.env[:klass]
+      assert_equal parse_type("singleton(::Object)"), falsy_result.env[:klass]
+    end
+  end
+
   def test_call_not
     with_checker(<<-RBS) do |checker|
       RBS

test/type_check_test.rb

@@ -1644,4 +1644,30 @@ class ConstantTest
       YAML
     )
   end
+
+  def test_class_narrowing
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+          module Foo
+            def self.foo: () -> void
+          end
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          klass = Class.new()
+
+          if klass < Foo
+            klass.foo()
+          end
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics: []
+      YAML
+    )
+  end
 end