\section{K-buckets}

K-buckets is a data structure for efficiently storing a set of nodes close to a
certain key called the base key.  The base key is constant throughout the
lifetime of a k-buckets instance.

\begin{code}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns             #-}
{-# LANGUAGE StrictData                 #-}
{-# LANGUAGE Trustworthy                #-}
module Network.Tox.DHT.KBuckets where

import           Control.Applicative           (Applicative, (<$>))
import           Data.Binary                   (Binary)
import           Data.Foldable                 (toList)
import           Data.List                     (sortBy)
import           Data.Map                      (Map)
import qualified Data.Map                      as Map
import           Data.Maybe                    (isJust)
import           Data.Ord                      (comparing)
import           Data.Traversable              (Traversable, mapAccumR,
                                                traverse)
import           Data.Word                     (Word8)
import           Test.QuickCheck.Arbitrary     (Arbitrary, arbitrary)
import           Test.QuickCheck.Gen           (Gen)
import qualified Test.QuickCheck.Gen           as Gen

import           Network.Tox.Crypto.Key        (PublicKey)
import           Network.Tox.DHT.ClientList    (ClientList)
import qualified Network.Tox.DHT.ClientList    as ClientList
import           Network.Tox.DHT.Distance      (Distance)
import qualified Network.Tox.DHT.Distance      as Distance
import           Network.Tox.NodeInfo.NodeInfo (NodeInfo)
import qualified Network.Tox.NodeInfo.NodeInfo as NodeInfo
import           Network.Tox.Time              (Timestamp)


{-------------------------------------------------------------------------------
 -
 - :: Implementation.
 -
 ------------------------------------------------------------------------------}

\end{code}

A k-buckets is a map from small integers \texttt{0 <= n < 256} to Client Lists
of maximum size $k$. Each Client List is called a (k-)bucket. A k-buckets is
equipped with a base key, and each bucket has this key as its base key.
\texttt{k} is called the bucket size.  The default bucket size is 8.
A large bucket size was chosen to increase the speed at which peers are found.

\begin{code}

data KBuckets = KBuckets
  { KBuckets -> Int
bucketSize :: Int
  , KBuckets -> Map KBucketIndex ClientList
buckets    :: Map KBucketIndex ClientList
  , KBuckets -> PublicKey
baseKey    :: PublicKey
  }
  deriving (KBuckets -> KBuckets -> Bool
(KBuckets -> KBuckets -> Bool)
-> (KBuckets -> KBuckets -> Bool) -> Eq KBuckets
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: KBuckets -> KBuckets -> Bool
$c/= :: KBuckets -> KBuckets -> Bool
== :: KBuckets -> KBuckets -> Bool
$c== :: KBuckets -> KBuckets -> Bool
Eq, ReadPrec [KBuckets]
ReadPrec KBuckets
Int -> ReadS KBuckets
ReadS [KBuckets]
(Int -> ReadS KBuckets)
-> ReadS [KBuckets]
-> ReadPrec KBuckets
-> ReadPrec [KBuckets]
-> Read KBuckets
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [KBuckets]
$creadListPrec :: ReadPrec [KBuckets]
readPrec :: ReadPrec KBuckets
$creadPrec :: ReadPrec KBuckets
readList :: ReadS [KBuckets]
$creadList :: ReadS [KBuckets]
readsPrec :: Int -> ReadS KBuckets
$creadsPrec :: Int -> ReadS KBuckets
Read, Int -> KBuckets -> ShowS
[KBuckets] -> ShowS
KBuckets -> String
(Int -> KBuckets -> ShowS)
-> (KBuckets -> String) -> ([KBuckets] -> ShowS) -> Show KBuckets
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [KBuckets] -> ShowS
$cshowList :: [KBuckets] -> ShowS
show :: KBuckets -> String
$cshow :: KBuckets -> String
showsPrec :: Int -> KBuckets -> ShowS
$cshowsPrec :: Int -> KBuckets -> ShowS
Show)


defaultBucketSize :: Int
defaultBucketSize :: Int
defaultBucketSize = Int
8


empty :: PublicKey -> KBuckets
empty :: PublicKey -> KBuckets
empty = Int -> Map KBucketIndex ClientList -> PublicKey -> KBuckets
KBuckets Int
defaultBucketSize Map KBucketIndex ClientList
forall k a. Map k a
Map.empty

\end{code}

The above number \texttt{n} is the bucket index.  It is a non-negative integer
with the range \texttt{[0, 255]}, i.e. the range of an 8 bit unsigned integer.

\begin{code}


newtype KBucketIndex = KBucketIndex Word8
  deriving (KBucketIndex -> KBucketIndex -> Bool
(KBucketIndex -> KBucketIndex -> Bool)
-> (KBucketIndex -> KBucketIndex -> Bool) -> Eq KBucketIndex
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: KBucketIndex -> KBucketIndex -> Bool
$c/= :: KBucketIndex -> KBucketIndex -> Bool
== :: KBucketIndex -> KBucketIndex -> Bool
$c== :: KBucketIndex -> KBucketIndex -> Bool
Eq, Eq KBucketIndex
Eq KBucketIndex
-> (KBucketIndex -> KBucketIndex -> Ordering)
-> (KBucketIndex -> KBucketIndex -> Bool)
-> (KBucketIndex -> KBucketIndex -> Bool)
-> (KBucketIndex -> KBucketIndex -> Bool)
-> (KBucketIndex -> KBucketIndex -> Bool)
-> (KBucketIndex -> KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex -> KBucketIndex)
-> Ord KBucketIndex
KBucketIndex -> KBucketIndex -> Bool
KBucketIndex -> KBucketIndex -> Ordering
KBucketIndex -> KBucketIndex -> KBucketIndex
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: KBucketIndex -> KBucketIndex -> KBucketIndex
$cmin :: KBucketIndex -> KBucketIndex -> KBucketIndex
max :: KBucketIndex -> KBucketIndex -> KBucketIndex
$cmax :: KBucketIndex -> KBucketIndex -> KBucketIndex
>= :: KBucketIndex -> KBucketIndex -> Bool
$c>= :: KBucketIndex -> KBucketIndex -> Bool
> :: KBucketIndex -> KBucketIndex -> Bool
$c> :: KBucketIndex -> KBucketIndex -> Bool
<= :: KBucketIndex -> KBucketIndex -> Bool
$c<= :: KBucketIndex -> KBucketIndex -> Bool
< :: KBucketIndex -> KBucketIndex -> Bool
$c< :: KBucketIndex -> KBucketIndex -> Bool
compare :: KBucketIndex -> KBucketIndex -> Ordering
$ccompare :: KBucketIndex -> KBucketIndex -> Ordering
$cp1Ord :: Eq KBucketIndex
Ord, ReadPrec [KBucketIndex]
ReadPrec KBucketIndex
Int -> ReadS KBucketIndex
ReadS [KBucketIndex]
(Int -> ReadS KBucketIndex)
-> ReadS [KBucketIndex]
-> ReadPrec KBucketIndex
-> ReadPrec [KBucketIndex]
-> Read KBucketIndex
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [KBucketIndex]
$creadListPrec :: ReadPrec [KBucketIndex]
readPrec :: ReadPrec KBucketIndex
$creadPrec :: ReadPrec KBucketIndex
readList :: ReadS [KBucketIndex]
$creadList :: ReadS [KBucketIndex]
readsPrec :: Int -> ReadS KBucketIndex
$creadsPrec :: Int -> ReadS KBucketIndex
Read, Int -> KBucketIndex -> ShowS
[KBucketIndex] -> ShowS
KBucketIndex -> String
(Int -> KBucketIndex -> ShowS)
-> (KBucketIndex -> String)
-> ([KBucketIndex] -> ShowS)
-> Show KBucketIndex
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [KBucketIndex] -> ShowS
$cshowList :: [KBucketIndex] -> ShowS
show :: KBucketIndex -> String
$cshow :: KBucketIndex -> String
showsPrec :: Int -> KBucketIndex -> ShowS
$cshowsPrec :: Int -> KBucketIndex -> ShowS
Show, Integer -> KBucketIndex
KBucketIndex -> KBucketIndex
KBucketIndex -> KBucketIndex -> KBucketIndex
(KBucketIndex -> KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex)
-> (Integer -> KBucketIndex)
-> Num KBucketIndex
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> KBucketIndex
$cfromInteger :: Integer -> KBucketIndex
signum :: KBucketIndex -> KBucketIndex
$csignum :: KBucketIndex -> KBucketIndex
abs :: KBucketIndex -> KBucketIndex
$cabs :: KBucketIndex -> KBucketIndex
negate :: KBucketIndex -> KBucketIndex
$cnegate :: KBucketIndex -> KBucketIndex
* :: KBucketIndex -> KBucketIndex -> KBucketIndex
$c* :: KBucketIndex -> KBucketIndex -> KBucketIndex
- :: KBucketIndex -> KBucketIndex -> KBucketIndex
$c- :: KBucketIndex -> KBucketIndex -> KBucketIndex
+ :: KBucketIndex -> KBucketIndex -> KBucketIndex
$c+ :: KBucketIndex -> KBucketIndex -> KBucketIndex
Num, Get KBucketIndex
[KBucketIndex] -> Put
KBucketIndex -> Put
(KBucketIndex -> Put)
-> Get KBucketIndex
-> ([KBucketIndex] -> Put)
-> Binary KBucketIndex
forall t. (t -> Put) -> Get t -> ([t] -> Put) -> Binary t
putList :: [KBucketIndex] -> Put
$cputList :: [KBucketIndex] -> Put
get :: Get KBucketIndex
$cget :: Get KBucketIndex
put :: KBucketIndex -> Put
$cput :: KBucketIndex -> Put
Binary, Int -> KBucketIndex
KBucketIndex -> Int
KBucketIndex -> [KBucketIndex]
KBucketIndex -> KBucketIndex
KBucketIndex -> KBucketIndex -> [KBucketIndex]
KBucketIndex -> KBucketIndex -> KBucketIndex -> [KBucketIndex]
(KBucketIndex -> KBucketIndex)
-> (KBucketIndex -> KBucketIndex)
-> (Int -> KBucketIndex)
-> (KBucketIndex -> Int)
-> (KBucketIndex -> [KBucketIndex])
-> (KBucketIndex -> KBucketIndex -> [KBucketIndex])
-> (KBucketIndex -> KBucketIndex -> [KBucketIndex])
-> (KBucketIndex -> KBucketIndex -> KBucketIndex -> [KBucketIndex])
-> Enum KBucketIndex
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: KBucketIndex -> KBucketIndex -> KBucketIndex -> [KBucketIndex]
$cenumFromThenTo :: KBucketIndex -> KBucketIndex -> KBucketIndex -> [KBucketIndex]
enumFromTo :: KBucketIndex -> KBucketIndex -> [KBucketIndex]
$cenumFromTo :: KBucketIndex -> KBucketIndex -> [KBucketIndex]
enumFromThen :: KBucketIndex -> KBucketIndex -> [KBucketIndex]
$cenumFromThen :: KBucketIndex -> KBucketIndex -> [KBucketIndex]
enumFrom :: KBucketIndex -> [KBucketIndex]
$cenumFrom :: KBucketIndex -> [KBucketIndex]
fromEnum :: KBucketIndex -> Int
$cfromEnum :: KBucketIndex -> Int
toEnum :: Int -> KBucketIndex
$ctoEnum :: Int -> KBucketIndex
pred :: KBucketIndex -> KBucketIndex
$cpred :: KBucketIndex -> KBucketIndex
succ :: KBucketIndex -> KBucketIndex
$csucc :: KBucketIndex -> KBucketIndex
Enum)


\end{code}

\subsection{Bucket Index}

The index of the bucket can be computed using the following function:
\texttt{bucketIndex(baseKey, nodeKey) = 255 - log\_2(distance(baseKey,
nodeKey))}.  This function is not defined when \texttt{baseKey == nodeKey},
meaning k-buckets will never contain a Node Info about the base node.

Thus, each k-bucket contains only Node Infos for whose keys the following
holds: if node with key \texttt{nodeKey} is in k-bucket with index \texttt{n},
then \texttt{bucketIndex(baseKey, nodeKey) == n}. Thus, n'th k-bucket consists
of nodes for which distance to the base node lies in range
\verb![2^n, 2^(n+1) - 1]!.

The bucket index can be efficiently computed by determining the first bit at
which the two keys differ, starting from the most significant bit.  So, if the
local DHT key starts with e.g. \texttt{0x80} and the bucketed node key starts
with \texttt{0x40}, then the bucket index for that node is 0.  If the second
bit differs, the bucket index is 1.  If the keys are almost exactly equal and
only the last bit differs, the bucket index is 255.

\begin{code}


bucketIndex :: PublicKey -> PublicKey -> Maybe KBucketIndex
bucketIndex :: PublicKey -> PublicKey -> Maybe KBucketIndex
bucketIndex PublicKey
pk1 PublicKey
pk2 =
  (Int -> KBucketIndex) -> Maybe Int -> Maybe KBucketIndex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
index -> KBucketIndex
255 KBucketIndex -> KBucketIndex -> KBucketIndex
forall a. Num a => a -> a -> a
- Int -> KBucketIndex
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
index) (Maybe Int -> Maybe KBucketIndex)
-> Maybe Int -> Maybe KBucketIndex
forall a b. (a -> b) -> a -> b
$ Distance -> Maybe Int
Distance.log2 (Distance -> Maybe Int) -> Distance -> Maybe Int
forall a b. (a -> b) -> a -> b
$ PublicKey -> PublicKey -> Distance
Distance.xorDistance PublicKey
pk1 PublicKey
pk2


\end{code}

\subsection{Manipulating k-buckets}

TODO: this is different from kademlia's least-recently-seen eviction policy; why
the existing solution was chosen, how does it affect security, performance and
resistance to poisoning? original paper claims that preference of old live nodes
results in better persistence and resistance to basic DDoS attacks;

Any update or lookup operation on a k-buckets instance that involves a single
node requires us to first compute the bucket index for that node.  An update
involving a Node Info with \texttt{nodeKey == baseKey} has no effect.  If the
update results in an empty bucket, that bucket is removed from the map.

\begin{code}


updateBucketForKey :: KBuckets -> PublicKey -> (ClientList -> ClientList) -> KBuckets
updateBucketForKey :: KBuckets -> PublicKey -> (ClientList -> ClientList) -> KBuckets
updateBucketForKey KBuckets
kBuckets PublicKey
key ClientList -> ClientList
f =
  case PublicKey -> PublicKey -> Maybe KBucketIndex
bucketIndex (KBuckets -> PublicKey
baseKey KBuckets
kBuckets) PublicKey
key of
    Maybe KBucketIndex
Nothing    -> KBuckets
kBuckets
    Just KBucketIndex
index -> KBuckets -> KBucketIndex -> (ClientList -> ClientList) -> KBuckets
updateBucketForIndex KBuckets
kBuckets KBucketIndex
index ClientList -> ClientList
f


updateBucketForIndex :: KBuckets -> KBucketIndex -> (ClientList -> ClientList) -> KBuckets
updateBucketForIndex :: KBuckets -> KBucketIndex -> (ClientList -> ClientList) -> KBuckets
updateBucketForIndex kBuckets :: KBuckets
kBuckets@KBuckets { Map KBucketIndex ClientList
buckets :: Map KBucketIndex ClientList
buckets :: KBuckets -> Map KBucketIndex ClientList
buckets, PublicKey
baseKey :: PublicKey
baseKey :: KBuckets -> PublicKey
baseKey, Int
bucketSize :: Int
bucketSize :: KBuckets -> Int
bucketSize } KBucketIndex
index ClientList -> ClientList
f =
  let
    -- Find the old bucket or create a new empty one.
    updatedBucket :: ClientList
updatedBucket = ClientList -> ClientList
f (ClientList -> ClientList) -> ClientList -> ClientList
forall a b. (a -> b) -> a -> b
$ ClientList
-> KBucketIndex -> Map KBucketIndex ClientList -> ClientList
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (PublicKey -> Int -> ClientList
ClientList.empty PublicKey
baseKey Int
bucketSize) KBucketIndex
index Map KBucketIndex ClientList
buckets
    -- Replace old bucket with updated bucket or delete if empty.
    updatedBuckets :: Map KBucketIndex ClientList
updatedBuckets =
      if ClientList -> Bool
ClientList.isEmpty ClientList
updatedBucket
      then KBucketIndex
-> Map KBucketIndex ClientList -> Map KBucketIndex ClientList
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete KBucketIndex
index Map KBucketIndex ClientList
buckets
      else KBucketIndex
-> ClientList
-> Map KBucketIndex ClientList
-> Map KBucketIndex ClientList
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert KBucketIndex
index ClientList
updatedBucket Map KBucketIndex ClientList
buckets
  in
  KBuckets
kBuckets { buckets :: Map KBucketIndex ClientList
buckets = Map KBucketIndex ClientList
updatedBuckets }


\end{code}

Adding a node to, or removing a node from, a k-buckets consists of performing
the corresponding operation on the Client List bucket whose index is that of
the node's public key, except that adding a new node to a full bucket has no
effect.  A node is considered \textit{viable} for entry if the corresponding
bucket is not full.

\begin{code}

addNode :: Timestamp -> NodeInfo -> KBuckets -> KBuckets
addNode :: Timestamp -> NodeInfo -> KBuckets -> KBuckets
addNode Timestamp
time NodeInfo
nodeInfo KBuckets
kBuckets =
  KBuckets -> PublicKey -> (ClientList -> ClientList) -> KBuckets
updateBucketForKey KBuckets
kBuckets PublicKey
publicKey ((ClientList -> ClientList) -> KBuckets)
-> (ClientList -> ClientList) -> KBuckets
forall a b. (a -> b) -> a -> b
$ \ClientList
clientList ->
    let
      full :: Bool
full = ClientList -> Bool
ClientList.full ClientList
clientList
      alreadyIn :: Bool
alreadyIn = Maybe NodeInfo -> Bool
forall a. Maybe a -> Bool
isJust (Maybe NodeInfo -> Bool) -> Maybe NodeInfo -> Bool
forall a b. (a -> b) -> a -> b
$ PublicKey -> ClientList -> Maybe NodeInfo
ClientList.lookup PublicKey
publicKey ClientList
clientList
    in
    if Bool -> Bool
not Bool
full Bool -> Bool -> Bool
|| Bool
alreadyIn
      then Timestamp -> NodeInfo -> ClientList -> ClientList
ClientList.addNode Timestamp
time NodeInfo
nodeInfo ClientList
clientList
      else ClientList
clientList
  where
    publicKey :: PublicKey
publicKey = NodeInfo -> PublicKey
NodeInfo.publicKey NodeInfo
nodeInfo

removeNode :: PublicKey -> KBuckets -> KBuckets
removeNode :: PublicKey -> KBuckets -> KBuckets
removeNode PublicKey
publicKey KBuckets
kBuckets =
  KBuckets -> PublicKey -> (ClientList -> ClientList) -> KBuckets
updateBucketForKey KBuckets
kBuckets PublicKey
publicKey ((ClientList -> ClientList) -> KBuckets)
-> (ClientList -> ClientList) -> KBuckets
forall a b. (a -> b) -> a -> b
$ PublicKey -> ClientList -> ClientList
ClientList.removeNode PublicKey
publicKey

viable :: NodeInfo -> KBuckets -> Bool
viable :: NodeInfo -> KBuckets -> Bool
viable NodeInfo
nodeInfo KBuckets{ PublicKey
baseKey :: PublicKey
baseKey :: KBuckets -> PublicKey
baseKey, Map KBucketIndex ClientList
buckets :: Map KBucketIndex ClientList
buckets :: KBuckets -> Map KBucketIndex ClientList
buckets } =
  case PublicKey -> PublicKey -> Maybe KBucketIndex
bucketIndex PublicKey
baseKey (PublicKey -> Maybe KBucketIndex)
-> PublicKey -> Maybe KBucketIndex
forall a b. (a -> b) -> a -> b
$ NodeInfo -> PublicKey
NodeInfo.publicKey NodeInfo
nodeInfo of
    Maybe KBucketIndex
Nothing    -> Bool
False
    Just KBucketIndex
index -> case KBucketIndex -> Map KBucketIndex ClientList -> Maybe ClientList
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup KBucketIndex
index Map KBucketIndex ClientList
buckets of
      Maybe ClientList
Nothing     -> Bool
True
      Just ClientList
bucket -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ClientList -> Bool
ClientList.full ClientList
bucket

\end{code}

Iteration order of a k-buckets instance is in order of distance from the base
key.  I.e. the first node seen in iteration is the closest, and the last node
is the furthest away in terms of the distance metric.

\begin{code}

traverseClientLists ::
    Applicative f => (ClientList -> f ClientList) -> KBuckets -> f KBuckets
traverseClientLists :: (ClientList -> f ClientList) -> KBuckets -> f KBuckets
traverseClientLists ClientList -> f ClientList
f kBuckets :: KBuckets
kBuckets@KBuckets{ Map KBucketIndex ClientList
buckets :: Map KBucketIndex ClientList
buckets :: KBuckets -> Map KBucketIndex ClientList
buckets } =
  (\Map KBucketIndex ClientList
x -> KBuckets
kBuckets{ buckets :: Map KBucketIndex ClientList
buckets = Map KBucketIndex ClientList
x }) (Map KBucketIndex ClientList -> KBuckets)
-> f (Map KBucketIndex ClientList) -> f KBuckets
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ClientList -> f ClientList)
-> Map KBucketIndex ClientList -> f (Map KBucketIndex ClientList)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ClientList -> f ClientList
f (Map KBucketIndex ClientList -> Map KBucketIndex ClientList
forall (t :: * -> *) a. Traversable t => t a -> t a
reverseT Map KBucketIndex ClientList
buckets)
  where
    reverseT :: (Traversable t) => t a -> t a
    reverseT :: t a -> t a
reverseT t a
t = ([a], t a) -> t a
forall a b. (a, b) -> b
snd (([a] -> a -> ([a], a)) -> [a] -> t a -> ([a], t a)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (\ (a
x:[a]
xs) a
_ -> ([a]
xs, a
x)) (t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t a
t) t a
t)

closeNodes :: PublicKey -> KBuckets -> [ (Distance, NodeInfo) ]
closeNodes :: PublicKey -> KBuckets -> [(Distance, NodeInfo)]
closeNodes PublicKey
publicKey KBuckets{ PublicKey
baseKey :: PublicKey
baseKey :: KBuckets -> PublicKey
baseKey, Map KBucketIndex ClientList
buckets :: Map KBucketIndex ClientList
buckets :: KBuckets -> Map KBucketIndex ClientList
buckets } =
  let
    (Map KBucketIndex ClientList
further, Maybe ClientList
at, Map KBucketIndex ClientList
nearer) = case PublicKey -> PublicKey -> Maybe KBucketIndex
bucketIndex PublicKey
baseKey PublicKey
publicKey of
      Maybe KBucketIndex
Nothing    -> (Map KBucketIndex ClientList
buckets, Maybe ClientList
forall a. Maybe a
Nothing, Map KBucketIndex ClientList
forall k a. Map k a
Map.empty)
      Just KBucketIndex
index -> KBucketIndex
-> Map KBucketIndex ClientList
-> (Map KBucketIndex ClientList, Maybe ClientList,
    Map KBucketIndex ClientList)
forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
Map.splitLookup KBucketIndex
index Map KBucketIndex ClientList
buckets
    clientClose :: ClientList -> [(Distance, NodeInfo)]
clientClose = PublicKey -> ClientList -> [(Distance, NodeInfo)]
ClientList.closeNodes PublicKey
publicKey
    bucketsClose :: [ClientList] -> [(Distance, NodeInfo)]
bucketsClose = ((Distance, NodeInfo) -> (Distance, NodeInfo) -> Ordering)
-> [(Distance, NodeInfo)] -> [(Distance, NodeInfo)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Distance, NodeInfo) -> Distance)
-> (Distance, NodeInfo) -> (Distance, NodeInfo) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Distance, NodeInfo) -> Distance
forall a b. (a, b) -> a
fst) ([(Distance, NodeInfo)] -> [(Distance, NodeInfo)])
-> ([ClientList] -> [(Distance, NodeInfo)])
-> [ClientList]
-> [(Distance, NodeInfo)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ClientList -> [(Distance, NodeInfo)])
-> [ClientList] -> [(Distance, NodeInfo)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ClientList -> [(Distance, NodeInfo)]
clientClose
  in
    [[(Distance, NodeInfo)]] -> [(Distance, NodeInfo)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ [(Distance, NodeInfo)]
-> (ClientList -> [(Distance, NodeInfo)])
-> Maybe ClientList
-> [(Distance, NodeInfo)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ClientList -> [(Distance, NodeInfo)]
clientClose Maybe ClientList
at
      , [ClientList] -> [(Distance, NodeInfo)]
bucketsClose ([ClientList] -> [(Distance, NodeInfo)])
-> [ClientList] -> [(Distance, NodeInfo)]
forall a b. (a -> b) -> a -> b
$ Map KBucketIndex ClientList -> [ClientList]
forall k a. Map k a -> [a]
Map.elems Map KBucketIndex ClientList
nearer
      , [ClientList] -> [(Distance, NodeInfo)]
bucketsClose ([ClientList] -> [(Distance, NodeInfo)])
-> [ClientList] -> [(Distance, NodeInfo)]
forall a b. (a -> b) -> a -> b
$ Map KBucketIndex ClientList -> [ClientList]
forall k a. Map k a -> [a]
Map.elems Map KBucketIndex ClientList
further
      ]


{-------------------------------------------------------------------------------
 -
 - :: Tests.
 -
 ------------------------------------------------------------------------------}


getAllNodes :: KBuckets -> [NodeInfo]
getAllNodes :: KBuckets -> [NodeInfo]
getAllNodes =
  (ClientList -> [NodeInfo]) -> [ClientList] -> [NodeInfo]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ClientList -> [NodeInfo]
ClientList.nodeInfos ([ClientList] -> [NodeInfo])
-> (KBuckets -> [ClientList]) -> KBuckets -> [NodeInfo]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map KBucketIndex ClientList -> [ClientList]
forall k a. Map k a -> [a]
Map.elems (Map KBucketIndex ClientList -> [ClientList])
-> (KBuckets -> Map KBucketIndex ClientList)
-> KBuckets
-> [ClientList]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KBuckets -> Map KBucketIndex ClientList
buckets


genKBuckets :: PublicKey -> Gen KBuckets
genKBuckets :: PublicKey -> Gen KBuckets
genKBuckets PublicKey
publicKey =
  (KBuckets -> (Timestamp, NodeInfo) -> KBuckets)
-> KBuckets -> [(Timestamp, NodeInfo)] -> KBuckets
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (((Timestamp, NodeInfo) -> KBuckets -> KBuckets)
-> KBuckets -> (Timestamp, NodeInfo) -> KBuckets
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((Timestamp, NodeInfo) -> KBuckets -> KBuckets)
 -> KBuckets -> (Timestamp, NodeInfo) -> KBuckets)
-> ((Timestamp, NodeInfo) -> KBuckets -> KBuckets)
-> KBuckets
-> (Timestamp, NodeInfo)
-> KBuckets
forall a b. (a -> b) -> a -> b
$ (Timestamp -> NodeInfo -> KBuckets -> KBuckets)
-> (Timestamp, NodeInfo) -> KBuckets -> KBuckets
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Timestamp -> NodeInfo -> KBuckets -> KBuckets
addNode) (PublicKey -> KBuckets
empty PublicKey
publicKey) ([(Timestamp, NodeInfo)] -> KBuckets)
-> Gen [(Timestamp, NodeInfo)] -> Gen KBuckets
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Timestamp, NodeInfo) -> Gen [(Timestamp, NodeInfo)]
forall a. Gen a -> Gen [a]
Gen.listOf Gen (Timestamp, NodeInfo)
forall a. Arbitrary a => Gen a
arbitrary


instance Arbitrary KBuckets where
  arbitrary :: Gen KBuckets
arbitrary = Gen PublicKey
forall a. Arbitrary a => Gen a
arbitrary Gen PublicKey -> (PublicKey -> Gen KBuckets) -> Gen KBuckets
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PublicKey -> Gen KBuckets
genKBuckets
\end{code}