1 files changed, 72 insertions, 23 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index f6f2807d..4fc9438c 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -1,30 +1,36 @@
 -- | This module can implement any query\/response protocol.  It was written
 -- with Kademlia implementations in mind.
 
-{-# LANGUAGE CPP #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE CPP                   #-}
+{-# LANGUAGE GADTs                 #-}
+{-# LANGUAGE LambdaCase            #-}
 {-# LANGUAGE PartialTypeSignatures #-}
-{-# LANGUAGE GADTs #-}
+{-# LANGUAGE RankNTypes            #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
 module Network.QueryResponse where
 
 #ifdef THREAD_DEBUG
 import Control.Concurrent.Lifted.Instrument
 #else
-import GHC.Conc (labelThread)
 import Control.Concurrent
+import GHC.Conc           (labelThread)
 #endif
 import Control.Concurrent.STM
-import System.Timeout
-import Data.Function
 import Control.Exception
 import Control.Monad
 import qualified Data.ByteString as B
          ;import Data.ByteString (ByteString)
+import Data.Function
+import qualified Data.IntMap     as IntMap
+         ;import Data.IntMap     (IntMap)
+import qualified Data.Map        as Map
+         ;import Data.Map        (Map)
+import Data.Maybe
+import Data.Typeable
 import Network.Socket
 import Network.Socket.ByteString as B
 import System.IO.Error
-import Data.Maybe
+import System.Timeout
 
 -- * Using a query\/response 'Client'.
 
@@ -184,12 +190,12 @@ layerTransport parse encode tr =
 
 
 -- | To dipatch responses to our outbound queries, we require three primitives.
--- See the 'transactionTableMethods' function to create these primitives out of a
+-- See the 'transactionMethods' function to create these primitives out of a
 -- lookup table and a generator for transaction ids.
 --
 -- The type variable /d/ is used to represent the current state of the
 -- transaction generator and the table of pending transactions.
-data TableMethods d tid x = TableMethods
+data TransactionMethods d tid x = TransactionMethods
     {
       -- | Before a query is sent, this function stores an 'MVar' to which the
       -- response will be written too.  The returned _tid_ is a transaction id
@@ -206,19 +212,55 @@ data TableMethods d tid x = TableMethods
     , dispatchCancel :: tid -> d -> d
     }
 
--- | Construct 'TableMethods' methods out of 3 lookup table primitives and a
+-- | The standard lookup table methods for use as input to 'transactionMethods'
+-- in lieu of directly implementing 'TransactionMethods'.
+data TableMethods t tid = TableMethods
+    { -- | Insert a new _tid_ entry into the transaction table.
+      tblInsert :: forall a. tid -> a -> t a -> t a
+      -- | Delete transaction _tid_ from the transaction table.
+    , tblDelete :: forall a. tid -> t a -> t a
+       -- | Lookup the value associated with transaction /tid/.
+    , tblLookup :: forall a. tid -> t a -> Maybe a
+    }
+
+-- | Methods for using 'Data.IntMap.
+intMapMethods :: TableMethods IntMap Int
+intMapMethods = TableMethods IntMap.insert IntMap.delete IntMap.lookup
+
+-- | Methods for using 'Data.Map'
+mapMethods :: Ord tid => TableMethods (Map tid) tid
+mapMethods = TableMethods Map.insert Map.delete Map.lookup
+
+-- | Change the key type for a lookup table implementation.
+--
+-- This can be used with 'intMapMethods' or 'mapMethods' to restrict lookups to
+-- only a part of the generated /tid/ value.  This is useful for /tid/ types
+-- that are especially large due their use for other purposes, such as secure
+-- nonces for encryption.
+contramapT f (TableMethods ins del lookup) =
+    TableMethods (\k v t -> ins (f k) v t)
+                 (\k t   -> del (f k) t)
+                 (\k t   -> lookup (f k) t)
+
+-- | Since 'Int' may be 32 or 64 bits, this function is provided as a
+-- convenience to test if an integral type, such as 'Data.Word.Word64', can be
+-- safely transformed into an 'Int' for use with 'IntMap'.
+--
+-- Returns 'True' if the proxied type can be losslessly convered to 'Int' using
+-- 'fromIntegral'.
+fitsInInt :: forall word. (Bounded word, Integral word) => Proxy word -> Bool
+fitsInInt Proxy = (original == casted)
+ where
+    original = div maxBound 2 :: word
+    casted   = fromIntegral (fromIntegral original :: Int) :: word
+
+-- | Construct 'TransactionMethods' methods out of 3 lookup table primitives and a
 -- function for generating unique transaction ids.
-transactionTableMethods ::
-    (forall a. tid -> a -> t a -> t a)
-    -- ^ Insert a new _tid_ entry into the transaction table.
-    -> (forall a. tid -> t a -> t a)
-    -- ^ Delete transaction _tid_ from the transaction table.
-    -> (forall a. tid -> t a -> Maybe a)
-    -- ^ Lookup the value associated with transaction _tid_.
-    -> (g -> (tid,g))
-    -- ^ Generate a new unique _tid_ value and update the generator state _g_.
-    -> TableMethods (g,t (MVar x)) tid x
-transactionTableMethods insert delete lookup generate = TableMethods
+transactionMethods ::
+    TableMethods t tid  -- ^ Table methods to lookup values by /tid/.
+    -> (g -> (tid,g))   -- ^ Generate a new unique /tid/ value and update the generator state /g/.
+    -> TransactionMethods (g,t (MVar x)) tid x
+transactionMethods (TableMethods insert delete lookup) generate = TransactionMethods
     { dispatchCancel = \tid (g,t) -> (g, delete tid t)
     , dispatchRegister = \v (g,t) ->
         let (tid,g') = generate g
@@ -238,7 +280,7 @@ data DispatchMethods tbl err meth tid addr x ctx = DispatchMethods
       -- | Lookup the handler for a inbound query.
     , lookupHandler :: meth -> Maybe (MethodHandler err tid addr x ctx)
       -- | Methods for handling incomming responses.
-    , tableMethods :: TableMethods tbl tid x
+    , tableMethods :: TransactionMethods tbl tid x
     }
 
 -- | These methods indicate what should be done upon various conditions.  Write
@@ -264,6 +306,13 @@ data ErrorReporter addr x meth tid err = ErrorReporter
     , reportTimeout :: meth -> tid -> addr -> IO ()
     }
 
+-- Change the /err/ type for an 'ErrorReporter'.
+contramapE f (ErrorReporter pe mh unk tim)
+    = ErrorReporter (\e -> pe (f e))
+                    mh
+                    (\addr x e -> unk addr x (f e))
+                    tim
+
 -- | Handle a single inbound packet and then invoke the given continuation.
 -- The 'forkListener' function is implemeneted by passing this function to
 -- 'fix' in a forked thread that loops until 'awaitMessage' returns 'Nothing'