From cce81b935c3a51a3feb9891e677ac8ecb3472812 Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Tue, 4 Jun 2019 21:41:02 -0700
Subject: [PATCH] Implement Complex#<=>

Implement Complex#<=> so that it is usable as an argument when
calling <=> on objects of other classes (since #coerce will coerce
such numbers to Complex).  If the complex number has a zero imaginary
part, and the other argument is a real number (or complex number with
zero imaginary part), return -1, 0, or 1.  Otherwise, return nil,
indicating the objects are not comparable.

Fixes [Bug #15857]
---
 NEWS                                     |  7 ++++
 complex.c                                | 49 +++++++++++++++++++++---
 spec/ruby/core/complex/spaceship_spec.rb | 27 +++++++++++++
 test/ruby/test_complex.rb                | 14 +++++--
 test/ruby/test_complexrational.rb        |  8 ++--
 5 files changed, 92 insertions(+), 13 deletions(-)
 create mode 100644 spec/ruby/core/complex/spaceship_spec.rb

diff --git a/NEWS b/NEWS
index c069f9e0eb..2346460914 100644
--- a/NEWS
+++ b/NEWS
@@ -50,6 +50,13 @@ sufficient information, see the ChangeLog file or Redmine
 
 === Core classes updates (outstanding ones only)
 
+Complex::
+
+  New method::
+
+    * Added Complex#<=>. So 0 <=> 0i will not raise NoMethodError.
+      [Bug #15857]
+
 Enumerable::
 
   New method::
diff --git a/complex.c b/complex.c
index aa905c1c88..58c1bcac51 100644
--- a/complex.c
+++ b/complex.c
@@ -36,7 +36,7 @@ static ID id_abs, id_arg,
     id_denominator, id_fdiv, id_numerator, id_quo,
     id_real_p, id_i_real, id_i_imag,
     id_finite_p, id_infinite_p, id_rationalize,
-    id_PI;
+    id_spaceship, id_PI;
 #define id_to_i idTo_i
 #define id_to_r idTo_r
 #define id_negate idUMinus
@@ -1021,14 +1021,50 @@ nucomp_eqeq_p(VALUE self, VALUE other)
     return f_boolcast(f_eqeq_p(other, self));
 }
 
+static VALUE
+nucomp_real_p(VALUE self)
+{
+    get_dat1(self);
+    return(f_zero_p(dat->imag) ? Qtrue : Qfalse);
+}
+
+/*
+ * call-seq:
+ *    cmp <=> object  ->  0, 1, -1, or nil
+ *
+ * If +cmp+'s imaginary part is zero, and +object+ is also a
+ * real number (or a Complex number where the imaginary part is zero),
+ * compare the real part of +cmp+ to object.  Otherwise, return nil.
+ *
+ *    Complex(2, 3)  <=> Complex(2, 3)   #=> nil
+ *    Complex(2, 3)  <=> 1               #=> nil
+ *    Complex(2)     <=> 1               #=> 1
+ *    Complex(2)     <=> 2               #=> 0
+ *    Complex(2)     <=> 3               #=> -1
+ */
+static VALUE
+nucomp_cmp(VALUE self, VALUE other)
+{
+    if (nucomp_real_p(self) && k_numeric_p(other)) {
+        if (RB_TYPE_P(other, T_COMPLEX) && nucomp_real_p(other)) {
+            get_dat2(self, other);
+            return rb_funcall(adat->real, id_spaceship, 1, bdat->real);
+        } else if (f_real_p(other)) {
+            get_dat1(self);
+            return rb_funcall(dat->real, id_spaceship, 1, other);
+        }
+    }
+    return Qnil;
+}
+
 /* :nodoc: */
 static VALUE
 nucomp_coerce(VALUE self, VALUE other)
 {
-    if (k_numeric_p(other) && f_real_p(other))
-	return rb_assoc_new(f_complex_new_bang1(CLASS_OF(self), other), self);
     if (RB_TYPE_P(other, T_COMPLEX))
 	return rb_assoc_new(other, self);
+    if (k_numeric_p(other) && f_real_p(other))
+	return rb_assoc_new(f_complex_new_bang1(CLASS_OF(self), other), self);
 
     rb_raise(rb_eTypeError, "%"PRIsVALUE" can't be coerced into %"PRIsVALUE,
 	     rb_obj_class(other), rb_obj_class(self));
@@ -1147,9 +1183,10 @@ rb_complex_conjugate(VALUE self)
 
 /*
  * call-seq:
- *    cmp.real?  ->  false
+ *    Complex(1).real?     ->  false
+ *    Complex(1, 2).real?  ->  false
  *
- * Returns false.
+ * Returns false, even if the complex number has no imaginary part.
  */
 static VALUE
 nucomp_false(VALUE self)
@@ -2211,6 +2248,7 @@ Init_Complex(void)
     id_finite_p = rb_intern("finite?");
     id_infinite_p = rb_intern("infinite?");
     id_rationalize = rb_intern("rationalize");
+    id_spaceship = rb_intern("<=>");
     id_PI = rb_intern("PI");
 
     rb_cComplex = rb_define_class("Complex", rb_cNumeric);
@@ -2254,6 +2292,7 @@ Init_Complex(void)
     rb_define_method(rb_cComplex, "**", rb_complex_pow, 1);
 
     rb_define_method(rb_cComplex, "==", nucomp_eqeq_p, 1);
+    rb_define_method(rb_cComplex, "<=>", nucomp_cmp, 1);
     rb_define_method(rb_cComplex, "coerce", nucomp_coerce, 1);
 
     rb_define_method(rb_cComplex, "abs", rb_complex_abs, 0);
diff --git a/spec/ruby/core/complex/spaceship_spec.rb b/spec/ruby/core/complex/spaceship_spec.rb
new file mode 100644
index 0000000000..7b2849a86a
--- /dev/null
+++ b/spec/ruby/core/complex/spaceship_spec.rb
@@ -0,0 +1,27 @@
+require_relative '../../spec_helper'
+
+describe "Complex#<=>" do
+  ruby_version_is '2.7' do
+    it "returns nil if either self or argument has imaginary part" do
+      (Complex(5, 1) <=> Complex(2)).should be_nil
+      (Complex(1) <=> Complex(2, 1)).should be_nil
+      (5 <=> Complex(2, 1)).should be_nil
+    end
+
+    it "returns nil if argument is not numeric" do
+      (Complex(5, 1) <=> "cmp").should be_nil
+      (Complex(1) <=> "cmp").should be_nil
+      (Complex(1) <=> Object.new).should be_nil
+    end
+
+    it "returns 0, 1, or -1 if self and argument do not have imaginary part" do
+      (Complex(5) <=> Complex(2)).should == 1
+      (Complex(2) <=> Complex(3)).should == -1
+      (Complex(2) <=> Complex(2)).should == 0
+
+      (Complex(5) <=> 2).should == 1
+      (Complex(2) <=> 3).should == -1
+      (Complex(2) <=> 2).should == 0
+    end
+  end
+end
diff --git a/test/ruby/test_complex.rb b/test/ruby/test_complex.rb
index 3da66890c1..60a46d737d 100644
--- a/test/ruby/test_complex.rb
+++ b/test/ruby/test_complex.rb
@@ -518,9 +518,16 @@ def test_expt
   end
 
   def test_cmp
-    assert_raise(NoMethodError){1 <=> Complex(1,1)}
-    assert_raise(NoMethodError){Complex(1,1) <=> 1}
-    assert_raise(NoMethodError){Complex(1,1) <=> Complex(1,1)}
+    assert_nil(Complex(5, 1) <=> Complex(2))
+    assert_nil(5 <=> Complex(2, 1))
+
+    assert_equal(1, Complex(5) <=> Complex(2))
+    assert_equal(-1, Complex(2) <=> Complex(3))
+    assert_equal(0, Complex(2) <=> Complex(2))
+
+    assert_equal(1, Complex(5) <=> 2)
+    assert_equal(-1, Complex(2) <=> 3)
+    assert_equal(0, Complex(2) <=> 2)
   end
 
   def test_eqeq
@@ -891,7 +898,6 @@ def o.to_c; raise; end
   def test_respond
     c = Complex(1,1)
     assert_not_respond_to(c, :%)
-    assert_not_respond_to(c, :<=>)
     assert_not_respond_to(c, :div)
     assert_not_respond_to(c, :divmod)
     assert_not_respond_to(c, :floor)
diff --git a/test/ruby/test_complexrational.rb b/test/ruby/test_complexrational.rb
index 0360f5ee42..bf4e2b1809 100644
--- a/test/ruby/test_complexrational.rb
+++ b/test/ruby/test_complexrational.rb
@@ -102,7 +102,7 @@ def test_comp_srat
     assert_equal(Complex(SimpleRat(4,3),SimpleRat(1,1)), c * 2)
     assert_equal(Complex(SimpleRat(1,3),SimpleRat(1,4)), c / 2)
     assert_equal(Complex(SimpleRat(7,36),SimpleRat(2,3)), c ** 2)
-    assert_raise(NoMethodError){c <=> 2}
+    assert_nil(c <=> 2)
 
     assert_equal(Complex(SimpleRat(8,3),SimpleRat(1,2)), 2 + c)
     assert_equal(Complex(SimpleRat(4,3),SimpleRat(-1,2)), 2 - c)
@@ -111,7 +111,7 @@ def test_comp_srat
     r = 2 ** c
     assert_in_delta(1.4940, r.real, 0.001)
     assert_in_delta(0.5392, r.imag, 0.001)
-    assert_raise(NoMethodError){2 <=> c}
+    assert_nil(2 <=> c)
 
     assert_equal(Complex(SimpleRat(13,6),SimpleRat(5,2)), c + cc)
     assert_equal(Complex(SimpleRat(-5,6),SimpleRat(-3,2)), c - cc)
@@ -120,7 +120,7 @@ def test_comp_srat
     r = c ** cc
     assert_in_delta(0.1732, r.real, 0.001)
     assert_in_delta(0.1186, r.imag, 0.001)
-    assert_raise(NoMethodError){c <=> cc}
+    assert_nil(c <=> cc)
 
     assert_equal(Complex(SimpleRat(13,6),SimpleRat(5,2)), cc + c)
     assert_equal(Complex(SimpleRat(5,6),SimpleRat(3,2)), cc - c)
@@ -129,7 +129,7 @@ def test_comp_srat
     r = cc ** c
     assert_in_delta(0.5498, r.real, 0.001)
     assert_in_delta(1.0198, r.imag, 0.001)
-    assert_raise(NoMethodError){cc <=> c}
+    assert_nil(cc <=> c)
 
     assert_equal([SimpleRat,SimpleRat],
 		 (+c).instance_eval{[real.class, imag.class]})
-- 
2.21.0

