#![allow(clippy::manual_div_ceil)]
use std::{
borrow::Borrow,
hash::{Hash, Hasher},
};
use libp2p_core::multihash::Multihash;
use libp2p_identity::PeerId;
use sha2::{
digest::generic_array::{typenum::U32, GenericArray},
Digest, Sha256,
};
use uint::*;
use crate::record;
construct_uint! {
pub struct U256(4);
}
#[derive(Clone, Copy, Debug)]
pub struct Key<T> {
preimage: T,
bytes: KeyBytes,
}
impl<T> Key<T> {
pub fn new(preimage: T) -> Key<T>
where
T: Borrow<[u8]>,
{
let bytes = KeyBytes::new(preimage.borrow());
Key { preimage, bytes }
}
pub fn preimage(&self) -> &T {
&self.preimage
}
pub fn into_preimage(self) -> T {
self.preimage
}
pub fn distance<U>(&self, other: &U) -> Distance
where
U: AsRef<KeyBytes>,
{
self.bytes.distance(other)
}
pub fn hashed_bytes(&self) -> &[u8] {
&self.bytes.0
}
pub fn for_distance(&self, d: Distance) -> KeyBytes {
self.bytes.for_distance(d)
}
}
impl<T> From<Key<T>> for KeyBytes {
fn from(key: Key<T>) -> KeyBytes {
key.bytes
}
}
impl<const S: usize> From<Multihash<S>> for Key<Multihash<S>> {
fn from(m: Multihash<S>) -> Self {
let bytes = KeyBytes(Sha256::digest(m.to_bytes()));
Key { preimage: m, bytes }
}
}
impl From<PeerId> for Key<PeerId> {
fn from(p: PeerId) -> Self {
let bytes = KeyBytes(Sha256::digest(p.to_bytes()));
Key { preimage: p, bytes }
}
}
impl From<Vec<u8>> for Key<Vec<u8>> {
fn from(b: Vec<u8>) -> Self {
Key::new(b)
}
}
impl From<record::Key> for Key<record::Key> {
fn from(k: record::Key) -> Self {
Key::new(k)
}
}
impl<T> AsRef<KeyBytes> for Key<T> {
fn as_ref(&self) -> &KeyBytes {
&self.bytes
}
}
impl<T, U> PartialEq<Key<U>> for Key<T> {
fn eq(&self, other: &Key<U>) -> bool {
self.bytes == other.bytes
}
}
impl<T> Eq for Key<T> {}
impl<T> Hash for Key<T> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.bytes.0.hash(state);
}
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub struct KeyBytes(GenericArray<u8, U32>);
impl KeyBytes {
pub fn new<T>(value: T) -> Self
where
T: Borrow<[u8]>,
{
KeyBytes(Sha256::digest(value.borrow()))
}
pub fn distance<U>(&self, other: &U) -> Distance
where
U: AsRef<KeyBytes>,
{
let a = U256::from_big_endian(self.0.as_slice());
let b = U256::from_big_endian(other.as_ref().0.as_slice());
Distance(a ^ b)
}
pub fn for_distance(&self, d: Distance) -> KeyBytes {
let key_int = U256::from_big_endian(self.0.as_slice()) ^ d.0;
KeyBytes(GenericArray::from(key_int.to_big_endian()))
}
}
impl AsRef<KeyBytes> for KeyBytes {
fn as_ref(&self) -> &KeyBytes {
self
}
}
#[derive(Copy, Clone, PartialEq, Eq, Default, PartialOrd, Ord, Debug)]
pub struct Distance(pub U256);
impl Distance {
pub fn ilog2(&self) -> Option<u32> {
(256 - self.0.leading_zeros()).checked_sub(1)
}
}
#[cfg(test)]
mod tests {
use quickcheck::*;
use super::*;
use crate::SHA_256_MH;
impl Arbitrary for Key<PeerId> {
fn arbitrary(_: &mut Gen) -> Key<PeerId> {
Key::from(PeerId::random())
}
}
impl Arbitrary for Key<Multihash<64>> {
fn arbitrary(g: &mut Gen) -> Key<Multihash<64>> {
let hash: [u8; 32] = core::array::from_fn(|_| u8::arbitrary(g));
Key::from(Multihash::wrap(SHA_256_MH, &hash).unwrap())
}
}
#[test]
fn identity() {
fn prop(a: Key<PeerId>) -> bool {
a.distance(&a) == Distance::default()
}
quickcheck(prop as fn(_) -> _)
}
#[test]
fn symmetry() {
fn prop(a: Key<PeerId>, b: Key<PeerId>) -> bool {
a.distance(&b) == b.distance(&a)
}
quickcheck(prop as fn(_, _) -> _)
}
#[test]
fn triangle_inequality() {
fn prop(a: Key<PeerId>, b: Key<PeerId>, c: Key<PeerId>) -> TestResult {
let ab = a.distance(&b);
let bc = b.distance(&c);
let (ab_plus_bc, overflow) = ab.0.overflowing_add(bc.0);
if overflow {
TestResult::discard()
} else {
TestResult::from_bool(a.distance(&c) <= Distance(ab_plus_bc))
}
}
quickcheck(prop as fn(_, _, _) -> _)
}
#[test]
fn unidirectionality() {
fn prop(a: Key<PeerId>, b: Key<PeerId>) -> bool {
let d = a.distance(&b);
(0..100).all(|_| {
let c = Key::from(PeerId::random());
a.distance(&c) != d || b == c
})
}
quickcheck(prop as fn(_, _) -> _)
}
}