2 files changed, 97 insertions, 88 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index 63c2b494..4a811fa2 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 -> STM ()
+      tblWrite :: R.BucketList ni -> STM ()
 
       -- | Read the routing table.  Typically 'readTVar'.
-    , tblRead :: STM (R.Table ni)
+    , tblRead :: STM (R.BucketList 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) -> (ni -> IO Bool) -> TableStateIO nid ni
+vanillaIO :: TVar (BucketList ni) -> (ni -> IO Bool) -> TableStateIO nid ni
 vanillaIO var ping = TableStateIO
     { tblRead = readTVar var
     , tblWrite = writeTVar var
@@ -152,7 +152,7 @@ data Kademlia nid ni = Kademlia (InsertionReporter ni)
 
 {-
 kademlia :: FiniteBits nid =>
-            TVar (Table nid nid) -> (nid -> IO Bool) -> Kademlia nid nid
+            TVar (BucketList nid nid) -> (nid -> IO Bool) -> Kademlia nid nid
 kademlia var ping = Kademlia quietInsertions
                              (KademliaSpace id testIdBit)
                              (vanillaIO var ping)
diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 0004ac09..1821ca1c 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -25,8 +25,8 @@
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 module Network.DHT.Routing
        {-
-       ( -- * Table
-         Table
+       ( -- * BucketList
+         BucketList
        , Info(..)
 
          -- * Attributes
@@ -86,8 +86,10 @@ import Text.PrettyPrint.HughesPJClass (pPrint,Pretty)
 import qualified Data.ByteString as BS
 import Data.Bits
 import Data.Ord
-
+import Data.Reflection
 import Network.Address
+import Data.Typeable
+import Data.Coerce
 
 -- | Last time the node was responding to our queries.
 --
@@ -167,6 +169,20 @@ bucketQ :: QueueMethods Identity ni (BucketQueue ni)
 bucketQ = seqQ
 
 
+newtype Compare a = Compare (a -> a -> Ordering)
+
+newtype Ordered s a = Ordered a
+
+-- | Hack to avoid UndecidableInstances
+newtype Shrink a = Shrink a
+
+instance Reifies s (Compare a) => Eq (Ordered s (Shrink a)) where
+    a == b   = (compare a b == EQ)
+
+instance Reifies s (Compare a) => Ord (Ordered s (Shrink a)) where
+    compare a b = cmp (coerce a) (coerce b)
+     where Compare cmp = reflect (Proxy :: Proxy s)
+
 -- | Bucket is also limited in its length — thus it's called k-bucket.
 --   When bucket becomes full, we should split it in two lists by
 --   current span bit. Span bit is defined by depth in the routing
@@ -332,14 +348,14 @@ split testNodeIdBit i b = (Bucket ns qs, Bucket ms rs)
 
 
 {-----------------------------------------------------------------------
---  Table
+--  BucketList
 -----------------------------------------------------------------------}
 
 defaultBucketCount :: Int
 defaultBucketCount = 20
 
 data Info ni nid = Info
-    { myBuckets :: Table ni
+    { myBuckets :: BucketList ni
     , myNodeId :: nid
     , myAddress :: SockAddr
     }
@@ -366,22 +382,13 @@ 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
-  -- most nearest bucket
-  = Tip  ni Int (Bucket ni)
-
-  -- left biased tree branch
-  | Zero (Table ni) (Bucket ni)
-
-  -- right biased tree branch
-  | One  (Bucket ni) (Table ni)
+data BucketList ni = BucketList !ni !Int ![Bucket ni]
     deriving Generic
 
-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)
+mapTable :: Ord ni => (a -> ni) -> BucketList a -> BucketList ni
+mapTable f (BucketList self n bs) = BucketList (f self) n (map (mapBucket f) bs)
 
-instance (Eq ni) => Eq (Table ni) where
+instance (Eq ni) => Eq (BucketList ni) where
   (==) = (==) `on` Network.DHT.Routing.toList
 
 #if 0
@@ -392,19 +399,19 @@ instance Serialize NominalDiffTime where
 
 #endif
 
-deriving instance (Show ni) => Show (Table ni)
+deriving instance (Show ni) => Show (BucketList ni)
 
 #if 0
 
 -- | Normally, routing table should be saved between invocations of
 -- the client software. Note that you don't need to store /this/
 -- 'NodeId' since it is already included in routing table.
-instance (Eq ip, Serialize ip, Ord (NodeId), Serialize (NodeId), Serialize u) => Serialize (Table)
+instance (Eq ip, Serialize ip, Ord (NodeId), Serialize (NodeId), Serialize u) => Serialize (BucketList)
 
 #endif
 
 -- | Shape of the table.
-instance Pretty (Table ni) where
+instance Pretty (BucketList ni) where
   pPrint t
     | bucketCount < 6 = hcat $ punctuate ", " $ L.map PP.int ss
     |    otherwise    = brackets $
@@ -415,8 +422,11 @@ instance Pretty (Table ni) where
       ss = shape t
 
 -- | Empty table with specified /spine/ node id.
-nullTable :: ni -> Int -> Table ni
-nullTable ni n = Tip ni (bucketCount (pred n)) (Bucket PSQ.empty (runIdentity $ emptyQueue bucketQ))
+nullTable :: ni -> Int -> BucketList ni
+nullTable ni n = BucketList
+                    ni
+                    (bucketCount (pred n))
+                    [Bucket PSQ.empty (runIdentity $ emptyQueue bucketQ)]
   where
     bucketCount x = max 0 (min 159 x)
 
@@ -424,19 +434,19 @@ nullTable ni n = Tip ni (bucketCount (pred n)) (Bucket PSQ.empty (runIdentity $
 
 -- | Test if table is empty. In this case DHT should start
 -- bootstrapping process until table becomes 'full'.
-null :: Table -> Bool
+null :: BucketList -> Bool
 null (Tip _ _ b) = PSQ.null $ bktNodes b
 null  _          = False
 
 -- | Test if table have maximum number of nodes. No more nodes can be
 -- 'insert'ed, except old ones becomes bad.
-full :: Table -> Bool
+full :: BucketList -> Bool
 full (Tip  _ n _) = n == 0
 full (Zero   t b) = PSQ.size (bktNodes b) == defaultBucketSize && full t
 full (One    b t) = PSQ.size (bktNodes b) == defaultBucketSize && full t
 
 -- | Get the /spine/ node id.
-thisId :: Table -> NodeId
+thisId :: BucketList -> NodeId
 thisId (Tip  nid _ _) = nid
 thisId (Zero table _) = thisId table
 thisId (One _  table) = thisId table
@@ -448,58 +458,57 @@ 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 -> [Int]
-shape = map (PSQ.size . bktNodes) . toBucketList
+shape :: BucketList ni -> [Int]
+shape = map (PSQ.size . bktNodes) . buckets
 
 #if 0
 
 -- | Get number of nodes in the table.
-size :: Table -> NodeCount
+size :: BucketList -> NodeCount
 size = L.sum . shape
 
 -- | Get number of buckets in the table.
-depth :: Table -> BucketCount
+depth :: BucketList -> BucketCount
 depth = L.length . shape
 
 #endif
 
 lookupBucket :: ( FiniteBits nid
                 , Ord nid
-                ) => nid -> Table ni -> [Bucket ni]
-lookupBucket nid = go 0 []
+                ) => (ni -> nid) -> nid -> BucketList ni -> [Bucket ni]
+lookupBucket nodeId nid (BucketList self _ bkts) = go 0 [] bkts
   where
-    go i bs (Zero table bucket)
-      |  testIdBit nid i     = bucket : toBucketList table ++ bs
-      |     otherwise        = go (succ i) (bucket:bs) table
-    go i bs (One  bucket table)
-      |  testIdBit nid i     = go (succ i) (bucket:bs) table
-      |     otherwise        = bucket : toBucketList table ++ bs
-    go _ bs (Tip _ _ bucket) = bucket : bs
+    d = complement (nid `xor` nodeId self)
+
+    go i bs (bucket : buckets)
+      |  testIdBit d i = bucket : buckets ++ bs
+      |     otherwise  = go (succ i) (bucket:bs) buckets
+    go _ bs []         = bs
 
 
 compatibleNodeId :: forall ni nid.
                     ( Serialize nid, FiniteBits nid) =>
-                    Table ni -> IO nid
-compatibleNodeId tbl = genBucketSample prefix br
+                    (ni -> nid) -> BucketList ni -> IO nid
+compatibleNodeId nodeId tbl = genBucketSample prefix br
  where
     br = bucketRange (L.length (shape tbl) - 1) True
     nodeIdSize = finiteBitSize (undefined :: nid) `div` 8
-    bs = BS.pack $ take nodeIdSize $ tablePrefix tbl ++ repeat 0
+    bs = BS.pack $ take nodeIdSize $ tablePrefix (testIdBit . nodeId) tbl ++ repeat 0
     prefix = either error id $ S.decode bs
 
-tablePrefix :: Table ni -> [Word8]
-tablePrefix = map (packByte . take 8 . (++repeat False))
-              . chunksOf 8
-              . tableBits
+tablePrefix :: (ni -> Word -> Bool) -> BucketList ni -> [Word8]
+tablePrefix testbit = map (packByte . take 8 . (++repeat False))
+                      . chunksOf 8
+                      . tableBits testbit
  where
     packByte = foldl1' (.|.) . zipWith bitmask [7,6 .. 0]
     bitmask ix True = bit ix
     bitmask _  _    = 0
 
-tableBits :: Table ni -> [Bool]
-tableBits (One _ tbl)  = True : tableBits tbl
-tableBits (Zero tbl _) = False : tableBits tbl
-tableBits (Tip _ _ _)  = []
+tableBits :: (ni -> Word -> Bool) -> BucketList ni -> [Bool]
+tableBits testbit (BucketList self _ bkts) =
+    zipWith const (map (testbit self) [0..])
+                  bkts
 
 chunksOf :: Int -> [e] -> [[e]]
 chunksOf i ls = map (take i) (build (splitter ls)) where
@@ -548,14 +557,14 @@ 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 -> [ni]
+            ) => (ni -> nid) -> Int -> nid -> BucketList ni -> [ni]
 kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
                               ++ rank nodeId nid (L.concat everyone)
  where
     (bucket,everyone) =
         L.splitAt 1
         . L.map (L.map PSQ.key . PSQ.toList . bktNodes)
-        . lookupBucket nid
+        . lookupBucket nodeId nid
         $ tbl
 
 
@@ -567,10 +576,10 @@ kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
 splitTip :: -- ( Eq ip , Ord (NodeId) , FiniteBits (NodeId)) =>
             Ord ni =>
             (ni -> Word -> Bool)
-            -> ni -> Int -> BitIx -> Bucket ni -> Table ni
+            -> ni -> Int -> BitIx -> Bucket ni -> [ Bucket 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)
+  | testNodeBit ni i = [zeros , ones ]
+  |    otherwise     = [ones , zeros ]
   where
     (ones, zeros) = split testNodeBit i bucket
 
@@ -581,23 +590,25 @@ splitTip testNodeBit ni n i bucket
 -- paper.  The rule requiring additional splits is in section 2.4.
 modifyBucket
   :: -- ( Eq ip , Ord (NodeId) , FiniteBits (NodeId)) =>
-    forall ni nid xs. Ord ni =>
+    forall ni nid xs. (Ord ni, Bits nid) =>
     (nid -> Word -> Bool)
-    -> (ni -> Word -> Bool)
-    -> nid -> (Bucket ni -> Maybe (xs, Bucket ni)) -> Table ni -> Maybe (xs,Table ni)
-modifyBucket testIdBit testNodeBit nodeId f = go (0 :: BitIx)
+    -> (ni -> nid)
+    -> nid -> (Bucket ni -> Maybe (xs, Bucket ni)) -> BucketList ni -> Maybe (xs,BucketList ni)
+modifyBucket testIdBit nodeId nid f (BucketList self n bkts)
+    = second (BucketList self n) <$> go (0 :: BitIx) bkts
   where
-    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
-    go !i (One  bucket table )
-      | testIdBit nodeId i  = second (One  bucket)   <$> go (succ i) table
-      |     otherwise       = second (`One` table)   <$> f bucket
-    go !i (Tip nid n bucket)
-      |        n == 0       = second (Tip nid n)    <$> f bucket
-      |     otherwise       = second (Tip nid n)    <$> f bucket
-                           <|>  go i (splitTip testNodeBit nid n i bucket)
+    d = nid `xor` nodeId self
+
+    go :: BitIx -> [Bucket ni] -> Maybe (xs, [Bucket ni])
+
+    go !i (bucket : buckets@(_:_))
+      | testIdBit d i   = second (bucket :)  <$> go (succ i) buckets
+      |       otherwise = second (: buckets) <$> f bucket
+
+    go !i [bucket]
+      |      (n == 0) = second (: []) <$> f bucket
+      |     otherwise = second (: []) <$> f bucket
+                        <|>  go i (splitTip (testIdBit . nodeId) self n i bucket)
 
 
 -- | Triggering event for atomic table update
@@ -642,15 +653,15 @@ deriving instance ( Show ip
 --                     may be added to a replacement queue and will be inserted if
 --                     one of the items in this list time out.
 --
---  [ /tbl'/ ]         The updated routing 'Table'.
+--  [ /tbl'/ ]         The updated routing 'BucketList'.
 --
-updateForInbound :: Ord ni =>
+updateForInbound :: (Ord ni, Bits nid) =>
     KademliaSpace nid ni
-    -> Timestamp -> ni -> Table ni -> (Bool, [ni], Table ni)
+    -> Timestamp -> ni -> BucketList ni -> (Bool, [ni], BucketList ni)
 updateForInbound space tm ni tbl =
     maybe (False, [],tbl) (\(ps,tbl') -> (True, ps, tbl'))
                    $ modifyBucket (kademliaTestBit space)
-                                   (\ni -> kademliaTestBit space $ kademliaLocation space ni)
+                                   (kademliaLocation space)
                                    (kademliaLocation space ni)
                                    (updateBucketForInbound tm ni)
                                    tbl
@@ -659,16 +670,16 @@ updateForInbound space tm ni tbl =
 --
 -- Each (a,(tm,b)) in the returned list indicates that the node /a/ was deleted from the
 -- routing table and the  node /b/, with timestamp /tm/, has taken its place.
-updateForPingResult :: Ord ni =>
+updateForPingResult :: (Ord ni, Bits nid) =>
     KademliaSpace nid ni
     -> ni           -- ^ The pinged node.
     -> Bool         -- ^ True if we got a reply, False if it timed out.
-    -> Table ni -- ^ The routing table.
-    -> ( [(ni,(Timestamp, ni))], Table ni )
+    -> BucketList ni -- ^ The routing table.
+    -> ( [(ni,(Timestamp, ni))], BucketList ni )
 updateForPingResult space ni got_reply tbl =
     fromMaybe ([],tbl)
                    $ modifyBucket (kademliaTestBit space)
-                                  (\ni -> kademliaTestBit space $ kademliaLocation space ni)
+                                  (kademliaLocation space)
                                   (kademliaLocation space ni)
                                   (updateBucketForPingResult ni got_reply)
                                   tbl
@@ -684,13 +695,11 @@ tableEntry :: NodeEntry ni -> TableEntry ni
 tableEntry (a :-> b) = (a, b)
 
 -- | Non-empty list of buckets.
-toBucketList :: Table ni -> [Bucket ni]
-toBucketList (Tip _ _ b) = [b]
-toBucketList (Zero  t b) = b : toBucketList t
-toBucketList (One   b t) = b : toBucketList t
+buckets :: BucketList ni -> [Bucket ni]
+buckets (BucketList _ _ bs) = bs
 
-toList :: Table ni -> [[TableEntry ni]]
-toList = L.map (L.map tableEntry . PSQ.toList . bktNodes) . toBucketList
+toList :: BucketList ni -> [[TableEntry ni]]
+toList = L.map (L.map tableEntry . PSQ.toList . bktNodes) . buckets
 
 data KademliaSpace nid ni = KademliaSpace
     { -- | Given a node record (probably including IP address), yields a