pipecircuit 0.0.22

use core::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
    num::{NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64},
};

use super::IndexKey;

const SEQUENTIAL_INDEX_MULTIPLIER: u64 = 0x9E37_79B9_7F4A_7C15;

/// A compact, sequential identifier scoped by its row type.
#[repr(transparent)]
pub struct WarpId<Row> {
    value: NonZeroU64,
    row: PhantomData<fn() -> Row>,
}

impl<Row> WarpId<Row> {
    pub const MIN: Self = Self::new(NonZeroU64::MIN);

    pub const fn new(value: NonZeroU64) -> Self {
        Self {
            value,
            row: PhantomData,
        }
    }

    pub const fn value(self) -> NonZeroU64 {
        self.value
    }
}

#[cfg(feature = "serde")]
impl<Row> serde::Serialize for WarpId<Row> {
    fn serialize<Serializer: serde::Serializer>(
        &self,
        serializer: Serializer,
    ) -> Result<Serializer::Ok, Serializer::Error> {
        serde::Serialize::serialize(&self.value, serializer)
    }
}

#[cfg(feature = "serde")]
impl<'message, Row> serde::Deserialize<'message> for WarpId<Row> {
    fn deserialize<Deserializer: serde::Deserializer<'message>>(
        deserializer: Deserializer,
    ) -> Result<Self, Deserializer::Error> {
        <NonZeroU64 as serde::Deserialize<'message>>::deserialize(deserializer).map(Self::new)
    }
}

#[cfg(all(feature = "postcard-bindings", not(target_arch = "wasm32")))]
impl<Row> postcard_bindgen::__private::GenBinding for WarpId<Row> {
    fn get_type() -> postcard_bindgen::__private::ValueType {
        <NonZeroU64 as postcard_bindgen::__private::GenBinding>::get_type()
    }
}

impl<Row> Clone for WarpId<Row> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<Row> Copy for WarpId<Row> {}

impl<Row> fmt::Debug for WarpId<Row> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_tuple("WarpId").field(&self.value).finish()
    }
}

impl<Row> PartialEq for WarpId<Row> {
    fn eq(&self, other: &Self) -> bool {
        self.value == other.value
    }
}

impl<Row> Eq for WarpId<Row> {}

impl<Row> PartialOrd for WarpId<Row> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<Row> Ord for WarpId<Row> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.value.cmp(&other.value)
    }
}

impl<Row> Hash for WarpId<Row> {
    fn hash<State: Hasher>(&self, state: &mut State) {
        self.value.hash(state);
    }
}

impl<Row> IndexKey for WarpId<Row> {
    fn index_word(self) -> u64 {
        sequential_index_word(self.value)
    }
}

/// Distributes a sequential nonzero integer across the complete `u64` range.
pub const fn sequential_index_word(sequence: NonZeroU64) -> u64 {
    integer_index_word(sequence.get())
}

/// Distributes an unsigned integer across the complete `u64` range.
pub const fn integer_index_word(value: u64) -> u64 {
    value.wrapping_mul(SEQUENTIAL_INDEX_MULTIPLIER)
}

impl IndexKey for u8 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u16 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u32 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u64 {
    fn index_word(self) -> u64 {
        integer_index_word(self)
    }
}

impl IndexKey for NonZeroU8 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU16 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU32 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU64 {
    fn index_word(self) -> u64 {
        sequential_index_word(self)
    }
}

#[cfg(test)]
mod tests {
    use core::{mem::size_of, num::NonZeroU64};

    use super::{WarpId, integer_index_word, sequential_index_word};
    use crate::storage::warp_table::IndexKey;

    struct Row;

    #[test]
    fn warp_id_is_one_word_and_preserves_its_sequence() {
        let id = WarpId::<Row>::new(NonZeroU64::MIN);

        assert_eq!(size_of::<WarpId<Row>>(), size_of::<u64>());
        assert_eq!(size_of::<Option<WarpId<Row>>>(), size_of::<u64>());
        assert_eq!(id.value(), NonZeroU64::MIN);
        assert_eq!(id.index_word(), sequential_index_word(NonZeroU64::MIN));
        assert_eq!(7_u64.index_word(), integer_index_word(7));
    }

    #[test]
    fn identity_traits_preserve_the_exact_nonzero_sequence() -> Result<(), &'static str> {
        use alloc::{format, vec::Vec};
        use core::hash::{Hash, Hasher};

        #[derive(Default)]
        struct RecordedHash(Vec<u8>);

        impl Hasher for RecordedHash {
            fn finish(&self) -> u64 {
                0
            }

            fn write(&mut self, bytes: &[u8]) {
                self.0.extend_from_slice(bytes);
            }
        }

        fn cloned<Value: Clone>(value: &Value) -> Value {
            value.clone()
        }

        for raw in [1, 7, u64::MAX] {
            let value = NonZeroU64::new(raw).ok_or("zero identity")?;
            let id = WarpId::<Row>::new(value);
            assert_eq!(cloned(&id).value(), value);
            assert_eq!(format!("{id:?}"), format!("WarpId({raw})"));
            let mut hash = RecordedHash::default();
            id.hash(&mut hash);
            assert_eq!(hash.0, raw.to_ne_bytes());

            for other_raw in [1, 7, u64::MAX] {
                let other = WarpId::<Row>::new(NonZeroU64::new(other_raw).ok_or("zero identity")?);
                assert_eq!(id == other, raw == other_raw);
                assert_eq!(id.cmp(&other), raw.cmp(&other_raw));
                assert_eq!(id.partial_cmp(&other), Some(raw.cmp(&other_raw)));
            }
        }
        Ok(())
    }

    #[test]
    fn integer_index_keys_preserve_width_and_wrapping_permutation() {
        use core::num::{NonZeroU8, NonZeroU16, NonZeroU32};

        for (actual, expected) in [
            (0_u8.index_word(), 0),
            (1_u8.index_word(), 0x9e37_79b9_7f4a_7c15),
            (u8::MAX.index_word(), 0x9942_3fc5_cb31_98eb),
            (0_u16.index_word(), 0),
            (1_u16.index_word(), 0x9e37_79b9_7f4a_7c15),
            (u16::MAX.index_word(), 0xdb82_0590_fcca_83eb),
            (0_u32.index_word(), 0),
            (1_u32.index_word(), 0x9e37_79b9_7f4a_7c15),
            (u32::MAX.index_word(), 0xe113_025b_80b5_83eb),
            (0_u64.index_word(), 0),
            (7_u64.index_word(), 0x5384_5412_7b09_6493),
            (u64::MAX.index_word(), 0x61c8_8646_80b5_83eb),
            (NonZeroU8::MIN.index_word(), 0x9e37_79b9_7f4a_7c15),
            (NonZeroU8::MAX.index_word(), 0x9942_3fc5_cb31_98eb),
            (NonZeroU16::MIN.index_word(), 0x9e37_79b9_7f4a_7c15),
            (NonZeroU16::MAX.index_word(), 0xdb82_0590_fcca_83eb),
            (NonZeroU32::MIN.index_word(), 0x9e37_79b9_7f4a_7c15),
            (NonZeroU32::MAX.index_word(), 0xe113_025b_80b5_83eb),
            (NonZeroU64::MIN.index_word(), 0x9e37_79b9_7f4a_7c15),
            (NonZeroU64::MAX.index_word(), 0x61c8_8646_80b5_83eb),
        ] {
            assert_eq!(actual, expected);
        }
    }

    #[cfg(feature = "serde")]
    #[test]
    fn warp_id_encoding_preserves_the_sequence_without_serializing_the_row()
    -> Result<(), postcard::Error> {
        for value in [NonZeroU64::MIN, NonZeroU64::MAX] {
            let id = WarpId::<Row>::new(value);
            let bytes = postcard::to_allocvec(&id)?;
            assert_eq!(bytes, postcard::to_allocvec(&value)?);
            assert_eq!(postcard::from_bytes::<WarpId<Row>>(&bytes)?, id);
        }
        assert!(postcard::from_bytes::<WarpId<Row>>(&[0]).is_err());
        Ok(())
    }

    #[cfg(all(feature = "postcard-bindings", not(target_arch = "wasm32")))]
    #[test]
    fn warp_id_bindings_preserve_the_nonzero_wire_bound() {
        use postcard_bindgen::__private::GenBinding;

        assert_eq!(WarpId::<Row>::get_type(), NonZeroU64::get_type());
    }

    #[cfg(feature = "serde")]
    proptest::proptest! {
        #[test]
        fn warp_id_round_trips_every_nonzero_sequence(raw in 1_u64..=u64::MAX) {
            let value = NonZeroU64::new(raw)
                .ok_or_else(|| proptest::test_runner::TestCaseError::fail("generated zero"))?;
            let id = WarpId::<Row>::new(value);
            let bytes = postcard::to_allocvec(&id)?;
            let decoded = postcard::from_bytes::<WarpId<Row>>(&bytes)?;
            proptest::prop_assert_eq!(decoded, id);
            proptest::prop_assert_eq!(decoded.index_word(), id.index_word());
        }
    }
}