2 files changed, 28 insertions, 21 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index b7a402b4..8c793a1a 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -22,6 +22,7 @@ import Control.Monad
 import qualified Data.ByteString    as B
          ;import Data.ByteString    (ByteString)
 import Data.Function
+import Data.Functor.Contravariant
 import qualified Data.IntMap.Strict as IntMap
          ;import Data.IntMap.Strict (IntMap)
 import qualified Data.Map.Strict    as Map
@@ -160,7 +161,7 @@ sendQuery (Client net d err pending whoami _) meth q addr = do
     mvar <- newEmptyMVar
     tid <- atomically $ do
         tbl <- readTVar pending
-        let (tid, tbl') = dispatchRegister (tableMethods d) mvar tbl
+        (tid, tbl') <- dispatchRegister (tableMethods d) mvar tbl
         writeTVar pending tbl'
         return tid
     self <- whoami (Just addr)
@@ -169,7 +170,7 @@ sendQuery (Client net d err pending whoami _) meth q addr = do
     case mres of
         Just x -> return $ Just $ unwrapResponse meth x
         Nothing -> do
-            atomically $ modifyTVar' pending (dispatchCancel (tableMethods d) tid)
+            atomically $ readTVar pending >>= dispatchCancel (tableMethods d) tid >>= writeTVar pending
             reportTimeout err (method meth) tid addr
             return Nothing
 
@@ -278,15 +279,15 @@ data TransactionMethods d tid x = TransactionMethods
       -- response will be written too.  The returned /tid/ is a transaction id
       -- that can be used to forget the 'MVar' if the remote peer is not
       -- responding.
-      dispatchRegister :: MVar x -> d -> (tid, d)
+      dispatchRegister :: MVar x -> d -> STM (tid, d)
       -- | This method is invoked when an incomming packet /x/ indicates it is
       -- a response to the transaction with id /tid/.  The returned IO action
       -- is will write the packet to the correct 'MVar' thus completing the
       -- dispatch.
-    , dispatchResponse :: tid -> x -> d -> (d, IO ())
+    , dispatchResponse :: tid -> x -> d -> STM (d, IO ())
       -- | When a timeout interval elapses, this method is called to remove the
       -- transaction from the table.
-    , dispatchCancel :: tid -> d -> d
+    , dispatchCancel :: tid -> d -> STM d
     }
 
 -- | The standard lookup table methods for use as input to 'transactionMethods'
@@ -314,10 +315,12 @@ mapMethods = TableMethods Map.insert Map.delete Map.lookup
 -- 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)
+instance Contravariant (TableMethods t) where
+    -- contramap :: (tid -> t1) -> TableMethods t t1 -> TableMethods t tid
+    contramap 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
@@ -338,16 +341,16 @@ transactionMethods ::
     -> (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)
+    { dispatchCancel = \tid (g,t) -> return (g, delete tid t)
     , dispatchRegister = \v (g,t) ->
         let (tid,g') = generate g
             t' = insert tid v t
-        in ( tid, (g',t') )
+        in return ( tid, (g',t') )
     , dispatchResponse = \tid x (g,t) ->
         case lookup tid t of
             Just v -> let t' = delete tid t
-                      in ((g,t'),void $ tryPutMVar v x)
-            Nothing -> ((g,t), return ())
+                      in return ((g,t'),void $ tryPutMVar v x)
+            Nothing -> return ((g,t), return ())
     }
 
 -- | A set of methods neccessary for dispatching incomming packets.
@@ -402,11 +405,13 @@ printErrors h = ErrorReporter
     }
 
 -- 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
+instance Contravariant (ErrorReporter addr x meth tid) where
+    -- contramap :: (t5 -> t4) -> ErrorReporter t3 t2 t1 t t4 -> ErrorReporter t3 t2 t1 t t5
+    contramap 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
@@ -436,7 +441,7 @@ handleMessage (Client net d err pending whoami responseID) addr plain = do
             IsResponse tid -> do
                 action <- atomically $ do
                     ts0 <- readTVar pending
-                    let (ts, action) = dispatchResponse (tableMethods d) tid plain ts0
+                    (ts, action) <- dispatchResponse (tableMethods d) tid plain ts0
                     writeTVar pending ts
                     return action
                 action
@@ -458,6 +463,7 @@ sockAddrFamily (SockAddrCan _        ) = AF_CAN
 -- | Packets with an empty payload may trigger eof exception.
 -- 'udpTransport' uses this function to avoid throwing in that
 -- case.
+ignoreEOF :: a -> IOError -> IO a
 ignoreEOF def e | isEOFError e = pure def
                 | otherwise    = throwIO e
 
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 7893d84a..8df1a09d 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -46,6 +46,7 @@ import qualified Data.ByteString.Char8        as C8
 import Data.ByteString.Lazy                   (toStrict)
 import Data.Char
 import Data.Data
+import Data.Functor.Contravariant
 import Data.Hashable
 import Data.IP
 import Data.Maybe
@@ -165,11 +166,11 @@ newClient drg net classify selfAddr handlers modifynet = do
     tblvar <-
       if fitsInInt (Proxy :: Proxy Word64)
         then do
-            let intmapT = transactionMethods (contramapT intKey intMapMethods) gen
+            let intmapT = transactionMethods (contramap intKey intMapMethods) gen
             intmap_var <- atomically $ newTVar (drg, mempty)
             return $ Right (intmapT,intmap_var)
          else do
-            let mapT = transactionMethods (contramapT nonceKey mapMethods) gen
+            let mapT = transactionMethods (contramap nonceKey mapMethods) gen
             map_var <- atomically $ newTVar (drg, mempty)
             return $ Left (mapT,map_var)
     let dispatch tbl var handlers = DispatchMethods