1 files changed, 88 insertions, 62 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 1821ca1c..58e0cbd3 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -16,10 +16,12 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE BangPatterns    #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ViewPatterns    #-}
 {-# LANGUAGE TypeOperators   #-}
 {-# LANGUAGE DeriveGeneric   #-}
 {-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving, FlexibleContexts, MultiParamTypeClasses, FlexibleInstances #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
@@ -171,15 +173,22 @@ bucketQ = seqQ
 
 newtype Compare a = Compare (a -> a -> Ordering)
 
-newtype Ordered s a = Ordered a
+contramapC :: (b -> a) -> Compare a -> Compare b
+contramapC f (Compare cmp) = Compare $ \a b -> cmp (f a) (f b)
+
+newtype Ordered' s a = Ordered a
+ deriving (Show)
 
 -- | Hack to avoid UndecidableInstances
 newtype Shrink a = Shrink a
+ deriving (Show)
+
+type Ordered s a = Ordered' s (Shrink a)
 
-instance Reifies s (Compare a) => Eq (Ordered s (Shrink a)) where
+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
+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)
 
@@ -189,16 +198,20 @@ instance Reifies s (Compare a) => Ord (Ordered s (Shrink a)) where
 --   table tree. Size of the bucket should be choosen such that it's
 --   very unlikely that all nodes in bucket fail within an hour of
 --   each other.
-data Bucket ni = Bucket
-    { bktNodes :: !(PSQ ni Timestamp)           -- current routing nodes
+data Bucket s ni = Bucket
+    { bktNodes :: !(PSQ (Ordered s ni) Timestamp)           -- current routing nodes
     , bktQ     :: !(BucketQueue (Timestamp,ni)) -- replacements pending time-outs
     } deriving (Generic)
 
-deriving instance Show ni => Show (Bucket ni)
+deriving instance Show ni => Show (Bucket s ni)
+
+bucketCompare :: forall p ni s. Reifies s (Compare ni) => p (Bucket s ni) -> Compare ni
+bucketCompare _ = reflect (Proxy :: Proxy s)
 
-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)
+mapBucket :: Reifies t (Compare ni) => (a -> ni) -> Bucket s a -> Bucket t ni
+mapBucket f (Bucket ns q) = Bucket (PSQ.fromList $ map (\(ni :-> tm) -> (f' ni :-> tm)) $ PSQ.toList ns)
                                    (fmap (second f) q)
+ where f' = coerce . f . coerce
 
 
 #if 0
@@ -255,10 +268,10 @@ updateBucketForInbound curTime info bucket
   -- timestamp updated here, since 'TryInsert' is called on every inbound packet,
   -- including ping results.
   | already_have
-    = pure ( [], map_ns $ PSQ.insertWith max info curTime )
+    = pure ( [], map_ns $ PSQ.insertWith max (coerce info) curTime )
   -- bucket is good, but not full => we can insert a new node
   | PSQ.size (bktNodes bucket) < defaultBucketSize
-    = pure ( [], map_ns $ PSQ.insert info curTime )
+    = pure ( [], map_ns $ PSQ.insert (coerce info) curTime )
   -- If there are any questionable nodes in the bucket have not been
   -- seen in the last 15 minutes, the least recently seen node is
   -- pinged. If any nodes in the bucket are known to have become bad,
@@ -267,7 +280,7 @@ updateBucketForInbound curTime info bucket
   | not (L.null stales)
     = pure ( stales
            , bucket { -- Update timestamps so that we don't redundantly ping.
-                      bktNodes = updateStamps curTime stales $ bktNodes bucket
+                      bktNodes = updateStamps curTime (coerce stales) $ bktNodes bucket
                       -- Update queue with the pending NodeInfo in case of ping fail.
                     , bktQ     = runIdentity $ pushBack bucketQ (curTime,info) $ bktQ bucket } )
   -- When the bucket is full of good nodes, the new node is simply discarded.
@@ -281,11 +294,11 @@ updateBucketForInbound curTime info bucket
     stales = -- One stale:
              do (n :-> t) <- maybeToList $ PSQ.findMin (bktNodes bucket)
                 guard (t < curTime - delta)
-                return n
+                return $ coerce n
              -- All stale:
              -- map key \$ PSQ.atMost (curTime - delta) $ bktNodes bucket
 
-    already_have = maybe False (const True) $ PSQ.lookup info (bktNodes bucket)
+    already_have = maybe False (const True) $ PSQ.lookup (coerce info) (bktNodes bucket)
 
     map_ns f = bucket { bktNodes = f (bktNodes bucket) }
     -- map_q f = bucket { bktQ = runIdentity \$ f (bktQ bucket) }
@@ -302,13 +315,13 @@ updateBucketForPingResult bad_node got_response bucket
     replacements | got_response     = [] -- Timestamp was already updated by TryInsert.
                  | Just info <- top = do
                                          -- Insert only if there's a removal.
-                                         _ <- maybeToList $ PSQ.lookup bad_node (bktNodes bucket)
+                                         _ <- maybeToList $ PSQ.lookup (coerce bad_node) (bktNodes bucket)
                                          return (bad_node, info)
                  | otherwise        = []
 
     replace (bad_node, (tm, info)) =
-                 PSQ.insert info tm
-                 . PSQ.delete bad_node
+                 PSQ.insert (coerce info) tm
+                 . PSQ.delete (coerce bad_node)
 
 
 updateStamps :: Ord ni => Timestamp -> [ni] -> PSQ ni Timestamp -> PSQ ni Timestamp
@@ -335,12 +348,12 @@ partitionQ imp test q0 = do
 
 
 split :: -- ( Eq ip , Ord (NodeId) , FiniteBits (NodeId)) =>
-    forall ni. Ord ni =>
+    forall ni s. ( Reifies s (Compare ni) ) =>
     (ni -> Word -> Bool)
-    -> BitIx -> Bucket ni -> (Bucket ni, Bucket ni)
+    -> BitIx -> Bucket s ni -> (Bucket s ni, Bucket s ni)
 split testNodeIdBit i b = (Bucket ns qs, Bucket ms rs)
   where
-    (ns,ms) = (PSQ.fromList *** PSQ.fromList) . partition (spanBit . key) . PSQ.toList $ bktNodes b
+    (ns,ms) = (PSQ.fromList *** PSQ.fromList) . partition (spanBit . coerce . key) . PSQ.toList $ bktNodes b
     (qs,rs) = runIdentity $ partitionQ bucketQ (spanBit . snd) $ bktQ b
 
     spanBit :: ni -> Bool
@@ -382,11 +395,23 @@ 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 BucketList ni = BucketList !ni !Int ![Bucket ni]
-    deriving Generic
-
-mapTable :: Ord ni => (a -> ni) -> BucketList a -> BucketList ni
-mapTable f (BucketList self n bs) = BucketList (f self) n (map (mapBucket f) bs)
+data BucketList ni = forall s. Reifies s (Compare ni) =>
+    BucketList { thisNode :: !ni
+               , bktCount :: !Int
+                 -- | Non-empty list of buckets.
+               , buckets :: [Bucket s ni]
+               }
+
+mapTable :: (b -> t) -> (t -> b) -> BucketList t -> BucketList b
+mapTable g f tbl@(BucketList self n bkts) = reify (contramapC g $ bucketCompare bkts)
+    $ \p -> BucketList
+        { thisNode = f self
+        , bktCount = n
+        , buckets = map (resolve p . mapBucket f) bkts
+        }
+ where
+    resolve :: Proxy s -> Bucket s ni -> Bucket s ni
+    resolve = const id
 
 instance (Eq ni) => Eq (BucketList ni) where
   (==) = (==) `on` Network.DHT.Routing.toList
@@ -422,14 +447,19 @@ instance Pretty (BucketList ni) where
       ss = shape t
 
 -- | Empty table with specified /spine/ node id.
-nullTable :: ni -> Int -> BucketList ni
-nullTable ni n = BucketList
+nullTable :: (ni -> ni -> Ordering) -> ni -> Int -> BucketList ni
+nullTable cmp ni n =
+    reify (Compare cmp)
+        $ \p -> BucketList
                     ni
                     (bucketCount (pred n))
-                    [Bucket PSQ.empty (runIdentity $ emptyQueue bucketQ)]
+                    [Bucket (empty p) (runIdentity $ emptyQueue bucketQ)]
   where
     bucketCount x = max 0 (min 159 x)
 
+    empty :: Proxy s -> PSQ (p s ni) t
+    empty = const $ PSQ.empty
+
 #if 0
 
 -- | Test if table is empty. In this case DHT should start
@@ -459,7 +489,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 :: BucketList ni -> [Int]
-shape = map (PSQ.size . bktNodes) . buckets
+shape (BucketList _ _ tbl) = map (PSQ.size . bktNodes) tbl
 
 #if 0
 
@@ -473,16 +503,18 @@ depth = L.length . shape
 
 #endif
 
-lookupBucket :: ( FiniteBits nid
+lookupBucket :: forall ni nid x.
+                ( FiniteBits nid
                 , Ord nid
-                ) => (ni -> nid) -> nid -> BucketList ni -> [Bucket ni]
-lookupBucket nodeId nid (BucketList self _ bkts) = go 0 [] bkts
+                ) => (ni -> nid) -> nid -> (forall s. [Bucket s ni] -> x) -> BucketList ni -> x
+lookupBucket nodeId nid kont (BucketList self _ bkts) = kont $ go 0 [] bkts
   where
-    d = complement (nid `xor` nodeId self)
+    d = nid `xor` nodeId self
 
+    go :: Word -> [Bucket s ni] -> [Bucket s ni] -> [Bucket s ni]
     go i bs (bucket : buckets)
-      |  testIdBit d i = bucket : buckets ++ bs
-      |     otherwise  = go (succ i) (bucket:bs) buckets
+      |  testIdBit d i = go (succ i) (bucket:bs) buckets
+      |     otherwise  = bucket : buckets ++ bs
     go _ bs []         = bs
 
 
@@ -563,8 +595,7 @@ kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
  where
     (bucket,everyone) =
         L.splitAt 1
-        . L.map (L.map PSQ.key . PSQ.toList . bktNodes)
-        . lookupBucket nodeId nid
+        . lookupBucket nodeId nid (L.map (coerce . L.map PSQ.key . PSQ.toList . bktNodes))
         $ tbl
 
 
@@ -574,9 +605,9 @@ kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
 -----------------------------------------------------------------------}
 
 splitTip :: -- ( Eq ip , Ord (NodeId) , FiniteBits (NodeId)) =>
-            Ord ni =>
+            ( Reifies s (Compare ni) ) =>
             (ni -> Word -> Bool)
-            -> ni -> Int -> BitIx -> Bucket ni -> [ Bucket ni ]
+            -> ni -> Int -> BitIx -> Bucket s ni -> [ Bucket s ni ]
 splitTip testNodeBit ni n i bucket
   | testNodeBit ni i = [zeros , ones ]
   |    otherwise     = [ones , zeros ]
@@ -590,25 +621,24 @@ 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, Bits nid) =>
-    (nid -> Word -> Bool)
-    -> (ni -> nid)
-    -> nid -> (Bucket ni -> Maybe (xs, Bucket ni)) -> BucketList ni -> Maybe (xs,BucketList ni)
-modifyBucket testIdBit nodeId nid f (BucketList self n bkts)
+    forall ni nid xs.
+    KademliaSpace nid ni
+    -> nid -> (forall s. Reifies s (Compare ni) => Bucket s ni -> Maybe (xs, Bucket s ni)) -> BucketList ni -> Maybe (xs,BucketList ni)
+modifyBucket space nid f (BucketList self n bkts)
     = second (BucketList self n) <$> go (0 :: BitIx) bkts
   where
-    d = nid `xor` nodeId self
+    d = kademliaXor space nid (kademliaLocation space self)
 
-    go :: BitIx -> [Bucket ni] -> Maybe (xs, [Bucket ni])
+    -- go :: BitIx -> [Bucket s ni] -> Maybe (xs, [Bucket s ni])
 
     go !i (bucket : buckets@(_:_))
-      | testIdBit d i   = second (bucket :)  <$> go (succ i) buckets
-      |       otherwise = second (: buckets) <$> f bucket
+      | kademliaTestBit space 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)
+                        <|>  go i (splitTip (kademliaTestBit space . kademliaLocation space) self n i bucket)
 
 
 -- | Triggering event for atomic table update
@@ -655,13 +685,12 @@ deriving instance ( Show ip
 --
 --  [ /tbl'/ ]         The updated routing 'BucketList'.
 --
-updateForInbound :: (Ord ni, Bits nid) =>
+updateForInbound ::
     KademliaSpace nid ni
     -> Timestamp -> ni -> BucketList ni -> (Bool, [ni], BucketList ni)
-updateForInbound space tm ni tbl =
+updateForInbound space tm ni tbl@(BucketList _ _ bkts) =
     maybe (False, [],tbl) (\(ps,tbl') -> (True, ps, tbl'))
-                   $ modifyBucket (kademliaTestBit space)
-                                   (kademliaLocation space)
+                   $ modifyBucket space
                                    (kademliaLocation space ni)
                                    (updateBucketForInbound tm ni)
                                    tbl
@@ -670,16 +699,15 @@ 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, Bits nid) =>
+updateForPingResult ::
     KademliaSpace nid ni
-    -> ni           -- ^ The pinged node.
-    -> Bool         -- ^ True if we got a reply, False if it timed out.
+    -> ni            -- ^ The pinged node.
+    -> Bool          -- ^ True if we got a reply, False if it timed out.
     -> BucketList ni -- ^ The routing table.
     -> ( [(ni,(Timestamp, ni))], BucketList ni )
 updateForPingResult space ni got_reply tbl =
     fromMaybe ([],tbl)
-                   $ modifyBucket (kademliaTestBit space)
-                                  (kademliaLocation space)
+                   $ modifyBucket space
                                   (kademliaLocation space ni)
                                   (updateBucketForPingResult ni got_reply)
                                   tbl
@@ -694,12 +722,8 @@ type TableEntry ni = (ni, Timestamp)
 tableEntry :: NodeEntry ni -> TableEntry ni
 tableEntry (a :-> b) = (a, b)
 
--- | Non-empty list of buckets.
-buckets :: BucketList ni -> [Bucket ni]
-buckets (BucketList _ _ bs) = bs
-
 toList :: BucketList ni -> [[TableEntry ni]]
-toList = L.map (L.map tableEntry . PSQ.toList . bktNodes) . buckets
+toList (BucketList _ _ bkts) = coerce $ L.map (L.map tableEntry . PSQ.toList . bktNodes) bkts
 
 data KademliaSpace nid ni = KademliaSpace
     { -- | Given a node record (probably including IP address), yields a
@@ -709,6 +733,8 @@ data KademliaSpace nid ni = KademliaSpace
       -- 'Data.Bits.testBit' except that the ordering of bits is reversed, so
       -- that 0 is the most significant bit.
     , kademliaTestBit :: nid -> Word -> Bool
+      -- | The Kademlia xor-metric.
+    , kademliaXor :: nid -> nid -> nid
     }
 
 contramapKS f ks = ks