1 files changed, 31 insertions, 20 deletions

diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index d48fa8ac..19f9fc9e 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -13,14 +13,16 @@
 --   See <http://www.bittorrent.org/beps/bep_0005.html#krpc-protocol>
 --
 {-# LANGUAGE CPP                    #-}
-{-# LANGUAGE OverloadedStrings      #-}
+{-# LANGUAGE DefaultSignatures      #-}
+{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE DeriveFunctor          #-}
+{-# LANGUAGE DeriveTraversable      #-}
 {-# LANGUAGE FlexibleContexts       #-}
 {-# LANGUAGE FlexibleInstances      #-}
-{-# LANGUAGE TypeSynonymInstances   #-}
-{-# LANGUAGE MultiParamTypeClasses  #-}
 {-# LANGUAGE FunctionalDependencies #-}
-{-# LANGUAGE DefaultSignatures      #-}
-{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE MultiParamTypeClasses  #-}
+{-# LANGUAGE OverloadedStrings      #-}
+{-# LANGUAGE TypeSynonymInstances   #-}
 module Network.KRPC.Message
        ( -- * Transaction
          TransactionId
@@ -32,7 +34,8 @@ module Network.KRPC.Message
        , unknownMessage
 
          -- * Query
-       , KQuery(..)
+       , KQueryOf(..)
+       , KQuery
 #ifndef VERSION_bencoding
        , queryArgs
        , queryMethod
@@ -41,11 +44,13 @@ module Network.KRPC.Message
        , MethodName
 
          -- * Response
-       , KResponse(..)
+       , KResponseOf(..)
+       , KResponse
        , ReflectedIP(..)
 
          -- * Message
-       , KMessage (..)
+       , KMessageOf (..)
+       , KMessage
        , KQueryArgs
 
        ) where
@@ -208,11 +213,13 @@ type KQueryArgs = ByteString      -- msgPayload
 -- callee and pass arguments in. Therefore query may be only sent from
 -- client to server but not in the opposite direction.
 --
-data KQuery = KQuery
-  { queryArgs   :: !KQueryArgs     -- ^ values to be passed to method;
+data KQueryOf a = KQuery
+  { queryArgs   :: !a              -- ^ values to be passed to method;
   , queryMethod :: !MethodName     -- ^ method to call;
   , queryId     :: !TransactionId  -- ^ one-time query token.
-  } deriving ( Show, Eq, Ord, Typeable, Read )
+  } deriving ( Show, Eq, Ord, Typeable, Read, Functor, Foldable, Traversable )
+
+type KQuery = KQueryOf KQueryArgs
 
 -- | Queries, or KRPC message dictionaries with a \"y\" value of
 -- \"q\", contain two additional keys; \"q\" and \"a\". Key \"q\" has
@@ -223,7 +230,7 @@ data KQuery = KQuery
 --
 --    > { "t" : "aa", "y" : "q", "q" : "ping", "a" : { "msg" : "hi!" } }
 --
-instance BEncode KQuery where
+instance (Typeable a, BEncode a) => BEncode (KQueryOf a) where
   toBEncode KQuery {..} = toDict $
        "a" .=! queryArgs
     .: "q" .=! queryMethod
@@ -288,11 +295,13 @@ encodeAddr _ = B.empty
 --   * KResponse can be only sent from server to client.
 --
 #ifdef VERSION_bencoding
-data KResponse = KResponse
-  { respVals :: KQueryArgs     -- ^ 'BDict' containing return values;
+data KResponseOf a = KResponse
+  { respVals :: a              -- ^ 'BDict' containing return values;
   , respId   :: TransactionId  -- ^ match to the corresponding 'queryId'.
   , respIP   :: Maybe ReflectedIP
-  } deriving (Show, Eq, Ord, Typeable)
+  } deriving (Show, Eq, Ord, Typeable, Functor, Foldable, Traversable)
+
+type KResponse = KResponseOf KQueryArgs
 
 -- | Responses, or KRPC message dictionaries with a \"y\" value of
 -- \"r\", contain one additional key \"r\". The value of \"r\" is a
@@ -302,7 +311,7 @@ data KResponse = KResponse
 --
 --   > { "t" : "aa", "y" : "r", "r" : { "msg" : "you've sent: hi!" } }
 --
-instance BEncode KResponse where
+instance (Typeable a, BEncode a) => BEncode (KResponseOf a) where
   toBEncode KResponse {..} = toDict $
        "ip" .=? respIP
     .: "r" .=! respVals
@@ -329,11 +338,13 @@ respIP         = Nothing :: Maybe ReflectedIP
 
 #ifdef VERSION_bencoding
 -- | Generic KRPC message.
-data KMessage
-  = Q KQuery
-  | R KResponse
+data KMessageOf a
+  = Q (KQueryOf a)
+  | R (KResponseOf a)
   | E KError
-    deriving (Show, Eq)
+    deriving (Show, Eq, Functor, Foldable, Traversable)
+
+type KMessage = KMessageOf KQueryArgs
 
 instance BEncode KMessage where
   toBEncode (Q q) = toBEncode q