1 files changed, 36 insertions, 12 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index d380c6f2..396c4b1d 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -93,6 +93,7 @@ import Data.Reflection
 import Network.Address
 import Data.Typeable
 import Data.Coerce
+import Data.Hashable
 
 -- | Last time the node was responding to our queries.
 --
@@ -172,10 +173,11 @@ bucketQ :: QueueMethods Identity ni (BucketQueue ni)
 bucketQ = seqQ
 
 
-newtype Compare a = Compare (a -> a -> Ordering)
+data Compare a = Compare (a -> a -> Ordering) (Int -> a -> Int)
 
 contramapC :: (b -> a) -> Compare a -> Compare b
-contramapC f (Compare cmp) = Compare $ \a b -> cmp (f a) (f b)
+contramapC f (Compare cmp hsh) = Compare (\a b -> cmp (f a) (f b))
+                                         (\s x -> hsh s (f x))
 
 newtype Ordered' s a = Ordered a
  deriving (Show)
@@ -191,7 +193,11 @@ instance Reifies s (Compare a) => Eq (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)
+     where Compare cmp _ = reflect (Proxy :: Proxy s)
+
+instance Reifies s (Compare a) => Hashable (Ordered' s (Shrink a)) where
+    hashWithSalt salt x = hash salt (coerce x)
+     where Compare _ hash = 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
@@ -204,12 +210,18 @@ data Bucket s ni = Bucket
     , bktQ     :: !(BucketQueue (Timestamp,ni)) -- replacements pending time-outs
     } deriving (Generic)
 
+#define CAN_SHOW_BUCKET 0
+
+#if CAN_SHOW_BUCKET
 deriving instance Show ni => Show (Bucket s ni)
+#endif
 
 bucketCompare :: forall p ni s. Reifies s (Compare ni) => p (Bucket s ni) -> Compare ni
 bucketCompare _ = reflect (Proxy :: Proxy s)
 
-mapBucket :: Reifies t (Compare ni) => (a -> ni) -> Bucket s a -> Bucket t ni
+mapBucket :: ( Reifies s (Compare a)
+             , 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
@@ -248,10 +260,10 @@ instance (Eq ip, Ord (NodeId), Serialize (NodeId), Serialize ip, Serialize u) =>
 
 #endif
 
-psqFromPairList :: (Ord p, Ord k) => [(k, p)] -> PSQ k p
+psqFromPairList :: (Ord p, PSQKey k) => [(k, p)] -> PSQ k p
 psqFromPairList xs = PSQ.fromList $ map (\(a,b) -> a :-> b) xs
 
-psqToPairList :: PSQ t t1 -> [(t, t1)]
+psqToPairList :: ( PSQKey t, Ord t1 ) => PSQ t t1 -> [(t, t1)]
 psqToPairList psq = map (\(a :-> b) -> (a,b)) $ PSQ.toList psq
 
 -- | Update interval, in seconds.
@@ -330,7 +342,7 @@ updateBucketForPingResult bad_node got_response bucket
                  . PSQ.delete (coerce bad_node)
 
 
-updateStamps :: Ord ni => Timestamp -> [ni] -> PSQ ni Timestamp -> PSQ ni Timestamp
+updateStamps :: PSQKey ni => Timestamp -> [ni] -> PSQ ni Timestamp -> PSQ ni Timestamp
 updateStamps curTime stales nodes = foldl' (\q n -> PSQ.insert n curTime q) nodes stales
 
 type BitIx = Word
@@ -430,7 +442,19 @@ instance Serialize NominalDiffTime where
 
 #endif
 
+#if CAN_SHOW_BUCKET
 deriving instance (Show ni) => Show (BucketList ni)
+#else
+instance Show ni => Show (BucketList ni) where
+    showsPrec d (BucketList self cnt bkts) =
+        mappend "BucketList "
+        . showsPrec (d+1) self
+        . mappend " "
+        . showsPrec (d+1) cnt
+        . mappend " (fromList "
+        . showsPrec (d+1) (L.map (L.map tableEntry . PSQ.toList . bktNodes) $ bkts)
+        . mappend ") "
+#endif
 
 #if 0
 
@@ -455,9 +479,9 @@ instance Pretty (BucketList ni) where
 -- | Empty table with specified /spine/ node id.
 --
 -- XXX: The comparison function argument is awkward here.
-nullTable :: (ni -> ni -> Ordering) -> ni -> Int -> BucketList ni
-nullTable cmp ni n =
-    reify (Compare cmp)
+nullTable :: (ni -> ni -> Ordering) -> (Int -> ni -> Int) -> ni -> Int -> BucketList ni
+nullTable cmp hsh ni n =
+    reify (Compare cmp hsh)
         $ \p -> BucketList
                     ni
                     (bucketCount (pred n))
@@ -465,7 +489,7 @@ nullTable cmp ni n =
   where
     bucketCount x = max 0 (min 159 x)
 
-    empty :: Proxy s -> PSQ (p s ni) t
+    empty :: Reifies s (Compare ni) => Proxy s -> PSQ (Ordered s ni) Timestamp
     empty = const $ PSQ.empty
 
 #if 0
@@ -514,7 +538,7 @@ depth = L.length . shape
 lookupBucket :: forall ni nid x.
                 ( FiniteBits nid
                 , Ord nid
-                ) => (ni -> nid) -> nid -> (forall s. [Bucket s ni] -> x) -> BucketList ni -> x
+                ) => (ni -> nid) -> nid -> (forall s. Reifies s (Compare ni) => [Bucket s ni] -> x) -> BucketList ni -> x
 lookupBucket nodeId nid kont (BucketList self _ bkts) = kont $ go 0 [] bkts
   where
     d = nid `xor` nodeId self