Move HashList to Torrent module

16 files changed, 213 insertions, 256 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 9d687d7d..9a86702d 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -46,7 +46,6 @@ library
   hs-source-dirs:      src
   exposed-modules:     Data.Torrent
                        Data.Torrent.Bitfield
-                       Data.Torrent.Piece
                        Data.Torrent.Progress
                        Data.Torrent.Tree
                        Network.BitTorrent
diff --git a/src/Data/Torrent.hs b/src/Data/Torrent.hs
index 701da9dd..98d6f94e 100644
--- a/src/Data/Torrent.hs
+++ b/src/Data/Torrent.hs
@@ -86,6 +86,34 @@ module Data.Torrent
          -- ** Internal
        , sizeInBase
 
+         -- * Pieces
+         -- ** Attributes
+       , PieceIx
+       , PieceCount
+       , PieceSize
+       , minPieceSize
+       , maxPieceSize
+       , defaultPieceSize
+       , PieceHash
+
+         -- ** Piece data
+       , Piece (..)
+       , pieceSize
+       , hashPiece
+
+         -- ** Piece control
+       , HashList (..)
+       , PieceInfo (..)
+       , pieceCount
+
+         -- ** Lens
+       , pieceLength
+       , pieceHashes
+
+         -- ** Validation
+       , pieceHash
+       , checkPieceLazy
+
          -- * Info dictionary
        , InfoDict (..)
        , infoDictionary
@@ -133,8 +161,11 @@ import Control.Exception
 import Control.Lens hiding (unsnoc)
 import Control.Monad
 import qualified Crypto.Hash.SHA1 as C
+import qualified Crypto.Hash.SHA1 as SHA1
 import Data.BEncode as BE
 import Data.BEncode.Types as BE
+import Data.Bits
+import Data.Bits.Extras
 import           Data.ByteString as BS
 import           Data.ByteString.Base16 as Base16
 import           Data.ByteString.Base32 as Base32
@@ -146,6 +177,7 @@ import           Data.Convertible
 import           Data.Default
 import           Data.Foldable as F
 import           Data.Hashable   as Hashable
+import           Data.Int
 import qualified Data.List as L
 import           Data.Map as M
 import           Data.Maybe
@@ -166,7 +198,6 @@ import Text.PrettyPrint.Class
 import System.FilePath
 import System.Posix.Types
 
-import Data.Torrent.Piece
 import Network.BitTorrent.Core.NodeInfo
 
 
@@ -527,6 +558,171 @@ sizeInBase n b = fromIntegral (n `div` fromIntegral b) + align
 {-# SPECIALIZE sizeInBase :: Integer -> Int -> Int #-}
 
 {-----------------------------------------------------------------------
+-- Piece attributes
+-----------------------------------------------------------------------}
+
+-- | Zero-based index of piece in torrent content.
+type PieceIx   = Int
+
+-- | Size of piece in bytes. Should be a power of 2.
+--
+--   NOTE: Have max and min size constrained to wide used
+--   semi-standard values. This bounds should be used to make decision
+--   about piece size for new torrents.
+--
+type PieceSize = Int
+
+-- | Number of pieces in torrent or a part of torrent.
+type PieceCount = Int
+
+defaultBlockSize :: Int
+defaultBlockSize = 16 * 1024
+
+-- | Optimal number of pieces in torrent.
+optimalPieceCount :: PieceCount
+optimalPieceCount = 1000
+{-# INLINE optimalPieceCount #-}
+
+-- | Piece size should not be less than this value.
+minPieceSize :: Int
+minPieceSize = defaultBlockSize * 4
+{-# INLINE minPieceSize #-}
+
+-- | To prevent transfer degradation piece size should not exceed this
+-- value.
+maxPieceSize :: Int
+maxPieceSize = 4 * 1024 * 1024
+{-# INLINE maxPieceSize #-}
+
+toPow2 :: Int -> Int
+toPow2 x = bit $ fromIntegral (leadingZeros (0 :: Int) - leadingZeros x)
+
+-- | Find the optimal piece size for a given torrent size.
+defaultPieceSize :: Int64 -> Int
+defaultPieceSize x = max minPieceSize $ min maxPieceSize $ toPow2 pc
+  where
+    pc = fromIntegral (x `div` fromIntegral optimalPieceCount)
+
+{-----------------------------------------------------------------------
+-- Piece data
+-----------------------------------------------------------------------}
+
+type PieceHash = ByteString
+
+hashsize :: Int
+hashsize = 20
+{-# INLINE hashsize #-}
+
+-- TODO check if pieceLength is power of 2
+-- | Piece payload should be strict or lazy bytestring.
+data Piece a = Piece
+  { -- | Zero-based piece index in torrent.
+    pieceIndex :: {-# UNPACK #-} !PieceIx
+
+    -- | Payload.
+  , pieceData  :: !a
+  } deriving (Show, Read, Eq, Functor, Typeable)
+
+instance NFData (Piece a)
+
+-- | Payload bytes are omitted.
+instance Pretty (Piece a) where
+  pretty Piece {..} = "Piece" <+> braces ("index" <+> "=" <+> int pieceIndex)
+
+-- | Get size of piece in bytes.
+pieceSize :: Piece BL.ByteString -> PieceSize
+pieceSize Piece {..} = fromIntegral (BL.length pieceData)
+
+-- | Get piece hash.
+hashPiece :: Piece BL.ByteString -> PieceHash
+hashPiece Piece {..} = SHA1.hashlazy pieceData
+
+{-----------------------------------------------------------------------
+-- Piece control
+-----------------------------------------------------------------------}
+
+-- | A flat array of SHA1 hash for each piece.
+newtype HashList = HashList { unHashList :: ByteString }
+  deriving (Show, Read, Eq, BEncode, Typeable)
+
+-- | Empty hash list.
+instance Default HashList where
+  def = HashList ""
+
+-- | Part of torrent file used for torrent content validation.
+data PieceInfo = PieceInfo
+  { piPieceLength  :: {-# UNPACK #-} !PieceSize
+    -- ^ Number of bytes in each piece.
+
+  , piPieceHashes  :: !HashList
+    -- ^ Concatenation of all 20-byte SHA1 hash values.
+  } deriving (Show, Read, Eq, Typeable)
+
+-- | Number of bytes in each piece.
+makeLensesFor [("piPieceLength", "pieceLength")] ''PieceInfo
+
+-- | Concatenation of all 20-byte SHA1 hash values.
+makeLensesFor [("piPieceHashes", "pieceHashes")] ''PieceInfo
+
+instance NFData PieceInfo
+
+instance Default PieceInfo where
+  def = PieceInfo 1 def
+
+class Lint a where
+   lint :: a -> Either String a
+
+instance Lint PieceInfo where
+  lint pinfo @ PieceInfo {..}
+    | BS.length (unHashList piPieceHashes) `rem` hashsize == 0
+    , piPieceLength >= 0 = return pinfo
+    |       otherwise    = Left undefined
+
+
+putPieceInfo :: Data.Torrent.Put PieceInfo
+putPieceInfo PieceInfo {..} cont =
+       "piece length" .=! piPieceLength
+    .: "pieces"       .=! piPieceHashes
+    .: cont
+
+getPieceInfo :: BE.Get PieceInfo
+getPieceInfo = do
+    PieceInfo <$>! "piece length"
+              <*>! "pieces"
+
+instance BEncode PieceInfo where
+  toBEncode   = toDict . (`putPieceInfo` endDict)
+  fromBEncode = fromDict getPieceInfo
+
+-- | Hashes are omitted.
+instance Pretty PieceInfo where
+  pretty  PieceInfo {..} = "Piece size: " <> int piPieceLength
+
+slice :: Int -> Int -> ByteString -> ByteString
+slice start len = BS.take len . BS.drop start
+{-# INLINE slice #-}
+
+-- | Extract validation hash by specified piece index.
+pieceHash :: PieceInfo -> PieceIx -> PieceHash
+pieceHash PieceInfo {..} i = slice (hashsize * i) hashsize (unHashList piPieceHashes)
+
+-- | Find count of pieces in the torrent. If torrent size is not a
+-- multiple of piece size then the count is rounded up.
+pieceCount :: PieceInfo -> PieceCount
+pieceCount PieceInfo {..} = BS.length (unHashList piPieceHashes) `quot` hashsize
+
+-- | Test if this is last piece in torrent content.
+isLastPiece :: PieceInfo -> PieceIx -> Bool
+isLastPiece ci i = pieceCount ci == succ i
+
+-- | Validate piece with metainfo hash.
+checkPieceLazy :: PieceInfo -> Piece BL.ByteString -> Bool
+checkPieceLazy pinfo @ PieceInfo {..} Piece {..}
+  = (fromIntegral (BL.length pieceData) == piPieceLength
+     || isLastPiece pinfo pieceIndex)
+  && SHA1.hashlazy pieceData == pieceHash pinfo pieceIndex
+
+{-----------------------------------------------------------------------
 --  Info dictionary
 -----------------------------------------------------------------------}
 
@@ -620,6 +816,7 @@ instance Pretty InfoDict where
 {-----------------------------------------------------------------------
 --  Torrent info
 -----------------------------------------------------------------------}
+-- TODO add torrent file validation
 
 -- | Metainfo about particular torrent.
 data Torrent = Torrent
diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index b65f058b..ff701d75 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -92,7 +92,7 @@ import           Data.List (foldl')
 import           Data.Monoid
 import           Data.Ratio
 
-import Data.Torrent.Piece
+import Data.Torrent
 
 -- TODO cache some operations
 
diff --git a/src/Data/Torrent/Piece.hs b/src/Data/Torrent/Piece.hs
deleted file mode 100644
index d4b2c399..00000000
--- a/src/Data/Torrent/Piece.hs
+++ /dev/null
@@ -1,232 +0,0 @@
--- |
---   Copyright   :  (c) Sam Truzjan 2013
---   License     :  BSD3
---   Maintainer  :  pxqr.sta@gmail.com
---   Stability   :  experimental
---   Portability :  portable
---
---   Pieces are used to validate torrent content.
---
-{-# LANGUAGE TemplateHaskell            #-}
-{-# LANGUAGE DeriveDataTypeable         #-}
-{-# LANGUAGE DeriveFunctor              #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-module Data.Torrent.Piece
-       ( -- * Piece attributes
-         PieceIx
-       , PieceCount
-       , PieceSize
-       , minPieceSize
-       , maxPieceSize
-       , defaultPieceSize
-       , PieceHash
-
-         -- * Piece data
-       , Piece (..)
-       , pieceSize
-       , hashPiece
-
-         -- * Piece control
-       , HashList (..)
-       , PieceInfo (..)
-       , pieceCount
-
-         -- * Lens
-       , pieceLength
-       , pieceHashes
-
-         -- * Validation
-       , pieceHash
-       , checkPieceLazy
-
-         -- * Internal
-       , getPieceInfo
-       , putPieceInfo
-       ) where
-
-import Control.DeepSeq
-import Control.Lens
-import qualified Crypto.Hash.SHA1 as SHA1
-import Data.BEncode
-import Data.BEncode.Types
-import Data.Bits
-import Data.Bits.Extras
-import           Data.ByteString as BS
-import qualified Data.ByteString.Lazy as BL
-import qualified Data.ByteString.Base64 as Base64
-import Data.Default
-import Data.Int
-import Data.Text.Encoding as T
-import Data.Typeable
-import Text.PrettyPrint
-import Text.PrettyPrint.Class
-
-
--- TODO add torrent file validation
-class Lint a where
-   lint :: a -> Either String a
-
---class Validation a where
---  validate :: PieceInfo -> Piece a -> Bool
-
-{-----------------------------------------------------------------------
--- Piece attributes
------------------------------------------------------------------------}
-
--- | Zero-based index of piece in torrent content.
-type PieceIx   = Int
-
--- | Size of piece in bytes. Should be a power of 2.
---
---   NOTE: Have max and min size constrained to wide used
---   semi-standard values. This bounds should be used to make decision
---   about piece size for new torrents.
---
-type PieceSize = Int
-
--- | Number of pieces in torrent or a part of torrent.
-type PieceCount = Int
-
-defaultBlockSize :: Int
-defaultBlockSize = 16 * 1024
-
--- | Optimal number of pieces in torrent.
-optimalPieceCount :: PieceCount
-optimalPieceCount = 1000
-{-# INLINE optimalPieceCount #-}
-
--- | Piece size should not be less than this value.
-minPieceSize :: Int
-minPieceSize = defaultBlockSize * 4
-{-# INLINE minPieceSize #-}
-
--- | To prevent transfer degradation piece size should not exceed this
--- value.
-maxPieceSize :: Int
-maxPieceSize = 4 * 1024 * 1024
-{-# INLINE maxPieceSize #-}
-
-toPow2 :: Int -> Int
-toPow2 x = bit $ fromIntegral (leadingZeros (0 :: Int) - leadingZeros x)
-
--- | Find the optimal piece size for a given torrent size.
-defaultPieceSize :: Int64 -> Int
-defaultPieceSize x = max minPieceSize $ min maxPieceSize $ toPow2 pc
-  where
-    pc = fromIntegral (x `div` fromIntegral optimalPieceCount)
-
-{-----------------------------------------------------------------------
--- Piece data
------------------------------------------------------------------------}
-
-type PieceHash = ByteString
-
-hashsize :: Int
-hashsize = 20
-{-# INLINE hashsize #-}
-
--- TODO check if pieceLength is power of 2
--- | Piece payload should be strict or lazy bytestring.
-data Piece a = Piece
-  { -- | Zero-based piece index in torrent.
-    pieceIndex :: {-# UNPACK #-} !PieceIx
-
-    -- | Payload.
-  , pieceData  :: !a
-  } deriving (Show, Read, Eq, Functor, Typeable)
-
-instance NFData (Piece a)
-
--- | Payload bytes are omitted.
-instance Pretty (Piece a) where
-  pretty Piece {..} = "Piece" <+> braces ("index" <+> "=" <+> int pieceIndex)
-
--- | Get size of piece in bytes.
-pieceSize :: Piece BL.ByteString -> PieceSize
-pieceSize Piece {..} = fromIntegral (BL.length pieceData)
-
--- | Get piece hash.
-hashPiece :: Piece BL.ByteString -> PieceHash
-hashPiece Piece {..} = SHA1.hashlazy pieceData
-
-{-----------------------------------------------------------------------
--- Piece control
------------------------------------------------------------------------}
-
--- | A flat array of SHA1 hash for each piece.
-newtype HashList = HashList { unHashList :: ByteString }
-  deriving (Show, Read, Eq, BEncode, Typeable)
-
--- | Empty hash list.
-instance Default HashList where
-  def = HashList ""
-
--- | Part of torrent file used for torrent content validation.
-data PieceInfo = PieceInfo
-  { piPieceLength  :: {-# UNPACK #-} !PieceSize
-    -- ^ Number of bytes in each piece.
-
-  , piPieceHashes  :: !HashList
-    -- ^ Concatenation of all 20-byte SHA1 hash values.
-  } deriving (Show, Read, Eq, Typeable)
-
--- | Number of bytes in each piece.
-makeLensesFor [("piPieceLength", "pieceLength")] ''PieceInfo
-
--- | Concatenation of all 20-byte SHA1 hash values.
-makeLensesFor [("piPieceHashes", "pieceHashes")] ''PieceInfo
-
-instance NFData PieceInfo
-
-instance Default PieceInfo where
-  def = PieceInfo 1 def
-
-instance Lint PieceInfo where
-  lint pinfo @ PieceInfo {..}
-    | BS.length (unHashList piPieceHashes) `rem` hashsize == 0
-    , piPieceLength >= 0 = return pinfo
-    |       otherwise    = Left undefined
-
-
-putPieceInfo :: PieceInfo -> BDict -> BDict
-putPieceInfo PieceInfo {..} cont =
-       "piece length" .=! piPieceLength
-    .: "pieces"       .=! piPieceHashes
-    .: cont
-
-getPieceInfo :: Get PieceInfo
-getPieceInfo = do
-    PieceInfo <$>! "piece length"
-              <*>! "pieces"
-
-instance BEncode PieceInfo where
-  toBEncode   = toDict . (`putPieceInfo` endDict)
-  fromBEncode = fromDict getPieceInfo
-
--- | Hashes are omitted.
-instance Pretty PieceInfo where
-  pretty  PieceInfo {..} = "Piece size: " <> int piPieceLength
-
-slice :: Int -> Int -> ByteString -> ByteString
-slice start len = BS.take len . BS.drop start
-{-# INLINE slice #-}
-
--- | Extract validation hash by specified piece index.
-pieceHash :: PieceInfo -> PieceIx -> PieceHash
-pieceHash PieceInfo {..} i = slice (hashsize * i) hashsize (unHashList piPieceHashes)
-
--- | Find count of pieces in the torrent. If torrent size is not a
--- multiple of piece size then the count is rounded up.
-pieceCount :: PieceInfo -> PieceCount
-pieceCount PieceInfo {..} = BS.length (unHashList piPieceHashes) `quot` hashsize
-
--- | Test if this is last piece in torrent content.
-isLastPiece :: PieceInfo -> PieceIx -> Bool
-isLastPiece ci i = pieceCount ci == succ i
-
--- | Validate piece with metainfo hash.
-checkPieceLazy :: PieceInfo -> Piece BL.ByteString -> Bool
-checkPieceLazy pinfo @ PieceInfo {..} Piece {..}
-  = (fromIntegral (BL.length pieceData) == piPieceLength
-     || isLastPiece pinfo pieceIndex)
-  && SHA1.hashlazy pieceData == pieceHash pinfo pieceIndex
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 87a6d4ea..8fe81abd 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -91,7 +91,7 @@ import System.Random (randomIO)
 import Text.PrettyPrint as PP hiding ((<>), ($$))
 import Text.PrettyPrint.Class
 
-import Data.Torrent
+import Data.Torrent as Torrent
 import Network.KRPC hiding (Options, def)
 import qualified Network.KRPC as KRPC (def)
 import Network.BitTorrent.Core
diff --git a/src/Network/BitTorrent/Exchange/Assembler.hs b/src/Network/BitTorrent/Exchange/Assembler.hs
index e5834948..e17dfbe2 100644
--- a/src/Network/BitTorrent/Exchange/Assembler.hs
+++ b/src/Network/BitTorrent/Exchange/Assembler.hs
@@ -67,7 +67,7 @@ import Data.Map as M
 import Data.Maybe
 import Data.IP
 
-import Data.Torrent.Piece
+import Data.Torrent
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Block as B
 
diff --git a/src/Network/BitTorrent/Exchange/Block.hs b/src/Network/BitTorrent/Exchange/Block.hs
index 16c124e9..ccc7a0a9 100644
--- a/src/Network/BitTorrent/Exchange/Block.hs
+++ b/src/Network/BitTorrent/Exchange/Block.hs
@@ -69,7 +69,7 @@ import Numeric
 import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.Class
 
-import Data.Torrent.Piece
+import Data.Torrent
 
 {-----------------------------------------------------------------------
 --  Block attributes
diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index bd5c6526..5ca7c97e 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -118,8 +118,8 @@ import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.Class
 
 import Data.Torrent.Bitfield
-import Data.Torrent
-import qualified Data.Torrent.Piece as P
+import Data.Torrent hiding (Piece (..))
+import qualified Data.Torrent as P (Piece (..))
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Block
 
@@ -864,7 +864,7 @@ instance PeerMessage ExtendedMetadata where
 
 -- | All 'Piece's in 'MetadataData' messages MUST have size equal to
 -- this value. The last trailing piece can be shorter.
-metadataPieceSize :: P.PieceSize
+metadataPieceSize :: PieceSize
 metadataPieceSize = 16 * 1024
 
 isLastPiece :: P.Piece a -> Int -> Bool
@@ -877,8 +877,8 @@ isLastPiece P.Piece {..} total = succ pieceIndex == pcnt
 -- length; otherwise serialization MUST fail.
 isValidPiece :: P.Piece BL.ByteString -> Int -> Bool
 isValidPiece p @ P.Piece {..} total
-  | isLastPiece p total = P.pieceSize p <= metadataPieceSize
-  |       otherwise     = P.pieceSize p == metadataPieceSize
+  | isLastPiece p total = pieceSize p <= metadataPieceSize
+  |       otherwise     = pieceSize p == metadataPieceSize
 
 setMetadataPayload :: BS.ByteString -> ExtendedMetadata -> ExtendedMetadata
 setMetadataPayload bs (MetadataData (P.Piece pix _) t) =
diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 0adb08c8..cae3a2d5 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -45,8 +45,7 @@ import Text.PrettyPrint.Class
 import System.Log.FastLogger (LogStr, ToLogStr (..))
 
 import Data.BEncode as BE
-import Data.Torrent as T
-import Data.Torrent.Piece as Torrent
+import Data.Torrent as Torrent
 import Data.Torrent.Bitfield as BF
 import Network.BitTorrent.Internal.Types
 import Network.BitTorrent.Core
diff --git a/src/Network/BitTorrent/Exchange/Session/Metadata.hs b/src/Network/BitTorrent/Exchange/Session/Metadata.hs
index bdd5b322..a4e54659 100644
--- a/src/Network/BitTorrent/Exchange/Session/Metadata.hs
+++ b/src/Network/BitTorrent/Exchange/Session/Metadata.hs
@@ -27,7 +27,6 @@ import Data.Tuple
 
 import Data.BEncode as BE
 import Data.Torrent as Torrent
-import Data.Torrent.Piece as Torrent
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Block   as Block
 import Network.BitTorrent.Exchange.Message as Message hiding (Status)
diff --git a/src/Network/BitTorrent/Exchange/Session/Status.hs b/src/Network/BitTorrent/Exchange/Session/Status.hs
index 565c3bf3..4feff8d6 100644
--- a/src/Network/BitTorrent/Exchange/Session/Status.hs
+++ b/src/Network/BitTorrent/Exchange/Session/Status.hs
@@ -28,7 +28,7 @@ import Data.Map as M
 import Data.Set as S
 import Data.Tuple
 
-import Data.Torrent.Piece
+import Data.Torrent
 import Data.Torrent.Bitfield as BF
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Block as Block
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index 697e3def..1123cea9 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -57,8 +57,6 @@ import Data.Typeable
 
 import Data.Torrent
 import Data.Torrent.Bitfield as BF
-import Data.Torrent
-import Data.Torrent.Piece
 import System.Torrent.FileMap as FM
 
 
diff --git a/tests/Data/Torrent/MetainfoSpec.hs b/tests/Data/Torrent/MetainfoSpec.hs
index 537b3f99..1a8f97c7 100644
--- a/tests/Data/Torrent/MetainfoSpec.hs
+++ b/tests/Data/Torrent/MetainfoSpec.hs
@@ -14,7 +14,6 @@ import Test.QuickCheck
 import Test.QuickCheck.Instances ()
 
 import Data.Torrent
-import Data.Torrent.Piece
 import Data.Torrent.LayoutSpec ()
 import Network.BitTorrent.Core.NodeInfoSpec ()
 
diff --git a/tests/Data/Torrent/PieceSpec.hs b/tests/Data/Torrent/PieceSpec.hs
index ef1f2938..d3933396 100644
--- a/tests/Data/Torrent/PieceSpec.hs
+++ b/tests/Data/Torrent/PieceSpec.hs
@@ -3,7 +3,7 @@ module Data.Torrent.PieceSpec (spec) where
 import Control.Applicative
 import Test.Hspec
 import Test.QuickCheck
-import Data.Torrent.Piece
+import Data.Torrent
 
 
 instance Arbitrary a => Arbitrary (Piece a) where
diff --git a/tests/Network/BitTorrent/Exchange/Session/MetadataSpec.hs b/tests/Network/BitTorrent/Exchange/Session/MetadataSpec.hs
index 975ceb5b..5392d74b 100644
--- a/tests/Network/BitTorrent/Exchange/Session/MetadataSpec.hs
+++ b/tests/Network/BitTorrent/Exchange/Session/MetadataSpec.hs
@@ -7,8 +7,7 @@ import Test.Hspec
 import Test.QuickCheck
 
 import Data.BEncode as BE
-import Data.Torrent
-import Data.Torrent.Piece as P
+import Data.Torrent as Torrent
 import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Message
 import Network.BitTorrent.Exchange.Session.Metadata
@@ -36,7 +35,7 @@ simulateFetch :: InfoDict -> Updates (Maybe InfoDict)
 simulateFetch dict = go
   where
     blocks = chunkBy metadataPieceSize (BL.toStrict (BE.encode dict))
-    packPiece ix = P.Piece ix (blocks !! ix)
+    packPiece ix = Torrent.Piece ix (blocks !! ix)
     ih     = idInfoHash dict
 
     go = do
diff --git a/tests/System/Torrent/StorageSpec.hs b/tests/System/Torrent/StorageSpec.hs
index ebf4fe3e..96f1b036 100644
--- a/tests/System/Torrent/StorageSpec.hs
+++ b/tests/System/Torrent/StorageSpec.hs
@@ -7,9 +7,8 @@ import System.Directory
 import System.IO.Unsafe
 import Test.Hspec
 
-import Data.Torrent.Bitfield as BF
 import Data.Torrent
-import Data.Torrent.Piece
+import Data.Torrent.Bitfield as BF
 import System.Torrent.Storage