From 6ab23ea9acf74a9e49018d121a7a0480b48555ec Mon Sep 17 00:00:00 2001
From: Akinori MUSHA <knu@idaemons.org>
Date: Sun, 3 Sep 2017 19:15:46 +0900
Subject: [PATCH] Add type coercion support to Prefix

This enables expressions like `32 - prefix` and `24 <= prefix`.
---
 lib/ipaddress/prefix.rb       | 26 +++++++++++++++++++++++---
 test/ipaddress/prefix_test.rb | 20 +++++++++++++++++++-
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/lib/ipaddress/prefix.rb b/lib/ipaddress/prefix.rb
index e819800..a6b9ce4 100644
--- a/lib/ipaddress/prefix.rb
+++ b/lib/ipaddress/prefix.rb
@@ -45,12 +45,30 @@ def to_s
     def to_i
       @prefix
     end
+    alias_method :to_int, :to_i
+
+    #
+    # Provides support for Ruby type coercion
+    #
+    def coerce(other)
+      case other
+      when Integer
+        [other, @prefix]
+      else
+        other.coerce(@prefix).reverse!
+      end
+    end
 
     # 
     # Compare the prefix
     #
     def <=>(oth)
-      @prefix <=> oth.to_i
+      if oth.is_a? Integer
+        @prefix <=> oth
+      else
+        x, y = oth.coerce(@prefix)
+        x <=> y
+      end
     end
 
     #
@@ -61,7 +79,8 @@ def +(oth)
       if oth.is_a? Integer
         self.prefix + oth
       else
-        self.prefix + oth.prefix
+        x, y = oth.coerce(@prefix)
+        x + y
       end
     end
 
@@ -74,7 +93,8 @@ def -(oth)
       if oth.is_a? Integer
         self.prefix - oth
       else
-        (self.prefix - oth.prefix).abs
+        x, y = oth.coerce(@prefix)
+        x - y
       end
     end
     
diff --git a/test/ipaddress/prefix_test.rb b/test/ipaddress/prefix_test.rb
index f0a3d8c..e6ac607 100644
--- a/test/ipaddress/prefix_test.rb
+++ b/test/ipaddress/prefix_test.rb
@@ -73,19 +73,37 @@ def test_method_to_u32
     end
   end
 
+  def test_method_cmp
+    p1 = @klass.new 8
+    p2 = @klass.new 24
+    assert_equal false, (p1 > 16)
+    assert_equal true,  (p2 > 16)
+    assert_equal true,  (16 >= p1)
+    assert_equal false, (16 >= p2)
+    assert_equal true,  (p1 < p2)
+    assert_equal true,  (p1 == 8r)
+    assert_equal true,  (24r == p2)
+  end
+
   def test_method_plus
     p1 = @klass.new 8
     p2 = @klass.new 10
     assert_equal 18, p1+p2
     assert_equal 12, p1+4
+    assert_equal 14r, p2+4r
+    assert_equal 12r, 4r+p1
   end
 
   def test_method_minus
     p1 = @klass.new 8
     p2 = @klass.new 24
-    assert_equal 16, p1-p2
+    assert_equal -16, p1-p2
     assert_equal 16, p2-p1
     assert_equal 20, p2-4
+    assert_equal 20r, p2-4r
+    assert_equal 24, 32-p1
+    assert_equal 8, 32-p2
+    assert_equal 8r, 32r-p2
   end
 
   def test_initialize