ckey 0.4.3

// Copyright (C) 2024 Christian Mauduit <ufoot@ufoot.org>

pub(crate) const DATA_U64_SIZE: usize = 4;
#[cfg(feature = "serde")]
pub(crate) const DATA_U8_SIZE: usize = 32;

#[cfg(feature = "serde")]
use bincode::Options;
#[cfg(feature = "rand")]
use rand_core::{OsRng, RngCore};
#[cfg(feature = "serde")]
use serde::Serialize;
#[cfg(feature = "serde")]
use sha2::Digest;
#[cfg(feature = "serde")]
use sha2::Sha256;
use std::cmp::Ordering;
use std::hash::Hash;

/// CKey implements a consistent hash key.
///
/// Theory here: <https://en.wikipedia.org/wiki/Consistent_hashing>
///
/// # Examples
///
/// ```
/// use ckey::CKey;
///
/// let k1 = CKey::from(0.1);
/// let k2 = k1.next();
/// let k3 = k2 + 10u8;
/// assert!(k2.inside(k1, k3));
/// let k4 = CKey::from(1000u16);
/// assert_eq!("0.015258789", format!("{}", k4));
/// ```
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Default, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CKey {
    pub(crate) data: [u64; DATA_U64_SIZE],
}

impl CKey {
    /// Make a digest from bytes and returns a new key from it.
    ///
    /// Typical usage is to create a a CKey from a string.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::digest("/a/unique/path");
    /// assert_eq!("3b818742b0f27cfc10f4de1e5ba5594c975d580c412a31011ad58d36a2b17cdc", format!("{:?}",k));
    /// ```
    #[cfg(feature = "serde")]
    pub fn digest(data: impl AsRef<[u8]>) -> Self {
        let mut hasher = Sha256::default();
        hasher.update(data);
        let hash = hasher.finalize();
        bincode::deserialize(&hash).unwrap()
    }

    /// Make a hash from serialized data and returns a new key from it.
    ///
    /// Typical usage is to create a a CKey from some object content.
    ///
    /// Note: this is not using the Hash trait, and instead serializes the complete
    /// data. You could achieve "build a CKey from a hashable object" easily
    /// by doing a digest on the hash. This is not recommended, as you could get many
    /// more collisions, and it defeats the purpose of having 256-bit keys.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    /// use serde::Serialize;
    ///
    /// #[derive(Serialize)]
    /// struct Obj {
    ///     a: usize,
    ///     b: usize,
    /// }
    ///
    /// let obj = Obj{a: 10, b: 100};
    /// let k = CKey::serial(obj);
    /// assert_eq!("fc6326e9af50c0ae08d34921f61afa012988998fac0269dbcf2dbe30748ed4eb", format!("{:?}",k));
    /// ```
    ///
    /// ```
    /// // DONT'DO THIS
    /// use ckey::CKey;
    /// use std::collections::hash_map::DefaultHasher;
    /// use bincode::serialize;
    /// use std::hash::{Hash, Hasher};
    ///
    /// #[derive(Hash)]
    /// struct Obj {
    ///     a: usize,
    ///     b: usize,
    /// }
    ///
    /// let obj = Obj{a: 10, b: 100};
    /// let mut s = DefaultHasher::new();
    /// obj.hash(&mut s);
    /// let h = s.finish();
    /// let data = serialize(&h).unwrap();
    /// let k = CKey::digest(data);
    /// // The code above works, but has 2 caveats:
    /// // 1 - since DefaultHasher is randomized, results are unpredictable;
    /// // 2 - risks of collision, as we are only using 64 bits.
    /// println!("{}", k);
    /// ```
    #[cfg(feature = "serde")]
    pub fn serial<T>(v: T) -> Self
    where
        T: Serialize,
    {
        let data = bincode::serialize(&v).unwrap();
        Self::digest(data)
    }

    /// Generate a random key.
    ///
    /// This can be time consuming as it is using
    /// the OS random generation, which is better from a cryptographic
    /// point of view, but possibly slower than a standard prng.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::rand();
    /// print!("k: {:?}", k);
    /// ```
    #[cfg(feature = "rand")]
    #[cfg(feature = "serde")]
    pub fn rand() -> Self {
        let mut hash = [0u8; DATA_U8_SIZE];
        OsRng.fill_bytes(&mut hash);
        bincode::deserialize(&hash).unwrap()
    }

    /// Generate a key with the value 0.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::zero();
    /// assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", format!("{:?}", k));
    /// ```
    pub fn zero() -> Self {
        Self::default()
    }

    /// Generate a key with the value 1, the smallest key possible just after zero.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::unit();
    /// assert_eq!("0000000000000000000000000000000000000000000000000000000000000001", format!("{:?}", k));
    /// ```
    pub fn unit() -> Self {
        Self { data: [0, 0, 0, 1] }
    }

    /// Generate a key with the last, highest possible value.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::last();
    /// assert_eq!("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", format!("{:?}", k));
    /// ```
    pub fn last() -> Self {
        Self {
            data: [
                0xffffffffffffffff,
                0xffffffffffffffff,
                0xffffffffffffffff,
                0xffffffffffffffff,
            ],
        }
    }

    /// Generate a key which is exactly in the middle of the ring.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::halfway();
    /// assert_eq!("8000000000000000000000000000000000000000000000000000000000000000", format!("{:?}", k));
    /// ```
    pub fn halfway() -> Self {
        Self {
            data: [
                0x8000000000000000,
                0x0000000000000000,
                0x0000000000000000,
                0x0000000000000000,
            ],
        }
    }

    /// Returns true if self is inside lower and upper.
    /// Lower limit is excluded, upper is included.
    /// This is typically used in a ring to know if a given
    /// key is handled by a node. You would ask if key
    /// is inside previous node (lower) and this node (upper)
    /// and if true -> yes, this is the right node to handle it.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k1 = CKey::from(0.1);
    /// let k2 = CKey::from(0.3);
    /// let k3 = CKey::from(0.9);
    /// assert!(k2.inside(k1, k3));
    /// ```
    pub fn inside(self, lower: CKey, upper: CKey) -> bool {
        if lower == upper {
            // If lower and upper are equal, consider it matches all the ring,
            // so any value is inside the interval.
            return true;
        }
        // Upper limit is included within the interval. Lower is not.
        if self == lower {
            return false;
        }
        if self == upper {
            return true;
        }
        // Now compare member by member, since we made a diff initially,
        // self and upper are correctly ordered, according to lower.
        let lower_to_self = self - lower;
        let lower_to_upper = upper - lower;
        for i in 0..DATA_U64_SIZE {
            if lower_to_self.data[i] > lower_to_upper.data[i] {
                return false;
            }
            if lower_to_self.data[i] < lower_to_upper.data[i] {
                return true;
            }
        }
        // This should be unreachable code as to get here, we need to
        // have self == upper and this is a special case was tested before.
        true
    }

    /// Increment current key by one.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let mut k = CKey::zero();
    /// k.incr();
    /// assert_eq!(CKey::unit(), k);
    /// ```
    pub fn incr(&mut self) {
        *self = self.next();
    }

    /// Decrement current key by one.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let mut k = CKey::zero();
    /// k.decr();
    /// assert_eq!(CKey::last(), k);
    /// ```
    pub fn decr(&mut self) {
        *self = self.prev();
    }

    /// Return current key plus one.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::zero();
    /// assert_eq!(CKey::unit(), k.next());
    /// ```
    pub fn next(self) -> CKey {
        self + Self::unit()
    }

    /// Return current key minus one.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::zero();
    /// assert_eq!(CKey::last(), k.prev());
    /// ```
    pub fn prev(self) -> CKey {
        self - Self::unit()
    }

    /// Try to convert from a &str.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::parse("12345678ffffffff01234567fffffffff00123456ffffffff0012345ffffffff").unwrap();
    /// assert_eq!("12345678ffffffff01234567fffffffff00123456ffffffff0012345ffffffff", format!("{:?}", k));
    /// ```
    pub fn parse(value: &str) -> Result<CKey, <Self as std::convert::TryFrom<&str>>::Error> {
        if value.len() != 64 {
            return Err(hex::FromHexError::InvalidStringLength);
        }
        Ok(CKey {
            data: [
                u64::from_str_radix(&value[0..16], 16).unwrap(),
                u64::from_str_radix(&value[16..32], 16).unwrap(),
                u64::from_str_radix(&value[32..48], 16).unwrap(),
                u64::from_str_radix(&value[48..64], 16).unwrap(),
            ],
        })
    }

    /// Set value from bytes, big endian.
    ///
    /// The data can be longer or shorter than 32 bytes.
    /// Extra bytes will be ignored.
    /// Missing bytes will be replaced by zeroes.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    /// let bytes: [u8; 3] = [4, 2, 1];
    /// let k = CKey::set(bytes);
    /// assert_eq!("0402010000000000000000000000000000000000000000000000000000000000", format!("{:?}", k));
    /// ```
    #[cfg(feature = "serde")]
    pub fn set(data: impl AsRef<[u8]>) -> Self {
        let unserializer = bincode::DefaultOptions::new()
            .with_fixint_encoding()
            .with_big_endian()
            .allow_trailing_bytes();
        if data.as_ref().len() < DATA_U8_SIZE {
            let mut buf: Vec<u8> = Vec::new();
            for byte in data.as_ref().iter() {
                buf.push(*byte);
            }
            while buf.len() < DATA_U8_SIZE {
                buf.push(0);
            }
            unserializer.deserialize(&buf).unwrap()
        } else {
            unserializer.deserialize(data.as_ref()).unwrap()
        }
    }

    /// Get value as bytes, big endian.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let k = CKey::from(42u8);
    /// let bytes = k.bytes();
    /// assert_eq!(32, bytes.len());
    /// assert_eq!(42, bytes[0]);
    /// for byte in &bytes[1..] {
    ///     assert_eq!(0, *byte);
    /// }
    /// ```
    #[cfg(feature = "serde")]
    pub fn bytes(&self) -> Vec<u8> {
        let serializer = bincode::DefaultOptions::new()
            .with_fixint_encoding()
            .with_big_endian()
            .allow_trailing_bytes();
        serializer.serialize(&self).unwrap()
    }

    /// Sort an array of keys, using a start reference.
    ///
    /// There is no way to sort keys in an absolute way,
    /// since they are virtually on a circle and wrapping,
    /// everything depends on where you start from.
    ///
    /// This is why there's a start parameter.
    ///
    /// # Examples
    /// ```
    /// use ckey::CKey;
    ///
    /// let start = CKey::from(0.7);
    /// let k1 = CKey::from(0.2);
    /// let k2 = CKey::from(0.4);
    /// let k3 = CKey::from(0.3);
    /// let k4 = CKey::from(0.9);
    /// let mut sorted = vec![k1, k2, k3, k4];
    /// CKey::sort(&mut sorted, start);
    /// assert_eq!(4, sorted.len());
    /// assert_eq!("0.900000000", format!("{}", sorted[0]));
    /// assert_eq!("0.200000000", format!("{}", sorted[1]));
    /// assert_eq!("0.300000000", format!("{}", sorted[2]));
    /// assert_eq!("0.400000000", format!("{}", sorted[3]));
    /// ```
    pub fn sort(v: &mut Vec<CKey>, start: CKey) {
        v.sort_by(|a, b| {
            // So here, the logic is to consider that everything is
            // "normalized" so that the start is "key 0" and then we
            // look at the order *after* this starting point.
            //
            // In practice, self is lower than other if self is between
            // start and other.
            //
            // And self is greater than other if other is between start
            // and self.
            if *a == *b {
                Ordering::Equal
            } else if a.inside(start, *b) {
                // Other key is after us, considering the "zero" is start.
                Ordering::Less
            } else {
                Ordering::Greater
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::CKey;
    #[cfg(feature = "serde")]
    use serde_json::json;

    #[test]
    fn test_ckey_inside() {
        let k_zero = CKey::zero();
        let k_unit = CKey::unit();
        let k_halfway = CKey::halfway();
        let k_last = CKey::last();
        let k2 = CKey::from(0.2);
        let k3 = CKey::from(0.3);
        let k9 = CKey::from(0.9);

        assert!(k3.inside(k2, k9));
        assert!(!k2.inside(k3, k9));
        assert!(k2.inside(k9, k3));
        assert!(!k3.inside(k9, k2));
        assert!(k3.inside(k2, k2));
        assert!(k3.inside(k2, k3), "upper bound is included");
        assert!(!k2.inside(k2, k3), "lower bound is excluded");
        assert!(k_unit.inside(k_zero, k2));
        assert!(k_zero.inside(k9, k_zero));
        assert!(k_last.inside(k9, k_zero));
        assert!(k_zero.inside(k_last, k_zero));
        assert!(!k_last.inside(k_last, k_zero));
        assert!(k_halfway.inside(k3, k9));
    }

    #[test]
    fn test_ckey_next() {
        assert_eq!(
            "0000000000000000000000000000000000000000000000000000000000000001",
            format!("{:?}", CKey::zero().next())
        );
        assert_eq!(
            "0000000000000000000000000000000000000000000000000000000000000002",
            format!("{:?}", CKey::unit().next())
        );
        assert_eq!(
            "8000000000000000000000000000000000000000000000000000000000000001",
            format!("{:?}", CKey::halfway().next())
        );
        assert_eq!(
            "0000000000000000000000000000000000000000000000000000000000000000",
            format!("{:?}", CKey::last().next())
        );
    }

    #[test]
    fn test_ckey_prev() {
        assert_eq!(
            "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
            format!("{:?}", CKey::zero().prev())
        );
        assert_eq!(
            "0000000000000000000000000000000000000000000000000000000000000000",
            format!("{:?}", CKey::unit().prev())
        );
        assert_eq!(
            "7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
            format!("{:?}", CKey::halfway().prev())
        );
        assert_eq!(
            "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe",
            format!("{:?}", CKey::last().prev())
        );
    }

    #[test]
    fn test_ckey_sort() {
        let k_zero = CKey::zero();
        let k_unit = CKey::unit();
        let k_halfway = CKey::halfway();
        let k_last = CKey::last();
        let k2 = CKey::from(0.2);
        let k3 = CKey::from(0.3);
        let k9 = CKey::from(0.9);

        let mut sorted = vec![k2, k9, k3, k_halfway, k_unit, k_last, k_unit, k_zero];
        CKey::sort(&mut sorted, k3);
        assert_eq!(8, sorted.len());
        assert_eq!(k_halfway, sorted[0]);
        assert_eq!(k9, sorted[1]);
        assert_eq!(k_last, sorted[2]);
        assert_eq!(k_zero, sorted[3]);
        assert_eq!(k_unit, sorted[4]);
        assert_eq!(k_unit, sorted[5]);
        assert_eq!(k2, sorted[6]);
        assert_eq!(k3, sorted[7]);
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_json() {
        let key_halfway = CKey::halfway();
        let js_halfway = json!(&key_halfway).to_string();
        assert_eq!("{\"data\":[9223372036854775808,0,0,0]}", js_halfway);

        let key_from_js_halfway: CKey = serde_json::from_str(js_halfway.as_str()).unwrap();
        assert_eq!(key_halfway, key_from_js_halfway);
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_set_bytes() {
        for n in 0..100usize {
            let mut v: Vec<u8> = Vec::new();
            for i in 0..n {
                v.push(i as u8);
            }
            let k = CKey::set(&v);
            let bytes = k.bytes();
            assert_eq!(32, bytes.len());
            for i in 0..std::cmp::min(n, bytes.len()) {
                assert_eq!(i, bytes[i] as usize);
            }
        }
    }
}