From 00afdb72c30d2d0c7d5b6de389eccf2900c53c0f Mon Sep 17 00:00:00 2001
From: Fumiaki MATSUSHIMA <mtsmfm@gmail.com>
Date: Tue, 21 May 2024 00:10:49 +0900
Subject: [PATCH] Treat `present?` as a special method that guarantees the
 receiver is not nil

---
 lib/steep/interface/builder.rb                | 14 ++++++++
 .../type_inference/logic_type_interpreter.rb  | 23 ++++++++++++
 test/type_construction_test.rb                | 35 +++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/lib/steep/interface/builder.rb b/lib/steep/interface/builder.rb
index 4d5933300..7976fc9ee 100644
--- a/lib/steep/interface/builder.rb
+++ b/lib/steep/interface/builder.rb
@@ -720,6 +720,20 @@ def replace_primitive_method(method_name, method_def, method_type)
             end
             end
 
+          when :present?
+            case defined_in
+            when RBS::BuiltinNames::Object.name,
+              AST::Builtin::NilClass.module_name,
+              RBS::BuiltinNames::Kernel.name
+              if member.instance?
+                return method_type.with(
+                  type: method_type.type.with(
+                    return_type: AST::Types::Logic::ReceiverIsNotNil.new(location: method_type.type.return_type.location)
+                  )
+                )
+            end
+            end
+
           when :!
             case defined_in
             when RBS::BuiltinNames::BasicObject.name,
diff --git a/lib/steep/type_inference/logic_type_interpreter.rb b/lib/steep/type_inference/logic_type_interpreter.rb
index 621bc4d1b..a4f335616 100644
--- a/lib/steep/type_inference/logic_type_interpreter.rb
+++ b/lib/steep/type_inference/logic_type_interpreter.rb
@@ -300,6 +300,29 @@ def evaluate_method_call(env:, type:, receiver:, arguments:)
             [truthy_result, falsy_result]
           end
 
+        when AST::Types::Logic::ReceiverIsNotNil
+          if receiver && arguments.size.zero?
+            receiver_type = typing.type_of(node: receiver)
+            unwrap = factory.unwrap_optional(receiver_type)
+            truthy_receiver = unwrap || receiver_type
+            falsy_receiver = AST::Builtin.nil_type
+
+            truthy_env, falsy_env = refine_node_type(
+              env: env,
+              node: receiver,
+              truthy_type: truthy_receiver,
+              falsy_type: falsy_receiver
+            )
+
+            truthy_result = Result.new(type: TRUE, env: truthy_env, unreachable: false)
+            truthy_result.unreachable! unless unwrap
+
+            falsy_result = Result.new(type: FALSE, env: falsy_env, unreachable: false)
+            falsy_result.unreachable! if no_subtyping?(sub_type: AST::Builtin.nil_type, super_type: receiver_type)
+
+            [truthy_result, falsy_result]
+          end
+
         when AST::Types::Logic::ReceiverIsArg
           if receiver && (arg = arguments[0])
             receiver_type = typing.type_of(node: receiver)
diff --git a/test/type_construction_test.rb b/test/type_construction_test.rb
index 2138b89d4..c3af635b6 100644
--- a/test/type_construction_test.rb
+++ b/test/type_construction_test.rb
@@ -6195,6 +6195,41 @@ def test_logic_receiver_is_nil
     end
   end
 
+  def test_logic_receiver_is_not_nil
+    with_checker(<<-RBS) do |checker|
+      class Object
+        def present?: () -> bool
+      end
+    RBS
+      source = parse_ruby(<<-RUBY)
+a = [1].first
+return unless a.present?
+a + 1
+
+b = [1].first
+return unless (x = b).present?
+x + 1
+
+c = [1].first
+return unless (y = c.present?)
+c + 1
+
+d = [1].first
+puts "nil!" and return unless d.present?
+d + 1
+
+e = [1].first
+nil or return unless e.present?
+e + 1
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        construction.synthesize(source.node)
+        assert_no_error typing
+      end
+    end
+  end
+
   def test_logic_receiver_is_arg
     with_checker(<<-RBS) do |checker|
     RBS