2 files changed, 41 insertions, 34 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index 7bffe4c1..63c2b494 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -114,10 +114,10 @@ contramapIR f ir = InsertionReporter
 -- | All the IO operations neccessary to maintain a Kademlia routing table.
 data TableStateIO nid ni = TableStateIO
     { -- | Write the routing table.  Typically 'writeTVar'.
-      tblWrite :: R.Table ni nid -> STM ()
+      tblWrite :: R.Table ni -> STM ()
 
       -- | Read the routing table.  Typically 'readTVar'.
-    , tblRead :: STM (R.Table ni nid)
+    , tblRead :: STM (R.Table ni)
 
       -- | Issue a ping to a remote node and report 'True' if the node
       -- responded within an acceptable time and 'False' otherwise.
@@ -136,7 +136,7 @@ data TableStateIO nid ni = TableStateIO
     , tblChanged :: RoutingTableChanged ni -> STM (IO ())
     }
 
-vanillaIO :: TVar (Table ni nid) -> (ni -> IO Bool) -> TableStateIO nid ni
+vanillaIO :: TVar (Table ni) -> (ni -> IO Bool) -> TableStateIO nid ni
 vanillaIO var ping = TableStateIO
     { tblRead = readTVar var
     , tblWrite = writeTVar var
diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index e2bc5aa9..0004ac09 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -19,6 +19,7 @@
 {-# LANGUAGE ViewPatterns    #-}
 {-# LANGUAGE TypeOperators   #-}
 {-# LANGUAGE DeriveGeneric   #-}
+{-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving, FlexibleContexts, MultiParamTypeClasses, FlexibleInstances #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
@@ -175,10 +176,13 @@ bucketQ = seqQ
 data Bucket ni = Bucket
     { bktNodes :: !(PSQ ni Timestamp)           -- current routing nodes
     , bktQ     :: !(BucketQueue (Timestamp,ni)) -- replacements pending time-outs
-    } deriving Generic
+    } deriving (Generic)
 
 deriving instance Show ni => Show (Bucket ni)
 
+mapBucket :: Ord ni => (a -> ni) -> Bucket a -> Bucket ni
+mapBucket f (Bucket ns q) = Bucket (PSQ.fromList $ map (\(ni :-> tm) -> (f ni :-> tm)) $ PSQ.toList ns)
+                                   (fmap (second f) q)
 
 
 #if 0
@@ -335,7 +339,7 @@ defaultBucketCount :: Int
 defaultBucketCount = 20
 
 data Info ni nid = Info
-    { myBuckets :: Table ni nid
+    { myBuckets :: Table ni
     , myNodeId :: nid
     , myAddress :: SockAddr
     }
@@ -362,18 +366,22 @@ deriving instance (Show ni, Show nid) => Show (Info ni nid)
 -- is always split into two new buckets covering the ranges @0..2 ^
 -- 159@ and @2 ^ 159..2 ^ 160@.
 --
-data Table ni nid
+data Table ni
   -- most nearest bucket
-  = Tip  nid Int (Bucket ni)
+  = Tip  ni Int (Bucket ni)
 
   -- left biased tree branch
-  | Zero (Table ni nid) (Bucket ni)
+  | Zero (Table ni) (Bucket ni)
 
   -- right biased tree branch
-  | One  (Bucket ni) (Table ni nid)
+  | One  (Bucket ni) (Table ni)
     deriving Generic
 
-instance (Eq ni, Eq nid) => Eq (Table ni nid) where
+mapTable f (One b t) = One (mapBucket f b) (mapTable f t)
+mapTable f (Zero t b) = Zero (mapTable f t) (mapBucket f b)
+mapTable f (Tip ni n b) = Tip (f ni) n (mapBucket f b)
+
+instance (Eq ni) => Eq (Table ni) where
   (==) = (==) `on` Network.DHT.Routing.toList
 
 #if 0
@@ -384,7 +392,7 @@ instance Serialize NominalDiffTime where
 
 #endif
 
-deriving instance (Show ni, Show nid) => Show (Table ni nid)
+deriving instance (Show ni) => Show (Table ni)
 
 #if 0
 
@@ -396,7 +404,7 @@ instance (Eq ip, Serialize ip, Ord (NodeId), Serialize (NodeId), Serialize u) =>
 #endif
 
 -- | Shape of the table.
-instance Pretty (Table ni nid) where
+instance Pretty (Table ni) where
   pPrint t
     | bucketCount < 6 = hcat $ punctuate ", " $ L.map PP.int ss
     |    otherwise    = brackets $
@@ -407,8 +415,8 @@ instance Pretty (Table ni nid) where
       ss = shape t
 
 -- | Empty table with specified /spine/ node id.
-nullTable :: nid -> Int -> Table ni nid
-nullTable nid n = Tip nid (bucketCount (pred n)) (Bucket PSQ.empty (runIdentity $ emptyQueue bucketQ))
+nullTable :: ni -> Int -> Table ni
+nullTable ni n = Tip ni (bucketCount (pred n)) (Bucket PSQ.empty (runIdentity $ emptyQueue bucketQ))
   where
     bucketCount x = max 0 (min 159 x)
 
@@ -440,7 +448,7 @@ type NodeCount   = Int
 
 -- | Internally, routing table is similar to list of buckets or a
 -- /matrix/ of nodes. This function returns the shape of the matrix.
-shape :: Table ni nid -> [Int]
+shape :: Table ni -> [Int]
 shape = map (PSQ.size . bktNodes) . toBucketList
 
 #if 0
@@ -457,7 +465,7 @@ depth = L.length . shape
 
 lookupBucket :: ( FiniteBits nid
                 , Ord nid
-                ) => nid -> Table ni nid -> [Bucket ni]
+                ) => nid -> Table ni -> [Bucket ni]
 lookupBucket nid = go 0 []
   where
     go i bs (Zero table bucket)
@@ -471,7 +479,7 @@ lookupBucket nid = go 0 []
 
 compatibleNodeId :: forall ni nid.
                     ( Serialize nid, FiniteBits nid) =>
-                    Table ni nid -> IO nid
+                    Table ni -> IO nid
 compatibleNodeId tbl = genBucketSample prefix br
  where
     br = bucketRange (L.length (shape tbl) - 1) True
@@ -479,7 +487,7 @@ compatibleNodeId tbl = genBucketSample prefix br
     bs = BS.pack $ take nodeIdSize $ tablePrefix tbl ++ repeat 0
     prefix = either error id $ S.decode bs
 
-tablePrefix :: Table ni nid -> [Word8]
+tablePrefix :: Table ni -> [Word8]
 tablePrefix = map (packByte . take 8 . (++repeat False))
               . chunksOf 8
               . tableBits
@@ -488,7 +496,7 @@ tablePrefix = map (packByte . take 8 . (++repeat False))
     bitmask ix True = bit ix
     bitmask _  _    = 0
 
-tableBits :: Table ni nid -> [Bool]
+tableBits :: Table ni -> [Bool]
 tableBits (One _ tbl)  = True : tableBits tbl
 tableBits (Zero tbl _) = False : tableBits tbl
 tableBits (Tip _ _ _)  = []
@@ -540,7 +548,7 @@ rank f nid = L.sortBy (comparing (distance nid . f))
 -- 'find_node' and 'get_peers' queries.
 kclosest :: ( FiniteBits nid
             , Ord nid
-            ) => (ni -> nid) -> Int -> nid -> Table ni nid -> [ni]
+            ) => (ni -> nid) -> Int -> nid -> Table ni -> [ni]
 kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
                               ++ rank nodeId nid (L.concat everyone)
  where
@@ -558,12 +566,11 @@ kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
 
 splitTip :: -- ( Eq ip , Ord (NodeId) , FiniteBits (NodeId)) =>
             Ord ni =>
-            (nid -> Word -> Bool)
-            -> (ni -> Word -> Bool)
-            -> nid -> Int -> BitIx -> Bucket ni -> Table ni nid
-splitTip testIdBit testNodeBit nid n i bucket
-  | testIdBit nid i = (One  zeros (Tip nid (pred n) ones))
-  |    otherwise    = (Zero (Tip nid (pred n) zeros) ones)
+            (ni -> Word -> Bool)
+            -> ni -> Int -> BitIx -> Bucket ni -> Table ni
+splitTip testNodeBit ni n i bucket
+  | testNodeBit ni i = (One  zeros (Tip ni (pred n) ones))
+  |    otherwise     = (Zero (Tip ni (pred n) zeros) ones)
   where
     (ones, zeros) = split testNodeBit i bucket
 
@@ -577,10 +584,10 @@ modifyBucket
     forall ni nid xs. Ord ni =>
     (nid -> Word -> Bool)
     -> (ni -> Word -> Bool)
-    -> nid -> (Bucket ni -> Maybe (xs, Bucket ni)) -> Table ni nid -> Maybe (xs,Table ni nid)
+    -> nid -> (Bucket ni -> Maybe (xs, Bucket ni)) -> Table ni -> Maybe (xs,Table ni)
 modifyBucket testIdBit testNodeBit nodeId f = go (0 :: BitIx)
   where
-    go :: BitIx -> Table ni nid -> Maybe (xs, Table ni nid)
+    go :: BitIx -> Table ni -> Maybe (xs, Table ni)
     go !i (Zero table  bucket)
       | testIdBit nodeId i  = second (Zero  table) <$> f bucket
       |     otherwise       = second (`Zero` bucket) <$> go (succ i) table
@@ -590,7 +597,7 @@ modifyBucket testIdBit testNodeBit nodeId f = go (0 :: BitIx)
     go !i (Tip nid n bucket)
       |        n == 0       = second (Tip nid n)    <$> f bucket
       |     otherwise       = second (Tip nid n)    <$> f bucket
-                           <|>  go i (splitTip testIdBit testNodeBit nid n i bucket)
+                           <|>  go i (splitTip testNodeBit nid n i bucket)
 
 
 -- | Triggering event for atomic table update
@@ -639,7 +646,7 @@ deriving instance ( Show ip
 --
 updateForInbound :: Ord ni =>
     KademliaSpace nid ni
-    -> Timestamp -> ni -> Table ni nid -> (Bool, [ni], Table ni nid)
+    -> Timestamp -> ni -> Table ni -> (Bool, [ni], Table ni)
 updateForInbound space tm ni tbl =
     maybe (False, [],tbl) (\(ps,tbl') -> (True, ps, tbl'))
                    $ modifyBucket (kademliaTestBit space)
@@ -656,8 +663,8 @@ updateForPingResult :: Ord ni =>
     KademliaSpace nid ni
     -> ni           -- ^ The pinged node.
     -> Bool         -- ^ True if we got a reply, False if it timed out.
-    -> Table ni nid -- ^ The routing table.
-    -> ( [(ni,(Timestamp, ni))], Table ni nid )
+    -> Table ni -- ^ The routing table.
+    -> ( [(ni,(Timestamp, ni))], Table ni )
 updateForPingResult space ni got_reply tbl =
     fromMaybe ([],tbl)
                    $ modifyBucket (kademliaTestBit space)
@@ -677,12 +684,12 @@ tableEntry :: NodeEntry ni -> TableEntry ni
 tableEntry (a :-> b) = (a, b)
 
 -- | Non-empty list of buckets.
-toBucketList :: Table ni nid -> [Bucket ni]
+toBucketList :: Table ni -> [Bucket ni]
 toBucketList (Tip _ _ b) = [b]
 toBucketList (Zero  t b) = b : toBucketList t
 toBucketList (One   b t) = b : toBucketList t
 
-toList :: Table ni nid -> [[TableEntry ni]]
+toList :: Table ni -> [[TableEntry ni]]
 toList = L.map (L.map tableEntry . PSQ.toList . bktNodes) . toBucketList
 
 data KademliaSpace nid ni = KademliaSpace