From 5e92eec501e0a1ca6d09a01e078cf54ff3277273 Mon Sep 17 00:00:00 2001
From: Sam T <pxqr.sta@gmail.com>
Date: Sat, 25 May 2013 21:41:15 +0400
Subject: ~ Use PortNumber instead of Word16 in Port message.

---
 bench/Main.hs                              | 5 +++++
 network-bittorrent.cabal                   | 1 +
 src/Network/BitTorrent/PeerWire/Message.hs | 6 +++---
 tests/Encoding.hs                          | 7 ++++++-
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/bench/Main.hs b/bench/Main.hs
index 269e60f5..6e8a0ce3 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -7,10 +7,15 @@ import Criterion.Main
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import Data.Serialize
+import Network
+
 import Network.BitTorrent as BT
 import Data.Bitfield as BT
 
 
+instance NFData PortNumber where
+  rnf = rnf . (fromIntegral :: PortNumber -> Int)
+
 instance NFData BlockIx where
   rnf (BlockIx a b c) = a `deepseq` b `deepseq` rnf c
 
diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index 339dc376..0a89603a 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -127,6 +127,7 @@ benchmark benchmarks
   build-depends:       base == 4.*
                      , bytestring >= 0.10.2.0
                      , cereal
+                     , network
 
                      , criterion
                      , deepseq
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 1c520221..5544dca7 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -6,13 +6,13 @@ module Network.BitTorrent.PeerWire.Message
 
 import Control.Applicative
 import qualified Data.ByteString as B
-
 import Data.Serialize
+import Network
 
 import Network.BitTorrent.PeerWire.Block
 import Data.Bitfield
 
-import Data.Array
+
 
 -- | Messages used in communication between peers.
 --
@@ -48,7 +48,7 @@ data Message = KeepAlive
                -- used during "End Game".
              | Cancel   BlockIx
 
-             | Port     Int
+             | Port     PortNumber
 
                -- | BEP 6: Then peer have all pieces it might send the
                --   'HaveAll' message instead of 'Bitfield'
diff --git a/tests/Encoding.hs b/tests/Encoding.hs
index fe0e1d7e..d43045bc 100644
--- a/tests/Encoding.hs
+++ b/tests/Encoding.hs
@@ -13,6 +13,8 @@ import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck
 
 import Network.URI
+import Network
+
 
 import Data.Bitfield
 import Data.Torrent
@@ -33,6 +35,9 @@ instance Arbitrary Block where
 
 deriving instance Arbitrary Bitfield
 
+instance Arbitrary PortNumber where
+  arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
+
 instance Arbitrary Message where
   arbitrary = oneof
     [ pure KeepAlive
@@ -45,7 +50,7 @@ instance Arbitrary Message where
     , Request <$> arbitrary
     , Piece <$> arbitrary
     , Cancel <$> arbitrary
-    , Port <$> choose (0, fromIntegral (maxBound :: Word16))
+    , Port <$> arbitrary
     ]
 
 instance Arbitrary PeerID where
-- 
cgit v1.2.3