From 0136af50253ea8d608367d0bc5514a65fb7b083a Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Fri, 28 Aug 2020 09:00:07 +0200
Subject: [PATCH] protocols/kad: Implement ilog2 for Distance (#1719)

In order to make use of the distances returned by `KBucketRef::range` in
a human readable format, one needs to be able to translate the 256 bit in
a `Distance` to a smaller space.
---
 protocols/kad/src/kbucket.rs     | 4 +---
 protocols/kad/src/kbucket/key.rs | 9 +++++++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/protocols/kad/src/kbucket.rs b/protocols/kad/src/kbucket.rs
index 6ab7e27a..be976884 100644
--- a/protocols/kad/src/kbucket.rs
+++ b/protocols/kad/src/kbucket.rs
@@ -106,9 +106,7 @@ impl BucketIndex {
     /// `local_key` is the `local_key` itself, which does not belong in any
     /// bucket.
     fn new(d: &Distance) -> Option<BucketIndex> {
-        (NUM_BUCKETS - d.0.leading_zeros() as usize)
-            .checked_sub(1)
-            .map(BucketIndex)
+        d.ilog2().map(|i| BucketIndex(i as usize))
     }
 
     /// Gets the index value as an unsigned integer.
diff --git a/protocols/kad/src/kbucket/key.rs b/protocols/kad/src/kbucket/key.rs
index 302d93a0..8d5ec616 100644
--- a/protocols/kad/src/kbucket/key.rs
+++ b/protocols/kad/src/kbucket/key.rs
@@ -169,6 +169,15 @@ impl AsRef<KeyBytes> for KeyBytes {
 #[derive(Copy, Clone, PartialEq, Eq, Default, PartialOrd, Ord, Debug)]
 pub struct Distance(pub(super) U256);
 
+impl Distance {
+    /// Returns the integer part of the base 2 logarithm of the [`Distance`].
+    ///
+    /// Returns `None` if the distance is zero.
+    pub fn ilog2(&self) -> Option<u32> {
+        (256 - self.0.leading_zeros()).checked_sub(1)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;